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

[stellaraccident]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-mlprogram

Author: Stella Laurenzo (stellaraccident)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/85479.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt (+1)
• (modified) mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp (+13-1)

diff --git a/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt b/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt
index 90c9c9aa787503..725bb5fd9da9ed 100644
--- a/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt
@@ -15,5 +15,6 @@ add_mlir_dialect_library(MLIRMLProgramDialect
   MLIRControlFlowInterfaces
   MLIRFunctionInterfaces
   MLIRInferTypeOpInterface
+  MLIRTransforms
   MLIRIR
   )
diff --git a/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp b/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp
index 1a8fe208d4099e..0b186a0d072aa6 100644
--- a/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp
+++ b/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,17 @@ 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 {
+    // 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 +65,5 @@ void ml_program::MLProgramDialect::initialize() {
 #include "mlir/Dialect/MLProgram/IR/MLProgramOps.cpp.inc"
       >();
 
-  addInterfaces<MLProgramOpAsmDialectInterface>();
+  addInterfaces<MLProgramInlinerInterface, MLProgramOpAsmDialectInterface>();
 }

[kuhar]

Not super familiar with this but would this be easy to test? I can see some other dialect have those (LLVM, SPIR-V, Vector, Linalg).

[stellaraccident]

Not super familiar with this but would this be easy to test? I can see some other dialect have those (LLVM, SPIR-V, Vector, Linalg).

Good point. Let me add a test.

[stellaraccident]

@kuhar ptal

[jpienaar]

Why is the utils needed for the interface?