[MLIR] Add ComplexToStandard conversion pass.

[akuegel]

Summary

So far, only a conversion for complex::AbsOp is done, but more will be added.

[pifon2a]

mlir/include/mlir/Conversion/Passes.td

106

is it a ModuleOp pass because you want to convert func args of ComplexType? If not, then would FunctionPass be enough?

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

38

nit: realSqr?

[akuegel]

Address review comments.

[akuegel]

mlir/include/mlir/Conversion/Passes.td

106

Thanks for the suggestion :)
FunctionPass is enough. Done.

[kiranchandramohan]

The conversion of AbsOp in this patch does not account for underflows or overflows. Since it does not use higher precisions for computing the intermediate results, it might not be precise as well. We are hitting an underflow issue in Flang due to using this conversion. https://github.com/llvm/llvm-project/issues/62011
What are the plans for this pass? Would you accept fixes to improve precision or is the assumption that we should only use this pass where preciseness is not important?
Thanks in advance.

[akuegel]

In D101442#4257833, @kiranchandramohan wrote:

The conversion of AbsOp in this patch does not account for underflows or overflows. Since it does not use higher precisions for computing the intermediate results, it might not be precise as well. We are hitting an underflow issue in Flang due to using this conversion. https://github.com/llvm/llvm-project/issues/62011
What are the plans for this pass? Would you accept fixes to improve precision or is the assumption that we should only use this pass where preciseness is not important?
Thanks in advance.

We would accept fixes to improve precision. Thanks in advance :)

mlir/include/mlir/Conversion/ComplexToStandard/ComplexToStandard.h

+//===- ComplexToStandard.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_COMPLEXTOSTANDARD_COMPLEXTOSTANDARD_H_
+#define MLIR_CONVERSION_COMPLEXTOSTANDARD_COMPLEXTOSTANDARD_H_
+
+#include <memory>
+
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+class FuncOp;
+class RewritePatternSet;
+template <typename T>
+class OperationPass;
+
+/// Populate the given list with patterns that convert from Complex to Standard.
+void populateComplexToStandardConversionPatterns(RewritePatternSet &patterns);
+
+/// Create a pass to convert Complex operations to the Standard dialect.
+std::unique_ptr<OperationPass<FuncOp>> createConvertComplexToStandardPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_COMPLEXTOSTANDARD_COMPLEXTOSTANDARD_H_

mlir/include/mlir/Conversion/Passes.h

 //===- Passes.h - Conversion Pass Construction and Registration -----------===//
 //
 // 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_PASSES_H
 #define MLIR_CONVERSION_PASSES_H
 
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"
+#include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"

mlir/include/mlir/Conversion/Passes.td

 
 def ConvertComplexToLLVM : Pass<"convert-complex-to-llvm", "ModuleOp"> {
   let summary = "Convert Complex dialect to LLVM dialect";
   let constructor = "mlir::createConvertComplexToLLVMPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
+// ComplexToStandard
+//===----------------------------------------------------------------------===//
+
+def ConvertComplexToStandard : FunctionPass<"convert-complex-to-standard"> {
+  let summary = "Convert Complex dialect to standard dialect";
+  let constructor = "mlir::createConvertComplexToStandardPass()";
+  let dependentDialects = [
+    "complex::ComplexDialect",
+    "math::MathDialect",
+    "StandardOpsDialect"
+  ];
+}
+
+//===----------------------------------------------------------------------===//
 // GPUCommon
 //===----------------------------------------------------------------------===//
 
 def GpuToLLVMConversionPass : Pass<"gpu-to-llvm", "ModuleOp"> {
   let summary = "Convert GPU dialect to LLVM dialect with GPU runtime calls";
   let constructor = "mlir::createGpuToLLVMConversionPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];
 }

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AsyncToLLVM)
 add_subdirectory(ComplexToLLVM)
+add_subdirectory(ComplexToStandard)
 add_subdirectory(GPUCommon)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LinalgToStandard)

mlir/lib/Conversion/ComplexToStandard/CMakeLists.txt

+add_mlir_conversion_library(MLIRComplexToStandard
+  ComplexToStandard.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ComplexToStandard
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRComplex
+  MLIRIR
+  MLIRMath
+  MLIRStandard
+  MLIRTransforms
+  )

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

