Address reviewer changes.

Add conversion for complex.exp.

Author: TIFitis

flang/lib/Optimizer/CodeGen/CodeGen.cpp

@@ -4146,24 +4146,24 @@ class FIRToLLVMLowering
     // conversions that affect the ModuleOp, e.g. create new
     // function operations in it. We have to run such conversions
     // as passes here.
-    mlir::OpPassManager mathConvertionPM("builtin.module");
+    mlir::OpPassManager mathConversionPM("builtin.module");
 
     bool isAMDGCN = fir::getTargetTriple(mod).isAMDGCN();
     // If compiling for AMD target some math operations must be lowered to AMD
     // GPU library calls, the rest can be converted to LLVM intrinsics, which
     // is handled in the mathToLLVM conversion. The lowering to libm calls is
     // not needed since all math operations are handled this way.
     if (isAMDGCN) {
-      mathConvertionPM.addPass(mlir::createConvertMathToROCDL());
-      mathConvertionPM.addPass(mlir::createConvertComplexToROCDL());
+      mathConversionPM.addPass(mlir::createConvertMathToROCDL());
+      mathConversionPM.addPass(mlir::createConvertComplexToROCDL());
     }
 
     // Convert math::FPowI operations to inline implementation
     // only if the exponent's width is greater than 32, otherwise,
     // it will be lowered to LLVM intrinsic operation by a later conversion.
     mlir::ConvertMathToFuncsOptions mathToFuncsOptions{};
     mathToFuncsOptions.minWidthOfFPowIExponent = 33;
-    mathConvertionPM.addPass(
+    mathConversionPM.addPass(
         mlir::createConvertMathToFuncs(mathToFuncsOptions));
 
     mlir::ConvertComplexToStandardPassOptions complexToStandardOptions{};
@@ -4176,15 +4176,15 @@ class FIRToLLVMLowering
       complexToStandardOptions.complexRange =
           mlir::complex::ComplexRangeFlags::improved;
     }
