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

[llvmlistbot at llvm.org]

Author: Okwan Kwon
Date: 2024-01-10T09:23:36-08:00
New Revision: 7cc9ae95512edd0b969823fdfa062b92cb3c4d4e

URL: https://github.com/llvm/llvm-project/commit/7cc9ae95512edd0b969823fdfa062b92cb3c4d4e
DIFF: https://github.com/llvm/llvm-project/commit/7cc9ae95512edd0b969823fdfa062b92cb3c4d4e.diff

LOG: [mlir] allow inlining complex ops (#77514)

Complex ops are pure ops just like the arithmetic ops so they can be
inlined.

Added: 
    

Modified: 
    mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp
    mlir/test/Transforms/inlining.mlir

Removed: 
    


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