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

[Stella Laurenzo]

https://github.com/stellaraccident updated https://github.com/llvm/llvm-project/pull/85479

>From 0efd79e0d28e54a297686317d36dbc371575a22f Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident at gmail.com>
Date: Fri, 15 Mar 2024 15:30:25 -0700
Subject: [PATCH 1/2] [mlir] Make the ml_program dialect allow all of its
 operations to be inlined.

---
 mlir/lib/Dialect/MLProgram/IR/CMakeLists.txt       |  1 +
 mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp | 14 +++++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

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>();
 }

>From 96ba583ce9048769bf255b1cb6af30f0c21d0a9d Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident at gmail.com>
Date: Fri, 15 Mar 2024 16:08:11 -0700
Subject: [PATCH 2/2] Add test

---
 .../Dialect/MLProgram/IR/MLProgramDialect.cpp |  3 ++-
 mlir/test/Dialect/MLProgram/inlining.mlir     | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/MLProgram/inlining.mlir

diff --git a/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp b/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp
index 0b186a0d072aa6..bda1032ed98847 100644
--- a/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp
+++ b/mlir/lib/Dialect/MLProgram/IR/MLProgramDialect.cpp
@@ -29,7 +29,8 @@ namespace {
 struct MLProgramInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
-  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const {
+  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;
diff --git a/mlir/test/Dialect/MLProgram/inlining.mlir b/mlir/test/Dialect/MLProgram/inlining.mlir
new file mode 100644
index 00000000000000..ac3677e1b92882
--- /dev/null
+++ b/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>
+}