[Mlir-commits] [mlir] [MLIR][UB] Add inliner interface for UB dialect (PR #67115)

[Christian Ulmann]

https://github.com/Dinistro created https://github.com/llvm/llvm-project/pull/67115

This revision adds an inliner interface to the UB dialect that allows inlining of `ub.poison` operations.

>From c518334c86082a867b5e963707e77a4ef7b0df9d Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Fri, 22 Sep 2023 11:16:55 +0000
Subject: [PATCH] [MLIR][UB] Add inliner interface for UB dialect

This revision adds an inliner interface to the UB dialect that allows
inlining of `ub.poison` operations.
---
 mlir/lib/Dialect/UB/IR/UBOps.cpp   | 15 +++++++++++++++
 mlir/test/Dialect/UB/inlining.mlir | 13 +++++++++++++
 2 files changed, 28 insertions(+)
 create mode 100644 mlir/test/Dialect/UB/inlining.mlir

diff --git a/mlir/lib/Dialect/UB/IR/UBOps.cpp b/mlir/lib/Dialect/UB/IR/UBOps.cpp
index d63da90e16195f1..e0cd5dafcfa61eb 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<UBInlinerInterface>();
 }
 
 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 000000000000000..cb76750044a3b8f
--- /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
+}