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

Author: ingomueller-net

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
 
 //===----------------------------------------------------------------------===//

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

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

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