+//===- ComplexToStandard.cpp - conversion from Complex to Standard 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
+
+#include <memory>
+
+#include "../PassDetail.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+struct AbsOpConversion : public OpConversionPattern<complex::AbsOp> {
+  using OpConversionPattern<complex::AbsOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::AbsOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    complex::AbsOp::Adaptor transformed(operands);
+    auto loc = op.getLoc();
+    auto type = op.getType();
+
+    Value real =
+        rewriter.create<complex::ReOp>(loc, type, transformed.complex());
+    Value imag =
+        rewriter.create<complex::ImOp>(loc, type, transformed.complex());
+    Value realSqr = rewriter.create<MulFOp>(loc, real, real);
+    Value imagSqr = rewriter.create<MulFOp>(loc, imag, imag);
+    Value sqNorm = rewriter.create<AddFOp>(loc, realSqr, imagSqr);
+
+    rewriter.replaceOpWithNewOp<math::SqrtOp>(op, sqNorm);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateComplexToStandardConversionPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<AbsOpConversion>(patterns.getContext());
+}
+
+namespace {
+struct ConvertComplexToStandardPass
+    : public ConvertComplexToStandardBase<ConvertComplexToStandardPass> {
+  void runOnFunction() override;
+};
+
+void ConvertComplexToStandardPass::runOnFunction() {
+  auto function = getFunction();
+
+  // Convert to the Standard dialect using the converter defined above.
+  RewritePatternSet patterns(&getContext());
+  populateComplexToStandardConversionPatterns(patterns);
+
+  ConversionTarget target(getContext());
+  target.addLegalDialect<StandardOpsDialect, math::MathDialect,
+                         complex::ComplexDialect>();
+  target.addIllegalOp<complex::AbsOp>();
+  if (failed(applyPartialConversion(function, target, std::move(patterns))))
+    signalPassFailure();
+}
+} // namespace
+
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createConvertComplexToStandardPass() {
+  return std::make_unique<ConvertComplexToStandardPass>();
+}

mlir/lib/Conversion/PassDetail.h

 namespace LLVM {
 class LLVMDialect;
 } // end namespace LLVM
 
 namespace NVVM {
 class NVVMDialect;
 } // end namespace NVVM
 
+namespace math {
+class MathDialect;
+} // end namespace math
+
 namespace memref {
 class MemRefDialect;
 } // end namespace memref
 
 namespace omp {
 class OpenMPDialect;
 } // end namespace omp
 

mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir

+// RUN: mlir-opt %s -convert-complex-to-standard | FileCheck %s
+
+// CHECK-LABEL: func @complex_abs
+// CHECK-SAME: %[[ARG:.*]]: complex<f32>
+func @complex_abs(%arg: complex<f32>) -> f32 {
+  %abs = complex.abs %arg: complex<f32>
+  return %abs : f32
+}
+// CHECK: %[[REAL:.*]] = complex.re %[[ARG]] : complex<f32>
+// CHECK: %[[IMAG:.*]] = complex.im %[[ARG]] : complex<f32>
+// CHECK-DAG: %[[REAL_SQ:.*]] = mulf %[[REAL]], %[[REAL]] : f32
+// CHECK-DAG: %[[IMAG_SQ:.*]] = mulf %[[IMAG]], %[[IMAG]] : f32
+// CHECK: %[[SQ_NORM:.*]] = addf %[[REAL_SQ]], %[[IMAG_SQ]] : f32
+// CHECK: %[[NORM:.*]] = math.sqrt %[[SQ_NORM]] : f32
+// CHECK: return %[[NORM]] : f32
+

mlir/test/Conversion/ComplexToStandard/full-conversion.mlir

+// RUN: mlir-opt %s -convert-complex-to-standard -convert-complex-to-llvm -convert-std-to-llvm | FileCheck %s
+
+// CHECK-LABEL: llvm.func @complex_abs
+// CHECK-SAME: %[[ARG:.*]]: ![[C_TY:.*]])
+func @complex_abs(%arg: complex<f32>) -> f32 {
+  %abs = complex.abs %arg: complex<f32>
+  return %abs : f32
+}
+// CHECK: %[[REAL:.*]] = llvm.extractvalue %[[ARG]][0] : ![[C_TY]]
+// CHECK: %[[IMAG:.*]] = llvm.extractvalue %[[ARG]][1] : ![[C_TY]]
+// CHECK-DAG: %[[REAL_SQ:.*]] = llvm.fmul %[[REAL]], %[[REAL]]  : f32
+// CHECK-DAG: %[[IMAG_SQ:.*]] = llvm.fmul %[[IMAG]], %[[IMAG]]  : f32
+// CHECK: %[[SQ_NORM:.*]] = llvm.fadd %[[REAL_SQ]], %[[IMAG_SQ]]  : f32
+// CHECK: %[[NORM:.*]] = "llvm.intr.sqrt"(%[[SQ_NORM]]) : (f32) -> f32
+// CHECK: llvm.return %[[NORM]] : f32
+