From df0296742b7fa82bb3285f2a4eac99d20ac5a6cc Mon Sep 17 00:00:00 2001 From: Thomas Joerg Date: Fri, 10 May 2019 04:26:58 -0700 Subject: [PATCH] Annotate outlined kernel functions with the attribute `gpu.kernel`. -- PiperOrigin-RevId: 247589560 --- mlir/include/mlir/GPU/GPUDialect.h | 3 +++ mlir/lib/GPU/Transforms/KernelOutlining.cpp | 4 ++++ mlir/test/GPU/outlining.mlir | 1 + 3 files changed, 8 insertions(+) diff --git a/mlir/include/mlir/GPU/GPUDialect.h b/mlir/include/mlir/GPU/GPUDialect.h index d03ebe7..86a1acf 100644 --- a/mlir/include/mlir/GPU/GPUDialect.h +++ b/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 diff --git a/mlir/lib/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/GPU/Transforms/KernelOutlining.cpp index d57046c..dd73c19 100644 --- a/mlir/lib/GPU/Transforms/KernelOutlining.cpp +++ b/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; diff --git a/mlir/test/GPU/outlining.mlir b/mlir/test/GPU/outlining.mlir index 4dda193..bd8bca3 100644 --- a/mlir/test/GPU/outlining.mlir +++ b/mlir/test/GPU/outlining.mlir @@ -26,6 +26,7 @@ func @launch() { } // CHECK: func @launch_kernel(%arg0: f32, %arg1: memref) +// 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 -- 2.7.4