From 3a43e68ed1cbd96bc70359f0d9dced517195ae9c Mon Sep 17 00:00:00 2001
From: =?utf8?q?Ingo=20M=C3=BCller?= <ingomueller@google.com>
Date: Wed, 14 Dec 2022 15:16:06 +0000
Subject: [PATCH] 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
---
 .../include/mlir/Interfaces/SideEffectInterfaces.h |  6 +++++
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp         | 22 ++++++++++++++++-
 mlir/lib/Interfaces/SideEffectInterfaces.cpp       |  6 +++++
 mlir/test/Dialect/LLVMIR/inlining.mlir             | 28 ++++++++++++++++++++++
 4 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/LLVMIR/inlining.mlir

diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index a6922df..306f4cf 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 0694f9b..2bb483c 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"
@@ -2584,6 +2585,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 a59d52d..bff1517 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 0000000..ffe03ee
--- /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>
+}
-- 
2.7.4