diff --git a/mlir/lib/Dialect/UB/IR/UBOps.cpp b/mlir/lib/Dialect/UB/IR/UBOps.cpp index d63da90e16195f..e0cd5dafcfa61e 100644 --- a/mlir/lib/Dialect/UB/IR/UBOps.cpp +++ b/mlir/lib/Dialect/UB/IR/UBOps.cpp @@ -7,6 +7,7 @@ //===----------------------------------------------------------------------===// #include "mlir/Dialect/UB/IR/UBOps.h" +#include "mlir/Transforms/InliningUtils.h" #include "mlir/IR/Builders.h" #include "mlir/IR/DialectImplementation.h" @@ -17,6 +18,19 @@ using namespace mlir; using namespace mlir::ub; +namespace { +/// This class defines the interface for handling inlining with UB +/// operations. +struct UBInlinerInterface : public DialectInlinerInterface { + using DialectInlinerInterface::DialectInlinerInterface; + + /// All UB ops can be inlined. + bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final { + return true; + } +}; +} // namespace + //===----------------------------------------------------------------------===// // UBDialect //===----------------------------------------------------------------------===// @@ -30,6 +44,7 @@ void UBDialect::initialize() { #define GET_ATTRDEF_LIST #include "mlir/Dialect/UB/IR/UBOpsAttributes.cpp.inc" >(); + addInterfaces(); } Operation *UBDialect::materializeConstant(OpBuilder &builder, Attribute value, diff --git a/mlir/test/Dialect/UB/inlining.mlir b/mlir/test/Dialect/UB/inlining.mlir new file mode 100644 index 00000000000000..cb76750044a3b8 --- /dev/null +++ b/mlir/test/Dialect/UB/inlining.mlir @@ -0,0 +1,13 @@ +// RUN: mlir-opt %s -inline -split-input-file | FileCheck %s + +func.func @func() -> i32 { + %0 = ub.poison : i32 + return %0 : i32 +} + +// CHECK-LABEL: func @test_inline +func.func @test_inline(%ptr : !llvm.ptr) -> i32 { +// CHECK-NOT: call + %0 = call @func() : () -> i32 + return %0 : i32 +}