[MLIR][Math] Emit llvm.func when converting to libm

[chelini]

Summary

Currently, conversion to libm does not play nicely with
convert-func-to-llvm='emit-c-wrappers=1'. Conversion from func to LLVM emits a
'_mlir_ciface_*' call which is not the intended behaviour here. Emit an
llvm.func instead of a func.func when converting to libm.

Example:

llvm.func @tanh_caller(%float: f32) -> (f32) {
 %float_result = math.tanh %float : f32
 llvm.return %float_result : f32
}

Will convert to:

llvm.func @tanhf(%arg0: f32) -> f32 attributes {sym_visibility = "private"} {
  %0 = llvm.call @_mlir_ciface_tanhf(%arg0) : (f32) -> f32
  llvm.return %0 : f32
}

llvm.func @_mlir_ciface_tanhf(f32) -> f32 attributes {sym_visibility = "private"}

llvm.func @tanh_caller(%arg0: f32) -> f32 {
  %0 = llvm.call @tanhf(%arg0) : (f32) -> f32
  llvm.return %0 : f32
}

This is not what you want in this situation.

[chelini]

I noticed that this exact behavior applies when converting complex to libm. I will follow up with a patch for complex if this is the right direction.

[ftynse]

This feels a bit like premature lowering. Do you have a flow that relies on the emit-c-wrappers option? We have a better alternative -- attach the llvm.emit_c_interface attributes to functions that need the interface -- that leaves most function alone. Maybe we should just remove the option instead and only have the attribute, anyone who wants blanket wrapper emission can just add a pass that attaches the attribute to all funcs.

[chelini]

In D126964#3563386, @ftynse wrote:

This feels a bit like premature lowering. Do you have a flow that relies on the emit-c-wrappers option? We have a better alternative -- attach the llvm.emit_c_interface attributes to functions that need the interface -- that leaves most function alone. Maybe we should just remove the option instead and only have the attribute, anyone who wants blanket wrapper emission can just add a pass that attaches the attribute to all funcs.

You are right; the flow I have now uses only the 'emit-c-wrapper' option, and I am hitting the same problem Nicolas mentioned in the revision https://reviews.llvm.org/D77314 about undefined symbols. It makes sense to drop the option if we have a better alternative.

[chelini]

See discussion here: https://reviews.llvm.org/D127952

