Load functions along with module.

This allows concurrent execution of different kernels (different function or different module).

See https://discourse.llvm.org/t/how-to-lower-the-combination-of-async-gpu-ops-in-gpu-dialect/72796/17.

Author: chsigg

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

@@ -14,7 +14,8 @@
 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 
-#include <stdio.h>
+#include <cstdio>
+#include <vector>
 
 #include "cuda.h"
 #include "cuda_bf16.h"
@@ -56,14 +57,10 @@
 
 thread_local static int32_t defaultDevice = 0;
 
-const char *kDebugEnvironmentVariable = "MLIR_CUDA_DEBUG";
-
 /// Helper method that checks environment value for debugging.
 bool isDebugEnabled() {
-  static bool isInitialized = false;
-  static bool isEnabled = false;
-  if (!isInitialized)
-    isEnabled = getenv(kDebugEnvironmentVariable) != nullptr;
+  const char *kDebugEnvironmentVariable = "MLIR_CUDA_DEBUG";
+  static bool isEnabled = getenv(kDebugEnvironmentVariable) != nullptr;
   return isEnabled;
 }
 
@@ -125,6 +122,16 @@ mgpuModuleLoad(void *data, size_t /*gpuBlobSize*/) {
   ScopedContext scopedContext;
   CUmodule module = nullptr;
   CUDA_REPORT_IF_ERROR(cuModuleLoadData(&module, data));
+  // Preload functions in the module so that the first call to
+  // cuModuleGetFunction below doesn't synchronize context.
+  unsigned numFunctions = 0;
+  CUDA_REPORT_IF_ERROR(cuModuleGetFunctionCount(&numFunctions, module));
+  std::vector<CUfunction> functions(numFunctions);
+  CUDA_REPORT_IF_ERROR(
+      cuModuleEnumerateFunctions(functions.data(), numFunctions, module));
+  for (CUfunction function : functions) {
+    CUDA_REPORT_IF_ERROR(cuFuncLoad(function));
+  }
   return module;
 }