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

[Nicolas Vasilache]

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

None

>From 41999e3ed52f763ca6a732183f23b0ee5c7b8d2d Mon Sep 17 00:00:00 2001
From: Nicolas Vasilache <nicolas.vasilache at gmail.com>
Date: Wed, 1 Nov 2023 13:53:29 +0000
Subject: [PATCH 1/2] [mlir][LLVM] Allow call_intrinsic to inline

LLVM::CallIntrinsicOp was previously overlooked from the isLegalToInline.
---
 mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp | 1 +
 mlir/test/Dialect/LLVMIR/inlining.mlir      | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index b40be73ff21f703..6063abdba7b9a1f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -694,6 +694,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
             LLVM::AtomicRMWOp,
             LLVM::AtomicCmpXchgOp,
             LLVM::CallOp,
+            LLVM::CallIntrinsicOp,
             LLVM::DbgDeclareOp,
             LLVM::DbgLabelOp,
             LLVM::DbgValueOp,
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index 3f14dc6de6b764c..1296b8e031c1330 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -28,6 +28,7 @@ func.func @inner_func_inlinable(%ptr : !llvm.ptr) -> i32 {
   llvm.unreachable
 ^bb2:
   llvm.intr.stackrestore %stack : !llvm.ptr
+  llvm.call_intrinsic "llvm.x86.sse41.round.ss"() : () -> (vector<8xf32>)
   return %1 : i32
 }
 
@@ -50,6 +51,7 @@ func.func @inner_func_inlinable(%ptr : !llvm.ptr) -> i32 {
 // CHECK: llvm.inline_asm has_side_effects "foo", "bar"
 // CHECK: llvm.unreachable
 // CHECK: llvm.intr.stackrestore %[[STACK]]
+// CHECK: llvm.call_intrinsic "llvm.x86.sse41.round.ss"(
 func.func @test_inline(%ptr : !llvm.ptr) -> i32 {
   %0 = call @inner_func_inlinable(%ptr) : (!llvm.ptr) -> i32
   return %0 : i32

>From 34e8d42f6f72102d7bff27f127442928753db142 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 2/2] [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>
+}