[Mlir-commits] [mlir] 3a43e68 - Allow inline of all pure ops from the LLVM dialect.

[Ingo Müller]

Author: Ingo Müller
Date: 2022-12-14T15:18:37Z
New Revision: 3a43e68ed1cbd96bc70359f0d9dced517195ae9c

URL: https://github.com/llvm/llvm-project/commit/3a43e68ed1cbd96bc70359f0d9dced517195ae9c
DIFF: https://github.com/llvm/llvm-project/commit/3a43e68ed1cbd96bc70359f0d9dced517195ae9c.diff

LOG: Allow inline of all pure ops from the LLVM dialect.

This allows to inline regions containing pure LLVM ops into their call
sites. (Note that this is not related to inlining of llvm.func but to
any the inlining of any Callable.) For now, we only allow the inlining
of Pure ops to be conservative but other ops may be considered inlinable
in the future.

Testing for purity of ops requires the C++ equivalent of the Pure trait
from SideEffectInterfaces.td, which this patch also provide. Its
implementation calls the C++ equivalents of the two traits that the Pure
trait is based on and, thus, has to be kept with the tablegen trait.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D139187

Added: 
    mlir/test/Dialect/LLVMIR/inlining.mlir

Modified: 
    mlir/include/mlir/Interfaces/SideEffectInterfaces.h
    mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
    mlir/lib/Interfaces/SideEffectInterfaces.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index a6922df8035f3..306f4cf973181 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -337,6 +337,12 @@ bool isMemoryEffectFree(Operation *op);
 /// getSpeculatability hook in the ConditionallySpeculatable op interface.
 bool isSpeculatable(Operation *op);
 
+/// Returns true if the given operation is pure, i.e., is speculatable that does
+/// not touch memory.
+///
+/// This function is the C++ equivalent of the `Pure` trait.
+bool isPure(Operation *op);
+
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 0694f9b68eaf7..2bb483c78dc15 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -20,6 +20,7 @@
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/Transforms/InliningUtils.h"
 
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/AsmParser/Parser.h"
@@ -2583,6 +2584,22 @@ struct LLVMOpAsmDialectInterface : public OpAsmDialectInterface {
 };
 } // namespace
 
+//===----------------------------------------------------------------------===//
+// DialectInlinerInterface
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct LLVMInlinerInterface : public DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+
+  /// Conservatively only allow inlining of pure ops.
+  bool isLegalToInline(Operation *op, Region *, bool,
+                       BlockAndValueMapping &) const final {
+    return isPure(op);
+  }
+};
+} // end anonymous namespace
+
 //===----------------------------------------------------------------------===//
 // LLVMDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
@@ -2611,7 +2628,10 @@ void LLVMDialect::initialize() {
 
   // Support unknown operations because not all LLVM operations are registered.
   allowUnknownOperations();
-  addInterfaces<LLVMOpAsmDialectInterface>();
+  // clang-format off
+  addInterfaces<LLVMOpAsmDialectInterface,
+                LLVMInlinerInterface>();
+  // clang-format on
 }
 
 #define GET_OP_CLASSES

diff  --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index a59d52d167654..bff1517bf61cf 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -202,3 +202,9 @@ bool mlir::isSpeculatable(Operation *op) {
 
   llvm_unreachable("Unhandled enum in mlir::isSpeculatable!");
 }
+
+/// The implementation of this function replicates the `def Pure : TraitList`
+/// in `SideEffectInterfaces.td` and has to be kept in sync manually.
+bool mlir::isPure(Operation *op) {
+  return isSpeculatable(op) && isMemoryEffectFree(op);
+}

diff  --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
new file mode 100644
index 0000000000000..ffe03ee0e9738
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -0,0 +1,28 @@
+// RUN: mlir-opt %s -inline | FileCheck %s
+
+// CHECK-LABEL: func.func @test_inline() -> i32 {
+// CHECK-NEXT: %[[RES:.*]] = llvm.mlir.constant(42 : i32) : i32
+// CHECK-NEXT: return %[[RES]] : i32
+func.func @test_inline() -> i32 {
+  %0 = call @inner_func_inlinable() : () -> i32
+  return %0 : i32
+}
+
+func.func @inner_func_inlinable() -> i32 {
+  %0 = llvm.mlir.constant(42 : i32) : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: func.func @test_not_inline() -> !llvm.ptr<f64> {
+// CHECK-NEXT: %[[RES:.*]] = call @inner_func_not_inlinable() : () -> !llvm.ptr<f64>
+// CHECK-NEXT: return %[[RES]] : !llvm.ptr<f64>
+func.func @test_not_inline() -> !llvm.ptr<f64> {
+  %0 = call @inner_func_not_inlinable() : () -> !llvm.ptr<f64>
+  return %0 : !llvm.ptr<f64>
+}
+
+func.func @inner_func_not_inlinable() -> !llvm.ptr<f64> {
+  %0 = llvm.mlir.constant(0 : i32) : i32
+  %1 = llvm.alloca %0 x f64 : (i32) -> !llvm.ptr<f64>
+  return %1 : !llvm.ptr<f64>
+}