Add inliner interface for GPU dialect

[bondhugula]

Summary

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

[bondhugula]

I couldn't immediately find anything that should be prevented from inlining. I can refine the condition if there are such situations.

[bondhugula]

Fix typo.

[ftynse]

Is it legal though? Should we have something like gpu.thread_id outside of a GPU function/module? I would rather think it is legal to inline GPU functions into GPU functions.

[herhut]

In D116889#3230744, @ftynse wrote:

Is it legal though? Should we have something like gpu.thread_id outside of a GPU function/module? I would rather think it is legal to inline GPU functions into GPU functions.

Calling a gpu function from a non-gpu function should be illegal anyway, so this situation cannot really arise. And the gpu.launch is not a call in the sense of the inlining interface.

[ftynse]

Would be nice to have a test to clarify this. LGTM otherwise.

[bondhugula]

Add a test case.

[bondhugula]

In D116889#3230930, @ftynse wrote:

Would be nice to have a test to clarify this. LGTM otherwise.

Done.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/Matchers.h"
 #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;
 using namespace mlir::gpu;
 
 #include "mlir/Dialect/GPU/GPUOpsDialect.cpp.inc"
 
 //===----------------------------------------------------------------------===//
enum GPUMemorySpace {
   kSharedMemorySpace = 3
 };
 
 bool GPUDialect::isKernel(Operation *op) {
   UnitAttr isKernelAttr = op->getAttrOfType<UnitAttr>(getKernelFuncAttrName());
   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>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
       >();
+  addInterfaces<GPUInlinerInterface>();
 }
 
 Type GPUDialect::parseType(DialectAsmParser &parser) const {
   // Parse the main keyword for the type.
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return Type();
   MLIRContext *context = getContext();

mlir/test/Transforms/inlining.mlir

func @inline_simplify() -> i32 {
   return %res : i32
 }
 
 // CHECK-LABEL: func @no_inline_invalid_call
 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>
+}