750e90e440
Many machine-learning applications (and most software written at AMD) expect the operation that truncates floats to 8-bit floats to be saturatinng. That is, they expect `truncf 256.0 : f32 to f8E4M3FNUZ` to yield `240.0`, not `NaN`, and similarly for negative numbers. However, the underlying hardware instruction that can be used for this truncation implements overflow-to-NaN semantics. To enable handling this usecase, we add the saturate-fp8-truncf option to ArithToAMDGPU (off by default), which causes the requisite clamping code to be emitted. Said clamping code ensures that Inf and NaN are passed through exactly (and thus trancate to NaN). Per review feedback, this commit efactors createScalarOrSplatConstant() to the Arith dialect utilities and uses it in this code. It also fixes naming of existing patterns and switches from vector.extractelement/insertelement to vector.extract/insert.
263 lines
11 KiB
C++
263 lines
11 KiB
C++
//===- ArithToAMDGPU.cpp - Arith to AMDGPU dialect conversion ---------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Conversion/ArithToAMDGPU/ArithToAMDGPU.h"
|
|
|
|
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Arith/Utils/Utils.h"
|
|
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/IR/TypeUtilities.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Support/LogicalResult.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
|
|
namespace mlir {
|
|
#define GEN_PASS_DEF_ARITHTOAMDGPUCONVERSIONPASS
|
|
#include "mlir/Conversion/Passes.h.inc"
|
|
} // namespace mlir
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
struct ArithToAMDGPUConversionPass final
|
|
: impl::ArithToAMDGPUConversionPassBase<ArithToAMDGPUConversionPass> {
|
|
using impl::ArithToAMDGPUConversionPassBase<
|
|
ArithToAMDGPUConversionPass>::ArithToAMDGPUConversionPassBase;
|
|
|
|
void runOnOperation() override;
|
|
};
|
|
|
|
struct ExtFOnFloat8RewritePattern final : OpRewritePattern<arith::ExtFOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult match(arith::ExtFOp op) const override;
|
|
void rewrite(arith::ExtFOp op, PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
struct TruncFToFloat8RewritePattern final : OpRewritePattern<arith::TruncFOp> {
|
|
bool saturateFP8 = false;
|
|
TruncFToFloat8RewritePattern(MLIRContext *ctx, bool saturateFP8)
|
|
: OpRewritePattern::OpRewritePattern(ctx), saturateFP8(saturateFP8) {}
|
|
|
|
LogicalResult match(arith::TruncFOp op) const override;
|
|
void rewrite(arith::TruncFOp op, PatternRewriter &rewriter) const override;
|
|
};
|
|
} // end namespace
|
|
|
|
static Value castF32To(Type elementType, Value f32, Location loc,
|
|
PatternRewriter &rewriter) {
|
|
if (elementType.isF32())
|
|
return f32;
|
|
if (elementType.getIntOrFloatBitWidth() < 32)
|
|
return rewriter.create<arith::TruncFOp>(loc, elementType, f32);
|
|
if (elementType.getIntOrFloatBitWidth() > 32)
|
|
return rewriter.create<arith::ExtFOp>(loc, elementType, f32);
|
|
llvm_unreachable("The only 32-bit float type is f32");
|
|
}
|
|
|
|
LogicalResult ExtFOnFloat8RewritePattern::match(arith::ExtFOp op) const {
|
|
Type inType = op.getIn().getType();
|
|
if (auto inVecType = inType.dyn_cast<VectorType>()) {
|
|
if (inVecType.isScalable())
|
|
return failure();
|
|
if (inVecType.getShape().size() > 1)
|
|
// Multi-dimensional vectors are currently unsupported.
|
|
return failure();
|
|
inType = inVecType.getElementType();
|
|
}
|
|
return success(inType.isFloat8E5M2FNUZ() || inType.isFloat8E4M3FNUZ());
|
|
}
|
|
|
|
void ExtFOnFloat8RewritePattern::rewrite(arith::ExtFOp op,
|
|
PatternRewriter &rewriter) const {
|
|
Location loc = op.getLoc();
|
|
Value in = op.getIn();
|
|
Type outElemType = getElementTypeOrSelf(op.getOut().getType());
|
|
if (!in.getType().isa<VectorType>()) {
|
|
Value asFloat = rewriter.create<amdgpu::ExtPackedFp8Op>(
|
|
loc, rewriter.getF32Type(), in, 0);
|
|
Value result = castF32To(outElemType, asFloat, loc, rewriter);
|
|
return rewriter.replaceOp(op, result);
|
|
}
|
|
VectorType inType = in.getType().cast<VectorType>();
|
|
int64_t numElements = inType.getNumElements();
|
|
Value zero = rewriter.createOrFold<arith::ConstantOp>(
|
|
loc, outElemType, rewriter.getFloatAttr(outElemType, 0.0));
|
|
Value result =
|
|
rewriter.createOrFold<vector::SplatOp>(loc, op.getOut().getType(), zero);
|
|
if (inType.getShape().empty()) {
|
|
Value scalarIn =
|
|
rewriter.create<vector::ExtractOp>(loc, in, ArrayRef<int64_t>{});
|
|
// Recurse to send the 0-D vector case to the 1-D vector case
|
|
Value scalarExt =
|
|
rewriter.create<arith::ExtFOp>(loc, outElemType, scalarIn);
|
|
result = rewriter.create<vector::InsertOp>(loc, scalarExt, zero,
|
|
ArrayRef<int64_t>{});
|
|
return rewriter.replaceOp(op, result);
|
|
}
|
|
for (int64_t i = 0; i < numElements; i += 4) {
|
|
int64_t elemsThisOp = std::min(numElements, i + 4) - i;
|
|
Value inSlice = rewriter.create<vector::ExtractStridedSliceOp>(
|
|
loc, in, i, elemsThisOp, 1);
|
|
for (int64_t j = 0; j < elemsThisOp; ++j) {
|
|
Value asFloat = rewriter.create<amdgpu::ExtPackedFp8Op>(
|
|
loc, rewriter.getF32Type(), inSlice, j);
|
|
Value asType = castF32To(outElemType, asFloat, loc, rewriter);
|
|
result = rewriter.create<vector::InsertOp>(loc, asType, result, i + j);
|
|
}
|
|
}
|
|
rewriter.replaceOp(op, result);
|
|
}
|
|
|
|
static Value castToF32(Value value, Location loc, PatternRewriter &rewriter) {
|
|
Type type = value.getType();
|
|
if (type.isF32())
|
|
return value;
|
|
if (type.getIntOrFloatBitWidth() < 32)
|
|
return rewriter.create<arith::ExtFOp>(loc, rewriter.getF32Type(), value);
|
|
if (type.getIntOrFloatBitWidth() > 32)
|
|
return rewriter.create<arith::TruncFOp>(loc, rewriter.getF32Type(), value);
|
|
llvm_unreachable("The only 32-bit float type is f32");
|
|
}
|
|
|
|
// If `in` is a finite value, clamp it between the maximum and minimum values
|
|
// of `outElemType` so that subsequent conversion instructions don't
|
|
// overflow those out-of-range values to NaN. These semantics are commonly
|
|
// used in machine-learning contexts where failure to clamp would lead to
|
|
// excessive NaN production.
|
|
static Value clampInput(PatternRewriter &rewriter, Location loc,
|
|
Type outElemType, Value source) {
|
|
Type sourceType = source.getType();
|
|
const llvm::fltSemantics &sourceSem =
|
|
cast<FloatType>(getElementTypeOrSelf(sourceType)).getFloatSemantics();
|
|
const llvm::fltSemantics &targetSem =
|
|
cast<FloatType>(outElemType).getFloatSemantics();
|
|
|
|
APFloat min = APFloat::getLargest(targetSem, /*Negative=*/true);
|
|
APFloat max = APFloat::getLargest(targetSem, /*Negative=*/false);
|
|
bool ignoredLosesInfo = false;
|
|
// We can ignore conversion failures here because this conversion promotes
|
|
// from a smaller type to a larger one - ex. there can be no loss of precision
|
|
// when casting fp8 to f16.
|
|
(void)min.convert(sourceSem, APFloat::rmNearestTiesToEven, &ignoredLosesInfo);
|
|
(void)max.convert(sourceSem, APFloat::rmNearestTiesToEven, &ignoredLosesInfo);
|
|
|
|
Value minCst = createScalarOrSplatConstant(rewriter, loc, sourceType, min);
|
|
Value maxCst = createScalarOrSplatConstant(rewriter, loc, sourceType, max);
|
|
|
|
Value inf = createScalarOrSplatConstant(
|
|
rewriter, loc, sourceType,
|
|
APFloat::getInf(sourceSem, /*Negative=*/false));
|
|
Value negInf = createScalarOrSplatConstant(
|
|
rewriter, loc, sourceType, APFloat::getInf(sourceSem, /*Negative=*/true));
|
|
Value isInf = rewriter.createOrFold<arith::CmpFOp>(
|
|
loc, arith::CmpFPredicate::OEQ, source, inf);
|
|
Value isNegInf = rewriter.createOrFold<arith::CmpFOp>(
|
|
loc, arith::CmpFPredicate::OEQ, source, negInf);
|
|
Value isNan = rewriter.createOrFold<arith::CmpFOp>(
|
|
loc, arith::CmpFPredicate::UNO, source, source);
|
|
Value isNonFinite = rewriter.create<arith::OrIOp>(
|
|
loc, rewriter.create<arith::OrIOp>(loc, isInf, isNegInf), isNan);
|
|
|
|
Value clampedBelow = rewriter.create<arith::MaximumFOp>(loc, source, minCst);
|
|
Value clamped = rewriter.create<arith::MinimumFOp>(loc, clampedBelow, maxCst);
|
|
Value res =
|
|
rewriter.create<arith::SelectOp>(loc, isNonFinite, source, clamped);
|
|
return res;
|
|
}
|
|
|
|
LogicalResult TruncFToFloat8RewritePattern::match(arith::TruncFOp op) const {
|
|
Type outType = op.getOut().getType();
|
|
if (auto outVecType = outType.dyn_cast<VectorType>()) {
|
|
if (outVecType.isScalable())
|
|
return failure();
|
|
if (outVecType.getShape().size() > 1)
|
|
// Multi-dimensional vectors are currently unsupported.
|
|
return failure();
|
|
outType = outVecType.getElementType();
|
|
}
|
|
auto inType = dyn_cast<FloatType>(getElementTypeOrSelf(op.getIn().getType()));
|
|
if (inType && inType.getWidth() <= 8 && saturateFP8)
|
|
// Conversion between 8-bit floats is not supported with truncation enabled.
|
|
return failure();
|
|
return success(outType.isFloat8E5M2FNUZ() || outType.isFloat8E4M3FNUZ());
|
|
}
|
|
|
|
void TruncFToFloat8RewritePattern::rewrite(arith::TruncFOp op,
|
|
PatternRewriter &rewriter) const {
|
|
Location loc = op.getLoc();
|
|
Value in = op.getIn();
|
|
Type outElemType = getElementTypeOrSelf(op.getOut().getType());
|
|
if (saturateFP8)
|
|
in = clampInput(rewriter, loc, outElemType, in);
|
|
VectorType truncResType = VectorType::get(4, outElemType);
|
|
if (!in.getType().isa<VectorType>()) {
|
|
Value asFloat = castToF32(in, loc, rewriter);
|
|
Value asF8s = rewriter.create<amdgpu::PackedTrunc2xFp8Op>(
|
|
loc, truncResType, asFloat, /*sourceB=*/nullptr, 0,
|
|
/*existing=*/nullptr);
|
|
Value result = rewriter.create<vector::ExtractOp>(loc, asF8s, 0);
|
|
return rewriter.replaceOp(op, result);
|
|
}
|
|
VectorType outType = op.getOut().getType().cast<VectorType>();
|
|
int64_t numElements = outType.getNumElements();
|
|
Value zero = rewriter.createOrFold<arith::ConstantOp>(
|
|
loc, outElemType, rewriter.getFloatAttr(outElemType, 0.0));
|
|
Value result = rewriter.createOrFold<vector::SplatOp>(loc, outType, zero);
|
|
if (outType.getShape().empty()) {
|
|
Value scalarIn =
|
|
rewriter.create<vector::ExtractOp>(loc, in, ArrayRef<int64_t>{});
|
|
// Recurse to send the 0-D vector case to the 1-D vector case
|
|
Value scalarTrunc =
|
|
rewriter.create<arith::TruncFOp>(loc, outElemType, scalarIn);
|
|
result = rewriter.create<vector::InsertOp>(loc, scalarTrunc, zero,
|
|
ArrayRef<int64_t>{});
|
|
return rewriter.replaceOp(op, result);
|
|
}
|
|
|
|
for (int64_t i = 0; i < numElements; i += 4) {
|
|
int64_t elemsThisOp = std::min(numElements, i + 4) - i;
|
|
Value thisResult = nullptr;
|
|
for (int64_t j = 0; j < elemsThisOp; j += 2) {
|
|
Value elemA = rewriter.create<vector::ExtractOp>(loc, in, i + j);
|
|
Value asFloatA = castToF32(elemA, loc, rewriter);
|
|
Value asFloatB = nullptr;
|
|
if (j + 1 < elemsThisOp) {
|
|
Value elemB = rewriter.create<vector::ExtractOp>(loc, in, i + j + 1);
|
|
asFloatB = castToF32(elemB, loc, rewriter);
|
|
}
|
|
thisResult = rewriter.create<amdgpu::PackedTrunc2xFp8Op>(
|
|
loc, truncResType, asFloatA, asFloatB, j / 2, thisResult);
|
|
}
|
|
if (elemsThisOp < 4)
|
|
thisResult = rewriter.create<vector::ExtractStridedSliceOp>(
|
|
loc, thisResult, 0, elemsThisOp, 1);
|
|
result = rewriter.create<vector::InsertStridedSliceOp>(loc, thisResult,
|
|
result, i, 1);
|
|
}
|
|
rewriter.replaceOp(op, result);
|
|
}
|
|
|
|
void mlir::arith::populateArithToAMDGPUConversionPatterns(
|
|
RewritePatternSet &patterns, bool saturateFP8TruncF) {
|
|
patterns.add<ExtFOnFloat8RewritePattern>(patterns.getContext());
|
|
patterns.add<TruncFToFloat8RewritePattern>(patterns.getContext(),
|
|
saturateFP8TruncF);
|
|
}
|
|
|
|
void ArithToAMDGPUConversionPass::runOnOperation() {
|
|
Operation *op = getOperation();
|
|
RewritePatternSet patterns(op->getContext());
|
|
arith::populateArithToAMDGPUConversionPatterns(patterns, saturateFP8Truncf);
|
|
if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
|
|
return signalPassFailure();
|
|
}
|