blob: 2c1c0e107a57890ec70c7f68e6001c9d2221b945 [file] [log] [blame]
//===- OpToFuncCallLowering.h - GPU ops lowering to custom calls *- C++ -*-===//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/GPU/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/IR/Builders.h"
namespace mlir {
/// Rewriting that replace SourceOp with a CallOp to `f32Func` or `f64Func`
/// depending on the element type that Op operates upon. The function
/// declaration is added in case it was not added before.
/// If the input values are of f16 type, the value is first casted to f32, the
/// function called and then the result casted back.
/// Example with NVVM:
/// %exp_f32 = std.exp %arg_f32 : f32
/// will be transformed into
/// @__nv_expf(%arg_f32) : (f32) -> f32
template <typename SourceOp>
struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
explicit OpToFuncCallLowering(LLVMTypeConverter &lowering_, StringRef f32Func,
StringRef f64Func)
: ConvertOpToLLVMPattern<SourceOp>(lowering_), f32Func(f32Func),
f64Func(f64Func) {}
matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
using LLVM::LLVMFuncOp;
std::is_base_of<OpTrait::OneResult<SourceOp>, SourceOp>::value,
"expected single result op");
"expected op with same operand and result types");
SmallVector<Value, 1> castedOperands;
for (Value operand : adaptor.getOperands())
castedOperands.push_back(maybeCast(operand, rewriter));
Type resultType = castedOperands.front().getType();
Type funcType = getFunctionType(resultType, castedOperands);
StringRef funcName = getFunctionName(
if (funcName.empty())
return failure();
LLVMFuncOp funcOp = appendOrGetFuncOp(funcName, funcType, op);
auto callOp = rewriter.create<LLVM::CallOp>(
op->getLoc(), resultType, SymbolRefAttr::get(funcOp), castedOperands);
if (resultType == adaptor.getOperands().front().getType()) {
rewriter.replaceOp(op, {callOp.getResult(0)});
return success();
Value truncated = rewriter.create<LLVM::FPTruncOp>(
op->getLoc(), adaptor.getOperands().front().getType(),
rewriter.replaceOp(op, {truncated});
return success();
Value maybeCast(Value operand, PatternRewriter &rewriter) const {
Type type = operand.getType();
if (!type.isa<Float16Type>())
return operand;
return rewriter.create<LLVM::FPExtOp>(
operand.getLoc(), Float32Type::get(rewriter.getContext()), operand);
Type getFunctionType(Type resultType, ValueRange operands) const {
SmallVector<Type> operandTypes(operands.getTypes());
return LLVM::LLVMFunctionType::get(resultType, operandTypes);
StringRef getFunctionName(Type type) const {
if (type.isa<Float32Type>())
return f32Func;
if (type.isa<Float64Type>())
return f64Func;
return "";
LLVM::LLVMFuncOp appendOrGetFuncOp(StringRef funcName, Type funcType,
Operation *op) const {
using LLVM::LLVMFuncOp;
auto funcAttr = StringAttr::get(op->getContext(), funcName);
Operation *funcOp = SymbolTable::lookupNearestSymbolFrom(op, funcAttr);
if (funcOp)
return cast<LLVMFuncOp>(*funcOp);
mlir::OpBuilder b(op->getParentOfType<LLVMFuncOp>());
return b.create<LLVMFuncOp>(op->getLoc(), funcName, funcType);
const std::string f32Func;
const std::string f64Func;
} // namespace mlir