mlir/include/mlir/Conversion/Passes.td

 def ConvertMathToLibm : Pass<"convert-math-to-libm", "ModuleOp"> {
   let summary = "Convert Math dialect to libm calls";
   let description = [{
     This pass converts supported Math ops to libm calls.
   }];
   let constructor = "mlir::createConvertMathToLibmPass()";
   let dependentDialects = [
     "arith::ArithmeticDialect",
-    "func::FuncDialect",
+    "LLVM::LLVMDialect",
     "vector::VectorDialect",
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // MathToLLVM
 //===----------------------------------------------------------------------===//
 

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

 //===-- MathToLibm.cpp - conversion from Math 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/MathToLibm/MathToLibm.h"
 
 #include "../PassDetail.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
PromoteOpToF32<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const {
   rewriter.replaceOpWithNewOp<arith::TruncFOp>(op, opType, newOp);
   return success();
 }
 
 template <typename Op>
 LogicalResult
 ScalarOpToLibmCall<Op>::matchAndRewrite(Op op,
                                         PatternRewriter &rewriter) const {
+  assert(op->getNumResults() == 1 && "expect 1 result only");
   auto module = SymbolTable::getNearestSymbolTable(op);
   auto type = op.getType();
   if (!type.template isa<Float32Type, Float64Type>())
     return failure();
 
   auto name = type.getIntOrFloatBitWidth() == 64 ? doubleFunc : floatFunc;
-  auto opFunc = dyn_cast_or_null<SymbolOpInterface>(
+  LLVM::LLVMFuncOp opFunc = dyn_cast_or_null<LLVM::LLVMFuncOp>(
       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,
+    auto opFunctionTy = LLVM::LLVMFunctionType::get(
+        op->getResultTypes()[0], llvm::to_vector<4>(op->getOperandTypes()));
+    opFunc = rewriter.create<LLVM::LLVMFuncOp>(rewriter.getUnknownLoc(), name,
                                                opFunctionTy);
     opFunc.setPrivate();
   }
   assert(isa<FunctionOpInterface>(SymbolTable::lookupSymbolIn(module, name)));
 
-  rewriter.replaceOpWithNewOp<func::CallOp>(op, name, op.getType(),
-                                            op->getOperands());
+  rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, opFunc, op->getOperands());
 
   return success();
 }
 
 void mlir::populateMathToLibmConversionPatterns(RewritePatternSet &patterns,
                                                 PatternBenefit benefit) {
   patterns.add<VecOpToScalarOp<math::Atan2Op>, VecOpToScalarOp<math::ExpM1Op>,
                VecOpToScalarOp<math::TanhOp>>(patterns.getContext(), benefit);

 void ConvertMathToLibmPass::runOnOperation() {
   auto module = getOperation();
 
   RewritePatternSet patterns(&getContext());
   populateMathToLibmConversionPatterns(patterns, /*benefit=*/1);
 
   ConversionTarget target(getContext());
   target.addLegalDialect<arith::ArithmeticDialect, BuiltinDialect,
-                         func::FuncDialect, vector::VectorDialect>();
+                         LLVM::LLVMDialect, vector::VectorDialect>();
   target.addIllegalDialect<math::MathDialect>();
   if (failed(applyPartialConversion(module, target, std::move(patterns))))
     signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertMathToLibmPass() {
   return std::make_unique<ConvertMathToLibmPass>();
 }

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

 // RUN: mlir-opt %s -convert-math-to-libm -canonicalize | FileCheck %s
 
 // CHECK-DAG: @erf(f64) -> f64
 // CHECK-DAG: @erff(f32) -> f32
 // CHECK-DAG: @expm1(f64) -> f64
 // CHECK-DAG: @expm1f(f32) -> f32
 // CHECK-DAG: @atan2(f64, f64) -> f64
 // CHECK-DAG: @atan2f(f32, f32) -> f32
 // CHECK-DAG: @tanh(f64) -> f64
 // CHECK-DAG: @tanhf(f32) -> f32
 
 // CHECK-LABEL: func @tanh_caller
 // CHECK-SAME: %[[FLOAT:.*]]: f32
 // CHECK-SAME: %[[DOUBLE:.*]]: f64
 func.func @tanh_caller(%float: f32, %double: f64) -> (f32, f64)  {
-  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @tanhf(%[[FLOAT]]) : (f32) -> f32
+  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = llvm.call @tanhf(%[[FLOAT]]) : (f32) -> f32
   %float_result = math.tanh %float : f32
-  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @tanh(%[[DOUBLE]]) : (f64) -> f64
+  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = llvm.call @tanh(%[[DOUBLE]]) : (f64) -> f64
   %double_result = math.tanh %double : f64
   // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
   return %float_result, %double_result : f32, f64
 }
 
 
 // CHECK-LABEL: func @atan2_caller
 // CHECK-SAME: %[[FLOAT:.*]]: f32
 // CHECK-SAME: %[[DOUBLE:.*]]: f64
 // CHECK-SAME: %[[HALF:.*]]: f16
 // CHECK-SAME: %[[BFLOAT:.*]]: bf16
 func.func @atan2_caller(%float: f32, %double: f64, %half: f16, %bfloat: bf16) -> (f32, f64, f16, bf16) {
-  // CHECK: %[[FLOAT_RESULT:.*]] = call @atan2f(%[[FLOAT]], %[[FLOAT]]) : (f32, f32) -> f32
+  // CHECK: %[[FLOAT_RESULT:.*]] = llvm.call @atan2f(%[[FLOAT]], %[[FLOAT]]) : (f32, f32) -> f32
   %float_result = math.atan2 %float, %float : f32
-  // CHECK: %[[DOUBLE_RESULT:.*]] = call @atan2(%[[DOUBLE]], %[[DOUBLE]]) : (f64, f64) -> f64
+  // CHECK: %[[DOUBLE_RESULT:.*]] = llvm.call @atan2(%[[DOUBLE]], %[[DOUBLE]]) : (f64, f64) -> f64
   %double_result = math.atan2 %double, %double : f64
   // CHECK: %[[HALF_PROMOTED1:.*]] = arith.extf %[[HALF]] : f16 to f32
   // CHECK: %[[HALF_PROMOTED2:.*]] = arith.extf %[[HALF]] : f16 to f32
-  // CHECK: %[[HALF_CALL:.*]] = call @atan2f(%[[HALF_PROMOTED1]], %[[HALF_PROMOTED2]]) : (f32, f32) -> f32
+  // CHECK: %[[HALF_CALL:.*]] = llvm.call @atan2f(%[[HALF_PROMOTED1]], %[[HALF_PROMOTED2]]) : (f32, f32) -> f32
   // CHECK: %[[HALF_RESULT:.*]] = arith.truncf %[[HALF_CALL]] : f32 to f16
   %half_result = math.atan2 %half, %half : f16
   // CHECK: %[[BFLOAT_PROMOTED1:.*]] = arith.extf %[[BFLOAT]] : bf16 to f32
   // CHECK: %[[BFLOAT_PROMOTED2:.*]] = arith.extf %[[BFLOAT]] : bf16 to f32
-  // CHECK: %[[BFLOAT_CALL:.*]] = call @atan2f(%[[BFLOAT_PROMOTED1]], %[[BFLOAT_PROMOTED2]]) : (f32, f32) -> f32
+  // CHECK: %[[BFLOAT_CALL:.*]] = llvm.call @atan2f(%[[BFLOAT_PROMOTED1]], %[[BFLOAT_PROMOTED2]]) : (f32, f32) -> f32
   // CHECK: %[[BFLOAT_RESULT:.*]] = arith.truncf %[[BFLOAT_CALL]] : f32 to bf16
   %bfloat_result = math.atan2 %bfloat, %bfloat : bf16
   // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]], %[[HALF_RESULT]], %[[BFLOAT_RESULT]]
   return %float_result, %double_result, %half_result, %bfloat_result : f32, f64, f16, bf16
 }
 
 // CHECK-LABEL: func @erf_caller
 // CHECK-SAME: %[[FLOAT:.*]]: f32
 // CHECK-SAME: %[[DOUBLE:.*]]: f64
 func.func @erf_caller(%float: f32, %double: f64) -> (f32, f64)  {
-  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @erff(%[[FLOAT]]) : (f32) -> f32
+  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = llvm.call @erff(%[[FLOAT]]) : (f32) -> f32
   %float_result = math.erf %float : f32
-  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @erf(%[[DOUBLE]]) : (f64) -> f64
+  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = llvm.call @erf(%[[DOUBLE]]) : (f64) -> f64
   %double_result = math.erf %double : f64
   // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
   return %float_result, %double_result : f32, f64
 }
 
 // CHECK-LABEL: func @expm1_caller
 // CHECK-SAME: %[[FLOAT:.*]]: f32
 // CHECK-SAME: %[[DOUBLE:.*]]: f64
 func.func @expm1_caller(%float: f32, %double: f64) -> (f32, f64) {
-  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @expm1f(%[[FLOAT]]) : (f32) -> f32
+  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = llvm.call @expm1f(%[[FLOAT]]) : (f32) -> f32
   %float_result = math.expm1 %float : f32
-  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @expm1(%[[DOUBLE]]) : (f64) -> f64
+  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = llvm.call @expm1(%[[DOUBLE]]) : (f64) -> f64
   %double_result = math.expm1 %double : f64
   // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
   return %float_result, %double_result : f32, f64
 }
 
 func.func @expm1_vec_caller(%float: vector<2xf32>, %double: vector<2xf64>) -> (vector<2xf32>, vector<2xf64>) {
   %float_result = math.expm1 %float : vector<2xf32>
   %double_result = math.expm1 %double : vector<2xf64>
   return %float_result, %double_result : vector<2xf32>, vector<2xf64>
 }
 // CHECK-LABEL:   func @expm1_vec_caller(
 // CHECK-SAME:                           %[[VAL_0:.*]]: vector<2xf32>,
 // CHECK-SAME:                           %[[VAL_1:.*]]: vector<2xf64>) -> (vector<2xf32>, vector<2xf64>) {
 // CHECK-DAG:       %[[CVF:.*]] = arith.constant dense<0.000000e+00> : vector<2xf32>
 // CHECK-DAG:       %[[CVD:.*]] = arith.constant dense<0.000000e+00> : vector<2xf64>
 // CHECK:           %[[IN0_F32:.*]] = vector.extract %[[VAL_0]][0] : vector<2xf32>
-// CHECK:           %[[OUT0_F32:.*]] = call @expm1f(%[[IN0_F32]]) : (f32) -> f32
+// CHECK:           %[[OUT0_F32:.*]] = llvm.call @expm1f(%[[IN0_F32]]) : (f32) -> f32
 // CHECK:           %[[VAL_8:.*]] = vector.insert %[[OUT0_F32]], %[[CVF]] [0] : f32 into vector<2xf32>
 // CHECK:           %[[IN1_F32:.*]] = vector.extract %[[VAL_0]][1] : vector<2xf32>
-// CHECK:           %[[OUT1_F32:.*]] = call @expm1f(%[[IN1_F32]]) : (f32) -> f32
+// CHECK:           %[[OUT1_F32:.*]] = llvm.call @expm1f(%[[IN1_F32]]) : (f32) -> f32
 // CHECK:           %[[VAL_11:.*]] = vector.insert %[[OUT1_F32]], %[[VAL_8]] [1] : f32 into vector<2xf32>
 // CHECK:           %[[IN0_F64:.*]] = vector.extract %[[VAL_1]][0] : vector<2xf64>
-// CHECK:           %[[OUT0_F64:.*]] = call @expm1(%[[IN0_F64]]) : (f64) -> f64
+// CHECK:           %[[OUT0_F64:.*]] = llvm.call @expm1(%[[IN0_F64]]) : (f64) -> f64
 // CHECK:           %[[VAL_14:.*]] = vector.insert %[[OUT0_F64]], %[[CVD]] [0] : f64 into vector<2xf64>
 // CHECK:           %[[IN1_F64:.*]] = vector.extract %[[VAL_1]][1] : vector<2xf64>
-// CHECK:           %[[OUT1_F64:.*]] = call @expm1(%[[IN1_F64]]) : (f64) -> f64
+// CHECK:           %[[OUT1_F64:.*]] = llvm.call @expm1(%[[IN1_F64]]) : (f64) -> f64
 // CHECK:           %[[VAL_17:.*]] = vector.insert %[[OUT1_F64]], %[[VAL_14]] [1] : f64 into vector<2xf64>
 // CHECK:           return %[[VAL_11]], %[[VAL_17]] : vector<2xf32>, vector<2xf64>
 // CHECK:         }
 
 func.func @expm1_multidim_vec_caller(%float: vector<2x2xf32>) -> (vector<2x2xf32>) {
   %result = math.expm1 %float : vector<2x2xf32>
   return %result : vector<2x2xf32>
 }
 // CHECK-LABEL:   func @expm1_multidim_vec_caller(
 // CHECK-SAME:                           %[[VAL:.*]]: vector<2x2xf32>
 // CHECK-DAG:       %[[CVF:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf32>
 // CHECK:           %[[IN0_0_F32:.*]] = vector.extract %[[VAL]][0, 0] : vector<2x2xf32>
-// CHECK:           %[[OUT0_0_F32:.*]] = call @expm1f(%[[IN0_0_F32]]) : (f32) -> f32
+// CHECK:           %[[OUT0_0_F32:.*]] = llvm.call @expm1f(%[[IN0_0_F32]]) : (f32) -> f32
 // CHECK:           %[[VAL_1:.*]] = vector.insert %[[OUT0_0_F32]], %[[CVF]] [0, 0] : f32 into vector<2x2xf32>
 // CHECK:           %[[IN0_1_F32:.*]] = vector.extract %[[VAL]][0, 1] : vector<2x2xf32>
-// CHECK:           %[[OUT0_1_F32:.*]] = call @expm1f(%[[IN0_1_F32]]) : (f32) -> f32
+// CHECK:           %[[OUT0_1_F32:.*]] = llvm.call @expm1f(%[[IN0_1_F32]]) : (f32) -> f32
 // CHECK:           %[[VAL_2:.*]] = vector.insert %[[OUT0_1_F32]], %[[VAL_1]] [0, 1] : f32 into vector<2x2xf32>
 // CHECK:           %[[IN1_0_F32:.*]] = vector.extract %[[VAL]][1, 0] : vector<2x2xf32>
-// CHECK:           %[[OUT1_0_F32:.*]] = call @expm1f(%[[IN1_0_F32]]) : (f32) -> f32
+// CHECK:           %[[OUT1_0_F32:.*]] = llvm.call @expm1f(%[[IN1_0_F32]]) : (f32) -> f32
 // CHECK:           %[[VAL_3:.*]] = vector.insert %[[OUT1_0_F32]], %[[VAL_2]] [1, 0] : f32 into vector<2x2xf32>
 // CHECK:           %[[IN1_1_F32:.*]] = vector.extract %[[VAL]][1, 1] : vector<2x2xf32>
-// CHECK:           %[[OUT1_1_F32:.*]] = call @expm1f(%[[IN1_1_F32]]) : (f32) -> f32
+// CHECK:           %[[OUT1_1_F32:.*]] = llvm.call @expm1f(%[[IN1_1_F32]]) : (f32) -> f32
 // CHECK:           %[[VAL_4:.*]] = vector.insert %[[OUT1_1_F32]], %[[VAL_3]] [1, 1] : f32 into vector<2x2xf32>
 // CHECK:           return %[[VAL_4]] : vector<2x2xf32>
 // CHECK:         }