issue/1031fxi avg_pool1d cross_entropy omp

Author: PanZezhong1725

src/infiniop/ops/avg_pool1d/cpu/avg_pool1d_cpu.cc

@@ -36,33 +36,35 @@ infiniStatus_t calculateAvgPool1d(const AvgPool1dInfo &info,
                                   const T *x) {
     const float inv_kernel = 1.0f / static_cast<float>(info.kernel_size);
 
-#pragma omp parallel for collapse(2)
-    for (ptrdiff_t b = 0; b < ptrdiff_t(info.batch); ++b) {
-        for (ptrdiff_t c = 0; c < ptrdiff_t(info.channels); ++c) {
+#pragma omp parallel for
+    for (ptrdiff_t bc = 0; bc < ptrdiff_t(info.batch * info.channels); ++bc) {
 
-            size_t y_base = b * info.y_stride_batch + c * info.y_stride_channel;
-            size_t x_base = b * info.x_stride_batch + c * info.x_stride_channel;
+        ptrdiff_t b = bc / info.channels;
+        ptrdiff_t c = bc % info.channels;
 
-            for (size_t ow = 0; ow < info.out_width; ++ow) {
-                size_t y_offset = y_base + ow * info.y_stride_width;
+        size_t y_base = b * info.y_stride_batch + c * info.y_stride_channel;
+        size_t x_base = b * info.x_stride_batch + c * info.x_stride_channel;
 
-                long long start_w = static_cast<long long>(ow * info.stride) - info.padding;
-                long long end_w = start_w + info.kernel_size;
+        for (size_t ow = 0; ow < info.out_width; ++ow) {
+            size_t y_offset = y_base + ow * info.y_stride_width;
 
-                long long valid_start = std::max(0LL, start_w);
-                long long valid_end = std::min(static_cast<long long>(info.in_width), end_w);
+            long long start_w = static_cast<long long>(ow * info.stride) - info.padding;
+            long long end_w = start_w + info.kernel_size;
 
-                float sum = 0.0f;
-                for (long long iw = valid_start; iw < valid_end; ++iw) {
-                    size_t x_offset = x_base + iw * info.x_stride_width;
-                    sum += utils::cast<float>(x[x_offset]);
-                }
+            long long valid_start = std::max(0LL, start_w);
+            long long valid_end = std::min(static_cast<long long>(info.in_width), end_w);
 
-                const float avg = sum * inv_kernel;
-                y[y_offset] = utils::cast<T>(avg);
+            float sum = 0.0f;
+            for (long long iw = valid_start; iw < valid_end; ++iw) {
+                size_t x_offset = x_base + iw * info.x_stride_width;
+                sum += utils::cast<float>(x[x_offset]);
             }
+
+            const float avg = sum * inv_kernel;
+            y[y_offset] = utils::cast<T>(avg);
         }
     }
+
     return INFINI_STATUS_SUCCESS;
 }
 

src/infiniop/ops/cross_entropy/cpu/cross_entropy_cpu.cc

@@ -41,7 +41,7 @@ infiniStatus_t cross_entropy_kernel(const CrossEntropyInfo *info,
     const Tidx *label = reinterpret_cast<const Tidx *>(target);
 
 #pragma omp parallel for
-    for (size_t i = 0; i < info->outer_size; ++i) {
+    for (ptrdiff_t i = 0; i < info->outer_size; ++i) {
         const T *row = x + i * info->x_stride;
         Tidx idx = label[i];