[Mlir-commits] [mlir] [mlir][mesh] Introduce `DialectInlinerInterface` for Mesh dialect (PR #108297)

[Matteo Franciolini]

https://github.com/mfrancio updated https://github.com/llvm/llvm-project/pull/108297

>From e812edc54f949f2b7205af01c5cca3280c4c3aad Mon Sep 17 00:00:00 2001
From: Matteo Franciolini <mfranciolini at tesla.com>
Date: Wed, 11 Sep 2024 14:31:57 -0700
Subject: [PATCH] [mlir][mesh] Introduce `DialectInlinerInterface` for the Mesh
 dialect

The inliner interface does not implement any restrictions for inlining.
---
 mlir/lib/Dialect/Mesh/IR/MeshOps.cpp | 22 ++++++++++++++++++++++
 mlir/test/Dialect/Mesh/inlining.mlir | 15 +++++++++++++++
 2 files changed, 37 insertions(+)
 create mode 100644 mlir/test/Dialect/Mesh/inlining.mlir

diff --git a/mlir/lib/Dialect/Mesh/IR/MeshOps.cpp b/mlir/lib/Dialect/Mesh/IR/MeshOps.cpp
index c35020b4c20ccc..19e9212157ae47 100644
--- a/mlir/lib/Dialect/Mesh/IR/MeshOps.cpp
+++ b/mlir/lib/Dialect/Mesh/IR/MeshOps.cpp
@@ -24,6 +24,7 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
@@ -74,6 +75,26 @@ static DimensionSize operator*(DimensionSize lhs, DimensionSize rhs) {
   return lhs.value() * rhs.value();
 }
 
+//===----------------------------------------------------------------------===//
+// Inliner
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct MeshInlinerInterface : public DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+  // Currently no restrictions are encoded for inlining.
+  bool isLegalToInline(Operation *, Operation *, bool) const final {
+    return true;
+  }
+  bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
+    return true;
+  }
+  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
+    return true;
+  }
+};
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // Mesh dialect
 //===----------------------------------------------------------------------===//
@@ -91,6 +112,7 @@ void MeshDialect::initialize() {
 #define GET_TYPEDEF_LIST
 #include "mlir/Dialect/Mesh/IR/MeshTypes.cpp.inc"
       >();
+  addInterface<MeshInlinerInterface>();
 }
 
 Operation *MeshDialect::materializeConstant(OpBuilder &builder, Attribute value,
diff --git a/mlir/test/Dialect/Mesh/inlining.mlir b/mlir/test/Dialect/Mesh/inlining.mlir
new file mode 100644
index 00000000000000..c41a709e1a4ebb
--- /dev/null
+++ b/mlir/test/Dialect/Mesh/inlining.mlir
@@ -0,0 +1,15 @@
+// RUN: mlir-opt -inline %s | FileCheck %s
+
+mesh.mesh @mesh0(shape = 4x?x2)
+
+func.func private @mesh_to_inline() -> (index, index) {
+  %0:2 = mesh.mesh_shape @mesh0 axes = [2, 1] : index, index
+  return %0#0, %0#1 : index, index
+}
+// CHECK-LABEL: func.func @main
+func.func @main() -> (index, index) {
+  // CHECK-NEXT: %[[AXIS_SIZE:.*]]:2 = mesh.mesh_shape @mesh0 axes = [2, 1] : index
+  %0:2 = func.call @mesh_to_inline() : () -> (index, index)
+  // CHECK-NEXT: return %[[AXIS_SIZE]]#0, %[[AXIS_SIZE]]#1
+  return %0#0, %0#1 : index, index
+}