[mlir] Make the ml_program dialect allow all of its operations to be inlined. (#85479)

Author: stellaraccident

mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt

@@ -15,5 +15,6 @@ add_mlir_dialect_library(MLIRMLProgramDialect
  MLIRControlFlowInterfaces
  MLIRFunctionInterfaces
  MLIRInferTypeOpInterface
+ MLIRTransforms
  MLIRIR
  )

mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp

@@ -8,6 +8,7 @@
 
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
@@ -24,6 +25,18 @@ using namespace mlir::ml_program;
 #include "mlir/Dialect/MLProgram/IR/MLProgramTypes.cpp.inc"
 
 namespace {
+
+struct MLProgramInlinerInterface : public DialectInlinerInterface {
+ using DialectInlinerInterface::DialectInlinerInterface;
+
+ bool isLegalToInline(Operation *, Region *, bool,
+ IRMapping &) const override {
+ // We have no specific opinion on whether ops defined in this dialect should
+ // be inlined.
+ return true;
+ }
+};
+
 struct MLProgramOpAsmDialectInterface : public OpAsmDialectInterface {
  using OpAsmDialectInterface::OpAsmDialectInterface;
 
@@ -53,5 +66,5 @@ void ml_program::MLProgramDialect::initialize() {
 #include "mlir/Dialect/MLProgram/IR/MLProgramOps.cpp.inc"
  >();
 
- addInterfaces<MLProgramOpAsmDialectInterface>();
+ addInterfaces<MLProgramInlinerInterface, MLProgramOpAsmDialectInterface>();
 }

mlir/test/Dialect/MLProgram/inlining.mlir

@@ -0,0 +1,19 @@
+// RUN: mlir-opt %s -inline | FileCheck %s
+
+// Verifies that regions with operations from the ml_program dialect can
+// be inlined.
+
+ml_program.global private @global(dense<4> : tensor<4xi32>) : tensor<4xi32>
+
+// CHECK: @inline_into
+func.func @inline_into() -> tensor<4xi32> {
+ // CHECK-NOT: @inline_from
+ // CHECK: ml_program.global_load_const
+ %0 = call @inline_from() : () -> tensor<4xi32>
+ return %0 : tensor<4xi32>
+}
+
+func.func @inline_from() -> tensor<4xi32> {
+ %0 = ml_program.global_load_const @global : tensor<4xi32>
+ return %0 : tensor<4xi32>
+}