[Mlir-commits] [mlir] [mlir] allow inlining complex ops (PR #77514)

[Okwan Kwon]

https://github.com/okkwon updated https://github.com/llvm/llvm-project/pull/77514

>From 7bcc59e41fab234b11a18cfecabc38002f917ddf Mon Sep 17 00:00:00 2001
From: Okwan Kwon <okkwon at gmail.com>
Date: Tue, 9 Jan 2024 11:01:53 -0800
Subject: [PATCH] [mlir] allow inlining complex ops

Complex ops are pure ops just like arithmetic ops so they can be inlined.
---
 .../lib/Dialect/Complex/IR/ComplexDialect.cpp | 14 +++++++++++++
 mlir/test/Transforms/inlining.mlir            | 20 +++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp b/mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp
index e54b3a71bbc33e..ca57171af156f9 100644
--- a/mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp
+++ b/mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp
@@ -11,6 +11,7 @@
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 
@@ -18,6 +19,18 @@ using namespace mlir;
 
 #include "mlir/Dialect/Complex/IR/ComplexOpsDialect.cpp.inc"
 
+namespace {
+/// This class defines the interface for handling inlining for complex
+/// dialect operations.
+struct ComplexInlinerInterface : public DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+  /// All complex dialect ops can be inlined.
+  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
+    return true;
+  }
+};
+} // namespace
+
 void complex::ComplexDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
@@ -28,6 +41,7 @@ void complex::ComplexDialect::initialize() {
 #include "mlir/Dialect/Complex/IR/ComplexAttributes.cpp.inc"
       >();
   declarePromisedInterface<ComplexDialect, ConvertToLLVMPatternInterface>();
+  addInterfaces<ComplexInlinerInterface>();
 }
 
 Operation *complex::ComplexDialect::materializeConstant(OpBuilder &builder,
diff --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index 9544f1eb09170e..2a08e625ba79e2 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -297,3 +297,23 @@ func.func @inline_convert_and_handle_attr_call(%arg0 : i16) -> (i16) {
   %res = "test.conversion_call_op"(%arg0) { callee=@handle_attr_callee_fn } : (i16) -> (i16)
   return %res : i16
 }
+
+// Check a function with complex ops is inlined.
+func.func @double_square_complex(%cplx: complex<f32>) -> complex<f32> {
+  %double = complex.add %cplx, %cplx : complex<f32>
+  %square = complex.mul %double, %double : complex<f32>
+  return %square : complex<f32>
+}
+
+// CHECK-LABEL: func @inline_with_complex_ops
+func.func @inline_with_complex_ops() -> complex<f32> {
+  %c1 = arith.constant 1.0 : f32
+  %c2 = arith.constant 2.0 : f32
+  %c = complex.create %c1, %c2 : complex<f32>
+
+  // CHECK: complex.add
+  // CHECK: complex.mul
+  // CHECK-NOT: call
+  %r = call @double_square_complex(%c) : (complex<f32>) -> (complex<f32>)
+  return %r : complex<f32>
+}