diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp --- 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 @@ return static_cast(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(); addTypes(); @@ -109,6 +124,7 @@ #define GET_OP_LIST #include "mlir/Dialect/GPU/GPUOps.cpp.inc" >(); + addInterfaces(); } Type GPUDialect::parseType(DialectAsmParser &parser) const { diff --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir --- a/mlir/test/Transforms/inlining.mlir +++ b/mlir/test/Transforms/inlining.mlir @@ -190,3 +190,15 @@ %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> +}