210 lines · c
1//===- OpToFuncCallLowering.h - GPU ops lowering to custom calls *- C++ -*-===//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-exception6//7//===----------------------------------------------------------------------===//8#ifndef MLIR_CONVERSION_GPUCOMMON_OPTOFUNCCALLLOWERING_H_9#define MLIR_CONVERSION_GPUCOMMON_OPTOFUNCCALLLOWERING_H_10 11#include "mlir/Conversion/LLVMCommon/Pattern.h"12#include "mlir/Dialect/Arith/IR/Arith.h"13#include "mlir/Dialect/GPU/IR/GPUDialect.h"14#include "mlir/Dialect/LLVMIR/LLVMDialect.h"15#include "mlir/IR/Builders.h"16 17namespace mlir {18 19namespace {20/// Detection trait tor the `getFastmath` instance method.21template <typename T>22using has_get_fastmath_t = decltype(std::declval<T>().getFastmath());23} // namespace24 25/// Rewriting that replaces SourceOp with a CallOp to `f32Func` or `f64Func` or26/// `f32ApproxFunc` or `f16Func` or `i32Type` depending on the element type and27/// the fastMathFlag of that Op, if present. The function declaration is added28/// in case it was not added before.29///30/// If the input values are of bf16 type (or f16 type if f16Func is empty), the31/// value is first casted to f32, the function called and then the result casted32/// back.33///34/// Example with NVVM:35/// %exp_f32 = math.exp %arg_f32 : f3236///37/// will be transformed into38/// llvm.call @__nv_expf(%arg_f32) : (f32) -> f3239///40/// If the fastMathFlag attribute of SourceOp is `afn` or `fast`, this Op lowers41/// to the approximate calculation function.42///43/// Also example with NVVM:44/// %exp_f32 = math.exp %arg_f32 fastmath<afn> : f3245///46/// will be transformed into47/// llvm.call @__nv_fast_expf(%arg_f32) : (f32) -> f3248///49/// Final example with NVVM:50/// %pow_f32 = math.fpowi %arg_f32, %arg_i3251///52/// will be transformed into53/// llvm.call @__nv_powif(%arg_f32, %arg_i32) : (f32, i32) -> f3254template <typename SourceOp>55struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {56public:57 explicit OpToFuncCallLowering(const LLVMTypeConverter &lowering,58 StringRef f32Func, StringRef f64Func,59 StringRef f32ApproxFunc, StringRef f16Func,60 StringRef i32Func = "",61 PatternBenefit benefit = 1)62 : ConvertOpToLLVMPattern<SourceOp>(lowering, benefit), f32Func(f32Func),63 f64Func(f64Func), f32ApproxFunc(f32ApproxFunc), f16Func(f16Func),64 i32Func(i32Func) {}65 66 LogicalResult67 matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,68 ConversionPatternRewriter &rewriter) const override {69 using LLVM::LLVMFuncOp;70 71 static_assert(72 std::is_base_of<OpTrait::OneResult<SourceOp>, SourceOp>::value,73 "expected single result op");74 75 bool isResultBool = op->getResultTypes().front().isInteger(1);76 if constexpr (!std::is_base_of<OpTrait::SameOperandsAndResultType<SourceOp>,77 SourceOp>::value) {78 assert(op->getNumOperands() > 0 &&79 "expected op to take at least one operand");80 assert((op->getResultTypes().front() == op->getOperand(0).getType() ||81 isResultBool) &&82 "expected op with same operand and result types");83 }84 85 if (!op->template getParentOfType<FunctionOpInterface>()) {86 return rewriter.notifyMatchFailure(87 op, "expected op to be within a function region");88 }89 90 SmallVector<Value, 1> castedOperands;91 for (Value operand : adaptor.getOperands())92 castedOperands.push_back(maybeCast(operand, rewriter));93 94 Type castedOperandType = castedOperands.front().getType();95 96 // At ABI level, booleans are treated as i32.97 Type resultType =98 isResultBool ? rewriter.getIntegerType(32) : castedOperandType;99 Type funcType = getFunctionType(resultType, castedOperands);100 StringRef funcName = getFunctionName(castedOperandType, op);101 if (funcName.empty())102 return failure();103 104 LLVMFuncOp funcOp = appendOrGetFuncOp(funcName, funcType, op);105 auto callOp =106 LLVM::CallOp::create(rewriter, op->getLoc(), funcOp, castedOperands);107 108 if (resultType == adaptor.getOperands().front().getType()) {109 rewriter.replaceOp(op, {callOp.getResult()});110 return success();111 }112 113 // Boolean result are mapping to i32 at the ABI level with zero values being114 // interpreted as false and non-zero values being interpreted as true. Since115 // there is no guarantee of a specific value being used to indicate true,116 // compare for inequality with zero (rather than truncate or shift).117 if (isResultBool) {118 Value zero = LLVM::ConstantOp::create(rewriter, op->getLoc(),119 rewriter.getIntegerType(32),120 rewriter.getI32IntegerAttr(0));121 Value truncated =122 LLVM::ICmpOp::create(rewriter, op->getLoc(), LLVM::ICmpPredicate::ne,123 callOp.getResult(), zero);124 rewriter.replaceOp(op, {truncated});125 return success();126 }127 128 assert(callOp.getResult().getType().isF32() &&129 "only f32 types are supposed to be truncated back");130 Value truncated = LLVM::FPTruncOp::create(131 rewriter, op->getLoc(), adaptor.getOperands().front().getType(),132 callOp.getResult());133 rewriter.replaceOp(op, {truncated});134 return success();135 }136 137 Value maybeCast(Value operand, PatternRewriter &rewriter) const {138 Type type = operand.getType();139 if (!isa<Float16Type, BFloat16Type>(type))140 return operand;141 142 // If there's an f16 function, no need to cast f16 values.143 if (!f16Func.empty() && isa<Float16Type>(type))144 return operand;145 146 return LLVM::FPExtOp::create(rewriter, operand.getLoc(),147 Float32Type::get(rewriter.getContext()),148 operand);149 }150 151 Type getFunctionType(Type resultType, ValueRange operands) const {152 SmallVector<Type> operandTypes(operands.getTypes());153 return LLVM::LLVMFunctionType::get(resultType, operandTypes);154 }155 156 LLVM::LLVMFuncOp appendOrGetFuncOp(StringRef funcName, Type funcType,157 Operation *op) const {158 using LLVM::LLVMFuncOp;159 160 auto funcAttr = StringAttr::get(op->getContext(), funcName);161 auto funcOp =162 SymbolTable::lookupNearestSymbolFrom<LLVMFuncOp>(op, funcAttr);163 if (funcOp)164 return funcOp;165 166 auto parentFunc = op->getParentOfType<FunctionOpInterface>();167 assert(parentFunc && "expected there to be a parent function");168 OpBuilder b(parentFunc);169 170 // Create a valid global location removing any metadata attached to the171 // location as debug info metadata inside of a function cannot be used172 // outside of that function.173 auto globalloc = op->getLoc()->findInstanceOfOrUnknown<FileLineColLoc>();174 return LLVMFuncOp::create(b, globalloc, funcName, funcType);175 }176 177 StringRef getFunctionName(Type type, SourceOp op) const {178 bool useApprox = false;179 if constexpr (llvm::is_detected<has_get_fastmath_t, SourceOp>::value) {180 arith::FastMathFlags flag = op.getFastmath();181 useApprox = ((uint32_t)arith::FastMathFlags::afn & (uint32_t)flag) &&182 !f32ApproxFunc.empty();183 }184 185 if (isa<Float16Type>(type))186 return f16Func;187 if (isa<Float32Type>(type)) {188 if (useApprox)189 return f32ApproxFunc;190 return f32Func;191 }192 if (isa<Float64Type>(type))193 return f64Func;194 195 if (type.isInteger(32))196 return i32Func;197 return "";198 }199 200 const std::string f32Func;201 const std::string f64Func;202 const std::string f32ApproxFunc;203 const std::string f16Func;204 const std::string i32Func;205};206 207} // namespace mlir208 209#endif // MLIR_CONVERSION_GPUCOMMON_OPTOFUNCCALLLOWERING_H_210