issue/130: use collapse

src/infiniop/ops/gemm/cpu/gemm_cpu.cc

@@ -42,34 +42,32 @@ void calculate(
     if (info.is_transed) {
         std::swap(a, b);
     }
-#pragma omp parallel for
-    for (ptrdiff_t index = 0; index < ptrdiff_t(info.batch * info.m * info.n); ++index) {
-        size_t ind = index;
-        size_t n_ = ind % info.n;
-        ind /= info.n;
-        size_t m_ = ind % info.m;
-        ind /= info.m;
-        size_t i = ind;
-        auto c_ = reinterpret_cast<Tdata *>(c) + i * info.c_matrix.stride + m_ * info.c_matrix.row_stride + n_ * info.c_matrix.col_stride;
-        float sum = 0;
-        for (size_t k_ = 0; k_ < info.k; ++k_) {
-            auto a_ = reinterpret_cast<const Tdata *>(a) + i * info.a_matrix.stride + m_ * info.a_matrix.row_stride + k_ * info.a_matrix.col_stride;
-            auto b_ = reinterpret_cast<const Tdata *>(b) + i * info.b_matrix.stride + n_ * info.b_matrix.col_stride + k_ * info.b_matrix.row_stride;
-            if constexpr (std::is_same<Tdata, fp16_t>::value) {
-                sum += utils::cast<float>(*a_) * utils::cast<float>(*b_);
-            } else {
-                sum += *a_ * (*b_);
+#pragma omp parallel for collapse(3)
+    for (size_t i = 0; i < info.batch; ++i) {
+        for (size_t m_ = 0; m_ < info.m; ++m_) {
+            for (size_t n_ = 0; n_ < info.n; ++n_) {
+                auto c_ = reinterpret_cast<Tdata *>(c) + i * info.c_matrix.stride + m_ * info.c_matrix.row_stride + n_ * info.c_matrix.col_stride;
+                float sum = 0;
+                for (size_t k_ = 0; k_ < info.k; ++k_) {
+                    auto a_ = reinterpret_cast<const Tdata *>(a) + i * info.a_matrix.stride + m_ * info.a_matrix.row_stride + k_ * info.a_matrix.col_stride;
+                    auto b_ = reinterpret_cast<const Tdata *>(b) + i * info.b_matrix.stride + n_ * info.b_matrix.col_stride + k_ * info.b_matrix.row_stride;
+                    if constexpr (std::is_same<Tdata, fp16_t>::value) {
+                        sum += utils::cast<float>(*a_) * utils::cast<float>(*b_);
+                    } else {
+                        sum += *a_ * (*b_);
+                    }
+                }
+                if constexpr (std::is_same<Tdata, fp16_t>::value) {
+                    if (beta == 0) {
+                        *c_ = utils::cast<fp16_t>(alpha * sum);
+                    } else {
+                        *c_ = utils::cast<fp16_t>(beta * utils::cast<float>(*c_) + alpha * sum);
+                    }
+                } else {
+                    *c_ = beta * (*c_) + alpha * sum;
+                }
             }
         }
-        if constexpr (std::is_same<Tdata, fp16_t>::value) {
-            if (beta == 0) {
-                *c_ = utils::cast<fp16_t>(alpha * sum);
-            } else {
-                *c_ = utils::cast<fp16_t>(beta * utils::cast<float>(*c_) + alpha * sum);
-            }
-        } else {
-            *c_ = beta * (*c_) + alpha * sum;
-        }
     }
 }