[mlir][complex] Add pow/sqrt/tanh ops and lowering to libm

[bkramer]

Summary

Lowering through libm gives us a baseline version, even though it's not
going to be particularly fast. This is similar to what we do for some
math dialect ops.

[DavidTruby]

I think there's an ABI issue here, MLIR complex numbers are always lowered to an LLVM struct of 2 elements, but that isn't necessarily the correct platform ABI for passing complex numbers/C structs (C complex numbers and C structs have the same ABI).
I see an issue with complex<f32> on x86_64: cpow expects both arguments to be passed in %xmm0 (i,e, [2 x float] at the LLVM level) but here they are passed in %xmm0 and %xmm1 ({float, float} at the LLVM level)

I'm not sure what the fix is here. Does MLIR have target awareness that we can switch on to lower to the correct thing when the ABI is different?

[bkramer]

In D125550#3774395, @DavidTruby wrote:

I think there's an ABI issue here, MLIR complex numbers are always lowered to an LLVM struct of 2 elements, but that isn't necessarily the correct platform ABI for passing complex numbers/C structs (C complex numbers and C structs have the same ABI).
I see an issue with complex<f32> on x86_64: cpow expects both arguments to be passed in %xmm0 (i,e, [2 x float] at the LLVM level) but here they are passed in %xmm0 and %xmm1 ({float, float} at the LLVM level)

I'm not sure what the fix is here. Does MLIR have target awareness that we can switch on to lower to the correct thing when the ABI is different?

I was afraid this would happen. As far as I know MLIR doesn't have a way of using the correct ABI. The only place that has this information in LLVM is in Clang.

We probably just have to add non-libm expansions for all the ops.

[DavidTruby]

I was afraid this would happen. As far as I know MLIR doesn't have a way of using the correct ABI. The only place that has this information in LLVM is in Clang.

We probably just have to add non-libm expansions for all the ops.

I posted about this on discourse too, maybe we could start a discussion there? https://discourse.llvm.org/t/complex-to-libm-conversion-abi-issues/65131

Thanks!

mlir/include/mlir/Conversion/ComplexToLibm/ComplexToLibm.h

+//===- ComplexToLibm.h - Utils to convert from the complex dialect --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_CONVERSION_COMPLEXTOLIBM_COMPLEXTOLIBM_H_
+#define MLIR_CONVERSION_COMPLEXTOLIBM_COMPLEXTOLIBM_H_
+
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+template <typename T>
+class OperationPass;
+
+/// Populate the given list with patterns that convert from Complex to Libm
+/// calls.
+void populateComplexToLibmConversionPatterns(RewritePatternSet &patterns,
+                                             PatternBenefit benefit);
+
+/// Create a pass to convert Complex operations to libm calls.
+std::unique_ptr<OperationPass<ModuleOp>> createConvertComplexToLibmPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_COMPLEXTOLIBM_COMPLEXTOLIBM_H_

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h"
 #include "mlir/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.h"
 #include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h"
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"
 #include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"
+#include "mlir/Conversion/ComplexToLibm/ComplexToLibm.h"
 #include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.h"
 #include "mlir/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRVPass.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
 #include "mlir/Conversion/FuncToSPIRV/FuncToSPIRVPass.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"

mlir/include/mlir/Conversion/Passes.td

 
 def ConvertComplexToLLVM : Pass<"convert-complex-to-llvm"> {
   let summary = "Convert Complex dialect to LLVM dialect";
   let constructor = "mlir::createConvertComplexToLLVMPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
+// ComplexToLibm
+//===----------------------------------------------------------------------===//
+
+def ConvertComplexToLibm : Pass<"convert-complex-to-libm", "ModuleOp"> {
+  let summary = "Convert Complex dialect to libm calls";
+  let description = [{
+    This pass converts supported Complex ops to libm calls.
+  }];
+  let constructor = "mlir::createConvertComplexToLibmPass()";
+  let dependentDialects = [
+    "func::FuncDialect",
+  ];
+}
+
+//===----------------------------------------------------------------------===//
 // ComplexToStandard
 //===----------------------------------------------------------------------===//
 
 def ConvertComplexToStandard : Pass<"convert-complex-to-standard"> {
   let summary = "Convert Complex dialect to standard dialect";
   let constructor = "mlir::createConvertComplexToStandardPass()";
   let dependentDialects = ["math::MathDialect"];
 }

mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td

def NotEqualOp : Complex_Op<"neq",
 
   let arguments = (ins Complex<AnyFloat>:$lhs, Complex<AnyFloat>:$rhs);
   let results = (outs I1:$result);
 
   let assemblyFormat = "$lhs `,` $rhs  attr-dict `:` type($lhs)";
 }
 
 //===----------------------------------------------------------------------===//
