[MetaX] Implement cooperative kernel launch for CINN plugin

Author: YuhanXu

backends/metax_gpu/cinn/cinn_interface.cc

@@ -67,6 +67,20 @@ extern C_Status MetaxLaunchKernel(void* dev_ptr,
                                   int shm,
                                   void* stream);
 
+// Launches a cooperative kernel function (grid-level sync)
+extern C_Status MetaxLaunchCooperativeKernel(void* dev_ptr,
+                                             void* func_ptr,
+                                             void** args,
+                                             int num_args,
+                                             int gx,
+                                             int gy,
+                                             int gz,
+                                             int bx,
+                                             int by,
+                                             int bz,
+                                             int shm,
+                                             void* stream);
+
 // --- From passes/pass_manager.cc ---
 // Applies custom graph optimization passes
 extern C_Status MetaxApplyCustomPass(void* dev_ptr, void* ir_module);
@@ -99,6 +113,7 @@ void InitCinnInterface(C_DeviceInterface* device_interface) {
   metax_cinn_impl.module_unload = MetaxModuleUnload;
   metax_cinn_impl.get_kernel_address = MetaxGetKernelAddress;
   metax_cinn_impl.launch_kernel = MetaxLaunchKernel;
+  metax_cinn_impl.launch_cooperative_kernel = MetaxLaunchCooperativeKernel;
 
   // 6. Register Compilation Strategy interface
   metax_cinn_impl.apply_custom_pass = MetaxApplyCustomPass;

backends/metax_gpu/cinn/runtime/cinn_runtime.cc

@@ -82,6 +82,33 @@ C_Status MetaxLaunchKernel(void* dev_ptr,
   return C_Status::C_SUCCESS;
 }
 
+// Launch cooperative kernel: equivalent to cuLaunchCooperativeKernel
+C_Status MetaxLaunchCooperativeKernel(void* dev_ptr,
+                                      void* func_ptr,
+                                      void** args,
+                                      int num_args,
+                                      int gx,
+                                      int gy,
+                                      int gz,
+                                      int bx,
+                                      int by,
+                                      int bz,
+                                      int shm,
+                                      void* stream) {
+  CUresult err = cuLaunchCooperativeKernel((CUfunction)func_ptr,
+                                           gx,
+                                           gy,
+                                           gz,
+                                           bx,
+                                           by,
+                                           bz,
+                                           shm,
+                                           (CUstream)stream,
+                                           args);
+  if (err != CUDA_SUCCESS) return C_Status::C_FAILED;
+  return C_Status::C_SUCCESS;
+}
+
 }  // namespace metax
 }  // namespace custom_device
 }  // namespace paddle