Use cooperative_groups::this_grid().sync() in CINN_GRID_REDUCE_IMPL. Add

CINN_GRID_REDUCE_FP16_MACRO.

Author: YuhanXu

backends/metax_gpu/cinn/compiler/compiler.cc

@@ -39,6 +39,7 @@ namespace metax {
 // ============================================================
 static const char* kMacaRuntimeSource = R"MACA_SOURCE(
 #pragma once
+#include <cooperative_groups.h>
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
@@ -812,7 +813,7 @@ __device__ inline void __cinn_grid_sync() {
 }
 
 #define CINN_GRID_REDUCE_IMPL(REDUCE_TYPE, init_value, DTYPE)               \
-  __cinn_grid_sync();                                                        \
+  cooperative_groups::this_grid().sync();                                    \
   DTYPE tmp_val = init_value;                                                \
   for (int y = 0; y < gridDim.y; y++) {                                      \
     tmp_val = cinn_##REDUCE_TYPE(tmp_val, mem[y * spatial_size + spatial_index]); \
@@ -830,7 +831,28 @@ EXPAND_REDUCE_INT64_MACRO(CINN_GRID_REDUCE_MACRO)
 EXPAND_REDUCE_FP32_MACRO(CINN_GRID_REDUCE_MACRO)
 EXPAND_REDUCE_FP64_MACRO(CINN_GRID_REDUCE_MACRO)
 EXPAND_REDUCE_BOOL_MACRO(CINN_GRID_REDUCE_MACRO)
-EXPAND_REDUCE_FP16_MACRO(CINN_GRID_REDUCE_MACRO)
+
+// FP16 grid reduce: accumulate in FP32 to avoid precision loss when summing
+// multiple FP16 block-level partial sums. Each partial sum can have magnitude
+// O(block_size * input_scale), and accumulating N such values in FP16 incurs
+// error proportional to N * magnitude * eps_fp16. Using FP32 for the inter-
+// block accumulation step keeps the error at FP16 quantization level only.
+#define CINN_GRID_REDUCE_FP16_MACRO(FP16_TYPE, FP32_FUNC, INIT_VAL)           \
+  __device__ inline float16 cinn_grid_reduce_##FP16_TYPE(                      \
+      const float16 *mem, int spatial_size, int spatial_index) {               \
+    cooperative_groups::this_grid().sync();                                     \
+    float tmp_val = (float)(INIT_VAL);                                          \
+    for (int y = 0; y < gridDim.y; y++) {                                       \
+      tmp_val = FP32_FUNC(                                                       \
+          tmp_val, __half2float(mem[y * spatial_size + spatial_index]));        \
+    }                                                                            \
+    return __float2half(tmp_val);                                                \
+  }
+
+CINN_GRID_REDUCE_FP16_MACRO(sum_fp16,  cinn_sum_fp32,  0.0f)
+CINN_GRID_REDUCE_FP16_MACRO(prod_fp16, cinn_prod_fp32, 1.0f)
+CINN_GRID_REDUCE_FP16_MACRO(max_fp16,  cinn_max_fp32,  -65504.0f)
+CINN_GRID_REDUCE_FP16_MACRO(min_fp16,  cinn_min_fp32,   65504.0f)
 
 __device__ inline bool cinn_grid_reduce_update_semaphore(int *semaphores) {
   __shared__ bool done;

backends/metax_gpu/cinn/runtime/cinn_runtime.cc

@@ -34,15 +34,15 @@ C_Status MetaxModuleLoad(void* dev_ptr, const char* path, void** mod_out) {
     return C_Status::C_FAILED;
   }
   *mod_out = reinterpret_cast<void*>(module);
-  std::cerr << "[MetaxModuleLoad] OK path=" << path << " module=" << module
-            << std::endl;
+  // std::cerr << "[MetaxModuleLoad] OK path=" << path << " module=" << module
+  //           << std::endl;
   return C_Status::C_SUCCESS;
 }
 
 // Unload module
 C_Status MetaxModuleUnload(void* dev_ptr, void* module_handle) {
   cuModuleUnload((CUmodule)module_handle);
-  std::cout << "YUHAN!!! [MetaxModuleUnload] module_handle=" << module_handle << std::endl;
+  // std::cout << "YUHAN!!! [MetaxModuleUnload] module_handle=" << module_handle << std::endl;
   return C_Status::C_SUCCESS;
 }
 
@@ -59,8 +59,8 @@ C_Status MetaxGetKernelAddress(void* dev_ptr,
     return C_Status::C_FAILED;
   }
   *func_out = reinterpret_cast<void*>(func);
-  std::cout << "YUHAN!!! [MetaxGetKernelAddress] OK func_name=" << func_name
-            << " func_ptr=" << func << " module_handle=" << module_handle << std::endl;
+  // std::cout << "YUHAN!!! [MetaxGetKernelAddress] OK func_name=" << func_name
+  //           << " func_ptr=" << func << " module_handle=" << module_handle << std::endl;
   return C_Status::C_SUCCESS;
 }
 
@@ -109,7 +109,7 @@ C_Status MetaxLaunchCooperativeKernel(void* dev_ptr,
                                       int bz,
                                       int shm,
                                       void* stream) {
-  std::cout << "YUHAN!!! [MetaxLaunchCooperativeKernel] func_ptr=" << func_ptr;
+  // std::cout << "YUHAN!!! [MetaxLaunchCooperativeKernel] func_ptr=" << func_ptr;
   CUmodule module;
   CUresult errModule = cuFuncGetModule(&module ,static_cast<CUfunction>(func_ptr));
   if (errModule != CUDA_SUCCESS) {