+// PowOp
+//===----------------------------------------------------------------------===//
+
+def PowOp : ComplexArithmeticOp<"pow"> {
+  let summary = "complex power function";
+  let description = [{
+    The `sqrt` operation takes a complex number raises it to the given complex
+    exponent.
+
+    Example:
+
+    ```mlir
+    %a = complex.pow %b, %c : complex<f32>
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // ReOp
 //===----------------------------------------------------------------------===//
 
 def ReOp : ComplexUnaryOp<"re",
     [TypesMatchWith<"complex element type matches result type",
                     "complex", "real",
                     "$_self.cast<ComplexType>().getElementType()">]> {
   let summary = "extracts the real part of a complex number";
let description = [{
     %a = complex.sin %b : complex<f32>
     ```
   }];
 
   let results = (outs Complex<AnyFloat>:$result);
 }
 
 //===----------------------------------------------------------------------===//
+// SqrtOp
+//===----------------------------------------------------------------------===//
+
+def SqrtOp : ComplexUnaryOp<"sqrt", [SameOperandsAndResultType]> {
+  let summary = "complex square root";
+  let description = [{
+    The `sqrt` operation takes a complex number and returns its square root.
+
+    Example:
+
+    ```mlir
+    %a = complex.sqrt %b : complex<f32>
+    ```
+  }];
+
+  let results = (outs Complex<AnyFloat>:$result);
+}
+
+//===----------------------------------------------------------------------===//
 // SubOp
 //===----------------------------------------------------------------------===//
 
 def SubOp : ComplexArithmeticOp<"sub"> {
   let summary = "complex subtraction";
   let description = [{
     The `sub` operation takes two complex numbers and returns their difference.
 
     Example:
 
     ```mlir
     %a = complex.sub %b, %c : complex<f32>
     ```
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// TanhOp
+//===----------------------------------------------------------------------===//
+
+def TanhOp : ComplexUnaryOp<"tanh", [SameOperandsAndResultType]> {
+  let summary = "complex hyperbolic tangent";
+  let description = [{
+    The `tanh` operation takes a complex number and returns its hyperbolic
+    tangent.
+
+    Example:
+
+    ```mlir
+    %a = complex.tanh %b : complex<f32>
+    ```
+  }];
+
+  let results = (outs Complex<AnyFloat>:$result);
+}
+
 #endif // COMPLEX_OPS

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AMDGPUToROCDL)
 add_subdirectory(ArithmeticToLLVM)
 add_subdirectory(ArithmeticToSPIRV)
 add_subdirectory(ArmNeon2dToIntr)
 add_subdirectory(AsyncToLLVM)
 add_subdirectory(BufferizationToMemRef)
 add_subdirectory(ComplexToLLVM)
+add_subdirectory(ComplexToLibm)
 add_subdirectory(ComplexToStandard)
 add_subdirectory(ControlFlowToLLVM)
 add_subdirectory(ControlFlowToSPIRV)
 add_subdirectory(FuncToLLVM)
 add_subdirectory(FuncToSPIRV)
 add_subdirectory(GPUCommon)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)

mlir/lib/Conversion/ComplexToLibm/CMakeLists.txt

+add_mlir_conversion_library(MLIRComplexToLibm
+  ComplexToLibm.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ComplexToLibm
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRDialectUtils
+  MLIRFunc
+  MLIRComplex
+  MLIRTransformUtils
+  )

mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp

+//===-- ComplexToLibm.cpp - conversion from Complex to libm calls ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/ComplexToLibm/ComplexToLibm.h"
+
+#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 {
+// Pattern to convert scalar complex operations to calls to libm functions.
+// Additionally the libm function signatures are declared.
+template <typename Op>
+struct ScalarOpToLibmCall : public OpRewritePattern<Op> {
+public:
+  using OpRewritePattern<Op>::OpRewritePattern;
+  ScalarOpToLibmCall<Op>(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 {
+  auto module = SymbolTable::getNearestSymbolTable(op);
+  auto type = op.getType().template cast<ComplexType>();
+  Type elementType = type.getElementType();
+  if (!elementType.isa<Float32Type, Float64Type>())
+    return failure();
+
+  auto name =
+      elementType.getIntOrFloatBitWidth() == 64 ? 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());
+
+  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);
+}
+
+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>();
+  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-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>
+  %double_result = complex.pow %double, %double : complex<f64>
+  // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
+  return %float_result, %double_result : complex<f32>, complex<f64>
+}
+
+// CHECK-LABEL: func @csqrt_caller
+// CHECK-SAME: %[[FLOAT:.*]]: complex<f32>
+// CHECK-SAME: %[[DOUBLE:.*]]: complex<f64>
+func.func @csqrt_caller(%float: complex<f32>, %double: complex<f64>) -> (complex<f32>, complex<f64>)  {
+  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @csqrtf(%[[FLOAT]]) : (complex<f32>) -> complex<f32>
+  %float_result = complex.sqrt %float : complex<f32>
+  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @csqrt(%[[DOUBLE]]) : (complex<f64>) -> complex<f64>
+  %double_result = complex.sqrt %double : complex<f64>
+  // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
+  return %float_result, %double_result : complex<f32>, complex<f64>
+}
+
+// CHECK-LABEL: func @ctanh_caller
+// CHECK-SAME: %[[FLOAT:.*]]: complex<f32>
+// CHECK-SAME: %[[DOUBLE:.*]]: complex<f64>
+func.func @ctanh_caller(%float: complex<f32>, %double: complex<f64>) -> (complex<f32>, complex<f64>)  {
+  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @ctanhf(%[[FLOAT]]) : (complex<f32>) -> complex<f32>
+  %float_result = complex.tanh %float : complex<f32>
+  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @ctanh(%[[DOUBLE]]) : (complex<f64>) -> complex<f64>
+  %double_result = complex.tanh %double : complex<f64>
+  // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
+  return %float_result, %double_result : complex<f32>, complex<f64>
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":AMDGPUToROCDL",
         ":AffineToStandard",
         ":ArithmeticToLLVM",
         ":ArithmeticToSPIRV",
         ":ArmNeon2dToIntr",
         ":AsyncToLLVM",
         ":BufferizationToMemRef",
         ":ComplexToLLVM",
+        ":ComplexToLibm",
         ":ComplexToStandard",
         ":ControlFlowToLLVM",
         ":ControlFlowToSPIRV",
         ":ConversionPassIncGen",
         ":FuncToLLVM",
         ":FuncToSPIRV",
         ":GPUToGPURuntimeTransforms",
         ":GPUToNVVMTransforms",
deps = [
         ":Support",
         ":TensorDialect",
         ":Transforms",
         ":VectorOps",
         "//llvm:Support",
     ],
 )
 
