1 //===-- MathToLibm.cpp - conversion from Math to libm calls ---------------===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #include "mlir/Conversion/MathToLibm/MathToLibm.h" 10 11 #include "mlir/Dialect/Arith/IR/Arith.h" 12 #include "mlir/Dialect/Func/IR/FuncOps.h" 13 #include "mlir/Dialect/LLVMIR/LLVMDialect.h" 14 #include "mlir/Dialect/Math/IR/Math.h" 15 #include "mlir/Dialect/Utils/IndexingUtils.h" 16 #include "mlir/Dialect/Vector/IR/VectorOps.h" 17 #include "mlir/IR/BuiltinDialect.h" 18 #include "mlir/IR/PatternMatch.h" 19 #include "mlir/Pass/Pass.h" 20 #include "mlir/Transforms/DialectConversion.h" 21 22 namespace mlir { 23 #define GEN_PASS_DEF_CONVERTMATHTOLIBM 24 #include "mlir/Conversion/Passes.h.inc" 25 } // namespace mlir 26 27 using namespace mlir; 28 29 namespace { 30 // Pattern to convert vector operations to scalar operations. This is needed as 31 // libm calls require scalars. 32 template <typename Op> 33 struct VecOpToScalarOp : public OpRewritePattern<Op> { 34 public: 35 using OpRewritePattern<Op>::OpRewritePattern; 36 37 LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final; 38 }; 39 // Pattern to promote an op of a smaller floating point type to F32. 40 template <typename Op> 41 struct PromoteOpToF32 : public OpRewritePattern<Op> { 42 public: 43 using OpRewritePattern<Op>::OpRewritePattern; 44 45 LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final; 46 }; 47 // Pattern to convert scalar math operations to calls to libm functions. 48 // Additionally the libm function signatures are declared. 49 template <typename Op> 50 struct ScalarOpToLibmCall : public OpRewritePattern<Op> { 51 public: 52 using OpRewritePattern<Op>::OpRewritePattern; 53 ScalarOpToLibmCall(MLIRContext *context, StringRef floatFunc, 54 StringRef doubleFunc) 55 : OpRewritePattern<Op>(context), floatFunc(floatFunc), 56 doubleFunc(doubleFunc){}; 57 58 LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final; 59 60 private: 61 std::string floatFunc, doubleFunc; 62 }; 63 64 template <typename OpTy> 65 void populatePatternsForOp(RewritePatternSet &patterns, MLIRContext *ctx, 66 StringRef floatFunc, StringRef doubleFunc) { 67 patterns.add<VecOpToScalarOp<OpTy>, PromoteOpToF32<OpTy>>(ctx); 68 patterns.add<ScalarOpToLibmCall<OpTy>>(ctx, floatFunc, doubleFunc); 69 } 70 71 } // namespace 72 73 template <typename Op> 74 LogicalResult 75 VecOpToScalarOp<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const { 76 auto opType = op.getType(); 77 auto loc = op.getLoc(); 78 auto vecType = dyn_cast<VectorType>(opType); 79 80 if (!vecType) 81 return failure(); 82 if (!vecType.hasRank()) 83 return failure(); 84 auto shape = vecType.getShape(); 85 int64_t numElements = vecType.getNumElements(); 86 87 Value result = rewriter.create<arith::ConstantOp>( 88 loc, DenseElementsAttr::get( 89 vecType, FloatAttr::get(vecType.getElementType(), 0.0))); 90 SmallVector<int64_t> strides = computeStrides(shape); 91 for (auto linearIndex = 0; linearIndex < numElements; ++linearIndex) { 92 SmallVector<int64_t> positions = delinearize(linearIndex, strides); 93 SmallVector<Value> operands; 94 for (auto input : op->getOperands()) 95 operands.push_back( 96 rewriter.create<vector::ExtractOp>(loc, input, positions)); 97 Value scalarOp = 98 rewriter.create<Op>(loc, vecType.getElementType(), operands); 99 result = 100 rewriter.create<vector::InsertOp>(loc, scalarOp, result, positions); 101 } 102 rewriter.replaceOp(op, {result}); 103 return success(); 104 } 105 106 template <typename Op> 107 LogicalResult 108 PromoteOpToF32<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const { 109 auto opType = op.getType(); 110 if (!isa<Float16Type, BFloat16Type>(opType)) 111 return failure(); 112 113 auto loc = op.getLoc(); 114 auto f32 = rewriter.getF32Type(); 115 auto extendedOperands = llvm::to_vector( 116 llvm::map_range(op->getOperands(), [&](Value operand) -> Value { 117 return rewriter.create<arith::ExtFOp>(loc, f32, operand); 118 })); 119 auto newOp = rewriter.create<Op>(loc, f32, extendedOperands); 120 rewriter.replaceOpWithNewOp<arith::TruncFOp>(op, opType, newOp); 121 return success(); 122 } 123 124 template <typename Op> 125 LogicalResult 126 ScalarOpToLibmCall<Op>::matchAndRewrite(Op op, 127 PatternRewriter &rewriter) const { 128 auto module = SymbolTable::getNearestSymbolTable(op); 129 auto type = op.getType(); 130 if (!isa<Float32Type, Float64Type>(type)) 131 return failure(); 132 133 auto name = type.getIntOrFloatBitWidth() == 64 ? doubleFunc : floatFunc; 134 auto opFunc = dyn_cast_or_null<SymbolOpInterface>( 135 SymbolTable::lookupSymbolIn(module, name)); 136 // Forward declare function if it hasn't already been 137 if (!opFunc) { 138 OpBuilder::InsertionGuard guard(rewriter); 139 rewriter.setInsertionPointToStart(&module->getRegion(0).front()); 140 auto opFunctionTy = FunctionType::get( 141 rewriter.getContext(), op->getOperandTypes(), op->getResultTypes()); 142 opFunc = rewriter.create<func::FuncOp>(rewriter.getUnknownLoc(), name, 143 opFunctionTy); 144 opFunc.setPrivate(); 145 146 // By definition Math dialect operations imply LLVM's "readnone" 147 // function attribute, so we can set it here to provide more 148 // optimization opportunities (e.g. LICM) for backends targeting LLVM IR. 149 // This will have to be changed, when strict FP behavior is supported 150 // by Math dialect. 151 opFunc->setAttr(LLVM::LLVMDialect::getReadnoneAttrName(), 152 UnitAttr::get(rewriter.getContext())); 153 } 154 assert(isa<FunctionOpInterface>(SymbolTable::lookupSymbolIn(module, name))); 155 156 rewriter.replaceOpWithNewOp<func::CallOp>(op, name, op.getType(), 157 op->getOperands()); 158 159 return success(); 160 } 161 162 void mlir::populateMathToLibmConversionPatterns(RewritePatternSet &patterns) { 163 MLIRContext *ctx = patterns.getContext(); 164 165 populatePatternsForOp<math::AbsFOp>(patterns, ctx, "fabsf", "fabs"); 166 populatePatternsForOp<math::AcosOp>(patterns, ctx, "acosf", "acos"); 167 populatePatternsForOp<math::AcoshOp>(patterns, ctx, "acoshf", "acosh"); 168 populatePatternsForOp<math::AsinOp>(patterns, ctx, "asinf", "asin"); 169 populatePatternsForOp<math::AsinhOp>(patterns, ctx, "asinhf", "asinh"); 170 populatePatternsForOp<math::Atan2Op>(patterns, ctx, "atan2f", "atan2"); 171 populatePatternsForOp<math::AtanOp>(patterns, ctx, "atanf", "atan"); 172 populatePatternsForOp<math::AtanhOp>(patterns, ctx, "atanhf", "atanh"); 173 populatePatternsForOp<math::CbrtOp>(patterns, ctx, "cbrtf", "cbrt"); 174 populatePatternsForOp<math::CeilOp>(patterns, ctx, "ceilf", "ceil"); 175 populatePatternsForOp<math::CosOp>(patterns, ctx, "cosf", "cos"); 176 populatePatternsForOp<math::CoshOp>(patterns, ctx, "coshf", "cosh"); 177 populatePatternsForOp<math::ErfOp>(patterns, ctx, "erff", "erf"); 178 populatePatternsForOp<math::ExpOp>(patterns, ctx, "expf", "exp"); 179 populatePatternsForOp<math::Exp2Op>(patterns, ctx, "exp2f", "exp2"); 180 populatePatternsForOp<math::ExpM1Op>(patterns, ctx, "expm1f", "expm1"); 181 populatePatternsForOp<math::FloorOp>(patterns, ctx, "floorf", "floor"); 182 populatePatternsForOp<math::FmaOp>(patterns, ctx, "fmaf", "fma"); 183 populatePatternsForOp<math::LogOp>(patterns, ctx, "logf", "log"); 184 populatePatternsForOp<math::Log2Op>(patterns, ctx, "log2f", "log2"); 185 populatePatternsForOp<math::Log10Op>(patterns, ctx, "log10f", "log10"); 186 populatePatternsForOp<math::Log1pOp>(patterns, ctx, "log1pf", "log1p"); 187 populatePatternsForOp<math::PowFOp>(patterns, ctx, "powf", "pow"); 188 populatePatternsForOp<math::RoundEvenOp>(patterns, ctx, "roundevenf", 189 "roundeven"); 190 populatePatternsForOp<math::RoundOp>(patterns, ctx, "roundf", "round"); 191 populatePatternsForOp<math::SinOp>(patterns, ctx, "sinf", "sin"); 192 populatePatternsForOp<math::SinhOp>(patterns, ctx, "sinhf", "sinh"); 193 populatePatternsForOp<math::SqrtOp>(patterns, ctx, "sqrtf", "sqrt"); 194 populatePatternsForOp<math::RsqrtOp>(patterns, ctx, "rsqrtf", "rsqrt"); 195 populatePatternsForOp<math::TanOp>(patterns, ctx, "tanf", "tan"); 196 populatePatternsForOp<math::TanhOp>(patterns, ctx, "tanhf", "tanh"); 197 populatePatternsForOp<math::TruncOp>(patterns, ctx, "truncf", "trunc"); 198 } 199 200 namespace { 201 struct ConvertMathToLibmPass 202 : public impl::ConvertMathToLibmBase<ConvertMathToLibmPass> { 203 void runOnOperation() override; 204 }; 205 } // namespace 206 207 void ConvertMathToLibmPass::runOnOperation() { 208 auto module = getOperation(); 209 210 RewritePatternSet patterns(&getContext()); 211 populateMathToLibmConversionPatterns(patterns); 212 213 ConversionTarget target(getContext()); 214 target.addLegalDialect<arith::ArithDialect, BuiltinDialect, func::FuncDialect, 215 vector::VectorDialect>(); 216 target.addIllegalDialect<math::MathDialect>(); 217 if (failed(applyPartialConversion(module, target, std::move(patterns)))) 218 signalPassFailure(); 219 } 220 221 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertMathToLibmPass() { 222 return std::make_unique<ConvertMathToLibmPass>(); 223 } 224