Allow inline of all pure ops from the LLVM dialect.

[ingomueller-net]

Summary

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.

[ingomueller-net]

Two comments:

1. If someone thinks that adding the isPure function should be a different patch, please let me know.

2. The limitation to pure ops comes from a discussion I had offline with several people, in which those were deemed to be safe to inline. This is probably a very conservative restriction and it may be useful to widen it. However; I don't know enough about LLVM and other things to be able to judge what would be safe and what could cause problems. If somebody can vouch for a wider definition, I'll happily change the patch.

[ingomueller-net]

I have investigate a bit more on the limitation of allowing only pure ops to be inlined. I have found two that I would like to be inlinable for my use case and where I don't see any reason that they shouldn't be:

• CallOp (though if the callee is inlined, this op will also go away, so this may not be that big a problem)
• LoadOp

There are probably quite a few more ops that aren't pure and could still be inlined. My possibly naive intuition would assume that this is the case for the following:

• StoreOp
• FreezeOp
• FenceOp + Atomic*Op
• InlineAsmOp

Approaching the question from the opposite end, I suspect the following ops to be not inlineable or potentially more difficult to inline:

• InvokeOp + LandingpadOp
• MetadataOp + *MetadataOp
• Global*Op
• FuncOp

[ftynse]

mlir/test/Dialect/LLVMIR/inlining.mlir

29

Please add a newline.

[ingomueller-net]

Add missing new line at end of file.

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

 
 /// Returns true if the given operation is speculatable, i.e. has no undefined
 /// behavior or other side effects.
 ///
 /// An operation can indicate that it is speculatable by implementing the
 /// 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
 
 //===----------------------------------------------------------------------===//
 // SideEffect Interfaces
 //===----------------------------------------------------------------------===//
 
 /// Include the definitions of the side effect interfaces.
 #include "mlir/Interfaces/SideEffectInterfaces.h.inc"
 
 #endif // MLIR_INTERFACES_SIDEEFFECTINTERFACES_H

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #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"
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Type.h"
return TypeSwitch<Attribute, AliasResult>(attr)
           return AliasResult::OverridableAlias;
         })
         .Default([](Attribute) { return AliasResult::NoAlias; });
   }
 };
 } // 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.
 //===----------------------------------------------------------------------===//
 
 void LLVMDialect::initialize() {
   registerAttributes();
 
   // clang-format off
   addTypes<LLVMVoidType,

 #include "mlir/Dialect/LLVMIR/LLVMOps.cpp.inc"
       ,
 #define GET_OP_LIST
 #include "mlir/Dialect/LLVMIR/LLVMIntrinsicOps.cpp.inc"
       >();
 
   // 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
 #include "mlir/Dialect/LLVMIR/LLVMOps.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMIntrinsicOps.cpp.inc"
 

mlir/lib/Interfaces/SideEffectInterfaces.cpp

case Speculation::Speculatable:
     return true;
 
   case Speculation::NotSpeculatable:
     return false;
   }
 
   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

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