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-exception
6 //
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/GPU/GPUDialect.h"
13 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
14 #include "mlir/Dialect/StandardOps/IR/Ops.h"
15 #include "mlir/IR/Builders.h"
16 
17 namespace mlir {
18 
19 /// Rewriting that replace SourceOp with a CallOp to `f32Func` or `f64Func`
20 /// depending on the element type that Op operates upon. The function
21 /// declaration is added in case it was not added before.
22 ///
23 /// If the input values are of f16 type, the value is first casted to f32, the
24 /// function called and then the result casted back.
25 ///
26 /// Example with NVVM:
27 ///   %exp_f32 = std.exp %arg_f32 : f32
28 ///
29 /// will be transformed into
30 ///   llvm.call @__nv_expf(%arg_f32) : (f32) -> f32
31 template <typename SourceOp>
32 struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
33 public:
34   explicit OpToFuncCallLowering(LLVMTypeConverter &lowering_, StringRef f32Func,
35                                 StringRef f64Func)
36       : ConvertOpToLLVMPattern<SourceOp>(lowering_), f32Func(f32Func),
37         f64Func(f64Func) {}
38 
39   LogicalResult
40   matchAndRewrite(SourceOp op, ArrayRef<Value> operands,
41                   ConversionPatternRewriter &rewriter) const override {
42     using LLVM::LLVMFuncOp;
43 
44     static_assert(
45         std::is_base_of<OpTrait::OneResult<SourceOp>, SourceOp>::value,
46         "expected single result op");
47 
48     static_assert(std::is_base_of<OpTrait::SameOperandsAndResultType<SourceOp>,
49                                   SourceOp>::value,
50                   "expected op with same operand and result types");
51 
52     SmallVector<Value, 1> castedOperands;
53     for (Value operand : operands)
54       castedOperands.push_back(maybeCast(operand, rewriter));
55 
56     Type resultType = castedOperands.front().getType();
57     Type funcType = getFunctionType(resultType, castedOperands);
58     StringRef funcName = getFunctionName(
59         funcType.cast<LLVM::LLVMFunctionType>().getReturnType());
60     if (funcName.empty())
61       return failure();
62 
63     LLVMFuncOp funcOp = appendOrGetFuncOp(funcName, funcType, op);
64     auto callOp = rewriter.create<LLVM::CallOp>(
65         op->getLoc(), resultType, SymbolRefAttr::get(funcOp), castedOperands);
66 
67     if (resultType == operands.front().getType()) {
68       rewriter.replaceOp(op, {callOp.getResult(0)});
69       return success();
70     }
71 
72     Value truncated = rewriter.create<LLVM::FPTruncOp>(
73         op->getLoc(), operands.front().getType(), callOp.getResult(0));
74     rewriter.replaceOp(op, {truncated});
75     return success();
76   }
77 
78 private:
79   Value maybeCast(Value operand, PatternRewriter &rewriter) const {
80     Type type = operand.getType();
81     if (!type.isa<Float16Type>())
82       return operand;
83 
84     return rewriter.create<LLVM::FPExtOp>(
85         operand.getLoc(), Float32Type::get(rewriter.getContext()), operand);
86   }
87 
88   Type getFunctionType(Type resultType, ArrayRef<Value> operands) const {
89     SmallVector<Type, 1> operandTypes;
90     for (Value operand : operands) {
91       operandTypes.push_back(operand.getType());
92     }
93     return LLVM::LLVMFunctionType::get(resultType, operandTypes);
94   }
95 
96   StringRef getFunctionName(Type type) const {
97     if (type.isa<Float32Type>())
98       return f32Func;
99     if (type.isa<Float64Type>())
100       return f64Func;
101     return "";
102   }
103 
104   LLVM::LLVMFuncOp appendOrGetFuncOp(StringRef funcName, Type funcType,
105                                      Operation *op) const {
106     using LLVM::LLVMFuncOp;
107 
108     auto funcAttr = StringAttr::get(op->getContext(), funcName);
109     Operation *funcOp = SymbolTable::lookupNearestSymbolFrom(op, funcAttr);
110     if (funcOp)
111       return cast<LLVMFuncOp>(*funcOp);
112 
113     mlir::OpBuilder b(op->getParentOfType<LLVMFuncOp>());
114     return b.create<LLVMFuncOp>(op->getLoc(), funcName, funcType);
115   }
116 
117   const std::string f32Func;
118   const std::string f64Func;
119 };
120 
121 } // namespace mlir
122 
123 #endif // MLIR_CONVERSION_GPUCOMMON_OPTOFUNCCALLLOWERING_H_
124