[Mlir-commits] [mlir] [mlir][Vector] Add initial support for inlining in the presence of vector ops (PR #70942)

[Nicolas Vasilache]

https://github.com/nicolasvasilache updated https://github.com/llvm/llvm-project/pull/70942

>From 16e7b59eca70d7e76345d7ba9158645e8a6fdb1b Mon Sep 17 00:00:00 2001
From: Nicolas Vasilache <nicolas.vasilache at gmail.com>
Date: Wed, 1 Nov 2023 14:43:45 +0000
Subject: [PATCH] [mlir][Vector] Add initial support for inlining in the
 presence of vector ops

---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 16 ++++++++++++++++
 mlir/test/Dialect/Vector/inlining.mlir   | 14 ++++++++++++++
 2 files changed, 30 insertions(+)
 create mode 100644 mlir/test/Dialect/Vector/inlining.mlir

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 4e34caa6d8aaba8..60416f550ee619d 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -33,6 +33,7 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
@@ -348,6 +349,19 @@ struct BitmaskEnumStorage : public AttributeStorage {
 // VectorDialect
 //===----------------------------------------------------------------------===//
 
+namespace {
+/// This class defines the interface for handling inlining with vector dialect
+/// operations.
+struct VectorInlinerInterface : public DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+
+  /// All vector dialect ops can be inlined.
+  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
+    return true;
+  }
+};
+} // namespace
+
 void VectorDialect::initialize() {
   addAttributes<
 #define GET_ATTRDEF_LIST
@@ -358,6 +372,8 @@ void VectorDialect::initialize() {
 #define GET_OP_LIST
 #include "mlir/Dialect/Vector/IR/VectorOps.cpp.inc"
       >();
+
+  addInterfaces<VectorInlinerInterface>();
 }
 
 /// Materialize a single constant operation from a given attribute value with
diff --git a/mlir/test/Dialect/Vector/inlining.mlir b/mlir/test/Dialect/Vector/inlining.mlir
new file mode 100644
index 000000000000000..053a115613ff641
--- /dev/null
+++ b/mlir/test/Dialect/Vector/inlining.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s -inline | FileCheck %s
+
+func.func @inner_func_inlinable(%v: f32) -> vector<4xf32> {
+  %1 = vector.broadcast %v : f32 to vector<4xf32>
+  return %1 : vector<4xf32>
+}
+
+// CHECK-LABEL: func.func @test_inline(
+//  CHECK-NOT:    func.call
+//  CHECK-NEXT:   vector.broadcast
+func.func @test_inline(%v: f32) -> vector<4xf32> {
+  %0 = call @inner_func_inlinable(%v) : (f32) -> vector<4xf32>
+  return %0 : vector<4xf32>
+}