[mlir][complex] Convert complex.abs to libm

[Lewuathe]

Summary

Convert complex.abs to libm library

[bixia]

mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp

69

This template class implementation is very similar to ScalarOpToLibmCall, with the only difference that one handle operations that produce a complex type as a result while the other handles operations that produce a float type as a result. It would be nice if we can somehow combine these two template classes into one.

[Lewuathe]

@bixia Thank you for the feedback. I tried to combine them into one template function but the following error keeps showing up. Looks like complex::AbsOp fails to be embedded into the function in ScalarOpToLibmCall. I think we may need to change the op definition somewhat. Do you have any idea about that?

Assertion failed: (isa<U>()), function cast, file Types.h, line 251.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.	Program arguments: /Users/sasakikai/dev/llvm-project/build/bin/mlir-opt /Users/sasakikai/dev/llvm-project/mlir/test/Conversion/ComplexToLibm/convert-to-libm.mlir -convert-complex-to-libm -canonicalize
Stack dump without symbol names (ensure you have llvm-symbolizer in your PATH or set the environment var `LLVM_SYMBOLIZER_PATH` to point to it):
0  mlir-opt                 0x000000010bd95a8a llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 42
1  mlir-opt                 0x000000010bd94928 llvm::sys::RunSignalHandlers() + 248
2  mlir-opt                 0x000000010bd96120 SignalHandler(int) + 288
3  libsystem_platform.dylib 0x00007ff80e9d2dfd _sigtramp + 29
4  libsystem_platform.dylib 0x00007f9225042b20 _sigtramp + 18446743635998735680
5  libsystem_c.dylib        0x00007ff80e908d24 abort + 123
6  libsystem_c.dylib        0x00007ff80e9080cb err + 0
7  mlir-opt                 0x000000010d2fc003 (anonymous namespace)::ScalarOpToLibmCall<mlir::complex::AbsOp>::matchAndRewrite(mlir::complex::AbsOp, mlir::PatternRewriter&) const (.cold.10) + 35
8  mlir-opt                 0x000000010c8c3d29 (anonymous namespace)::ScalarOpToLibmCall<mlir::complex::AbsOp>::matchAndRewrite(mlir::complex::AbsOp, mlir::PatternRewriter&) const + 1257
9  mlir-opt                 0x000000010cf667d0 mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>) + 1552
10 mlir-opt                 0x000000010cd51e92 (anonymous namespace)::OperationLegalizer::legalize(mlir::Operation*, mlir::ConversionPatternRewriter&) + 2290
11 mlir-opt                 0x000000010cd4a43d (anonymous namespace)::OperationConverter::convertOperations(llvm::ArrayRef<mlir::Operation*>, llvm::function_ref<void (mlir::Diagnostic&)>) + 1117
12 mlir-opt                 0x000000010cd4ce9b mlir::applyPartialConversion(mlir::Operation*, mlir::ConversionTarget&, mlir::FrozenRewritePatternSet const&, llvm::DenseSet<mlir::Operation*, llvm::DenseMapInfo<mlir::Operation*, void>>*) + 75
13 mlir-opt                 0x000000010c8c410d (anonymous namespace)::ConvertComplexToLibmPass::runOnOperation() + 669
14 mlir-opt                 0x000000010cd0da95 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) + 485
15 mlir-opt                 0x000000010cd0e147 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) + 295
16 mlir-opt                 0x000000010cd0fb72 mlir::PassManager::run(mlir::Operation*) + 642
17 mlir-opt                 0x000000010cce4377 performActions(llvm::raw_ostream&, bool, bool, llvm::SourceMgr&, mlir::MLIRContext*, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>) + 503
18 mlir-opt                 0x000000010cce20c1 processBuffer(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, bool, bool, bool, bool, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, mlir::DialectRegistry&, llvm::ThreadPool*) + 641
19 mlir-opt                 0x000000010cce1d66 mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (mlir::PassManager&)>, mlir::DialectRegistry&, bool, bool, bool, bool, bool) + 182
20 mlir-opt                 0x000000010cce2bcd mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&, bool) + 2525
21 mlir-opt                 0x000000010bc2a9e6 main + 198
22 dyld                     0x0000000111ccf51e start + 462

