Skip to content

Commit 4df28af

Browse files
ftynsewsmoses
ftynse
and
wsmoses
authored
[mlir] Add lowering of absi and fpowi to libdevice (llvm#123644)
More concise version of llvm#123422. --------- Co-authored-by: William S. Moses <[email protected]>
1 parent 9c464e6 commit 4df28af

File tree

3 files changed

+97
-28
lines changed

3 files changed

+97
-28
lines changed

‎mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h‎

Lines changed: 54 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,16 @@
1616

1717
namespace mlir {
1818

19-
/// Rewriting that replace SourceOp with a CallOp to `f32Func` or `f64Func` or
20-
/// `f32ApproxFunc` or `f16Func` depending on the element type and the
21-
/// fastMathFlag of that Op. The function declaration is added in case it was
22-
/// not added before.
19+
namespace {
20+
/// Detection trait tor the `getFastmath` instance method.
21+
template <typename T>
22+
using has_get_fastmath_t = decltype(std::declval<T>().getFastmath());
23+
} // namespace
24+
25+
/// Rewriting that replaces SourceOp with a CallOp to `f32Func` or `f64Func` or
26+
/// `f32ApproxFunc` or `f16Func` or `i32Type` depending on the element type and
27+
/// the fastMathFlag of that Op, if present. The function declaration is added
28+
/// in case it was not added before.
2329
///
2430
/// If the input values are of bf16 type (or f16 type if f16Func is empty), the
2531
/// value is first casted to f32, the function called and then the result casted
@@ -39,14 +45,22 @@ namespace mlir {
3945
///
4046
/// will be transformed into
4147
/// llvm.call @__nv_fast_expf(%arg_f32) : (f32) -> f32
48+
///
49+
/// Final example with NVVM:
50+
/// %pow_f32 = math.fpowi %arg_f32, %arg_i32
51+
///
52+
/// will be transformed into
53+
/// llvm.call @__nv_powif(%arg_f32, %arg_i32) : (f32, i32) -> f32
4254
template <typename SourceOp>
4355
struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
4456
public:
4557
explicit OpToFuncCallLowering(const LLVMTypeConverter &lowering,
4658
StringRef f32Func, StringRef f64Func,
47-
StringRef f32ApproxFunc, StringRef f16Func)
59+
StringRef f32ApproxFunc, StringRef f16Func,
60+
StringRef i32Func = "")
4861
: ConvertOpToLLVMPattern<SourceOp>(lowering), f32Func(f32Func),
49-
f64Func(f64Func), f32ApproxFunc(f32ApproxFunc), f16Func(f16Func) {}
62+
f64Func(f64Func), f32ApproxFunc(f32ApproxFunc), f16Func(f16Func),
63+
i32Func(i32Func) {}
5064

5165
LogicalResult
5266
matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,
@@ -76,9 +90,8 @@ struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
7690

7791
Type resultType = castedOperands.front().getType();
7892
Type funcType = getFunctionType(resultType, castedOperands);
79-
StringRef funcName =
80-
getFunctionName(cast<LLVM::LLVMFunctionType>(funcType).getReturnType(),
81-
op.getFastmath());
93+
StringRef funcName = getFunctionName(
94+
cast<LLVM::LLVMFunctionType>(funcType).getReturnType(), op);
8295
if (funcName.empty())
8396
return failure();
8497

@@ -91,14 +104,15 @@ struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
91104
return success();
92105
}
93106

107+
assert(callOp.getResult().getType().isF32() &&
108+
"only f32 types are supposed to be truncated back");
94109
Value truncated = rewriter.create<LLVM::FPTruncOp>(
95110
op->getLoc(), adaptor.getOperands().front().getType(),
96111
callOp.getResult());
97112
rewriter.replaceOp(op, {truncated});
98113
return success();
99114
}
100115

101-
private:
102116
Value maybeCast(Value operand, PatternRewriter &rewriter) const {
103117
Type type = operand.getType();
104118
if (!isa<Float16Type, BFloat16Type>(type))
@@ -117,38 +131,50 @@ struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
117131
return LLVM::LLVMFunctionType::get(resultType, operandTypes);
118132
}
119133

120-
StringRef getFunctionName(Type type, arith::FastMathFlags flag) const {
121-
if (isa<Float16Type>(type))
122-
return f16Func;
123-
if (isa<Float32Type>(type)) {
124-
if (((uint32_t)arith::FastMathFlags::afn & (uint32_t)flag) &&
125-
!f32ApproxFunc.empty())
126-
return f32ApproxFunc;
127-
else
128-
return f32Func;
129-
}
130-
if (isa<Float64Type>(type))
131-
return f64Func;
132-
return "";
133-
}
134-
135134
LLVM::LLVMFuncOp appendOrGetFuncOp(StringRef funcName, Type funcType,
136135
Operation *op) const {
137136
using LLVM::LLVMFuncOp;
138137

139138
auto funcAttr = StringAttr::get(op->getContext(), funcName);
140-
Operation *funcOp = SymbolTable::lookupNearestSymbolFrom(op, funcAttr);
139+
auto funcOp =
140+
SymbolTable::lookupNearestSymbolFrom<LLVMFuncOp>(op, funcAttr);
141141
if (funcOp)
142-
return cast<LLVMFuncOp>(*funcOp);
142+
return funcOp;
143143

144-
mlir::OpBuilder b(op->getParentOfType<FunctionOpInterface>());
144+
auto parentFunc = op->getParentOfType<FunctionOpInterface>();
145+
assert(parentFunc && "expected there to be a parent function");
146+
OpBuilder b(parentFunc);
145147
return b.create<LLVMFuncOp>(op->getLoc(), funcName, funcType);
146148
}
147149