-    mathConvertionPM.addPass(
+    mathConversionPM.addPass(
         mlir::createConvertComplexToStandardPass(complexToStandardOptions));
 
     // Convert Math dialect operations into LLVM dialect operations.
     // There is no way to prefer MathToLLVM patterns over MathToLibm
     // patterns (applied below), so we have to run MathToLLVM conversion here.
-    mathConvertionPM.addNestedPass<mlir::func::FuncOp>(
+    mathConversionPM.addNestedPass<mlir::func::FuncOp>(
         mlir::createConvertMathToLLVMPass());
-    if (mlir::failed(runPipeline(mathConvertionPM, mod)))
+    if (mlir::failed(runPipeline(mathConversionPM, mod)))
       return signalPassFailure();
 
     std::optional<mlir::DataLayout> dl =

mlir/include/mlir/Conversion/ComplexToROCDL/ComplexToROCDL.h

@@ -20,8 +20,7 @@ class RewritePatternSet;
 
 /// Populate the given list with patterns that convert from Complex to ROCDL
 /// calls.
-void populateComplexToROCDLConversionPatterns(RewritePatternSet &patterns,
-                                              PatternBenefit benefit);
+void populateComplexToROCDLConversionPatterns(RewritePatternSet &patterns);
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_COMPLEXTOROCDL_COMPLEXTOROCDL_H_

mlir/lib/Conversion/ComplexToROCDL/CMakeLists.txt

@@ -7,9 +7,6 @@ add_mlir_conversion_library(MLIRComplexToROCDL
   DEPENDS
   MLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRComplexDialect
   MLIRFuncDialect

mlir/lib/Conversion/ComplexToROCDL/ComplexToROCDL.cpp

@@ -11,7 +11,6 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
-#include <optional>
 
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTCOMPLEXTOROCDL
@@ -21,36 +20,38 @@ namespace mlir {
 using namespace mlir;
 
 namespace {
-struct FloatTypeResolver {
-  std::optional<bool> operator()(Type type) const {
-    auto elementType = cast<FloatType>(type);
-    if (!isa<Float32Type, Float64Type>(elementType))
-      return {};
-    return elementType.getIntOrFloatBitWidth() == 64;
-  }
-};
 
-template <typename Op, typename TypeResolver = FloatTypeResolver>
-struct ScalarOpToROCDLCall : public OpRewritePattern<Op> {
+template <typename Op>
+// Pattern to convert Complex ops to ROCDL function calls.
+struct ComplexOpToROCDLCall : public OpRewritePattern<Op> {
   using OpRewritePattern<Op>::OpRewritePattern;
-  ScalarOpToROCDLCall(MLIRContext *context, StringRef floatFunc,
-                      StringRef doubleFunc, PatternBenefit benefit)
+  ComplexOpToROCDLCall(MLIRContext *context, StringRef floatFunc,
+                       StringRef doubleFunc, PatternBenefit benefit = 1)
       : OpRewritePattern<Op>(context, benefit), floatFunc(floatFunc),
         doubleFunc(doubleFunc) {}
 
   LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final {
-    auto module = SymbolTable::getNearestSymbolTable(op);
-    auto isDouble = TypeResolver()(op.getType());
-    if (!isDouble.has_value())
+    Operation *symTable = SymbolTable::getNearestSymbolTable(op);
+    Type resType = op.getType();
+    if (auto complexType = dyn_cast<ComplexType>(resType))
+      resType = complexType.getElementType();
+    FloatType floatTy = dyn_cast<FloatType>(resType);
+    if (!floatTy)
       return failure();
 
-    auto name = *isDouble ? doubleFunc : floatFunc;
+    StringRef name;
+    if (floatTy.isF64())
+      name = doubleFunc;
+    else if (floatTy.isF32())
+      name = floatFunc;
+    else
+      return failure();
 
     auto opFunc = dyn_cast_or_null<SymbolOpInterface>(
-        SymbolTable::lookupSymbolIn(module, name));
+        SymbolTable::lookupSymbolIn(symTable, name));
     if (!opFunc) {
       OpBuilder::InsertionGuard guard(rewriter);
-      rewriter.setInsertionPointToStart(&module->getRegion(0).front());
+      rewriter.setInsertionPointToStart(&symTable->getRegion(0).front());
       auto funcTy = FunctionType::get(
           rewriter.getContext(), op->getOperandTypes(), op->getResultTypes());
       opFunc =
@@ -67,10 +68,12 @@ struct ScalarOpToROCDLCall : public OpRewritePattern<Op> {
 };
 } // namespace
 
-void mlir::populateComplexToROCDLConversionPatterns(RewritePatternSet &patterns,
-                                                    PatternBenefit benefit) {
-  patterns.add<ScalarOpToROCDLCall<complex::AbsOp>>(
-      patterns.getContext(), "__ocml_cabs_f32", "__ocml_cabs_f64", benefit);
+void mlir::populateComplexToROCDLConversionPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<ComplexOpToROCDLCall<complex::AbsOp>>(
+      patterns.getContext(), "__ocml_cabs_f32", "__ocml_cabs_f64");
+  patterns.add<ComplexOpToROCDLCall<complex::ExpOp>>(
+      patterns.getContext(), "__ocml_cexp_f32", "__ocml_cexp_f64");
 }
 
 namespace {
@@ -81,14 +84,14 @@ struct ConvertComplexToROCDLPass
 } // namespace
 
 void ConvertComplexToROCDLPass::runOnOperation() {
-  auto module = getOperation();
+  Operation *op = getOperation();
 
   RewritePatternSet patterns(&getContext());
-  populateComplexToROCDLConversionPatterns(patterns, /*benefit=*/1);
+  populateComplexToROCDLConversionPatterns(patterns);
 
   ConversionTarget target(getContext());
   target.addLegalDialect<func::FuncDialect>();
-  target.addIllegalOp<complex::AbsOp>();
-  if (failed(applyPartialConversion(module, target, std::move(patterns))))
+  target.addIllegalOp<complex::AbsOp, complex::ExpOp>();
+  if (failed(applyPartialConversion(op, target, std::move(patterns))))
     signalPassFailure();
 }

mlir/test/Conversion/ComplexToROCDL/complex-to-rocdl.mlir

@@ -1,13 +1,26 @@
-// RUN: mlir-opt %s -convert-complex-to-rocdl -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -convert-complex-to-rocdl | FileCheck %s
 
 // CHECK-DAG: @__ocml_cabs_f32(complex<f32>) -> f32
 // CHECK-DAG: @__ocml_cabs_f64(complex<f64>) -> f64
+// CHECK-DAG: @__ocml_cexp_f32(complex<f32>) -> complex<f32>
+// CHECK-DAG: @__ocml_cexp_f64(complex<f64>) -> complex<f64>
 
+//CHECK-LABEL: @abs_caller
 func.func @abs_caller(%f: complex<f32>, %d: complex<f64>) -> (f32, f64) {
-  // CHECK: %[[RF:.*]] = call @__ocml_cabs_f32(%[[F:.*]])
+  // CHECK: %[[RF:.*]] = call @__ocml_cabs_f32(%{{.*}})
   %rf = complex.abs %f : complex<f32>
-  // CHECK: %[[RD:.*]] = call @__ocml_cabs_f64(%[[D:.*]])
+  // CHECK: %[[RD:.*]] = call @__ocml_cabs_f64(%{{.*}})
   %rd = complex.abs %d : complex<f64>
   // CHECK: return %[[RF]], %[[RD]]
   return %rf, %rd : f32, f64
 }
+
+//CHECK-LABEL: @exp_caller
+func.func @exp_caller(%f: complex<f32>, %d: complex<f64>) -> (complex<f32>, complex<f64>) {
+  // CHECK: %[[EF:.*]] = call @__ocml_cexp_f32(%{{.*}})
+  %ef = complex.exp %f : complex<f32>
+  // CHECK: %[[ED:.*]] = call @__ocml_cexp_f64(%{{.*}})
+  %ed = complex.exp %d : complex<f64>
+  // CHECK: return %[[EF]], %[[ED]]
+  return %ef, %ed : complex<f32>, complex<f64>
+}