[Mlir-commits] [mlir] fc61d07 - Add inliner interface for GPU dialect

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2022-01-12T12:55:02+05:30
New Revision: fc61d07dc1e1ecd26ff39b7256af9d1b81f1b609

URL: https://github.com/llvm/llvm-project/commit/fc61d07dc1e1ecd26ff39b7256af9d1b81f1b609
DIFF: https://github.com/llvm/llvm-project/commit/fc61d07dc1e1ecd26ff39b7256af9d1b81f1b609.diff

LOG: Add inliner interface for GPU dialect

Add inliner interface for GPU dialect. The interface marks all GPU
dialect ops legal to inline anywhere.

Differential Revision: https://reviews.llvm.org/D116889

Added: 
    

Modified: 
    mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
    mlir/test/Transforms/inlining.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index ea8ce177848ef..75249e7427f31 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -26,6 +26,7 @@
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
@@ -102,6 +103,20 @@ bool GPUDialect::isKernel(Operation *op) {
   return static_cast<bool>(isKernelAttr);
 }
 
+namespace {
+/// This class defines the interface for handling inlining with gpu
+/// operations.
+struct GPUInlinerInterface : public DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+
+  /// All gpu dialect ops can be inlined.
+  bool isLegalToInline(Operation *, Region *, bool,
+                       BlockAndValueMapping &) const final {
+    return true;
+  }
+};
+} // namespace
+
 void GPUDialect::initialize() {
   addTypes<AsyncTokenType>();
   addTypes<MMAMatrixType>();
@@ -109,6 +124,7 @@ void GPUDialect::initialize() {
 #define GET_OP_LIST
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
       >();
+  addInterfaces<GPUInlinerInterface>();
 }
 
 Type GPUDialect::parseType(DialectAsmParser &parser) const {

diff  --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index 00a0445f5a88d..a81050f16ecb0 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -190,3 +190,15 @@ func @no_inline_invalid_call() -> i32 {
   %res = "test.conversion_call_op"() { callee=@convert_callee_fn_multiblock, noinline } : () -> (i32)
   return %res : i32
 }
+
+func @gpu_alloc() -> memref<1024xf32> {
+  %m = gpu.alloc [] () : memref<1024xf32>
+  return %m : memref<1024xf32>
+}
+
+// CHECK-LABEL: func @inline_gpu_ops
+func @inline_gpu_ops() -> memref<1024xf32> {
+  // CHECK-NEXT: gpu.alloc
+  %m = call @gpu_alloc() : () -> memref<1024xf32>
+  return %m : memref<1024xf32>
+}