-
 cc_library(
     name = "TensorToSPIRV",
     srcs = glob([
         "lib/Conversion/TensorToSPIRV/*.cpp",
         "lib/Conversion/TensorToSPIRV/*.h",
     ]) + [":ConversionPassDetail"],
     hdrs = glob([
         "include/mlir/Conversion/TensorToSPIRV/*.h",
deps = [
         ":Async",
         ":AsyncPassIncGen",
         ":AsyncToLLVM",
         ":AsyncTransforms",
         ":BufferizationDialect",
         ":BufferizationTransforms",
         ":ComplexDialect",
         ":ComplexToLLVM",
+        ":ComplexToLibm",
         ":ControlFlowOps",
         ":ConversionPasses",
         ":DLTIDialect",
         ":EmitC",
         ":FuncDialect",
         ":FuncToLLVM",
         ":FuncToSPIRV",
         ":FuncTransforms",
deps = [
         ":Support",
         ":Transforms",
         "//llvm:Core",
         "//llvm:Support",
     ],
 )
 
 cc_library(
+    name = "ComplexToLibm",
+    srcs = glob([
+        "lib/Conversion/ComplexToLibm/*.cpp",
+        "lib/Conversion/ComplexToLibm/*.h",
+    ]) + [":ConversionPassDetail"],
+    hdrs = glob([
+        "include/mlir/Conversion/ComplexToLibm/*.h",
+    ]),
+    includes = ["include"],
+    deps = [
+        ":ComplexDialect",
+        ":ConversionPassIncGen",
+        ":DialectUtils",
+        ":FuncDialect",
+        ":IR",
+        ":Pass",
+        ":Support",
+        ":Transforms",
+        "//llvm:Core",
+        "//llvm:Support",
+    ],
+)
+
+cc_library(
     name = "ComplexToStandard",
     srcs = glob([
         "lib/Conversion/ComplexToStandard/*.cpp",
         "lib/Conversion/ComplexToStandard/*.h",
     ]) + [":ConversionPassDetail"],
     hdrs = glob([
         "include/mlir/Conversion/ComplexToStandard/*.h",
     ]),