[bixia]

How did you combine the two template? I was thinking the current template only take one parameter Op. The new template will take one additional parameter, that is, a functor that returns the return-type from Op.

[Lewuathe]

The result type is used for defining the function type.

auto opFunctionTy = FunctionType::get(
    rewriter.getContext(), op->getOperandTypes(), op->getResultTypes());

I thought abs op returns the float-like type so that we do not need to pass the functor to the template explicitly. Is this assumption wrong?

[Lewuathe]

Reuse single template function for complex.abs lowering.

[Lewuathe]

Remove unnecessary type declaration.

[Lewuathe]

Rename functors.

[bixia]

mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp

46–48

Would you please update the document to describe the new parameter, such as it must correspond to the result type of the operation?
You may also want to rename the parameter and assign a default value to it, such as TypeResolver = ComplexTypeResolver.

65

Would this help make it clear that the struct is used for the result type of the op?
s/given complex type/given complex result type/

67–68

This can be simplified to return a boolean, such as true for float and false for double.
If you make this change, you may also want to rename the struct, such as ComplexTypeResolver, FloatTypeResolver.

78

Similar to the above.

[Lewuathe]

mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp

67–68

We should return some value specifying the error for the incompatible element type. If we change the signature to bool, we can only return true or false which does not keep error information.

[Lewuathe]

Enhance docs to clarify the usage of FuncResolver.

[bixia]

mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp

67–68

can you use std::optional<bool> for the return type

[Lewuathe]

Good idea! Thanks let me try.

[Lewuathe]

Use llvm::Optional since MLIR looks mostly using llvm::Optional.

[pifon2a]

@Lewuathe, thank you for all your contributions to MLIR and CompexOps dialect! Would you be interested also in adding a dialect-specific attribute for complex numbers so that complex.constant would accept a dedicated attribute, e.g. F32ComplexAttr instead of ArrayAttr. That would make the code cleaner and also enable more opportunities for folding.

[Lewuathe]

@pifon2a Sounds interesting. Let me try to work on that. Do you have any specific idea on the interface of the custom attribute for complex types? (e.g. complex.constant<1.0, 2.0> ?)

[pifon2a]

In D127476#3637501, @Lewuathe wrote:

@pifon2a Sounds interesting. Let me try to work on that. Do you have any specific idea on the interface of the custom attribute for complex types? (e.g. complex.constant<1.0, 2.0> ?)

Thank you! I think if we can print/parse #complex.number<1.0, 2.0>: f32, then it would be pretty readable and tidy. Implementation-wise, this attribute should be pretty close to Builtin_FloatAttr, but with 3 parameters instead of 2.

[Lewuathe]

@pifon2a Got it. Thanks for sharing your idea.

mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp

 #include "../PassDetail.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
 
 namespace {
+// Functor to resolve the function name corresponding to the given complex
+// result type.
+struct ComplexTypeResolver {
+  llvm::Optional<bool> operator()(Type type) const {
+    auto complexType = type.cast<ComplexType>();
+    auto elementType = complexType.getElementType();
+    if (!elementType.isa<Float32Type, Float64Type>())
+      return {};
+
+    return elementType.getIntOrFloatBitWidth() == 64;
+  }
+};
+
+// Functor to resolve the function name corresponding to the given float result
+// type.
+struct FloatTypeResolver {
+  llvm::Optional<bool> operator()(Type type) const {
+    auto elementType = type.cast<FloatType>();
+    if (!elementType.isa<Float32Type, Float64Type>())
+      return {};
+
+    return elementType.getIntOrFloatBitWidth() == 64;
+  }
+};
+
 // Pattern to convert scalar complex operations to calls to libm functions.
 // Additionally the libm function signatures are declared.
-template <typename Op>
+// TypeResolver is a functor returning the libm function name according to the
+// expected type double or float.
+template <typename Op, typename TypeResolver = ComplexTypeResolver>
 struct ScalarOpToLibmCall : public OpRewritePattern<Op> {
 public:
   using OpRewritePattern<Op>::OpRewritePattern;
-  ScalarOpToLibmCall<Op>(MLIRContext *context, StringRef floatFunc,
-                         StringRef doubleFunc, PatternBenefit benefit)
+  ScalarOpToLibmCall<Op, TypeResolver>(MLIRContext *context,
+                                       StringRef floatFunc,
+                                       StringRef doubleFunc,
+                                       PatternBenefit benefit)
       : OpRewritePattern<Op>(context, benefit), floatFunc(floatFunc),
         doubleFunc(doubleFunc){};
 
   LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final;
 
 private:
   std::string floatFunc, doubleFunc;
 };
 } // namespace
 
