Metax MetaxLaunchCooperativeKernel, print more debug info

Author: YuhanXu

backends/metax_gpu/cinn/compiler/compiler.cc

@@ -1269,6 +1269,10 @@ C_Status MetaxCompile(void* dev_ptr,
     src_file << code;
     src_file.close();
   }
+  // std::cout << "[MetaX] src_file content written to: " << src_path
+  //             << "\n--- BEGIN src_file ---\n"
+  //             << kMacaRuntimeSource << "\n" << code
+  //             << "\n--- END src_file ---" << std::endl;
 
   // 2. Resolve compiler binary path
   const char* maca_path_env = std::getenv("MACA_PATH");

backends/metax_gpu/cinn/runtime/cinn_runtime.cc

@@ -42,6 +42,7 @@ C_Status MetaxModuleLoad(void* dev_ptr, const char* path, void** mod_out) {
 // 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;
   return C_Status::C_SUCCESS;
 }
 
@@ -58,8 +59,8 @@ C_Status MetaxGetKernelAddress(void* dev_ptr,
     return C_Status::C_FAILED;
   }
   *func_out = reinterpret_cast<void*>(func);
-  std::cerr << "[MetaxGetKernelAddress] OK func_name=" << func_name
-            << " func_ptr=" << func << 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;
 }
 
@@ -108,17 +109,17 @@ C_Status MetaxLaunchCooperativeKernel(void* dev_ptr,
                                       int bz,
                                       int shm,
                                       void* stream) {
-  std::cout << "YUHAN!!! [MetaxLaunchCooperativeKernel] func_ptr=" << func_ptr
-            << " grid=(" << gx << "," << gy << "," << gz << ")"
-            << " block=(" << bx << "," << by << "," << bz << ")"
-            << " shm=" << shm << std::endl;
+  std::cout << "YUHAN!!! [MetaxLaunchCooperativeKernel] func_ptr=" << func_ptr;
+  CUmodule module;
+  CUresult errModule = cuFuncGetModule(&module ,static_cast<CUfunction>(func_ptr));
+  if (errModule != CUDA_SUCCESS) {
+    std::cerr << "[MetaxLaunchCooperativeKernel] FAILED Module error=" << errModule
+              << std::endl;
+    return C_Status::C_FAILED;
+  }
   CUresult err = cuLaunchCooperativeKernel(static_cast<CUfunction>(func_ptr),
-                                           gx,
-                                           gy,
-                                           gz,
-                                           bx,
-                                           by,
-                                           bz,
+                                           gx, gy, gz,
+                                           bx, by, bz,
                                            shm,
                                            static_cast<CUstream>(stream),
                                            args);