[mlir][builtin] Make unrealized_conversion_cast inlineable

[matthias-springer]

Until now, builtin.unrealized_conversion_cast ops could not be inlined by the Inliner pass.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Matthias Springer (matthias-springer)

Changes

Until now, builtin.unrealized_conversion_cast ops could not be inlined by the Inliner pass.

---

Full diff: https://github.com/llvm/llvm-project/pull/139722.diff

4 Files Affected:

• (modified) mlir/include/mlir/InitAllDialects.h (+2)
• (modified) mlir/include/mlir/Transforms/InliningUtils.h (+6)
• (modified) mlir/lib/Transforms/Utils/InliningUtils.cpp (+21)
• (modified) mlir/test/Transforms/inlining.mlir (+7-3)

diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h
index e83be7b40eded..afae5232d1fb0 100644
--- a/mlir/include/mlir/InitAllDialects.h
+++ b/mlir/include/mlir/InitAllDialects.h
@@ -102,6 +102,7 @@
 #include "mlir/Target/LLVM/NVVM/Target.h"
 #include "mlir/Target/LLVM/ROCDL/Target.h"
 #include "mlir/Target/SPIRV/Target.h"
+#include "mlir/Transforms/InliningUtils.h"
 
 namespace mlir {
 
@@ -166,6 +167,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   builtin::registerCastOpInterfaceExternalModels(registry);
+  builtin::registerBuiltinDialectInlinerInterfaceExternalModel(registry);
   cf::registerBufferizableOpInterfaceExternalModels(registry);
   cf::registerBufferDeallocationOpInterfaceExternalModels(registry);
   gpu::registerBufferDeallocationOpInterfaceExternalModels(registry);
diff --git a/mlir/include/mlir/Transforms/InliningUtils.h b/mlir/include/mlir/Transforms/InliningUtils.h
index ed6413d8cd44c..48c5af2ca7585 100644
--- a/mlir/include/mlir/Transforms/InliningUtils.h
+++ b/mlir/include/mlir/Transforms/InliningUtils.h
@@ -305,6 +305,12 @@ inlineCall(InlinerInterface &interface,
            CallOpInterface call, CallableOpInterface callable, Region *src,
            bool shouldCloneInlinedRegion = true);
 
+namespace builtin {
+/// Register the builtin dialect inliner interface external model.
+void registerBuiltinDialectInlinerInterfaceExternalModel(
+    DialectRegistry &registry);
+} // namespace builtin
+
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_INLININGUTILS_H
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 8594136570a76..227d6bc86a75e 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -13,6 +13,7 @@
 #include "mlir/Transforms/InliningUtils.h"
 
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/CallInterfaces.h"
@@ -557,3 +558,23 @@ LogicalResult mlir::inlineCall(
     return cleanupState();
   return success();
 }
+
+namespace {
+/// This class defines the interface for handling inlining with builtin
+/// operations.
+struct BuiltinInlinerInterface : public DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+
+  /// All builtin ops can be inlined.
+  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
+    return true;
+  }
+};
+} // namespace
+
+void mlir::builtin::registerBuiltinDialectInlinerInterfaceExternalModel(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
+    dialect->addInterfaces<BuiltinInlinerInterface>();
+  });
+}
diff --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index 1ed08878430b5..d8e10aa4212ba 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -5,14 +5,18 @@
 // RUN: mlir-opt %s -inline='op-pipelines=func.func(canonicalize,cse)' | FileCheck %s --check-prefix INLINE_SIMPLIFY
 
 // Inline a function that takes an argument.
-func.func @func_with_arg(%c : i32) -> i32 {
-  %b = arith.addi %c, %c : i32
-  return %b : i32
+func.func @func_with_arg(%arg0 : i32) -> i32 {
+  %b = arith.addi %arg0, %arg0 : i32
+  %c = builtin.unrealized_conversion_cast %b : i32 to i64
+  %d = builtin.unrealized_conversion_cast %c : i64 to i32
+  return %d : i32
 }
 
 // CHECK-LABEL: func @inline_with_arg
 func.func @inline_with_arg(%arg0 : i32) -> i32 {
   // CHECK-NEXT: arith.addi
+  // CHECK-NEXT: unrealized_conversion_cast
+  // CHECK-NEXT: unrealized_conversion_cast
   // CHECK-NEXT: return
 
   %0 = call @func_with_arg(%arg0) : (i32) -> i32

[j2kun]

This also makes builtin.module inlineable, which is a weird thing to say, but, for example, the following IR verifies and presumably would also be inlined by this change.

module {
  func.func @foo() {
    builtin.module {
      func.func @bar() {
        return
      }
    }
    return
  }
  func.func @caller() {
    call @foo() : () -> ()
    return
  }
}

[matthias-springer]

Oh right, module is also a builtin ops. I wanted to make just unrealized_conversion_cast inlineable.

[javedabsar1]

LGTM.
One nitpick comment.

[javedabsar1]

maybe add .. ' e.g. unrealized_conversion_cast'