150+
StringRef getFunctionName(Type type, SourceOp op) const {
151+
bool useApprox = false;
152+
if constexpr (llvm::is_detected<has_get_fastmath_t, SourceOp>::value) {
153+
arith::FastMathFlags flag = op.getFastmath();
154+
useApprox = ((uint32_t)arith::FastMathFlags::afn & (uint32_t)flag) &&
155+
!f32ApproxFunc.empty();
156+
}
157+
158+
if (isa<Float16Type>(type))
159+
return f16Func;
160+
if (isa<Float32Type>(type)) {
161+
if (useApprox)
162+
return f32ApproxFunc;
163+
return f32Func;
164+
}
165+
if (isa<Float64Type>(type))
166+
return f64Func;
167+
168+
if (type.isInteger(32))
169+
return i32Func;
170+
return "";
171+
}
172+
148173
const std::string f32Func;
149174
const std::string f64Func;
150175
const std::string f32ApproxFunc;
151176
const std::string f16Func;
177+
const std::string i32Func;
152178
};
153179

154180
} // namespace mlir

‎mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -446,6 +446,22 @@ static void populateOpPatterns(const LLVMTypeConverter &converter,
446446
f32ApproxFunc, f16Func);
447447
}
448448

449+
template <typename OpTy>
450+
static void populateIntOpPatterns(const LLVMTypeConverter &converter,
451+
RewritePatternSet &patterns,
452+
StringRef i32Func) {
453+
patterns.add<ScalarizeVectorOpLowering<OpTy>>(converter);
454+
patterns.add<OpToFuncCallLowering<OpTy>>(converter, "", "", "", "", i32Func);
455+
}
456+
457+
template <typename OpTy>
458+
static void populateFloatIntOpPatterns(const LLVMTypeConverter &converter,
459+
RewritePatternSet &patterns,
460+
StringRef f32Func, StringRef f64Func) {
461+
patterns.add<ScalarizeVectorOpLowering<OpTy>>(converter);
462+
patterns.add<OpToFuncCallLowering<OpTy>>(converter, f32Func, f64Func, "", "");
463+
}
464+
449465
void mlir::populateGpuSubgroupReduceOpLoweringPattern(
450466
const LLVMTypeConverter &converter, RewritePatternSet &patterns) {
451467
patterns.add<GPUSubgroupReduceOpLowering>(converter);
@@ -509,6 +525,7 @@ void mlir::populateGpuToNVVMConversionPatterns(
509525

510526
populateOpPatterns<arith::RemFOp>(converter, patterns, "__nv_fmodf",
511527
"__nv_fmod");
528+
populateIntOpPatterns<math::AbsIOp>(converter, patterns, "__nv_abs");
512529
populateOpPatterns<math::AbsFOp>(converter, patterns, "__nv_fabsf",
513530
"__nv_fabs");
514531
populateOpPatterns<math::AcosOp>(converter, patterns, "__nv_acosf",
@@ -555,6 +572,8 @@ void mlir::populateGpuToNVVMConversionPatterns(
555572
"__nv_log2", "__nv_fast_log2f");
556573
populateOpPatterns<math::PowFOp>(converter, patterns, "__nv_powf", "__nv_pow",
557574
"__nv_fast_powf");
575+
populateFloatIntOpPatterns<math::FPowIOp>(converter, patterns, "__nv_powif",
576+
"__nv_powi");
558577
populateOpPatterns<math::RoundOp>(converter, patterns, "__nv_roundf",
559578
"__nv_round");
560579
populateOpPatterns<math::RoundEvenOp>(converter, patterns, "__nv_rintf",

‎mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1033,3 +1033,27 @@ module attributes {transform.with_named_sequence} {
10331033
transform.yield
10341034
}
10351035
}
1036+
1037+
1038+
gpu.module @test_module_52 {
1039+
// CHECK: llvm.func @__nv_abs(i32) -> i32
1040+
// CHECK-LABEL: func @gpu_abs
1041+
func.func @gpu_abs(%arg_i32 : i32) -> (i32) {
1042+
%result32 = math.absi %arg_i32 : i32
1043+
// CHECK: llvm.call @__nv_abs(%{{.*}}) : (i32) -> i32
1044+
func.return %result32 : i32
1045+
}
1046+
}
1047+
1048+
gpu.module @test_module_53 {
1049+
// CHECK: llvm.func @__nv_powif(f32, i32) -> f32
1050+
// CHECK: llvm.func @__nv_powi(f64, i32) -> f64
1051+
// CHECK-LABEL: func @gpu_powi
1052+
func.func @gpu_powi(%arg_f32 : f32, %arg_f64 : f64, %arg_i32 : i32) -> (f32, f64) {
1053+
%result32 = math.fpowi %arg_f32, %arg_i32 : f32, i32
1054+
// CHECK: llvm.call @__nv_powif(%{{.*}}, %{{.*}}) : (f32, i32) -> f32
1055+
%result64 = math.fpowi %arg_f64, %arg_i32 : f64, i32
1056+
// CHECK: llvm.call @__nv_powi(%{{.*}}, %{{.*}}) : (f64, i32) -> f64
1057+
func.return %result32, %result64 : f32, f64
1058+
}
1059+
}

0 commit comments

Comments
 (0)