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

Author: nicolasvasilache

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

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>
+}