More tuning improvements

external/llvm-project/mlir/lib/IR/MLIRContext.cpp

@@ -76,7 +76,7 @@ struct MLIRContextOptions {
 static llvm::ManagedStatic<MLIRContextOptions> clOptions;
 
 static bool isThreadingGloballyDisabled() {
-#if MLIR_ENABLE_THREADS != 0
+#if LLVM_ENABLE_THREADS != 0
   return clOptions.isConstructed() && clOptions->disableThreading;
 #else
   return true;

mlir/tools/rocmlir-tuning-driver/rocmlir-tuning-driver.cpp

@@ -47,6 +47,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/ThreadPool.h"
 
 #include <atomic>
 #include <cassert>
@@ -88,6 +89,24 @@ void pArgs(const std::tuple<Ts...> &formals, void **_vargs) {
 using namespace mlir;
 using namespace rocmlir::tuningdriver;
 
+//===----------------------------------------------------------------------===//
+// Shared Resources for Multi-threaded Compilation
+//===----------------------------------------------------------------------===//
+
+/// Returns a shared dialect registry, initialized exactly once.
+static DialectRegistry &getSharedDialectRegistry() {
+  static std::once_flag initFlag;
+  static DialectRegistry registry;
+  std::call_once(initFlag, []() { registerRocMLIRDialects(registry); });
+  return registry;
+}
+
+/// Returns a shared LLVM ThreadPool for all MLIR contexts.
+static llvm::DefaultThreadPool &getSharedThreadPool() {
+  static llvm::DefaultThreadPool pool;
+  return pool;
+}
+
 static llvm::cl::opt<std::string> inputFilename{
     llvm::cl::Positional, llvm::cl::desc("<input file>"), llvm::cl::init("-")};
 
@@ -316,9 +335,13 @@ struct ThreadResources {
                   const rock::KernelOptions &applicabilityOpts,
                   const rock::KernelOptions &compilationKernOpts,
                   const rock::BackendOptions &backendOpts) {
-    DialectRegistry registry;
-    registerRocMLIRDialects(registry);
-    ctx = std::make_unique<MLIRContext>(registry);
+    // Use the shared dialect registry (initialized exactly once)
+    DialectRegistry &registry = getSharedDialectRegistry();
+    // Create context with threading disabled internally, attach shared pool
+    ctx = std::make_unique<MLIRContext>(registry,
+                                        MLIRContext::Threading::DISABLED);
+    ctx->setThreadPool(getSharedThreadPool());
+    ctx->loadAllAvailableDialects();
     ctx->getDiagEngine().registerHandler([](Diagnostic &) {});
 
     // Pre-build pipelines once per thread
@@ -423,12 +446,12 @@ measureLargeKernel(unsigned iterations, hipStream_t stream,
 }
 
 // In order to match rocprof, returns time in nanoseconds
-static FailureOr<double>
-benchmarkKernels(ArrayRef<std::string> binaries,
-                 ArrayRef<std::string> funcNames, ArrayRef<uint32_t> blockSizes,
-                 ArrayRef<uint32_t> gridSizes, ArrayRef<void *> hostBuffers,
-                 MutableArrayRef<void *> gpuBuffers,
-                 ArrayRef<size_t> bufferSizes, const BenchmarkParams &params) {
+static FailureOr<double> benchmarkKernels(ArrayRef<std::string> binaries,
+                                          ArrayRef<std::string> funcNames,
+                                          ArrayRef<uint32_t> blockSizes,
+                                          ArrayRef<uint32_t> gridSizes,
+                                          MutableArrayRef<void *> gpuBuffers,
+                                          const BenchmarkParams &params) {
   bool benchmarkMode = !params.benchmarkConfig.empty();
   hipStream_t stream;
   HIPCHECK(hipStreamCreate(&stream));
@@ -440,12 +463,6 @@ benchmarkKernels(ArrayRef<std::string> binaries,
     }
   });
 
-  // Initialize device buffers
-  for (size_t i = 0; i < bufferSizes.size(); i++) {
-    HIPCHECK(hipMemcpyAsync(gpuBuffers[i], hostBuffers[i], bufferSizes[i],
-                            hipMemcpyHostToDevice, stream));
-  }
-
   // HIP wants an array of pointers to each argument
   std::vector<void *> argPointers;
   for (void *&item : gpuBuffers) {
@@ -731,6 +748,12 @@ static LogicalResult runTuningLoop(ModuleOp source) {
     gpuBuffers.push_back(gpuBuffer);
   }
 
+  // Copy host buffers to GPU once (reused across all config benchmarks)
+  for (size_t i = 0; i < bufferLengths.size(); i++) {
+    HIPCHECK(hipMemcpy(gpuBuffers[i], hostBuffers[i], bufferLengths[i],
+                       hipMemcpyHostToDevice));
+  }
+
   // 4. Multi-iteration tuning loop
   SmallString<64> bestConfigOverall;
   float bestTimeOverall = std::numeric_limits<float>::max();
@@ -975,10 +998,9 @@ static LogicalResult runTuningLoop(ModuleOp source) {
       assert(result.status == CompilationStatus::Success &&
              "Unexpected compilation status in benchmarking phase");
 
-      FailureOr<double> timing =
-          benchmarkKernels(result.hipModules, kernelFuncNames,
-                           result.blockSizes, result.gridSizes, hostBuffers,
-                           gpuBuffers, bufferLengths, benchmarkParams);
+      FailureOr<double> timing = benchmarkKernels(
+          result.hipModules, kernelFuncNames, result.blockSizes,
+          result.gridSizes, gpuBuffers, benchmarkParams);
 
       if (failed(timing)) {
         llvm::errs() << "Kernel execution failed\n";