[Mlir-commits] [mlir] [mlir][gpu] Add optional attributes of kernelModule and kernelFunc for outlining kernels. (PR #118861)

[Valentin Clement バレンタイン クレメン]

================
@@ -395,16 +403,35 @@ class GpuKernelOutliningPass
 
 private:
   /// Returns a gpu.module containing kernelFunc and all callees (recursive).
-  gpu::GPUModuleOp createKernelModule(gpu::GPUFuncOp kernelFunc,
+  gpu::GPUModuleOp createKernelModule(gpu::LaunchOp gpuLaunchOp,
+                                      gpu::GPUFuncOp kernelFunc,
                                       const SymbolTable &parentSymbolTable) {
     // TODO: This code cannot use an OpBuilder because it must be inserted into
     // a SymbolTable by the caller. SymbolTable needs to be refactored to
     // prevent manual building of Ops with symbols in code using SymbolTables
     // and then this needs to use the OpBuilder.
     auto *context = getOperation().getContext();
     OpBuilder builder(context);
-    auto kernelModule = builder.create<gpu::GPUModuleOp>(kernelFunc.getLoc(),
-                                                         kernelFunc.getName());
+    std::string kernelModuleName;
+    if (gpuLaunchOp.hasKernelModuleName()) {
+      kernelModuleName =
+          gpuLaunchOp->getAttrOfType<mlir::SymbolRefAttr>("kernelModule")
+              .getRootReference()
+              .str();
+    } else {
+      kernelModuleName = kernelFunc.getName();
+    }
+
+    gpu::GPUModuleOp kernelModule;
+    // Check if the module already exists in the symbol table
+    if (auto existingModule =
+            parentSymbolTable.lookup<gpu::GPUModuleOp>(kernelModuleName)) {
+      kernelModule = existingModule;
+    } else {
+      // If not found, create a new GPU module
+      kernelModule = builder.create<gpu::GPUModuleOp>(kernelFunc.getLoc(),
+                                                      kernelModuleName);
+    }
----------------
clementval wrote:

```suggestion
    // Check if the module already exists in the symbol table
    gpu::GPUModuleOp kernelModule = parentSymbolTable.lookup<gpu::GPUModuleOp>(kernelModuleName);
    if (!kernelModule) {
      // If not found, create a new GPU module
      kernelModule = builder.create<gpu::GPUModuleOp>(kernelFunc.getLoc(),
                                                      kernelModuleName);
    }
```

https://github.com/llvm/llvm-project/pull/118861