Annotate outlined kernel functions with the attribute `gpu.kernel`.

--

PiperOrigin-RevId: 247589560

mlir/include/mlir/GPU/GPUDialect.h

@@ -38,6 +38,9 @@ public:
 
   /// Get the canonical string name of the dialect.
   static StringRef getDialectName();
+
+  /// Get the name of the attribute used to annotate outlined kernel functions.
+  static StringRef getKernelFuncAttrName() { return "gpu.kernel"; }
 };
 
 /// Utility class for the GPU dialect to represent triples of `Value`s

mlir/lib/GPU/Transforms/KernelOutlining.cpp

@@ -71,6 +71,10 @@ Function *outlineKernelFunc(Module &module, gpu::LaunchOp &launchOp) {
   mlir::BlockAndValueMapping mapper;
   Function *outlinedFunc = new mlir::Function(loc, kernelFuncName, type);
   outlinedFunc->getBody().takeBody(launchOp.getBody());
+  Builder builder(&module);
+  outlinedFunc->getAttrList().set(
+      builder.getIdentifier(gpu::GPUDialect::getKernelFuncAttrName()),
+      builder.getUnitAttr());
   injectGpuIndexOperations(module, loc, outlinedFunc);
   module.getFunctions().push_back(outlinedFunc);
   return outlinedFunc;

mlir/test/GPU/outlining.mlir

@@ -26,6 +26,7 @@ func @launch() {
 }
 
 // CHECK: func @launch_kernel(%arg0: f32, %arg1: memref<?xf32, 1>)
+// CHECK-NEXT: attributes {gpu.kernel}
 // CHECK-NEXT: %0 = "gpu.block_id"() {dimension: "x"} : () -> index
 // CHECK-NEXT: %1 = "gpu.block_id"() {dimension: "y"} : () -> index
 // CHECK-NEXT: %2 = "gpu.block_id"() {dimension: "z"} : () -> index