Add optional attribute outline_module to gpu.launch

Author: wangzpgi

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

@@ -803,7 +803,8 @@ def GPU_LaunchOp : GPU_Op<"launch", [
                Optional<Index>:$clusterSizeX,
                Optional<Index>:$clusterSizeY,
                Optional<Index>:$clusterSizeZ,
-               Optional<I32>:$dynamicSharedMemorySize)>,
+               Optional<I32>:$dynamicSharedMemorySize,
+               OptionalAttr<SymbolRefAttr>:$outlineModule)>,
     Results<(outs Optional<GPU_AsyncToken>:$asyncToken)> {
   let summary = "GPU kernel launch operation";
 
@@ -837,6 +838,10 @@ def GPU_LaunchOp : GPU_Op<"launch", [
     -   a variadic number of Workgroup memory attributions.
     -   a variadic number of Private memory attributions.
 
+    The `outline_module` attribute is optional and specifies a module in which 
+    the kernel should be outlined. When this attribute is present, the kernel is
+    outlined into the specified module instead of the default behavior.
+
     Syntax:
 
     ```
@@ -1030,6 +1035,11 @@ def GPU_LaunchOp : GPU_Op<"launch", [
     static StringRef getNumWorkgroupAttributionsAttrName() {
       return "workgroup_attributions";
     }
+
+    /// Checks if the outline_module attribute is present.
+    bool hasOutlineModule() {
+      return getOutlineModule().has_value();
+    }
   }];
 
   let hasCanonicalizer = 1;

mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp

@@ -364,9 +364,16 @@ class GpuKernelOutliningPass
       Block::iterator insertPt(func->getNextNode());
       auto funcWalkResult = func.walk([&](gpu::LaunchOp op) {
         SetVector<Value> operands;
-        std::string kernelFnName =
-            Twine(op->getParentOfType<SymbolOpInterface>().getName(), "_kernel")
-                .str();
+        std::string kernelFnName;
+        if (auto outlineModuleAttr = op->getAttrOfType<SymbolRefAttr>("outline_module")) {
+          kernelFnName = outlineModuleAttr.getRootReference().str();
+          llvm::errs() << "outlined module name = " << kernelFnName << "\n";
+        } else {
+          kernelFnName =
+              Twine(op->getParentOfType<SymbolOpInterface>().getName(), "_kernel")
+                  .str();
+          llvm::errs() << "original module name = " << kernelFnName << "\n";
+        }
 
         gpu::GPUFuncOp outlinedFunc =
             outlineKernelFuncImpl(op, kernelFnName, operands);