-template <typename Op>
-LogicalResult
-ScalarOpToLibmCall<Op>::matchAndRewrite(Op op,
-                                        PatternRewriter &rewriter) const {
+template <typename Op, typename TypeResolver>
+LogicalResult ScalarOpToLibmCall<Op, TypeResolver>::matchAndRewrite(
+    Op op, PatternRewriter &rewriter) const {
   auto module = SymbolTable::getNearestSymbolTable(op);
-  auto type = op.getType().template cast<ComplexType>();
-  Type elementType = type.getElementType();
-  if (!elementType.isa<Float32Type, Float64Type>())
+  auto isDouble = TypeResolver()(op.getType());
+  if (!isDouble.hasValue())
     return failure();
 
-  auto name =
-      elementType.getIntOrFloatBitWidth() == 64 ? doubleFunc : floatFunc;
+  auto name = isDouble.value() ? doubleFunc : floatFunc;
+
   auto opFunc = dyn_cast_or_null<SymbolOpInterface>(
       SymbolTable::lookupSymbolIn(module, name));
   // Forward declare function if it hasn't already been
   if (!opFunc) {
     OpBuilder::InsertionGuard guard(rewriter);
     rewriter.setInsertionPointToStart(&module->getRegion(0).front());
     auto opFunctionTy = FunctionType::get(
         rewriter.getContext(), op->getOperandTypes(), op->getResultTypes());
     opFunc = rewriter.create<func::FuncOp>(rewriter.getUnknownLoc(), name,
                                            opFunctionTy);
     opFunc.setPrivate();
   }
   assert(isa<FunctionOpInterface>(SymbolTable::lookupSymbolIn(module, name)));
 
-  rewriter.replaceOpWithNewOp<func::CallOp>(op, name, type, op->getOperands());
+  rewriter.replaceOpWithNewOp<func::CallOp>(op, name, op.getType(),
+                                            op->getOperands());
 
   return success();
 }
 
 void mlir::populateComplexToLibmConversionPatterns(RewritePatternSet &patterns,
                                                    PatternBenefit benefit) {
   patterns.add<ScalarOpToLibmCall<complex::PowOp>>(patterns.getContext(),
                                                    "cpowf", "cpow", benefit);
   patterns.add<ScalarOpToLibmCall<complex::SqrtOp>>(patterns.getContext(),
                                                     "csqrtf", "csqrt", benefit);
   patterns.add<ScalarOpToLibmCall<complex::TanhOp>>(patterns.getContext(),
                                                     "ctanhf", "ctanh", benefit);
   patterns.add<ScalarOpToLibmCall<complex::CosOp>>(patterns.getContext(),
                                                    "ccosf", "ccos", benefit);
   patterns.add<ScalarOpToLibmCall<complex::SinOp>>(patterns.getContext(),
                                                    "csinf", "csin", benefit);
   patterns.add<ScalarOpToLibmCall<complex::ConjOp>>(patterns.getContext(),
                                                     "conjf", "conj", benefit);
+  patterns.add<ScalarOpToLibmCall<complex::AbsOp, FloatTypeResolver>>(
+      patterns.getContext(), "cabsf", "cabs", benefit);
 }
 
 namespace {
 struct ConvertComplexToLibmPass
     : public ConvertComplexToLibmBase<ConvertComplexToLibmPass> {
   void runOnOperation() override;
 };
 } // namespace
 
 void ConvertComplexToLibmPass::runOnOperation() {
   auto module = getOperation();
 
   RewritePatternSet patterns(&getContext());
   populateComplexToLibmConversionPatterns(patterns, /*benefit=*/1);
 
   ConversionTarget target(getContext());
   target.addLegalDialect<func::FuncDialect>();
-  target.addIllegalOp<complex::PowOp, complex::SqrtOp, complex::TanhOp>();
+  target.addIllegalOp<complex::PowOp, complex::SqrtOp, complex::TanhOp,
+                      complex::AbsOp>();
   if (failed(applyPartialConversion(module, target, std::move(patterns))))
     signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertComplexToLibmPass() {
   return std::make_unique<ConvertComplexToLibmPass>();
 }

mlir/test/Conversion/ComplexToLibm/convert-to-libm.mlir

 // RUN: mlir-opt %s -convert-complex-to-libm -canonicalize | FileCheck %s
 
 // CHECK-DAG: @cpowf(complex<f32>, complex<f32>) -> complex<f32>
 // CHECK-DAG: @cpow(complex<f64>, complex<f64>) -> complex<f64>
 // CHECK-DAG: @csqrtf(complex<f32>) -> complex<f32>
 // CHECK-DAG: @csqrt(complex<f64>) -> complex<f64>
 // CHECK-DAG: @ctanhf(complex<f32>) -> complex<f32>
 // CHECK-DAG: @ctanh(complex<f64>) -> complex<f64>
 // CHECK-DAG: @ccos(complex<f64>) -> complex<f64>
 // CHECK-DAG: @csin(complex<f64>) -> complex<f64>
 // CHECK-DAG: @conj(complex<f64>) -> complex<f64>
+// CHECK-DAG: @cabs(complex<f64>) -> f64
 
 // CHECK-LABEL: func @cpow_caller
 // CHECK-SAME: %[[FLOAT:.*]]: complex<f32>
 // CHECK-SAME: %[[DOUBLE:.*]]: complex<f64>
 func.func @cpow_caller(%float: complex<f32>, %double: complex<f64>) -> (complex<f32>, complex<f64>)  {
   // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @cpowf(%[[FLOAT]], %[[FLOAT]]) : (complex<f32>, complex<f32>) -> complex<f32>
   %float_result = complex.pow %float, %float : complex<f32>
   // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @cpow(%[[DOUBLE]], %[[DOUBLE]]) : (complex<f64>, complex<f64>) -> complex<f64>

 // CHECK-SAME: %[[DOUBLE:.*]]: complex<f64>
 func.func @conj_caller(%float: complex<f32>, %double: complex<f64>) -> (complex<f32>, complex<f64>)  {
   // CHECK: %[[FLOAT_RESULT:.*]] = call @conjf(%[[FLOAT]])
   %float_result = complex.conj %float : complex<f32>
   // CHECK: %[[DOUBLE_RESULT:.*]] = call @conj(%[[DOUBLE]])
   %double_result = complex.conj %double : complex<f64>
   // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
   return %float_result, %double_result : complex<f32>, complex<f64>
+}
+
+// CHECK-LABEL: func @cabs_caller
+// CHECK-SAME: %[[FLOAT:.*]]: complex<f32>
+// CHECK-SAME: %[[DOUBLE:.*]]: complex<f64>
+func.func @cabs_caller(%float: complex<f32>, %double: complex<f64>) -> (f32, f64)  {
+  // CHECK: %[[FLOAT_RESULT:.*]] = call @cabsf(%[[FLOAT]])
+  %float_result = complex.abs %float : complex<f32>
+  // CHECK: %[[DOUBLE_RESULT:.*]] = call @cabs(%[[DOUBLE]])
+  %double_result = complex.abs %double : complex<f64>
+  // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
+  return %float_result, %double_result : f32, f64
 }
 No newline at end of file