Skip to content

[MLIR][Math][GPU] Add lowering of absi and fpowi to libdevice #123422

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

[MLIR][Math][GPU] Add lowering of absi and fpowi to libdevice #123422

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7c849d7
[MLIR][Math][GPU] Add lowering of absi and fpowi to libdevice
wsmoses Jan 17, 2025
d6b89f0
Apply suggestions from code review
wsmoses Jan 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
190 changes: 140 additions & 50 deletions mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,37 +16,11 @@

namespace mlir {

/// Rewriting that replace SourceOp with a CallOp to `f32Func` or `f64Func` or
/// `f32ApproxFunc` or `f16Func` depending on the element type and the
/// fastMathFlag of that Op. The function declaration is added in case it was
/// not added before.
///
/// If the input values are of bf16 type (or f16 type if f16Func is empty), the
/// value is first casted to f32, the function called and then the result casted
/// back.
///
/// Example with NVVM:
/// %exp_f32 = math.exp %arg_f32 : f32
///
/// will be transformed into
/// llvm.call @__nv_expf(%arg_f32) : (f32) -> f32
///
/// If the fastMathFlag attribute of SourceOp is `afn` or `fast`, this Op lowers
/// to the approximate calculation function.
///
/// Also example with NVVM:
/// %exp_f32 = math.exp %arg_f32 fastmath<afn> : f32
///
/// will be transformed into
/// llvm.call @__nv_fast_expf(%arg_f32) : (f32) -> f32
template <typename SourceOp>
struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
template <typename SourceOp, typename DerivedTy>
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add some documentation. I also think the implicit convention for CRTP is to have the derived class the as the leading template argument.

struct OpToFuncCallLoweringBase : public ConvertOpToLLVMPattern<SourceOp> {
public:
explicit OpToFuncCallLowering(const LLVMTypeConverter &lowering,
StringRef f32Func, StringRef f64Func,
StringRef f32ApproxFunc, StringRef f16Func)
: ConvertOpToLLVMPattern<SourceOp>(lowering), f32Func(f32Func),
f64Func(f64Func), f32ApproxFunc(f32ApproxFunc), f16Func(f16Func) {}
explicit OpToFuncCallLoweringBase(const LLVMTypeConverter &lowering)
: ConvertOpToLLVMPattern<SourceOp>(lowering) {}

LogicalResult
matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,
Expand All @@ -72,13 +46,15 @@ struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {

SmallVector<Value, 1> castedOperands;
for (Value operand : adaptor.getOperands())
castedOperands.push_back(maybeCast(operand, rewriter));
castedOperands.push_back(
static_cast<const DerivedTy *>(this)->maybeCast(operand, rewriter));

Type resultType = castedOperands.front().getType();
Type funcType = getFunctionType(resultType, castedOperands);
StringRef funcName =
getFunctionName(cast<LLVM::LLVMFunctionType>(funcType).getReturnType(),
op.getFastmath());
static_cast<const DerivedTy *>(this)
->getFunctionName(
cast<LLVM::LLVMFunctionType>(funcType).getReturnType(), op);
if (funcName.empty())
return failure();

Expand All @@ -99,6 +75,63 @@ struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
}

private:
Type getFunctionType(Type resultType, ValueRange operands) const {
SmallVector<Type> operandTypes(operands.getTypes());
return LLVM::LLVMFunctionType::get(resultType, operandTypes);
}

LLVM::LLVMFuncOp appendOrGetFuncOp(StringRef funcName, Type funcType,
Operation *op) const {
using LLVM::LLVMFuncOp;

auto funcAttr = StringAttr::get(op->getContext(), funcName);
auto funcOp = SymbolTable::lookupNearestSymbolFrom<LLVMFuncOp>(op, funcAttr);
if (funcOp)
return funcOp;

auto parentFunc = op->getParentOfType<FunctionOpInterface>();
assert(parentFunc && "expected there to be a parent function");
OpBuilder b(parentFunc);
return b.create<LLVMFuncOp>(op->getLoc(), funcName, funcType);
}
};

/// Rewriting that replaces SourceOp with a CallOp to `f32Func` or `f64Func` or
/// `f32ApproxFunc` or `f16Func` depending on the element type and the
/// fastMathFlag of that Op. The function declaration is added in case it was
/// not added before.
///
/// If the input values are of bf16 type (or f16 type if f16Func is empty), the
/// value is first casted to f32, the function called and then the result casted
/// back.
///
/// Example with NVVM:
/// %exp_f32 = math.exp %arg_f32 : f32
///
/// will be transformed into
/// llvm.call @__nv_expf(%arg_f32) : (f32) -> f32
///
/// If the fastMathFlag attribute of SourceOp is `afn` or `fast`, this Op lowers
/// to the approximate calculation function.
///
/// Also example with NVVM:
/// %exp_f32 = math.exp %arg_f32 fastmath<afn> : f32
///
/// will be transformed into
/// llvm.call @__nv_fast_expf(%arg_f32) : (f32) -> f32
template <typename SourceOp>
struct OpToFuncCallLowering
: public OpToFuncCallLoweringBase<SourceOp,
OpToFuncCallLowering<SourceOp>> {
public:
explicit OpToFuncCallLowering(const LLVMTypeConverter &lowering,
StringRef f32Func, StringRef f64Func,
StringRef f32ApproxFunc, StringRef f16Func)
: OpToFuncCallLoweringBase<SourceOp, OpToFuncCallLowering<SourceOp>>(
lowering),
f32Func(f32Func), f64Func(f64Func), f32ApproxFunc(f32ApproxFunc),
f16Func(f16Func) {}

Value maybeCast(Value operand, PatternRewriter &rewriter) const {
Type type = operand.getType();
if (!isa<Float16Type, BFloat16Type>(type))
Expand All @@ -112,12 +145,8 @@ struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
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, arith::FastMathFlags flag) const {
StringRef getFunctionName(Type type, SourceOp op) const {
arith::FastMathFlags flag = op.getFastmath();
if (isa<Float16Type>(type))
return f16Func;
if (isa<Float32Type>(type)) {
Expand All @@ -132,23 +161,84 @@ struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
return "";
}

LLVM::LLVMFuncOp appendOrGetFuncOp(StringRef funcName, Type funcType,
Operation *op) const {
using LLVM::LLVMFuncOp;
const std::string f32Func;
const std::string f64Func;
const std::string f32ApproxFunc;
const std::string f16Func;
};

auto funcAttr = StringAttr::get(op->getContext(), funcName);
Operation *funcOp = SymbolTable::lookupNearestSymbolFrom(op, funcAttr);
if (funcOp)
return cast<LLVMFuncOp>(*funcOp);
/// Rewriting that replace SourceOp with a CallOp to `i32Func`
/// The function declaration is added in case it was not added before.
/// This assumes that all types integral.
///
/// Example with NVVM:
/// %abs_i32 = math.iabs %arg_i32 : i32
///
/// will be transformed into
/// llvm.call @__nv_abs(%arg_i32) : (i32) -> i32
///
template <typename SourceOp>
struct IntOpToFuncCallLowering
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can this just be folded into the the OpToFuncCallLowering as another option with an empty-by-default function name?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The maybeCast/ getFunctionName option in the floating point one currently does checks for floating point types whereas here we check if it's a 32bit int

: public OpToFuncCallLoweringBase<SourceOp,
IntOpToFuncCallLowering<SourceOp>> {
public:
explicit IntOpToFuncCallLowering(const LLVMTypeConverter &lowering,
StringRef i32Func)
: OpToFuncCallLoweringBase<SourceOp, IntOpToFuncCallLowering<SourceOp>>(
lowering),
i32Func(i32Func) {}

mlir::OpBuilder b(op->getParentOfType<FunctionOpInterface>());
return b.create<LLVMFuncOp>(op->getLoc(), funcName, funcType);
Value maybeCast(Value operand, PatternRewriter &rewriter) const {
return operand;
}

StringRef getFunctionName(Type type, SourceOp op) const {
IntegerType itype = dyn_cast<IntegerType>(type);
if (!itype || itype.getWidth() != 32)
return "";
return i32Func;
}

const std::string i32Func;
};

/// Rewriting that replaces SourceOp with a CallOp to `f32Func` or `f64Func`,
/// depending on the type of the result. This assumes that the first argument is
/// a floating type and the second argument is an integer type.
///
/// Example with NVVM:
/// %result32 = math.fpowi %arg_f32, %arg_i32 : f32, i32
///
/// will be transformed into
/// llvm.call @__nv_powf(%arg_f32, %arg_i32) : (f32, i32) -> f32
///
template <typename SourceOp>
struct FloatIntOpToFuncCallLowering
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this needed? I don't see anything specific in the logic here that differs from OpToFuncCallLowering.

Copy link
Member Author

@wsmoses wsmoses Jan 19, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah the optofunccalllowering checks if input types are valid floats (saying it doesnt apply otherwise) whereas here we need to permit a 32bit integer for the second arg

: public OpToFuncCallLoweringBase<SourceOp,
FloatIntOpToFuncCallLowering<SourceOp>> {
public:
explicit FloatIntOpToFuncCallLowering(const LLVMTypeConverter &lowering,
StringRef f32Func, StringRef f64Func)
: OpToFuncCallLoweringBase<SourceOp,
FloatIntOpToFuncCallLowering<SourceOp>>(
lowering),
f32Func(f32Func), f64Func(f64Func) {}

Value maybeCast(Value operand, PatternRewriter &rewriter) const {
return operand;
}

StringRef getFunctionName(Type type, SourceOp op) const {
if (isa<Float32Type>(type)) {
return f32Func;
}
if (isa<Float64Type>(type))
return f64Func;
return "";
}

const std::string f32Func;
const std::string f64Func;
const std::string f32ApproxFunc;
const std::string f16Func;
};

} // namespace mlir
Expand Down
19 changes: 19 additions & 0 deletions mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -446,6 +446,22 @@ static void populateOpPatterns(const LLVMTypeConverter &converter,
f32ApproxFunc, f16Func);
}

template <typename OpTy>
static void populateIntOpPatterns(const LLVMTypeConverter &converter,
RewritePatternSet &patterns,
StringRef i32Func) {
patterns.add<ScalarizeVectorOpLowering<OpTy>>(converter);
patterns.add<IntOpToFuncCallLowering<OpTy>>(converter, i32Func);
}

template <typename OpTy>
static void populateFloatIntOpPatterns(const LLVMTypeConverter &converter,
RewritePatternSet &patterns,
StringRef f32Func, StringRef f64Func) {
patterns.add<ScalarizeVectorOpLowering<OpTy>>(converter);
patterns.add<FloatIntOpToFuncCallLowering<OpTy>>(converter, f32Func, f64Func);
}

void mlir::populateGpuSubgroupReduceOpLoweringPattern(
const LLVMTypeConverter &converter, RewritePatternSet &patterns) {
patterns.add<GPUSubgroupReduceOpLowering>(converter);
Expand Down Expand Up @@ -509,6 +525,7 @@ void mlir::populateGpuToNVVMConversionPatterns(

populateOpPatterns<arith::RemFOp>(converter, patterns, "__nv_fmodf",
"__nv_fmod");
populateIntOpPatterns<math::AbsIOp>(converter, patterns, "__nv_abs");
populateOpPatterns<math::AbsFOp>(converter, patterns, "__nv_fabsf",
"__nv_fabs");
populateOpPatterns<math::AcosOp>(converter, patterns, "__nv_acosf",
Expand Down Expand Up @@ -555,6 +572,8 @@ void mlir::populateGpuToNVVMConversionPatterns(
"__nv_log2", "__nv_fast_log2f");
populateOpPatterns<math::PowFOp>(converter, patterns, "__nv_powf", "__nv_pow",
"__nv_fast_powf");
populateFloatIntOpPatterns<math::FPowIOp>(converter, patterns, "__nv_powif",
"__nv_powi");
populateOpPatterns<math::RoundOp>(converter, patterns, "__nv_roundf",
"__nv_round");
populateOpPatterns<math::RoundEvenOp>(converter, patterns, "__nv_rintf",
Expand Down
24 changes: 24 additions & 0 deletions mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1033,3 +1033,27 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}


gpu.module @test_module_52 {
// CHECK: llvm.func @__nv_abs(i32) -> i32
// CHECK-LABEL: func @gpu_abs
func.func @gpu_abs(%arg_i32 : i32) -> (i32) {
%result32 = math.absi %arg_i32 : i32
// CHECK: llvm.call @__nv_abs(%{{.*}}) : (i32) -> i32
func.return %result32 : i32
}
}

gpu.module @test_module_53 {
// CHECK: llvm.func @__nv_powif(f32, i32) -> f32
// CHECK: llvm.func @__nv_powi(f64, i32) -> f64
// CHECK-LABEL: func @gpu_powi
func.func @gpu_powi(%arg_f32 : f32, %arg_f64 : f64, %arg_i32 : i32) -> (f32, f64) {
%result32 = math.fpowi %arg_f32, %arg_i32 : f32, i32
// CHECK: llvm.call @__nv_powif(%{{.*}}, %{{.*}}) : (f32, i32) -> f32
%result64 = math.fpowi %arg_f64, %arg_i32 : f64, i32
// CHECK: llvm.call @__nv_powi(%{{.*}}, %{{.*}}) : (f64, i32) -> f64
func.return %result32, %result64 : f32, f64
}
}
Loading