[MLIR][NVGPU] Add truncf and extf Ops (#199700)
This change adds the `extf` and `truncf` ops to the NVGPU dialect to support floating-point conversion operations. These ops support scalar and vector inputs and lower to the corresponding NVVM ops after padding and chunking the input into `i32` registers.
This commit is contained in:
Srinivasa Ravi
GitHub
parent
96dcb03caa
commit
16094a86a6
11 changed files with 2103 additions and 1 deletions
|
|
@ -1106,7 +1106,8 @@ def ConvertNVGPUToNVVMPass : Pass<"convert-nvgpu-to-nvvm"> {
|
|||
"arith::ArithDialect",
|
||||
"LLVM::LLVMDialect",
|
||||
"memref::MemRefDialect",
|
||||
"NVVM::NVVMDialect"
|
||||
"NVVM::NVVMDialect",
|
||||
"vector::VectorDialect"
|
||||
];
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -671,4 +671,79 @@ def NVGPU_RcpOp : NVGPU_Op<"rcp", [Pure,
|
|||
let hasVerifier = 1;
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// NVGPU Conversion Ops
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
def AnyI32Like : TypeOrValueSemanticsContainer<I32, "scalar i32 or vector of i32">;
|
||||
|
||||
def NVGPU_TruncfOp : NVGPU_Op<"truncf", [Pure]> {
|
||||
let summary = "Truncate floating-point to narrower floating-point";
|
||||
let description = [{
|
||||
Truncate a floating-point value to a smaller floating-point type.
|
||||
Destination must be strictly narrower than source.
|
||||
|
||||
Supported paths:
|
||||
f16 -> f8 (f8E4M3FN, f8E5M2, f8E8M0FNU), f6 (f6E2M3FN, f6E3M2FN),
|
||||
f4 (f4E2M1FN)
|
||||
bf16 -> f8 (f8E4M3FN, f8E5M2, f8E8M0FNU), f6 (f6E2M3FN, f6E3M2FN),
|
||||
f4 (f4E2M1FN)
|
||||
f32 -> f16, bf16, f8 (f8E4M3FN, f8E5M2, f8E8M0FNU),
|
||||
f6 (f6E2M3FN, f6E3M2FN), f4 (f4E2M1FN)
|
||||
f64 -> f32, f16, bf16
|
||||
|
||||
The `random_bits` operand enables stochastic rounding (RS mode) for
|
||||
f32->f16/bf16 conversions; when provided, `rnd` must be RS.
|
||||
|
||||
Example:
|
||||
```mlir
|
||||
%r = nvgpu.truncf %in : vector<8xf32> to vector<8xf8E4M3FN>
|
||||
%r = nvgpu.truncf %in : vector<8xf32> to vector<8xf6E2M3FN>
|
||||
%r = nvgpu.truncf %in : f32 to f16
|
||||
%r = nvgpu.truncf %in : vector<2x4xf16> to vector<2x4xf8E5M2>
|
||||
```
|
||||
}];
|
||||
let arguments = (ins AnyTypeOf<[AnyFloat, VectorOfNonZeroRankOf<[AnyFloat]>]>:$in,
|
||||
DefaultValuedAttr<FPRoundingModeAttr, "NVVM::FPRoundingMode::RN">:$rnd,
|
||||
DefaultValuedAttr<ConfinedAttr<SaturationModeAttr,
|
||||
[EnumAttrIsOneOf<SaturationModeAttr,
|
||||
[SaturationModeNone, SaturationModeFinite]>]>,
|
||||
"NVVM::SaturationMode::SATFINITE">:$sat,
|
||||
DefaultValuedAttr<BoolAttr, "false">:$relu,
|
||||
Optional<I32>:$random_bits
|
||||
);
|
||||
let results = (outs AnyTypeOf<[AnyFloat, VectorOfNonZeroRankOf<[AnyFloat]>]>:$out);
|
||||
let assemblyFormat = "$in (`,` $random_bits^)? attr-dict `:` type($in) `to` type($out)";
|
||||
let hasVerifier = 1;
|
||||
}
|
||||
|
||||
def NVGPU_ExtfOp : NVGPU_Op<"extf", [Pure]> {
|
||||
let summary = "Extend floating-point to wider floating-point";
|
||||
let description = [{
|
||||
Extend a floating-point value to a wider floating-point type.
|
||||
Destination must be strictly wider than source.
|
||||
|
||||
Supported paths:
|
||||
f4 (f4E2M1FN) -> f16, bf16
|
||||
f6 (f6E2M3FN, f6E3M2FN) -> f16, bf16
|
||||
f8 (f8E4M3FN, f8E5M2, f8E8M0FNU) -> f16, bf16
|
||||
f16 -> f32
|
||||
bf16 -> f32
|
||||
|
||||
Example:
|
||||
```mlir
|
||||
%r = nvgpu.extf %in : vector<8xf8E4M3FN> to vector<8xf16>
|
||||
%r = nvgpu.extf %in : vector<8xf6E3M2FN> to vector<8xf16>
|
||||
%r = nvgpu.extf %in : vector<4xf8E5M2> to vector<4xf32>
|
||||
%r = nvgpu.extf %in : f8E4M3FN to f32
|
||||
```
|
||||
}];
|
||||
let arguments = (ins AnyTypeOf<[AnyFloat, VectorOfNonZeroRankOf<[AnyFloat]>]>:$in,
|
||||
DefaultValuedAttr<FPRoundingModeAttr, "NVVM::FPRoundingMode::RN">:$rnd,
|
||||
DefaultValuedAttr<BoolAttr, "false">:$relu);
|
||||
let results = (outs AnyTypeOf<[AnyFloat, VectorOfNonZeroRankOf<[AnyFloat]>]>:$out);
|
||||
let assemblyFormat = "$in attr-dict `:` type($in) `to` type($out)";
|
||||
let hasVerifier = 1;
|
||||
}
|
||||
|
||||
#endif // MLIR_DIALECT_NVGPU_IR_NVGPUOPS_TD
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ add_mlir_conversion_library(MLIRNVGPUToNVVM
|
|||
MLIRNVGPUDialect
|
||||
MLIRNVVMDialect
|
||||
MLIRArithDialect
|
||||
MLIRVectorDialect
|
||||
MLIRPass
|
||||
MLIRSCFTransforms
|
||||
MLIRTransforms
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@
|
|||
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
||||
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
|
||||
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
|
||||
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
||||
#include "mlir/IR/BuiltinTypes.h"
|
||||
#include "mlir/IR/PatternMatch.h"
|
||||
#include "mlir/IR/TypeUtilities.h"
|
||||
|
|
@ -461,6 +462,7 @@ struct ConvertNVGPUToNVVMPass
|
|||
target.addLegalDialect<::mlir::arith::ArithDialect>();
|
||||
target.addLegalDialect<::mlir::memref::MemRefDialect>();
|
||||
target.addLegalDialect<::mlir::NVVM::NVVMDialect>();
|
||||
target.addLegalDialect<::mlir::vector::VectorDialect>();
|
||||
mlir::scf::populateSCFStructuralTypeConversionsAndLegality(
|
||||
converter, patterns, target);
|
||||
if (failed(applyPartialConversion(getOperation(), target,
|
||||
|
|
@ -1709,6 +1711,668 @@ struct NVGPURcpOpLowering : public ConvertOpToLLVMPattern<nvgpu::RcpOp> {
|
|||
rewriter);
|
||||
}
|
||||
};
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// NVGPUTruncfOp Lowering
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
enum class FPKind { F32, BF16, F16, F8, F6, F4 };
|
||||
|
||||
/// Get the effective bit width of a floating-point type.
|
||||
/// f6 types are 6-bit but NVVM Ops expect 8-bit (i8) containers.
|
||||
static int getEffectiveBitWidth(int bitWidth) {
|
||||
return bitWidth == 6 ? 8 : bitWidth;
|
||||
}
|
||||
|
||||
static std::optional<FPKind> classifyFPType(Type t) {
|
||||
static constexpr auto isConvertibleF8Type = [](Type t) {
|
||||
return isa<Float8E4M3FNType, Float8E5M2Type, Float8E8M0FNUType>(t);
|
||||
};
|
||||
static constexpr auto isConvertibleF6Type = [](Type t) {
|
||||
return isa<Float6E2M3FNType, Float6E3M2FNType>(t);
|
||||
};
|
||||
static constexpr auto isConvertibleF4Type = [](Type t) {
|
||||
return isa<Float4E2M1FNType>(t);
|
||||
};
|
||||
|
||||
if (t.isF32())
|
||||
return FPKind::F32;
|
||||
if (t.isBF16())
|
||||
return FPKind::BF16;
|
||||
if (t.isF16())
|
||||
return FPKind::F16;
|
||||
if (isConvertibleF8Type(t))
|
||||
return FPKind::F8;
|
||||
if (isConvertibleF6Type(t))
|
||||
return FPKind::F6;
|
||||
if (isConvertibleF4Type(t))
|
||||
return FPKind::F4;
|
||||
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
/// Conversion op identifier for nvgpu.truncf lowering dispatch table.
|
||||
enum class FPTruncConvOp {
|
||||
F32x2_TO_F16x2,
|
||||
F32x2_TO_BF16x2,
|
||||
F32x2_TO_F8x2,
|
||||
F32x2_TO_F6x2,
|
||||
F32x2_TO_F4x2,
|
||||
F16x2_TO_F8x2,
|
||||
F16x2_TO_F6x2,
|
||||
F16x2_TO_F4x2,
|
||||
BF16x2_TO_F8x2,
|
||||
BF16x2_TO_F6x2,
|
||||
BF16x2_TO_F4x2,
|
||||
};
|
||||
|
||||
struct FPTruncTableEntry {
|
||||
FPKind src;
|
||||
FPKind dst;
|
||||
FPTruncConvOp convOp;
|
||||
};
|
||||
|
||||
static constexpr FPTruncTableEntry kFPTruncTable[] = {
|
||||
// f32 source
|
||||
{FPKind::F32, FPKind::F16, FPTruncConvOp::F32x2_TO_F16x2},
|
||||
{FPKind::F32, FPKind::BF16, FPTruncConvOp::F32x2_TO_BF16x2},
|
||||
{FPKind::F32, FPKind::F8, FPTruncConvOp::F32x2_TO_F8x2},
|
||||
{FPKind::F32, FPKind::F6, FPTruncConvOp::F32x2_TO_F6x2},
|
||||
{FPKind::F32, FPKind::F4, FPTruncConvOp::F32x2_TO_F4x2},
|
||||
// f16 source
|
||||
{FPKind::F16, FPKind::F8, FPTruncConvOp::F16x2_TO_F8x2},
|
||||
{FPKind::F16, FPKind::F6, FPTruncConvOp::F16x2_TO_F6x2},
|
||||
{FPKind::F16, FPKind::F4, FPTruncConvOp::F16x2_TO_F4x2},
|
||||
// bf16 source
|
||||
{FPKind::BF16, FPKind::F8, FPTruncConvOp::BF16x2_TO_F8x2},
|
||||
{FPKind::BF16, FPKind::F6, FPTruncConvOp::BF16x2_TO_F6x2},
|
||||
{FPKind::BF16, FPKind::F4, FPTruncConvOp::BF16x2_TO_F4x2},
|
||||
};
|
||||
|
||||
/// Find the conversion table entry whose source/destination `FPKind`s match the
|
||||
/// given element types.
|
||||
template <typename TableEntry, size_t N>
|
||||
static std::optional<TableEntry>
|
||||
lookupConvOp(const TableEntry (&table)[N], Type srcElemType, Type dstElemType) {
|
||||
std::optional<FPKind> srcKind = classifyFPType(srcElemType);
|
||||
std::optional<FPKind> dstKind = classifyFPType(dstElemType);
|
||||
if (!srcKind || !dstKind)
|
||||
return std::nullopt;
|
||||
for (const TableEntry &entry : table) {
|
||||
if (entry.src == *srcKind && entry.dst == *dstKind)
|
||||
return entry;
|
||||
}
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
/// Extract a single element from a vector.
|
||||
static Value extractElement(ImplicitLocOpBuilder &b, Value srcVec, int idx) {
|
||||
auto vecTy = cast<VectorType>(srcVec.getType());
|
||||
assert(idx >= 0 && idx < vecTy.getNumElements() &&
|
||||
"extractElement: index out of bounds");
|
||||
IntegerType i64Ty = b.getI64Type();
|
||||
return b.create<LLVM::ExtractElementOp>(
|
||||
srcVec, b.create<LLVM::ConstantOp>(i64Ty, b.getI64IntegerAttr(idx)));
|
||||
}
|
||||
|
||||
/// Extract a pair of f32 values from an i32 vector at the given base index.
|
||||
static std::pair<Value, Value> extractF32Pair(ImplicitLocOpBuilder &b,
|
||||
Value srcI32Vec, int baseIdx) {
|
||||
FloatType f32Ty = b.getF32Type();
|
||||
Value elem0 = extractElement(b, srcI32Vec, baseIdx);
|
||||
Value elem1 = extractElement(b, srcI32Vec, baseIdx + 1);
|
||||
return {b.create<LLVM::BitcastOp>(f32Ty, elem0),
|
||||
b.create<LLVM::BitcastOp>(f32Ty, elem1)};
|
||||
}
|
||||
|
||||
/// Extract a vector of elements of size i32 from an i32 vector and bitcast to
|
||||
/// the specified vector type.
|
||||
static Value extractAndBitcast(ImplicitLocOpBuilder &b, Value srcI32Vec,
|
||||
int idx, VectorType vecTy) {
|
||||
Value elem = extractElement(b, srcI32Vec, idx);
|
||||
return b.create<LLVM::BitcastOp>(vecTy, elem);
|
||||
}
|
||||
|
||||
/// Create a sub-byte conversion from an f32 pair source and return the native
|
||||
/// result.
|
||||
template <typename ConvertOp, typename... Args>
|
||||
static Value convertFromF32Pair(ImplicitLocOpBuilder &b, Value srcI32Vec,
|
||||
int srcBaseIdx, Type resultTy, Args &&...args) {
|
||||
auto [lo, hi] = extractF32Pair(b, srcI32Vec, srcBaseIdx);
|
||||
return b.create<ConvertOp>(resultTy, hi, lo, std::forward<Args>(args)...);
|
||||
}
|
||||
|
||||
/// Create a sub-byte conversion from a packed f16x2/bf16x2 source and return
|
||||
/// the native result.
|
||||
template <typename ConvertOp, typename... Args>
|
||||
static Value convertFromPacked(ImplicitLocOpBuilder &b, Value srcI32Vec,
|
||||
int srcBaseIdx, Type srcElemTy, Type resultTy,
|
||||
Args &&...args) {
|
||||
Value src = extractAndBitcast(b, srcI32Vec, srcBaseIdx,
|
||||
VectorType::get(2, srcElemTy));
|
||||
return b.create<ConvertOp>(resultTy, src, std::forward<Args>(args)...);
|
||||
}
|
||||
|
||||
/// Create a typed NVVM truncation conversion.
|
||||
static Value createTruncConversion(
|
||||
ImplicitLocOpBuilder &b, MLIRContext *ctx, FPTruncConvOp convOp,
|
||||
Value srcI32Vec, int srcBaseIdx, NVVM::FPRoundingModeAttr rndAttr,
|
||||
NVVM::SaturationModeAttr satAttr, BoolAttr reluAttr, Type dstElemType,
|
||||
Type actualDstFloatType, Value randomBits = Value()) {
|
||||
IntegerType i8Ty = b.getI8Type();
|
||||
IntegerType i16Ty = b.getI16Type();
|
||||
IntegerType i32Ty = b.getI32Type();
|
||||
TypeAttr dstTyAttr = TypeAttr::get(dstElemType);
|
||||
TypeAttr actualDstTyAttr = TypeAttr::get(actualDstFloatType);
|
||||
|
||||
switch (convOp) {
|
||||
case FPTruncConvOp::F32x2_TO_F16x2: {
|
||||
auto [lo, hi] = extractF32Pair(b, srcI32Vec, srcBaseIdx);
|
||||
Value r = b.create<NVVM::ConvertF32x2ToF16x2Op>(
|
||||
VectorType::get(2, b.getF16Type()), hi, lo, randomBits, rndAttr,
|
||||
satAttr, reluAttr);
|
||||
return b.create<LLVM::BitcastOp>(i32Ty, r);
|
||||
}
|
||||
case FPTruncConvOp::F32x2_TO_BF16x2: {
|
||||
auto [lo, hi] = extractF32Pair(b, srcI32Vec, srcBaseIdx);
|
||||
Value r = b.create<NVVM::ConvertF32x2ToBF16x2Op>(
|
||||
VectorType::get(2, b.getBF16Type()), hi, lo, randomBits, rndAttr,
|
||||
satAttr, reluAttr);
|
||||
return b.create<LLVM::BitcastOp>(i32Ty, r);
|
||||
}
|
||||
case FPTruncConvOp::F32x2_TO_F8x2:
|
||||
return convertFromF32Pair<NVVM::ConvertF32x2ToF8x2Op>(
|
||||
b, srcI32Vec, srcBaseIdx, i16Ty, rndAttr, satAttr, reluAttr, dstTyAttr);
|
||||
case FPTruncConvOp::F32x2_TO_F6x2:
|
||||
return convertFromF32Pair<NVVM::ConvertF32x2ToF6x2Op>(
|
||||
b, srcI32Vec, srcBaseIdx, i16Ty, reluAttr, actualDstTyAttr);
|
||||
case FPTruncConvOp::F32x2_TO_F4x2:
|
||||
return convertFromF32Pair<NVVM::ConvertF32x2ToF4x2Op>(
|
||||
b, srcI32Vec, srcBaseIdx, i8Ty, reluAttr, dstTyAttr);
|
||||
case FPTruncConvOp::F16x2_TO_F8x2:
|
||||
return convertFromPacked<NVVM::ConvertF16x2ToF8x2Op>(
|
||||
b, srcI32Vec, srcBaseIdx, b.getF16Type(), i16Ty, reluAttr, dstTyAttr);
|
||||
case FPTruncConvOp::F16x2_TO_F6x2:
|
||||
return convertFromPacked<NVVM::ConvertF16x2ToF6x2Op>(
|
||||
b, srcI32Vec, srcBaseIdx, b.getF16Type(), i16Ty, reluAttr,
|
||||
actualDstTyAttr);
|
||||
case FPTruncConvOp::F16x2_TO_F4x2:
|
||||
return convertFromPacked<NVVM::ConvertF16x2ToF4x2Op>(
|
||||
b, srcI32Vec, srcBaseIdx, b.getF16Type(), i8Ty, reluAttr,
|
||||
actualDstTyAttr);
|
||||
case FPTruncConvOp::BF16x2_TO_F8x2:
|
||||
return convertFromPacked<NVVM::ConvertBF16x2ToF8x2Op>(
|
||||
b, srcI32Vec, srcBaseIdx, b.getBF16Type(), i16Ty, rndAttr, satAttr,
|
||||
reluAttr, dstTyAttr);
|
||||
case FPTruncConvOp::BF16x2_TO_F6x2:
|
||||
return convertFromPacked<NVVM::ConvertBF16x2ToF6x2Op>(
|
||||
b, srcI32Vec, srcBaseIdx, b.getBF16Type(), i16Ty, reluAttr,
|
||||
actualDstTyAttr);
|
||||
case FPTruncConvOp::BF16x2_TO_F4x2:
|
||||
return convertFromPacked<NVVM::ConvertBF16x2ToF4x2Op>(
|
||||
b, srcI32Vec, srcBaseIdx, b.getBF16Type(), i8Ty, reluAttr,
|
||||
actualDstTyAttr);
|
||||
}
|
||||
llvm_unreachable("unhandled FPTruncConvOp");
|
||||
}
|
||||
|
||||
static LogicalResult lowerTruncf(nvgpu::TruncfOp op,
|
||||
nvgpu::TruncfOp::Adaptor adaptor,
|
||||
ConversionPatternRewriter &rewriter,
|
||||
const LLVMTypeConverter *typeConverter) {
|
||||
MLIRContext *ctx = op.getContext();
|
||||
ImplicitLocOpBuilder b(op->getLoc(), rewriter);
|
||||
IntegerType i32Ty = b.getI32Type();
|
||||
IntegerType i64Ty = b.getI64Type();
|
||||
static constexpr int regBits = 32;
|
||||
|
||||
auto srcType = llvm::dyn_cast<VectorType>(op.getIn().getType());
|
||||
auto dstType = llvm::dyn_cast<VectorType>(op.getOut().getType());
|
||||
if (!srcType || srcType.getRank() != 1 || !dstType || dstType.getRank() != 1)
|
||||
return rewriter.notifyMatchFailure(
|
||||
op, "expected 1-D vector; canonicalize pattern handles other shapes");
|
||||
|
||||
auto srcElemType = srcType.getElementType();
|
||||
auto dstElemType = dstType.getElementType();
|
||||
int srcBW = srcType.getElementTypeBitWidth();
|
||||
int dstBW = dstType.getElementTypeBitWidth();
|
||||
int numElems = srcType.getNumElements();
|
||||
|
||||
NVVM::FPRoundingModeAttr rndModeAttr = op.getRndAttr();
|
||||
NVVM::SaturationModeAttr satModeAttr = op.getSatAttr();
|
||||
auto reluBoolAttr = op.getReluAttr();
|
||||
Value randomBits = adaptor.getRandomBits();
|
||||
Type actualDstFloatType = dstElemType;
|
||||
|
||||
// STEP 1: bitcast input vector to i32 vector type.
|
||||
// f64 -> f32/f16/bf16 lowers to a single direct LLVM fptrunc
|
||||
// f64 -> f8/f6/f4 first truncates to f32 and then reuses the narrow
|
||||
// conversion path below.
|
||||
Value input = adaptor.getIn();
|
||||
if (srcBW == 64) {
|
||||
if (dstBW >= 16) {
|
||||
Type convertedType = typeConverter->convertType(dstType);
|
||||
assert(convertedType && "failed to convert type");
|
||||
Value result = b.create<LLVM::FPTruncOp>(convertedType, input);
|
||||
rewriter.replaceOp(op, result);
|
||||
return success();
|
||||
}
|
||||
auto f32VecTy = VectorType::get(srcType.getShape(), b.getF32Type());
|
||||
input = b.create<LLVM::FPTruncOp>(f32VecTy, input);
|
||||
srcType = f32VecTy;
|
||||
srcElemType = b.getF32Type();
|
||||
srcBW = 32;
|
||||
}
|
||||
|
||||
// f6 types are 6-bit in MLIR but NVVM uses 8-bit containers.
|
||||
int effectiveDstBW = getEffectiveBitWidth(dstBW);
|
||||
|
||||
int srcI32Elems = numElems * srcBW / regBits;
|
||||
int dstI32Elems = numElems * effectiveDstBW / regBits;
|
||||
Value srcI32Vec =
|
||||
b.create<LLVM::BitcastOp>(VectorType::get(srcI32Elems, i32Ty), input);
|
||||
Value dstI32Vec =
|
||||
b.create<LLVM::UndefOp>(VectorType::get(dstI32Elems, i32Ty));
|
||||
|
||||
// STEP 2: look up the conversion op from the (srcType, dstType) table.
|
||||
auto convEntry = lookupConvOp(kFPTruncTable, srcElemType, dstElemType);
|
||||
if (!convEntry)
|
||||
return rewriter.notifyMatchFailure(
|
||||
op, "unsupported type combination for truncation");
|
||||
FPTruncConvOp convOp = convEntry->convOp;
|
||||
|
||||
// Number of source-side i32 register slots consumed by each NVVM convert Op.
|
||||
auto getNumSrcI32PerConvert = [](FPKind src) {
|
||||
return src == FPKind::F32 ? 2 : 1;
|
||||
};
|
||||
int numSrcI32PerConv = getNumSrcI32PerConvert(convEntry->src);
|
||||
|
||||
// STEP 3: pack conversion results into destination i32 vector.
|
||||
const int srcStep = srcBW / effectiveDstBW;
|
||||
const int resultBW =
|
||||
effectiveDstBW * 2; // each conversion produces 2 (packed) elements
|
||||
const int numConvsPerI32 = regBits / resultBW;
|
||||
|
||||
for (int srcIdx = 0, dstIdx = 0; dstIdx < dstI32Elems;
|
||||
srcIdx += srcStep, dstIdx++) {
|
||||
Value dstIdxConst =
|
||||
b.create<LLVM::ConstantOp>(i64Ty, b.getI64IntegerAttr(dstIdx));
|
||||
Value dstValue;
|
||||
|
||||
if (numConvsPerI32 == 1) {
|
||||
// f16/bf16 destinations
|
||||
dstValue = createTruncConversion(
|
||||
b, ctx, convOp, srcI32Vec, srcIdx, rndModeAttr, satModeAttr,
|
||||
reluBoolAttr, dstElemType, actualDstFloatType, randomBits);
|
||||
} else {
|
||||
// f8/f6/f4 destinations: pack sub-results via vector insert + bitcast.
|
||||
auto subResultType = IntegerType::get(ctx, resultBW);
|
||||
auto subVecTy = VectorType::get(numConvsPerI32, subResultType);
|
||||
Value subVec = b.create<LLVM::UndefOp>(subVecTy);
|
||||
|
||||
int insertIdx = numConvsPerI32 - 1;
|
||||
int curStep = srcStep;
|
||||
while (curStep > 0) {
|
||||
curStep -= numSrcI32PerConv;
|
||||
Value subResult = createTruncConversion(
|
||||
b, ctx, convOp, srcI32Vec, srcIdx + curStep, rndModeAttr,
|
||||
satModeAttr, reluBoolAttr, dstElemType, actualDstFloatType,
|
||||
/*randomBits=*/Value());
|
||||
subVec = b.create<LLVM::InsertElementOp>(
|
||||
subVec, subResult,
|
||||
b.create<LLVM::ConstantOp>(i64Ty, b.getI64IntegerAttr(insertIdx)));
|
||||
insertIdx--;
|
||||
}
|
||||
|
||||
dstValue = b.create<LLVM::BitcastOp>(i32Ty, subVec);
|
||||
}
|
||||
|
||||
dstI32Vec =
|
||||
b.create<LLVM::InsertElementOp>(dstI32Vec, dstValue, dstIdxConst);
|
||||
}
|
||||
|
||||
// STEP 4: produce final result.
|
||||
Type convertedType = typeConverter->convertType(dstType);
|
||||
assert(convertedType && "failed to convert type");
|
||||
if (convEntry->dst == FPKind::F6) {
|
||||
IntegerType i8Ty = b.getI8Type();
|
||||
auto i8VecTy = VectorType::get(numElems, i8Ty);
|
||||
Value i8Vec = b.create<LLVM::BitcastOp>(i8VecTy, dstI32Vec);
|
||||
Value truncVec = b.create<LLVM::TruncOp>(convertedType, i8Vec);
|
||||
rewriter.replaceOp(op, truncVec);
|
||||
} else {
|
||||
auto dstVec = b.create<LLVM::BitcastOp>(convertedType, dstI32Vec);
|
||||
rewriter.replaceOp(op, dstVec);
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
struct NVGPUTruncfOpLowering : public ConvertOpToLLVMPattern<nvgpu::TruncfOp> {
|
||||
using ConvertOpToLLVMPattern<nvgpu::TruncfOp>::ConvertOpToLLVMPattern;
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(nvgpu::TruncfOp op, OpAdaptor adaptor,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return lowerTruncf(op, adaptor, rewriter, getTypeConverter());
|
||||
}
|
||||
};
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// NVGPUExtfOp Lowering
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
/// Conversion op identifier for nvgpu.extf lowering dispatch table.
|
||||
enum class FPExtConvOp {
|
||||
F8x2_TO_F16x2,
|
||||
F8x2_TO_BF16x2,
|
||||
F6x2_TO_F16x2,
|
||||
F6x2_TO_BF16x2,
|
||||
F4x2_TO_F16x2,
|
||||
F4x2_TO_BF16x2,
|
||||
};
|
||||
|
||||
struct FPExtTableEntry {
|
||||
FPKind src;
|
||||
FPKind dst;
|
||||
FPExtConvOp convOp;
|
||||
};
|
||||
|
||||
static constexpr FPExtTableEntry kFPExtTable[] = {
|
||||
{FPKind::F8, FPKind::F16, FPExtConvOp::F8x2_TO_F16x2},
|
||||
{FPKind::F8, FPKind::BF16, FPExtConvOp::F8x2_TO_BF16x2},
|
||||
{FPKind::F6, FPKind::F16, FPExtConvOp::F6x2_TO_F16x2},
|
||||
{FPKind::F6, FPKind::BF16, FPExtConvOp::F6x2_TO_BF16x2},
|
||||
{FPKind::F4, FPKind::F16, FPExtConvOp::F4x2_TO_F16x2},
|
||||
{FPKind::F4, FPKind::BF16, FPExtConvOp::F4x2_TO_BF16x2},
|
||||
};
|
||||
|
||||
/// Create a typed NVVM extension conversion.
|
||||
/// For f8/f6: src is vector<2xi8>. For f4: src is i8.
|
||||
/// Returns i32 (bitcast from vector<2xf16> or vector<2xbf16>).
|
||||
static Value createExtConversion(ImplicitLocOpBuilder &b, MLIRContext *ctx,
|
||||
FPExtConvOp convOp, Value src,
|
||||
BoolAttr reluAttr, Type actualSrcFloatType,
|
||||
Value extScaleFactor = Value()) {
|
||||
IntegerType i32Ty = b.getI32Type();
|
||||
auto srcTyAttr = TypeAttr::get(actualSrcFloatType);
|
||||
|
||||
switch (convOp) {
|
||||
case FPExtConvOp::F8x2_TO_F16x2: {
|
||||
Value r = NVVM::ConvertF8x2ToF16x2Op::create(
|
||||
b, VectorType::get(2, b.getF16Type()), src, srcTyAttr, reluAttr);
|
||||
return b.create<LLVM::BitcastOp>(i32Ty, r);
|
||||
}
|
||||
case FPExtConvOp::F8x2_TO_BF16x2: {
|
||||
Value r = NVVM::ConvertF8x2ToBF16x2Op::create(
|
||||
b, VectorType::get(2, b.getBF16Type()), src, extScaleFactor, srcTyAttr);
|
||||
return b.create<LLVM::BitcastOp>(i32Ty, r);
|
||||
}
|
||||
case FPExtConvOp::F6x2_TO_F16x2: {
|
||||
Value r = NVVM::ConvertF6x2ToF16x2Op::create(
|
||||
b, VectorType::get(2, b.getF16Type()), src, srcTyAttr, reluAttr);
|
||||
return b.create<LLVM::BitcastOp>(i32Ty, r);
|
||||
}
|
||||
case FPExtConvOp::F6x2_TO_BF16x2: {
|
||||
Value r = NVVM::ConvertF6x2ToBF16x2Op::create(
|
||||
b, VectorType::get(2, b.getBF16Type()), src, extScaleFactor, srcTyAttr);
|
||||
return b.create<LLVM::BitcastOp>(i32Ty, r);
|
||||
}
|
||||
case FPExtConvOp::F4x2_TO_F16x2: {
|
||||
Value r = NVVM::ConvertF4x2ToF16x2Op::create(
|
||||
b, VectorType::get(2, b.getF16Type()), src, srcTyAttr, reluAttr);
|
||||
return b.create<LLVM::BitcastOp>(i32Ty, r);
|
||||
}
|
||||
case FPExtConvOp::F4x2_TO_BF16x2: {
|
||||
Value r = NVVM::ConvertF4x2ToBF16x2Op::create(
|
||||
b, VectorType::get(2, b.getBF16Type()), src, extScaleFactor, srcTyAttr);
|
||||
return b.create<LLVM::BitcastOp>(i32Ty, r);
|
||||
}
|
||||
}
|
||||
llvm_unreachable("unhandled FPExtConvOp");
|
||||
}
|
||||
|
||||
static LogicalResult lowerExtf(nvgpu::ExtfOp op, nvgpu::ExtfOp::Adaptor adaptor,
|
||||
ConversionPatternRewriter &rewriter,
|
||||
const LLVMTypeConverter *typeConverter) {
|
||||
MLIRContext *ctx = op.getContext();
|
||||
ImplicitLocOpBuilder b(op->getLoc(), rewriter);
|
||||
IntegerType i8Ty = b.getI8Type();
|
||||
IntegerType i16Ty = b.getI16Type();
|
||||
IntegerType i32Ty = b.getI32Type();
|
||||
IntegerType i64Ty = b.getI64Type();
|
||||
|
||||
static constexpr int regBits = 32;
|
||||
auto srcType = llvm::dyn_cast<VectorType>(op.getIn().getType());
|
||||
auto dstType = llvm::dyn_cast<VectorType>(op.getOut().getType());
|
||||
if (!srcType || srcType.getRank() != 1 || !dstType || dstType.getRank() != 1)
|
||||
return rewriter.notifyMatchFailure(
|
||||
op, "expected 1-D vector; canonicalize pattern handles other shapes");
|
||||
|
||||
auto srcElemType = srcType.getElementType();
|
||||
auto dstElemType = dstType.getElementType();
|
||||
int srcBW = srcType.getElementTypeBitWidth();
|
||||
int dstBW = dstType.getElementTypeBitWidth();
|
||||
int numElems = srcType.getNumElements();
|
||||
|
||||
auto reluBoolAttr = op.getReluAttr();
|
||||
Type actualSrcFloatType = srcElemType;
|
||||
|
||||
assert(dstBW == 16 || dstBW == 32 || dstBW == 64);
|
||||
|
||||
// Wide source (f16/bf16/f32) to wide destination (f32/f64): single FPExt.
|
||||
if (srcBW >= 16 && dstBW >= 32) {
|
||||
Value result = adaptor.getIn();
|
||||
if (srcElemType != dstElemType) {
|
||||
Type convertedType = typeConverter->convertType(dstType);
|
||||
assert(convertedType && "failed to convert type");
|
||||
result = b.create<LLVM::FPExtOp>(convertedType, result);
|
||||
}
|
||||
rewriter.replaceOp(op, result);
|
||||
return success();
|
||||
}
|
||||
|
||||
// Narrow source (f8/f6/f4): NVVM typed op produces f16/bf16; optionally
|
||||
// followed by FPExt to the final f32/f64 destination.
|
||||
bool needsFinalFPExt = (dstBW >= 32);
|
||||
Type intermediateDstElem = dstElemType;
|
||||
if (needsFinalFPExt && llvm::isa<Float8E8M0FNUType>(srcElemType))
|
||||
intermediateDstElem = b.getBF16Type();
|
||||
else if (needsFinalFPExt)
|
||||
intermediateDstElem = b.getF16Type();
|
||||
int intermediateDstBW = needsFinalFPExt ? 16 : dstBW;
|
||||
|
||||
// f6 types are 6-bit in MLIR but NVVM uses 8-bit containers.
|
||||
int effectiveSrcBW = getEffectiveBitWidth(srcBW);
|
||||
|
||||
// STEP 1: prepare input as i32 register vector.
|
||||
// For f6: zext from vector<Nxi6> to vector<Nxi8>, then bitcast to i32s.
|
||||
Value inputVec = adaptor.getIn();
|
||||
if (srcBW == 6) {
|
||||
auto i8VecTy = VectorType::get(numElems, i8Ty);
|
||||
inputVec = b.create<LLVM::ZExtOp>(i8VecTy, inputVec);
|
||||
}
|
||||
|
||||
int srcI32Elems = numElems * effectiveSrcBW / regBits;
|
||||
int dstI32Elems = numElems * intermediateDstBW / regBits;
|
||||
Value srcI32Vec =
|
||||
b.create<LLVM::BitcastOp>(VectorType::get(srcI32Elems, i32Ty), inputVec);
|
||||
Value dstI32Vec =
|
||||
b.create<LLVM::UndefOp>(VectorType::get(dstI32Elems, i32Ty));
|
||||
|
||||
// STEP 2: look up the conversion op from the (srcType, dstType) table.
|
||||
auto convEntry = lookupConvOp(kFPExtTable, srcElemType, intermediateDstElem);
|
||||
if (!convEntry)
|
||||
return rewriter.notifyMatchFailure(
|
||||
op, "unsupported type combination for extension");
|
||||
FPExtConvOp convOp = convEntry->convOp;
|
||||
Value extScaleFactor;
|
||||
|
||||
// STEP 3: iterate over source i32 elements, producing destination i32s.
|
||||
for (int srcIdx = 0, dstIdx = 0; srcIdx < srcI32Elems; srcIdx++) {
|
||||
Value srcI32 = b.create<LLVM::ExtractElementOp>(
|
||||
srcI32Vec,
|
||||
b.create<LLVM::ConstantOp>(i64Ty, b.getI64IntegerAttr(srcIdx)));
|
||||
|
||||
if (effectiveSrcBW == 8) {
|
||||
// f8/f6: one i32 holds 4 bytes -> split into 2 pairs of i16 -> 2 convs.
|
||||
Value i16Vec =
|
||||
b.create<LLVM::BitcastOp>(VectorType::get(2, i16Ty), srcI32);
|
||||
for (int half = 0; half < 2; half++) {
|
||||
Value halfI16 = b.create<LLVM::ExtractElementOp>(
|
||||
i16Vec,
|
||||
b.create<LLVM::ConstantOp>(i64Ty, b.getI64IntegerAttr(half)));
|
||||
Value src =
|
||||
b.create<LLVM::BitcastOp>(VectorType::get(2, i8Ty), halfI16);
|
||||
Value dstValue =
|
||||
createExtConversion(b, ctx, convOp, src, reluBoolAttr,
|
||||
actualSrcFloatType, extScaleFactor);
|
||||
Value dstIdxConst =
|
||||
b.create<LLVM::ConstantOp>(i64Ty, b.getI64IntegerAttr(dstIdx));
|
||||
dstI32Vec =
|
||||
b.create<LLVM::InsertElementOp>(dstI32Vec, dstValue, dstIdxConst);
|
||||
dstIdx++;
|
||||
}
|
||||
} else {
|
||||
// f4: one i32 holds 4 bytes -> each byte is one conversion input.
|
||||
Value i8Vec = b.create<LLVM::BitcastOp>(VectorType::get(4, i8Ty), srcI32);
|
||||
for (int byteIdx = 0; byteIdx < 4; byteIdx++) {
|
||||
Value src = b.create<LLVM::ExtractElementOp>(
|
||||
i8Vec,
|
||||
b.create<LLVM::ConstantOp>(i64Ty, b.getI64IntegerAttr(byteIdx)));
|
||||
Value dstValue =
|
||||
createExtConversion(b, ctx, convOp, src, reluBoolAttr,
|
||||
actualSrcFloatType, extScaleFactor);
|
||||
Value dstIdxConst =
|
||||
b.create<LLVM::ConstantOp>(i64Ty, b.getI64IntegerAttr(dstIdx));
|
||||
dstI32Vec =
|
||||
b.create<LLVM::InsertElementOp>(dstI32Vec, dstValue, dstIdxConst);
|
||||
dstIdx++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// STEP 4: produce final result.
|
||||
Type convertedType = typeConverter->convertType(dstType);
|
||||
assert(convertedType && "failed to convert type");
|
||||
Value result;
|
||||
if (needsFinalFPExt) {
|
||||
auto intermediateVecTy = VectorType::get(numElems, intermediateDstElem);
|
||||
Value intermediateVec =
|
||||
b.create<LLVM::BitcastOp>(intermediateVecTy, dstI32Vec);
|
||||
result = b.create<LLVM::FPExtOp>(convertedType, intermediateVec);
|
||||
} else {
|
||||
result = b.create<LLVM::BitcastOp>(convertedType, dstI32Vec);
|
||||
}
|
||||
rewriter.replaceOp(op, result);
|
||||
return success();
|
||||
}
|
||||
|
||||
struct NVGPUExtfOpLowering : public ConvertOpToLLVMPattern<nvgpu::ExtfOp> {
|
||||
using ConvertOpToLLVMPattern<nvgpu::ExtfOp>::ConvertOpToLLVMPattern;
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(nvgpu::ExtfOp op, OpAdaptor adaptor,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return lowerExtf(op, adaptor, rewriter, getTypeConverter());
|
||||
}
|
||||
};
|
||||
|
||||
static int64_t computePaddedElems(int64_t numElems, int srcBW, int dstBW,
|
||||
int step) {
|
||||
static constexpr int regBits = 32;
|
||||
int effSrcBW = getEffectiveBitWidth(srcBW);
|
||||
int effDstBW = getEffectiveBitWidth(dstBW);
|
||||
auto ceilDiv = [](int64_t x, int64_t y) { return (x + y - 1) / y; };
|
||||
int64_t padded =
|
||||
std::max(ceilDiv(numElems * effSrcBW, regBits) * regBits / effSrcBW,
|
||||
ceilDiv(numElems * effDstBW, regBits) * regBits / effDstBW);
|
||||
return ceilDiv(padded, step) * step;
|
||||
}
|
||||
|
||||
/// Canonicalization pattern for nvgpu.truncf / nvgpu.extf:
|
||||
/// handles scalar inputs, non-32-bit-aligned vectors, and multi-rank vectors.
|
||||
/// Runs as an OpRewritePattern on MLIR types before LLVM type conversion.
|
||||
template <typename CvtOp, bool IsTrunc>
|
||||
struct NVGPUFPCanonicalizePattern : public OpRewritePattern<CvtOp> {
|
||||
using OpRewritePattern<CvtOp>::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(CvtOp op,
|
||||
PatternRewriter &rewriter) const override {
|
||||
Type inType = op.getIn().getType();
|
||||
Type outType = op.getOut().getType();
|
||||
|
||||
Type srcElemTy = getElementTypeOrSelf(inType);
|
||||
Type dstElemTy = getElementTypeOrSelf(outType);
|
||||
int srcBW = srcElemTy.getIntOrFloatBitWidth();
|
||||
int dstBW = dstElemTy.getIntOrFloatBitWidth();
|
||||
int effSrcBW = getEffectiveBitWidth(srcBW);
|
||||
int effDstBW = getEffectiveBitWidth(dstBW);
|
||||
|
||||
bool isScalar = !isa<VectorType>(inType);
|
||||
auto srcVecTy = dyn_cast<VectorType>(inType);
|
||||
bool isMultiRank = srcVecTy && srcVecTy.getRank() > 1;
|
||||
int64_t numElems = isScalar ? 1 : srcVecTy.getNumElements();
|
||||
int step = IsTrunc ? effSrcBW / effDstBW : effDstBW / effSrcBW;
|
||||
int64_t paddedElems = computePaddedElems(numElems, srcBW, dstBW, step);
|
||||
bool needsPad = (paddedElems != numElems);
|
||||
|
||||
if (!isScalar && !isMultiRank && !needsPad)
|
||||
return failure();
|
||||
|
||||
ImplicitLocOpBuilder b(op->getLoc(), rewriter);
|
||||
Value input = op.getIn();
|
||||
|
||||
if (isScalar)
|
||||
input = vector::BroadcastOp::create(b, VectorType::get({1}, srcElemTy),
|
||||
input);
|
||||
if (isMultiRank)
|
||||
input = vector::ShapeCastOp::create(
|
||||
b, VectorType::get({numElems}, srcElemTy), input);
|
||||
|
||||
if (needsPad) {
|
||||
auto paddedTy = VectorType::get({paddedElems}, srcElemTy);
|
||||
Value zero = arith::ConstantOp::create(
|
||||
b, DenseElementsAttr::get(paddedTy, b.getZeroAttr(srcElemTy)));
|
||||
input = vector::InsertStridedSliceOp::create(
|
||||
b, input, zero, SmallVector<int64_t>{0}, SmallVector<int64_t>{1});
|
||||
}
|
||||
|
||||
auto cvtDstTy =
|
||||
VectorType::get({needsPad ? paddedElems : numElems}, dstElemTy);
|
||||
Value cvt;
|
||||
if constexpr (IsTrunc) {
|
||||
cvt = CvtOp::create(b, cvtDstTy, input, op.getRndAttr(), op.getSatAttr(),
|
||||
op.getReluAttr(), op.getRandomBits());
|
||||
} else {
|
||||
cvt =
|
||||
CvtOp::create(b, cvtDstTy, input, op.getRndAttr(), op.getReluAttr());
|
||||
}
|
||||
Value result = cvt;
|
||||
|
||||
if (needsPad) {
|
||||
result = vector::ExtractStridedSliceOp::create(
|
||||
b, result, SmallVector<int64_t>{0}, SmallVector<int64_t>{numElems},
|
||||
SmallVector<int64_t>{1});
|
||||
}
|
||||
|
||||
if (isMultiRank) {
|
||||
result =
|
||||
vector::ShapeCastOp::create(b, cast<VectorType>(outType), result);
|
||||
}
|
||||
|
||||
if (isScalar) {
|
||||
result = vector::ExtractOp::create(b, result, SmallVector<int64_t>{0});
|
||||
}
|
||||
|
||||
rewriter.replaceOp(op, result);
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
using NVGPUTruncfCanonicalizePattern =
|
||||
NVGPUFPCanonicalizePattern<nvgpu::TruncfOp, true>;
|
||||
using NVGPUExtfCanonicalizePattern =
|
||||
NVGPUFPCanonicalizePattern<nvgpu::ExtfOp, false>;
|
||||
} // namespace
|
||||
|
||||
void mlir::nvgpu::populateCommonGPUTypeAndAttributeConversions(
|
||||
|
|
@ -1753,7 +2417,12 @@ void mlir::populateNVGPUToNVVMConversionPatterns(
|
|||
NVGPUWarpgroupMmaOpLowering, // nvgpu.warpgroup.mma
|
||||
NVGPUWarpgroupMmaStoreOpLowering, // nvgpu.warpgroup.mma.store
|
||||
NVGPUWarpgroupMmaInitAccumulatorOpLowering, // nvgpu.warpgroup.mma.init.accumulator
|
||||
NVGPUTruncfOpLowering, // nvgpu.truncf
|
||||
NVGPUExtfOpLowering, // nvgpu.extf
|
||||
MmaSyncOptoNVVM, MmaLdMatrixOpToNVVM, NVGPUAsyncCopyLowering,
|
||||
NVGPUAsyncCreateGroupLowering, NVGPUAsyncWaitLowering,
|
||||
NVGPUMmaSparseSyncLowering, NVGPURcpOpLowering>(converter);
|
||||
|
||||
patterns.add<NVGPUTruncfCanonicalizePattern, NVGPUExtfCanonicalizePattern>(
|
||||
patterns.getContext());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -700,6 +700,122 @@ LogicalResult RcpOp::verify() {
|
|||
return success();
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// NVGPU_TruncfOp
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
static LogicalResult verifyConversionShapes(Operation *op, Type inType,
|
||||
Type outType) {
|
||||
bool srcIsVector = llvm::isa<VectorType>(inType);
|
||||
bool dstIsVector = llvm::isa<VectorType>(outType);
|
||||
if (srcIsVector != dstIsVector)
|
||||
return op->emitOpError("input and output must both be scalars or both be "
|
||||
"vectors, got ")
|
||||
<< inType << " and " << outType;
|
||||
if (srcIsVector) {
|
||||
auto srcVector = llvm::cast<VectorType>(inType);
|
||||
auto dstVector = llvm::cast<VectorType>(outType);
|
||||
if (srcVector.getShape() != dstVector.getShape())
|
||||
return op->emitOpError("input and output shapes must match, got ")
|
||||
<< inType << " and " << outType;
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
LogicalResult TruncfOp::verify() {
|
||||
Type inType = getIn().getType();
|
||||
Type outType = getType();
|
||||
Type srcType = getElementTypeOrSelf(inType);
|
||||
Type dstType = getElementTypeOrSelf(outType);
|
||||
int srcBitWidth = srcType.getIntOrFloatBitWidth();
|
||||
int dstBitWidth = dstType.getIntOrFloatBitWidth();
|
||||
auto rnd = getRnd();
|
||||
|
||||
if (auto result = verifyConversionShapes(getOperation(), inType, outType);
|
||||
failed(result))
|
||||
return result;
|
||||
|
||||
if (srcBitWidth <= dstBitWidth)
|
||||
return emitOpError("result type ")
|
||||
<< dstType << " must be narrower than operand type " << srcType;
|
||||
|
||||
if (!(srcBitWidth == 64 || srcBitWidth == 32 || srcBitWidth == 16))
|
||||
return emitOpError("input type must be 64/32/16 bitwidth, but got ")
|
||||
<< srcBitWidth;
|
||||
|
||||
if (llvm::isa<Float8E8M0FNUType>(dstType)) {
|
||||
if (rnd != mlir::NVVM::FPRoundingMode::RZ &&
|
||||
rnd != mlir::NVVM::FPRoundingMode::RP)
|
||||
return emitOpError("expects RZ or RP rounding mode when result type is "
|
||||
"e8m0, but got ")
|
||||
<< getRndAttr();
|
||||
} else if (rnd == mlir::NVVM::FPRoundingMode::RS) {
|
||||
// TODO: Currently, we only support conversions which fit into a single i32
|
||||
// register. Support f32->f8/f6/f4 conversions with RS rounding.
|
||||
if (!(srcBitWidth == 32 && dstBitWidth == 16))
|
||||
return emitOpError("RS (stochastic) rounding is only supported for "
|
||||
"f32->f16/bf16, got ")
|
||||
<< srcType << " -> " << dstType;
|
||||
if (!getRandomBits())
|
||||
return emitOpError("random_bits operand is required with RS rounding");
|
||||
} else if (srcType.isF64() && dstBitWidth >= 16) {
|
||||
if (rnd != mlir::NVVM::FPRoundingMode::RN)
|
||||
return emitOpError("expects RN rounding mode for f64 input, but got ")
|
||||
<< getRndAttr();
|
||||
} else if (srcBitWidth == 32 && dstBitWidth == 16) {
|
||||
if (rnd != mlir::NVVM::FPRoundingMode::RN &&
|
||||
rnd != mlir::NVVM::FPRoundingMode::RZ)
|
||||
return emitOpError("expects RN or RZ rounding mode for f32 to f16/bf16, "
|
||||
"but got ")
|
||||
<< getRndAttr();
|
||||
} else if (rnd != mlir::NVVM::FPRoundingMode::RN) {
|
||||
return emitOpError("expects RN rounding mode, but got ") << getRndAttr();
|
||||
}
|
||||
|
||||
if (getRandomBits() && rnd != mlir::NVVM::FPRoundingMode::RS)
|
||||
return emitOpError("random_bits can only be used with RS rounding mode");
|
||||
|
||||
return success();
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// NVGPU_ExtfOp
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
LogicalResult ExtfOp::verify() {
|
||||
Type inType = getIn().getType();
|
||||
Type outType = getType();
|
||||
Type srcType = getElementTypeOrSelf(inType);
|
||||
Type dstType = getElementTypeOrSelf(outType);
|
||||
int srcBitWidth = srcType.getIntOrFloatBitWidth();
|
||||
int dstBitWidth = dstType.getIntOrFloatBitWidth();
|
||||
auto rnd = getRnd();
|
||||
|
||||
if (failed(verifyConversionShapes(getOperation(), inType, outType)))
|
||||
return failure();
|
||||
|
||||
if (srcBitWidth >= dstBitWidth)
|
||||
return emitOpError("result type ")
|
||||
<< dstType << " must be wider than operand type " << srcType;
|
||||
|
||||
if (dstBitWidth != 16 && dstBitWidth != 32 && dstBitWidth != 64)
|
||||
return emitOpError("result type must be 16, 32, or 64 bitwidth, but got ")
|
||||
<< dstBitWidth;
|
||||
|
||||
if (llvm::isa<Float8E8M0FNUType>(srcType) &&
|
||||
!llvm::isa<BFloat16Type>(dstType) && !dstType.isF32())
|
||||
return emitOpError("expects bf16 or f32 output type when input type is "
|
||||
"e8m0.");
|
||||
|
||||
if (rnd != mlir::NVVM::FPRoundingMode::RN)
|
||||
return emitOpError("expects RN rounding mode, but got ") << getRndAttr();
|
||||
|
||||
if (getRelu() && llvm::isa<BFloat16Type>(dstType))
|
||||
return emitOpError("relu is not supported for bf16 destination");
|
||||
|
||||
return success();
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// TableGen'd dialect, type, and op definitions
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
|
|
|||
78
mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf-large.mlir
Normal file
78
mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf-large.mlir
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
// RUN: mlir-opt %s -convert-nvgpu-to-nvvm | FileCheck %s
|
||||
|
||||
// Large-vector smoke tests for nvgpu.extf
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f8_to_f16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<400xf8E4M3FN>
|
||||
func.func @cvt_large_f8_to_f16(%in : vector<400xf8E4M3FN>) -> vector<400xf16> {
|
||||
// CHECK: builtin.unrealized_conversion_cast %[[IN]] : vector<400xf8E4M3FN> to vector<400xi8>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<400xi8> to vector<100xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-NOT: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<200xi32> to vector<400xf16>
|
||||
%out = nvgpu.extf %in : vector<400xf8E4M3FN> to vector<400xf16>
|
||||
return %out : vector<400xf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_e8m0_to_bf16(
|
||||
func.func @cvt_large_e8m0_to_bf16(%in : vector<400xf8E8M0FNU>) -> vector<400xbf16> {
|
||||
// CHECK: builtin.unrealized_conversion_cast %{{.*}} : vector<400xf8E8M0FNU> to vector<400xi8>
|
||||
// CHECK-COUNT-200: nvvm.convert.f8x2.to.bf16x2
|
||||
// CHECK-NOT: nvvm.convert.f8x2.to.bf16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<200xi32> to vector<400xbf16>
|
||||
%out = nvgpu.extf %in : vector<400xf8E8M0FNU> to vector<400xbf16>
|
||||
return %out : vector<400xbf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f6_to_f16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<400xf6E2M3FN>
|
||||
func.func @cvt_large_f6_to_f16(%in : vector<400xf6E2M3FN>) -> vector<400xf16> {
|
||||
// CHECK: builtin.unrealized_conversion_cast %[[IN]] : vector<400xf6E2M3FN> to vector<400xi6>
|
||||
// CHECK: llvm.zext {{.*}} : vector<400xi6> to vector<400xi8>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<400xi8> to vector<100xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK-NOT: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<200xi32> to vector<400xf16>
|
||||
%out = nvgpu.extf %in : vector<400xf6E2M3FN> to vector<400xf16>
|
||||
return %out : vector<400xf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f4_to_f16(
|
||||
func.func @cvt_large_f4_to_f16(%in : vector<400xf4E2M1FN>) -> vector<400xf16> {
|
||||
// CHECK: builtin.unrealized_conversion_cast %{{.*}} : vector<400xf4E2M1FN> to vector<400xi4>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<400xi4> to vector<50xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK-NOT: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<200xi32> to vector<400xf16>
|
||||
%out = nvgpu.extf %in : vector<400xf4E2M1FN> to vector<400xf16>
|
||||
return %out : vector<400xf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f8_to_f32(
|
||||
func.func @cvt_large_f8_to_f32(%in : vector<400xf8E5M2>) -> vector<400xf32> {
|
||||
// CHECK-COUNT-200: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-NOT: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<200xi32> to vector<400xf16>
|
||||
// CHECK: llvm.fpext {{.*}} : vector<400xf16> to vector<400xf32>
|
||||
%out = nvgpu.extf %in : vector<400xf8E5M2> to vector<400xf32>
|
||||
return %out : vector<400xf32>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f6_to_f32(
|
||||
func.func @cvt_large_f6_to_f32(%in : vector<400xf6E2M3FN>) -> vector<400xf32> {
|
||||
// CHECK: llvm.zext {{.*}} : vector<400xi6> to vector<400xi8>
|
||||
// CHECK-COUNT-200: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK-NOT: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK: llvm.fpext {{.*}} : vector<400xf16> to vector<400xf32>
|
||||
%out = nvgpu.extf %in : vector<400xf6E2M3FN> to vector<400xf32>
|
||||
return %out : vector<400xf32>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f16_to_f32(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<400xf16>
|
||||
func.func @cvt_large_f16_to_f32(%in : vector<400xf16>) -> vector<400xf32> {
|
||||
// CHECK-NOT: nvvm.convert
|
||||
// CHECK: llvm.fpext %[[IN]] : vector<400xf16> to vector<400xf32>
|
||||
%out = nvgpu.extf %in : vector<400xf16> to vector<400xf32>
|
||||
return %out : vector<400xf32>
|
||||
}
|
||||
506
mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf.mlir
Normal file
506
mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf.mlir
Normal file
|
|
@ -0,0 +1,506 @@
|
|||
// RUN: mlir-opt %s -convert-nvgpu-to-nvvm | FileCheck %s
|
||||
// RUN: mlir-opt %s -convert-nvgpu-to-nvvm -convert-vector-to-llvm | FileCheck %s --check-prefix=CHECK-E2E
|
||||
|
||||
// Basic aligned vector extension to f16/bf16.
|
||||
|
||||
// CHECK-LABEL: @cvt_float_e4m3fn_to_f16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf8E4M3FN>
|
||||
func.func @cvt_float_e4m3fn_to_f16(%in : vector<8xf8E4M3FN>) {
|
||||
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<8xf8E4M3FN> to vector<8xi8>
|
||||
// CHECK: llvm.bitcast %[[CAST]] : vector<8xi8> to vector<2xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<4xi32>
|
||||
// CHECK: llvm.bitcast {{.*}} : i32 to vector<2xi16>
|
||||
// CHECK: llvm.bitcast {{.*}} : i16 to vector<2xi8>
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f8E4M3FN)
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f8E4M3FN)
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<4xi32> to vector<8xf16>
|
||||
%out = nvgpu.extf %in : vector<8xf8E4M3FN> to vector<8xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @cvt_float_e5m2_to_f16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf8E5M2>
|
||||
func.func @cvt_float_e5m2_to_f16(%in : vector<8xf8E5M2>) {
|
||||
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<8xf8E5M2> to vector<8xi8>
|
||||
// CHECK: %[[IN_I32:.+]] = llvm.bitcast %[[CAST]] : vector<8xi8> to vector<2xi32>
|
||||
// CHECK: %[[OUT_I32:.+]] = llvm.mlir.undef : vector<4xi32>
|
||||
// CHECK: llvm.extractelement %[[IN_I32]]
|
||||
// CHECK: llvm.bitcast {{.*}} : i32 to vector<2xi16>
|
||||
// CHECK: llvm.extractelement {{.*}} : vector<2xi16>
|
||||
// CHECK: llvm.bitcast {{.*}} : i16 to vector<2xi8>
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f8E5M2)
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xf16> to i32
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<4xi32>
|
||||
// CHECK: llvm.extractelement {{.*}} : vector<2xi16>
|
||||
// CHECK: llvm.bitcast {{.*}} : i16 to vector<2xi8>
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f8E5M2)
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xf16> to i32
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<4xi32>
|
||||
// CHECK: llvm.bitcast {{.*}} : i32 to vector<2xi16>
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<4xi32> to vector<8xf16>
|
||||
%out = nvgpu.extf %in : vector<8xf8E5M2> to vector<8xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @cvt_float_e8m0_to_bf16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf8E8M0FNU>
|
||||
func.func @cvt_float_e8m0_to_bf16(%in : vector<8xf8E8M0FNU>) {
|
||||
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<8xf8E8M0FNU> to vector<8xi8>
|
||||
// CHECK: llvm.bitcast %[[CAST]] : vector<8xi8> to vector<2xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<4xi32>
|
||||
// CHECK: llvm.bitcast {{.*}} : i32 to vector<2xi16>
|
||||
// CHECK: llvm.bitcast {{.*}} : i16 to vector<2xi8>
|
||||
// CHECK: nvvm.convert.f8x2.to.bf16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f8E8M0FNU)
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xbf16> to i32
|
||||
// CHECK: nvvm.convert.f8x2.to.bf16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f8E8M0FNU)
|
||||
// CHECK: nvvm.convert.f8x2.to.bf16x2
|
||||
// CHECK: nvvm.convert.f8x2.to.bf16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<4xi32> to vector<8xbf16>
|
||||
%out = nvgpu.extf %in : vector<8xf8E8M0FNU> to vector<8xbf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @cvt_float_e2m3_to_f16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf6E2M3FN>
|
||||
func.func @cvt_float_e2m3_to_f16(%in : vector<8xf6E2M3FN>) {
|
||||
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<8xf6E2M3FN> to vector<8xi6>
|
||||
// CHECK: llvm.zext %[[CAST]] : vector<8xi6> to vector<8xi8>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<8xi8> to vector<2xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<4xi32>
|
||||
// CHECK: llvm.bitcast {{.*}} : i32 to vector<2xi16>
|
||||
// CHECK: llvm.bitcast {{.*}} : i16 to vector<2xi8>
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f6E2M3FN)
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f6E2M3FN)
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<4xi32> to vector<8xf16>
|
||||
%out = nvgpu.extf %in : vector<8xf6E2M3FN> to vector<8xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @cvt_float_e3m2_to_f16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf6E3M2FN>
|
||||
func.func @cvt_float_e3m2_to_f16(%in : vector<8xf6E3M2FN>) {
|
||||
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<8xf6E3M2FN> to vector<8xi6>
|
||||
// CHECK: llvm.zext %[[CAST]] : vector<8xi6> to vector<8xi8>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<8xi8> to vector<2xi32>
|
||||
// CHECK: %[[OUT_I32:.+]] = llvm.mlir.undef : vector<4xi32>
|
||||
// CHECK: llvm.bitcast {{.*}} : i32 to vector<2xi16>
|
||||
// CHECK: llvm.bitcast {{.*}} : i16 to vector<2xi8>
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f6E3M2FN)
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xf16> to i32
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<4xi32>
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f6E3M2FN)
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<4xi32>
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<4xi32> to vector<8xf16>
|
||||
%out = nvgpu.extf %in : vector<8xf6E3M2FN> to vector<8xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @cvt_float_e2m1_to_f16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<16xf4E2M1FN>
|
||||
func.func @cvt_float_e2m1_to_f16(%in : vector<16xf4E2M1FN>) {
|
||||
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<16xf4E2M1FN> to vector<16xi4>
|
||||
// CHECK: %[[IN_I32:.+]] = llvm.bitcast %[[CAST]] : vector<16xi4> to vector<2xi32>
|
||||
// CHECK: %[[OUT_I32:.+]] = llvm.mlir.undef : vector<8xi32>
|
||||
// CHECK: llvm.extractelement %[[IN_I32]]
|
||||
// CHECK: llvm.bitcast {{.*}} : i32 to vector<4xi8>
|
||||
// CHECK: llvm.extractelement {{.*}} : vector<4xi8>
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK-SAME: : i8(f4E2M1FN)
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xf16> to i32
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<8xi32>
|
||||
// CHECK: llvm.extractelement {{.*}} : vector<4xi8>
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK-SAME: : i8(f4E2M1FN)
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<8xi32>
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<8xi32>
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<8xi32>
|
||||
// CHECK: llvm.bitcast {{.*}} : i32 to vector<4xi8>
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<8xi32> to vector<16xf16>
|
||||
%out = nvgpu.extf %in : vector<16xf4E2M1FN> to vector<16xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @cvt_float_e2m3_to_bf16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf6E2M3FN>
|
||||
func.func @cvt_float_e2m3_to_bf16(%in : vector<8xf6E2M3FN>) {
|
||||
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<8xf6E2M3FN> to vector<8xi6>
|
||||
// CHECK: llvm.zext %[[CAST]] : vector<8xi6> to vector<8xi8>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<8xi8> to vector<2xi32>
|
||||
// CHECK: nvvm.convert.f6x2.to.bf16x2
|
||||
// CHECK-SAME: : vector<2xi8>(f6E2M3FN)
|
||||
// CHECK: nvvm.convert.f6x2.to.bf16x2
|
||||
// CHECK: nvvm.convert.f6x2.to.bf16x2
|
||||
// CHECK: nvvm.convert.f6x2.to.bf16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<4xi32> to vector<8xbf16>
|
||||
%out = nvgpu.extf %in : vector<8xf6E2M3FN> to vector<8xbf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @cvt_float_e2m1_to_bf16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<16xf4E2M1FN>
|
||||
func.func @cvt_float_e2m1_to_bf16(%in : vector<16xf4E2M1FN>) {
|
||||
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<16xf4E2M1FN> to vector<16xi4>
|
||||
// CHECK: llvm.bitcast %[[CAST]] : vector<16xi4> to vector<2xi32>
|
||||
// CHECK: nvvm.convert.f4x2.to.bf16x2
|
||||
// CHECK-SAME: : i8(f4E2M1FN)
|
||||
// CHECK: nvvm.convert.f4x2.to.bf16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<8xi32> to vector<16xbf16>
|
||||
%out = nvgpu.extf %in : vector<16xf4E2M1FN> to vector<16xbf16>
|
||||
return
|
||||
}
|
||||
|
||||
// Extension to f32 (two-step: narrow -> f16/bf16 -> f32).
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_e4m3fn_to_f32(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<4xf8E4M3FN>
|
||||
func.func @fpext_e4m3fn_to_f32(%in : vector<4xf8E4M3FN>) -> vector<4xf32> {
|
||||
// CHECK: builtin.unrealized_conversion_cast %[[IN]] : vector<4xf8E4M3FN> to vector<4xi8>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<4xi8> to vector<1xi32>
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<4xf16>
|
||||
// CHECK: llvm.fpext {{.*}} : vector<4xf16> to vector<4xf32>
|
||||
%out = nvgpu.extf %in : vector<4xf8E4M3FN> to vector<4xf32>
|
||||
return %out : vector<4xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_e5m2_to_f32(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<4xf8E5M2>
|
||||
func.func @fpext_e5m2_to_f32(%in : vector<4xf8E5M2>) -> vector<4xf32> {
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<4xf16>
|
||||
// CHECK: llvm.fpext {{.*}} : vector<4xf16> to vector<4xf32>
|
||||
%out = nvgpu.extf %in : vector<4xf8E5M2> to vector<4xf32>
|
||||
return %out : vector<4xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_e8m0_to_f32(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<4xf8E8M0FNU>
|
||||
func.func @fpext_e8m0_to_f32(%in : vector<4xf8E8M0FNU>) -> vector<4xf32> {
|
||||
// CHECK: nvvm.convert.f8x2.to.bf16x2
|
||||
// CHECK: nvvm.convert.f8x2.to.bf16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<4xbf16>
|
||||
// CHECK: llvm.fpext {{.*}} : vector<4xbf16> to vector<4xf32>
|
||||
%out = nvgpu.extf %in : vector<4xf8E8M0FNU> to vector<4xf32>
|
||||
return %out : vector<4xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_e2m3_to_f32(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<4xf6E2M3FN>
|
||||
func.func @fpext_e2m3_to_f32(%in : vector<4xf6E2M3FN>) -> vector<4xf32> {
|
||||
// CHECK: builtin.unrealized_conversion_cast %[[IN]] : vector<4xf6E2M3FN> to vector<4xi6>
|
||||
// CHECK: llvm.zext {{.*}} : vector<4xi6> to vector<4xi8>
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<4xf16>
|
||||
// CHECK: llvm.fpext {{.*}} : vector<4xf16> to vector<4xf32>
|
||||
%out = nvgpu.extf %in : vector<4xf6E2M3FN> to vector<4xf32>
|
||||
return %out : vector<4xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_e2m1_to_f32(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf4E2M1FN>
|
||||
func.func @fpext_e2m1_to_f32(%in : vector<8xf4E2M1FN>) -> vector<8xf32> {
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<4xi32> to vector<8xf16>
|
||||
// CHECK: llvm.fpext {{.*}} : vector<8xf16> to vector<8xf32>
|
||||
%out = nvgpu.extf %in : vector<8xf4E2M1FN> to vector<8xf32>
|
||||
return %out : vector<8xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_f16_to_f32(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<4xf16>
|
||||
func.func @fpext_f16_to_f32(%in : vector<4xf16>) -> vector<4xf32> {
|
||||
// CHECK: llvm.fpext %[[IN]] : vector<4xf16> to vector<4xf32>
|
||||
%out = nvgpu.extf %in : vector<4xf16> to vector<4xf32>
|
||||
return %out : vector<4xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_bf16_to_f32(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<4xbf16>
|
||||
func.func @fpext_bf16_to_f32(%in : vector<4xbf16>) -> vector<4xf32> {
|
||||
// CHECK: llvm.fpext %[[IN]] : vector<4xbf16> to vector<4xf32>
|
||||
%out = nvgpu.extf %in : vector<4xbf16> to vector<4xf32>
|
||||
return %out : vector<4xf32>
|
||||
}
|
||||
|
||||
// Extension to f64.
|
||||
|
||||
// CHECK-LABEL: @fpext_f16_to_f64
|
||||
func.func @fpext_f16_to_f64(%arg0: vector<4xf16>) -> vector<4xf64> {
|
||||
// CHECK: llvm.fpext %{{.*}} : vector<4xf16> to vector<4xf64>
|
||||
// CHECK-NOT: llvm.fpext
|
||||
%out = nvgpu.extf %arg0 : vector<4xf16> to vector<4xf64>
|
||||
return %out : vector<4xf64>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fpext_bf16_to_f64
|
||||
func.func @fpext_bf16_to_f64(%arg0: vector<4xbf16>) -> vector<4xf64> {
|
||||
// CHECK: llvm.fpext %{{.*}} : vector<4xbf16> to vector<4xf64>
|
||||
// CHECK-NOT: llvm.fpext
|
||||
%out = nvgpu.extf %arg0 : vector<4xbf16> to vector<4xf64>
|
||||
return %out : vector<4xf64>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fpext_f32_to_f64
|
||||
func.func @fpext_f32_to_f64(%arg0: vector<4xf32>) -> vector<4xf64> {
|
||||
// CHECK-NOT: nvvm
|
||||
// CHECK: llvm.fpext %{{.*}} : vector<4xf32> to vector<4xf64>
|
||||
%out = nvgpu.extf %arg0 : vector<4xf32> to vector<4xf64>
|
||||
return %out : vector<4xf64>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fpext_f8_to_f64
|
||||
func.func @fpext_f8_to_f64(%arg0: vector<4xf8E4M3FN>) -> vector<4xf64> {
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.fpext {{.*}} to {{.*}}f64
|
||||
// CHECK-NOT: llvm.fpext {{.*}} to {{.*}}f32
|
||||
// CHECK: vector.extract_strided_slice
|
||||
%out = nvgpu.extf %arg0 : vector<4xf8E4M3FN> to vector<4xf64>
|
||||
return %out : vector<4xf64>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fpext_f4_to_f64
|
||||
func.func @fpext_f4_to_f64(%arg0: vector<8xf4E2M1FN>) -> vector<8xf64> {
|
||||
// CHECK: nvvm.convert.f4x2.to.f16x2
|
||||
// CHECK: llvm.fpext {{.*}} to {{.*}}f64
|
||||
// CHECK-NOT: llvm.fpext {{.*}} to {{.*}}f32
|
||||
%out = nvgpu.extf %arg0 : vector<8xf4E2M1FN> to vector<8xf64>
|
||||
return %out : vector<8xf64>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fpext_e2m3_to_f64
|
||||
func.func @fpext_e2m3_to_f64(%arg0: vector<4xf6E2M3FN>) -> vector<4xf64> {
|
||||
// CHECK: llvm.zext {{.*}} : vector<8xi6> to vector<8xi8>
|
||||
// CHECK: nvvm.convert.f6x2.to.f16x2
|
||||
// CHECK: llvm.fpext {{.*}} to {{.*}}f64
|
||||
// CHECK-NOT: llvm.fpext {{.*}} to {{.*}}f32
|
||||
%out = nvgpu.extf %arg0 : vector<4xf6E2M3FN> to vector<4xf64>
|
||||
return %out : vector<4xf64>
|
||||
}
|
||||
|
||||
// Scalar inputs (canonicalize: broadcast + pad + extract).
|
||||
|
||||
// CHECK-LABEL: @fpext_scalar_f8_to_f16
|
||||
// CHECK-SAME: %[[IN:.+]]: f8E4M3FN
|
||||
func.func @fpext_scalar_f8_to_f16(%in : f8E4M3FN) -> f16 {
|
||||
// CHECK: vector.broadcast %[[IN]] : f8E4M3FN to vector<1xf8E4M3FN>
|
||||
// CHECK: vector.insert_strided_slice
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: vector.extract_strided_slice
|
||||
// CHECK: vector.extract {{.*}}[0] : f16 from vector<1xf16>
|
||||
%out = nvgpu.extf %in : f8E4M3FN to f16
|
||||
return %out : f16
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_scalar_f8_to_f32
|
||||
// CHECK-SAME: %[[IN:.+]]: f8E4M3FN
|
||||
func.func @fpext_scalar_f8_to_f32(%in : f8E4M3FN) -> f32 {
|
||||
// CHECK: vector.broadcast
|
||||
// CHECK: vector.insert_strided_slice
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.fpext
|
||||
// CHECK: vector.extract_strided_slice
|
||||
// CHECK: vector.extract {{.*}}[0] : f32 from vector<1xf32>
|
||||
%out = nvgpu.extf %in : f8E4M3FN to f32
|
||||
return %out : f32
|
||||
}
|
||||
|
||||
// Multi-rank vectors (canonicalize: shape_cast flatten).
|
||||
|
||||
// CHECK-LABEL: @fpext_v2x4_f8_to_f16
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<2x4xf8E4M3FN>
|
||||
func.func @fpext_v2x4_f8_to_f16(%in : vector<2x4xf8E4M3FN>) -> vector<2x4xf16> {
|
||||
// CHECK: vector.shape_cast %[[IN]] : vector<2x4xf8E4M3FN> to vector<8xf8E4M3FN>
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: vector.shape_cast {{.*}} to vector<2x4xf16>
|
||||
%out = nvgpu.extf %in : vector<2x4xf8E4M3FN> to vector<2x4xf16>
|
||||
return %out : vector<2x4xf16>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_v2x4_f8_to_f32
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<2x4xf8E4M3FN>
|
||||
func.func @fpext_v2x4_f8_to_f32(%in : vector<2x4xf8E4M3FN>) -> vector<2x4xf32> {
|
||||
// CHECK: vector.shape_cast {{.*}} to vector<8xf8E4M3FN>
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.fpext
|
||||
// CHECK: vector.shape_cast {{.*}} to vector<2x4xf32>
|
||||
%out = nvgpu.extf %in : vector<2x4xf8E4M3FN> to vector<2x4xf32>
|
||||
return %out : vector<2x4xf32>
|
||||
}
|
||||
|
||||
// Non-aligned 1-D vectors (canonicalize: pad via insert/extract_strided_slice).
|
||||
|
||||
// CHECK-LABEL: @fpext_v3f8_to_v3f16
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<3xf8E5M2>
|
||||
func.func @fpext_v3f8_to_v3f16(%in : vector<3xf8E5M2>) -> vector<3xf16> {
|
||||
// CHECK: vector.insert_strided_slice %[[IN]]
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: vector.extract_strided_slice
|
||||
%out = nvgpu.extf %in : vector<3xf8E5M2> to vector<3xf16>
|
||||
return %out : vector<3xf16>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @fpext_v3_f8_to_f32
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<3xf8E5M2>
|
||||
func.func @fpext_v3_f8_to_f32(%in : vector<3xf8E5M2>) -> vector<3xf32> {
|
||||
// CHECK: vector.insert_strided_slice
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: llvm.fpext
|
||||
// CHECK: vector.extract_strided_slice
|
||||
%out = nvgpu.extf %in : vector<3xf8E5M2> to vector<3xf32>
|
||||
return %out : vector<3xf32>
|
||||
}
|
||||
|
||||
// Multi-rank + padding combined.
|
||||
|
||||
// CHECK-LABEL: @fpext_v3x1_f8_to_f16
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<3x1xf8E4M3FN>
|
||||
func.func @fpext_v3x1_f8_to_f16(%in : vector<3x1xf8E4M3FN>) -> vector<3x1xf16> {
|
||||
// CHECK: vector.shape_cast %[[IN]] : vector<3x1xf8E4M3FN> to vector<3xf8E4M3FN>
|
||||
// CHECK: vector.insert_strided_slice
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK: vector.extract_strided_slice
|
||||
// CHECK: vector.shape_cast {{.*}} to vector<3x1xf16>
|
||||
%out = nvgpu.extf %in : vector<3x1xf8E4M3FN> to vector<3x1xf16>
|
||||
return %out : vector<3x1xf16>
|
||||
}
|
||||
|
||||
// Relu attribute.
|
||||
|
||||
// CHECK-LABEL: @fpext_f8_to_f16_relu
|
||||
func.func @fpext_f8_to_f16_relu(%in : vector<8xf8E4M3FN>) {
|
||||
// CHECK: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-SAME: relu = true
|
||||
%out = nvgpu.extf %in {relu = true}
|
||||
: vector<8xf8E4M3FN> to vector<8xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// End-to-end: no residual vector ops after full lowering.
|
||||
|
||||
// CHECK-E2E-LABEL: @e2e_scalar_f8_to_f16
|
||||
// CHECK-E2E-NOT: vector.broadcast
|
||||
// CHECK-E2E-NOT: vector.insert_strided_slice
|
||||
// CHECK-E2E-NOT: vector.extract_strided_slice
|
||||
// CHECK-E2E-NOT: vector.extract
|
||||
// CHECK-E2E-NOT: vector.shape_cast
|
||||
// CHECK-E2E: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-E2E: return
|
||||
func.func @e2e_scalar_f8_to_f16(%in : f8E4M3FN) -> f16 {
|
||||
%out = nvgpu.extf %in : f8E4M3FN to f16
|
||||
return %out : f16
|
||||
}
|
||||
|
||||
// CHECK-E2E-LABEL: @e2e_v2x4_f8_to_f16
|
||||
// CHECK-E2E-NOT: vector.shape_cast
|
||||
// CHECK-E2E: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-E2E: return
|
||||
func.func @e2e_v2x4_f8_to_f16(%in : vector<2x4xf8E4M3FN>) -> vector<2x4xf16> {
|
||||
%out = nvgpu.extf %in : vector<2x4xf8E4M3FN> to vector<2x4xf16>
|
||||
return %out : vector<2x4xf16>
|
||||
}
|
||||
|
||||
// CHECK-E2E-LABEL: @e2e_v3f8_to_v3f16
|
||||
// CHECK-E2E-NOT: vector.insert_strided_slice
|
||||
// CHECK-E2E-NOT: vector.extract_strided_slice
|
||||
// CHECK-E2E: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-E2E: return
|
||||
func.func @e2e_v3f8_to_v3f16(%in : vector<3xf8E5M2>) -> vector<3xf16> {
|
||||
%out = nvgpu.extf %in : vector<3xf8E5M2> to vector<3xf16>
|
||||
return %out : vector<3xf16>
|
||||
}
|
||||
|
||||
// CHECK-E2E-LABEL: @e2e_scalar_f8_to_f32
|
||||
// CHECK-E2E-NOT: vector.broadcast
|
||||
// CHECK-E2E-NOT: vector.insert_strided_slice
|
||||
// CHECK-E2E-NOT: vector.extract_strided_slice
|
||||
// CHECK-E2E-NOT: vector.extract
|
||||
// CHECK-E2E-NOT: vector.shape_cast
|
||||
// CHECK-E2E: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-E2E: llvm.fpext
|
||||
// CHECK-E2E: return
|
||||
func.func @e2e_scalar_f8_to_f32(%in : f8E4M3FN) -> f32 {
|
||||
%out = nvgpu.extf %in : f8E4M3FN to f32
|
||||
return %out : f32
|
||||
}
|
||||
|
||||
// CHECK-E2E-LABEL: @e2e_v2x4_f8_to_f32
|
||||
// CHECK-E2E-NOT: vector.shape_cast
|
||||
// CHECK-E2E: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-E2E: llvm.fpext
|
||||
// CHECK-E2E: return
|
||||
func.func @e2e_v2x4_f8_to_f32(%in : vector<2x4xf8E4M3FN>) -> vector<2x4xf32> {
|
||||
%out = nvgpu.extf %in : vector<2x4xf8E4M3FN> to vector<2x4xf32>
|
||||
return %out : vector<2x4xf32>
|
||||
}
|
||||
|
||||
// CHECK-E2E-LABEL: @e2e_v3f8_to_v3f32
|
||||
// CHECK-E2E-NOT: vector.insert_strided_slice
|
||||
// CHECK-E2E-NOT: vector.extract_strided_slice
|
||||
// CHECK-E2E: nvvm.convert.f8x2.to.f16x2
|
||||
// CHECK-E2E: llvm.fpext
|
||||
// CHECK-E2E: return
|
||||
func.func @e2e_v3f8_to_v3f32(%in : vector<3xf8E5M2>) -> vector<3xf32> {
|
||||
%out = nvgpu.extf %in : vector<3xf8E5M2> to vector<3xf32>
|
||||
return %out : vector<3xf32>
|
||||
}
|
||||
76
mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf-large.mlir
Normal file
76
mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf-large.mlir
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
// RUN: mlir-opt %s -convert-nvgpu-to-nvvm | FileCheck %s
|
||||
|
||||
// Large-vector smoke tests for nvgpu.truncf.
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f32_to_f16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<400xf32>
|
||||
func.func @cvt_large_f32_to_f16(%in : vector<400xf32>) -> vector<400xf16> {
|
||||
// CHECK: llvm.bitcast %[[IN]] : vector<400xf32> to vector<400xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<200xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-NOT: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<200xi32> to vector<400xf16>
|
||||
%out = nvgpu.truncf %in : vector<400xf32> to vector<400xf16>
|
||||
return %out : vector<400xf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f32_to_bf16(
|
||||
func.func @cvt_large_f32_to_bf16(%in : vector<400xf32>) -> vector<400xbf16> {
|
||||
// CHECK: llvm.bitcast %{{.*}} : vector<400xf32> to vector<400xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.f32x2.to.bf16x2
|
||||
// CHECK-NOT: nvvm.convert.f32x2.to.bf16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<200xi32> to vector<400xbf16>
|
||||
%out = nvgpu.truncf %in : vector<400xf32> to vector<400xbf16>
|
||||
return %out : vector<400xbf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f32_to_f8(
|
||||
func.func @cvt_large_f32_to_f8(%in : vector<400xf32>) -> vector<400xf8E4M3FN> {
|
||||
// CHECK: llvm.bitcast %{{.*}} : vector<400xf32> to vector<400xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<100xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK-NOT: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<100xi32> to vector<400xi8>
|
||||
%out = nvgpu.truncf %in : vector<400xf32> to vector<400xf8E4M3FN>
|
||||
return %out : vector<400xf8E4M3FN>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f32_to_f6(
|
||||
func.func @cvt_large_f32_to_f6(%in : vector<400xf32>) -> vector<400xf6E2M3FN> {
|
||||
// CHECK: llvm.bitcast %{{.*}} : vector<400xf32> to vector<400xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.f32x2.to.f6x2
|
||||
// CHECK-NOT: nvvm.convert.f32x2.to.f6x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<100xi32> to vector<400xi8>
|
||||
// CHECK: llvm.trunc {{.*}} : vector<400xi8> to vector<400xi6>
|
||||
%out = nvgpu.truncf %in : vector<400xf32> to vector<400xf6E2M3FN>
|
||||
return %out : vector<400xf6E2M3FN>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f32_to_f4(
|
||||
func.func @cvt_large_f32_to_f4(%in : vector<400xf32>) -> vector<400xf4E2M1FN> {
|
||||
// CHECK: llvm.bitcast %{{.*}} : vector<400xf32> to vector<400xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<50xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.f32x2.to.f4x2
|
||||
// CHECK-NOT: nvvm.convert.f32x2.to.f4x2
|
||||
%out = nvgpu.truncf %in : vector<400xf32> to vector<400xf4E2M1FN>
|
||||
return %out : vector<400xf4E2M1FN>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_f16_to_f8(
|
||||
func.func @cvt_large_f16_to_f8(%in : vector<400xf16>) -> vector<400xf8E4M3FN> {
|
||||
// CHECK: llvm.bitcast %{{.*}} : vector<400xf16> to vector<200xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.f16x2.to.f8x2
|
||||
// CHECK-NOT: nvvm.convert.f16x2.to.f8x2
|
||||
%out = nvgpu.truncf %in : vector<400xf16> to vector<400xf8E4M3FN>
|
||||
return %out : vector<400xf8E4M3FN>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_large_bf16_to_f6(
|
||||
func.func @cvt_large_bf16_to_f6(%in : vector<400xbf16>) -> vector<400xf6E3M2FN> {
|
||||
// CHECK: llvm.bitcast %{{.*}} : vector<400xbf16> to vector<200xi32>
|
||||
// CHECK-COUNT-200: nvvm.convert.bf16x2.to.f6x2
|
||||
// CHECK-NOT: nvvm.convert.bf16x2.to.f6x2
|
||||
// CHECK: llvm.trunc {{.*}} : vector<400xi8> to vector<400xi6>
|
||||
%out = nvgpu.truncf %in : vector<400xbf16> to vector<400xf6E3M2FN>
|
||||
return %out : vector<400xf6E3M2FN>
|
||||
}
|
||||
401
mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf.mlir
Normal file
401
mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf.mlir
Normal file
|
|
@ -0,0 +1,401 @@
|
|||
// RUN: mlir-opt %s -convert-nvgpu-to-nvvm | FileCheck %s
|
||||
// RUN: mlir-opt %s -convert-nvgpu-to-nvvm -convert-vector-to-llvm | FileCheck %s --check-prefix=CHECK-E2E
|
||||
|
||||
// Basic aligned vector inputs.
|
||||
|
||||
// CHECK-LABEL: @cvt_float_f32_to_f16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<4xf32>
|
||||
func.func @cvt_float_f32_to_f16(%in : vector<4xf32>) {
|
||||
// CHECK: llvm.bitcast %[[IN]] : vector<4xf32> to vector<4xi32>
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rn>
|
||||
// CHECK-SAME: : vector<2xf16>
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rn>
|
||||
// CHECK-SAME: : vector<2xf16>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<4xf16>
|
||||
%out = nvgpu.truncf %in : vector<4xf32> to vector<4xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_float_f32_to_f16_v8(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf32>
|
||||
func.func @cvt_float_f32_to_f16_v8(%in : vector<8xf32>) {
|
||||
// CHECK: llvm.bitcast %[[IN]] : vector<8xf32> to vector<8xi32>
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<4xi32> to vector<8xf16>
|
||||
%out = nvgpu.truncf %in : vector<8xf32> to vector<8xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_float_f32_to_bf16(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<4xf32>
|
||||
func.func @cvt_float_f32_to_bf16(%in : vector<4xf32>) {
|
||||
// CHECK: llvm.bitcast %[[IN]] : vector<4xf32> to vector<4xi32>
|
||||
// CHECK: nvvm.convert.f32x2.to.bf16x2
|
||||
// CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rn>
|
||||
// CHECK-SAME: : vector<2xbf16>
|
||||
// CHECK: nvvm.convert.f32x2.to.bf16x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<4xbf16>
|
||||
%out = nvgpu.truncf %in : vector<4xf32> to vector<4xbf16>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_float_f32_to_e4m3(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf32>
|
||||
func.func @cvt_float_f32_to_e4m3(%in : vector<8xf32>) {
|
||||
// CHECK: %[[IN_I32:.+]] = llvm.bitcast %[[IN]] : vector<8xf32> to vector<8xi32>
|
||||
// CHECK: %[[OUT_I32:.+]] = llvm.mlir.undef : vector<2xi32>
|
||||
// CHECK: %[[IDX_0:.+]] = llvm.mlir.constant(0 : i64) : i64
|
||||
// CHECK: %[[SUB_VEC_0:.+]] = llvm.mlir.undef : vector<2xi16>
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK-SAME: sat = #nvvm.sat_mode<satfinite>
|
||||
// CHECK-SAME: : i16(f8E4M3FN)
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi16>
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK-SAME: : i16(f8E4M3FN)
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi16>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi16> to i32
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<2xi16>
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi16> to i32
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi32>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<8xi8>
|
||||
%out = nvgpu.truncf %in : vector<8xf32> to vector<8xf8E4M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_float_f16_to_e2m3(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf16>
|
||||
func.func @cvt_float_f16_to_e2m3(%in : vector<8xf16>) {
|
||||
// CHECK: llvm.bitcast %[[IN]] : vector<8xf16> to vector<4xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<2xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<2xi16>
|
||||
// CHECK: nvvm.convert.f16x2.to.f6x2
|
||||
// CHECK-SAME: : vector<2xf16> -> i16(f6E2M3FN)
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi16>
|
||||
// CHECK: nvvm.convert.f16x2.to.f6x2
|
||||
// CHECK-SAME: : vector<2xf16> -> i16(f6E2M3FN)
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi16>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi16> to i32
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<8xi8>
|
||||
// CHECK: llvm.trunc {{.*}} : vector<8xi8> to vector<8xi6>
|
||||
%out = nvgpu.truncf %in : vector<8xf16> to vector<8xf6E2M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_float_bf16_to_e3m2(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xbf16>
|
||||
func.func @cvt_float_bf16_to_e3m2(%in : vector<8xbf16>) {
|
||||
// CHECK: llvm.bitcast %[[IN]] : vector<8xbf16> to vector<4xi32>
|
||||
// CHECK: nvvm.convert.bf16x2.to.f6x2
|
||||
// CHECK-SAME: : vector<2xbf16> -> i16(f6E3M2FN)
|
||||
// CHECK: nvvm.convert.bf16x2.to.f6x2
|
||||
// CHECK-SAME: : vector<2xbf16> -> i16(f6E3M2FN)
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<8xi8>
|
||||
// CHECK: llvm.trunc {{.*}} : vector<8xi8> to vector<8xi6>
|
||||
%out = nvgpu.truncf %in : vector<8xbf16> to vector<8xf6E3M2FN>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_float_f32_to_e2m3(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf32>
|
||||
func.func @cvt_float_f32_to_e2m3(%in : vector<8xf32>) {
|
||||
// CHECK: llvm.bitcast %[[IN]] : vector<8xf32> to vector<8xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<2xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<2xi16>
|
||||
// CHECK: nvvm.convert.f32x2.to.f6x2
|
||||
// CHECK-SAME: : i16(f6E2M3FN)
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi16>
|
||||
// CHECK: nvvm.convert.f32x2.to.f6x2
|
||||
// CHECK-SAME: : i16(f6E2M3FN)
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi16>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi16> to i32
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi32>
|
||||
// CHECK: llvm.mlir.undef : vector<2xi16>
|
||||
// CHECK: nvvm.convert.f32x2.to.f6x2
|
||||
// CHECK: nvvm.convert.f32x2.to.f6x2
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi16> to i32
|
||||
// CHECK: llvm.insertelement {{.*}} : vector<2xi32>
|
||||
// CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<8xi8>
|
||||
// CHECK: llvm.trunc {{.*}} : vector<8xi8> to vector<8xi6>
|
||||
%out = nvgpu.truncf %in : vector<8xf32> to vector<8xf6E2M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_float_f32_to_e8m0_rz(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf32>
|
||||
func.func @cvt_float_f32_to_e8m0_rz(%in : vector<8xf32>) {
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rz>
|
||||
// CHECK-SAME: : i16(f8E8M0FNU)
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rz>}
|
||||
: vector<8xf32> to vector<8xf8E8M0FNU>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @cvt_float_f32_to_e8m0_rp(
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<8xf32>
|
||||
func.func @cvt_float_f32_to_e8m0_rp(%in : vector<8xf32>) {
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rp>
|
||||
// CHECK-SAME: : i16(f8E8M0FNU)
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rp>}
|
||||
: vector<8xf32> to vector<8xf8E8M0FNU>
|
||||
return
|
||||
}
|
||||
|
||||
// Scalar inputs (canonicalize: broadcast + pad + extract).
|
||||
|
||||
// CHECK-LABEL: @fptrunc_scalar_f32_to_f16
|
||||
// CHECK-SAME: %[[IN:.+]]: f32
|
||||
func.func @fptrunc_scalar_f32_to_f16(%in : f32) -> f16 {
|
||||
// CHECK: vector.broadcast %[[IN]] : f32 to vector<1xf32>
|
||||
// CHECK: vector.insert_strided_slice
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK: vector.extract_strided_slice
|
||||
// CHECK: vector.extract {{.*}}[0] : f16 from vector<1xf16>
|
||||
%out = nvgpu.truncf %in : f32 to f16
|
||||
return %out : f16
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_scalar_f32_to_bf16
|
||||
// CHECK-SAME: %[[IN:.+]]: f32
|
||||
func.func @fptrunc_scalar_f32_to_bf16(%in : f32) -> bf16 {
|
||||
// CHECK: vector.broadcast %[[IN]]
|
||||
// CHECK: nvvm.convert.f32x2.to.bf16x2
|
||||
// CHECK: vector.extract
|
||||
%out = nvgpu.truncf %in : f32 to bf16
|
||||
return %out : bf16
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_scalar_f64_to_f32
|
||||
func.func @fptrunc_scalar_f64_to_f32(%arg0: f64) -> f32 {
|
||||
// CHECK: vector.broadcast
|
||||
// CHECK: llvm.fptrunc
|
||||
// CHECK: vector.extract
|
||||
%out = nvgpu.truncf %arg0 : f64 to f32
|
||||
return %out : f32
|
||||
}
|
||||
|
||||
// Multi-rank vectors (canonicalize: shape_cast flatten).
|
||||
|
||||
// CHECK-LABEL: @fptrunc_v2x4_f32_to_f16
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<2x4xf32>
|
||||
func.func @fptrunc_v2x4_f32_to_f16(%in : vector<2x4xf32>) -> vector<2x4xf16> {
|
||||
// CHECK: vector.shape_cast %[[IN]] : vector<2x4xf32> to vector<8xf32>
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK: vector.shape_cast {{.*}} to vector<2x4xf16>
|
||||
%out = nvgpu.truncf %in : vector<2x4xf32> to vector<2x4xf16>
|
||||
return %out : vector<2x4xf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_v4x2_f32_to_f8
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<4x2xf32>
|
||||
func.func @fptrunc_v4x2_f32_to_f8(%in : vector<4x2xf32>) -> vector<4x2xf8E4M3FN> {
|
||||
// CHECK: vector.shape_cast %[[IN]] : vector<4x2xf32> to vector<8xf32>
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK: vector.shape_cast {{.*}} to vector<4x2xf8E4M3FN>
|
||||
%out = nvgpu.truncf %in : vector<4x2xf32> to vector<4x2xf8E4M3FN>
|
||||
return %out : vector<4x2xf8E4M3FN>
|
||||
}
|
||||
|
||||
// Non-aligned 1-D vectors (canonicalize: pad via insert/extract_strided_slice).
|
||||
|
||||
// CHECK-LABEL: @fptrunc_v1f32_to_v1f16
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<1xf32>
|
||||
func.func @fptrunc_v1f32_to_v1f16(%in : vector<1xf32>) -> vector<1xf16> {
|
||||
// CHECK: vector.insert_strided_slice %[[IN]]
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK: vector.extract_strided_slice
|
||||
%out = nvgpu.truncf %in : vector<1xf32> to vector<1xf16>
|
||||
return %out : vector<1xf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_v3f16_to_v3f8
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<3xf16>
|
||||
func.func @fptrunc_v3f16_to_v3f8(%in : vector<3xf16>) -> vector<3xf8E4M3FN> {
|
||||
// CHECK: vector.insert_strided_slice %[[IN]]
|
||||
// CHECK: nvvm.convert.f16x2.to.f8x2
|
||||
// CHECK: vector.extract_strided_slice
|
||||
%out = nvgpu.truncf %in : vector<3xf16> to vector<3xf8E4M3FN>
|
||||
return %out : vector<3xf8E4M3FN>
|
||||
}
|
||||
|
||||
// Multi-rank + padding combined.
|
||||
|
||||
// CHECK-LABEL: @fptrunc_v3x1_f32_to_f16
|
||||
// CHECK-SAME: %[[IN:.+]]: vector<3x1xf32>
|
||||
func.func @fptrunc_v3x1_f32_to_f16(%in : vector<3x1xf32>) -> vector<3x1xf16> {
|
||||
// CHECK: vector.shape_cast %[[IN]] : vector<3x1xf32> to vector<3xf32>
|
||||
// CHECK: vector.insert_strided_slice
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK: vector.extract_strided_slice
|
||||
// CHECK: vector.shape_cast {{.*}} to vector<3x1xf16>
|
||||
%out = nvgpu.truncf %in : vector<3x1xf32> to vector<3x1xf16>
|
||||
return %out : vector<3x1xf16>
|
||||
}
|
||||
|
||||
// f64 source truncation.
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f64_to_f32
|
||||
func.func @fptrunc_f64_to_f32(%arg0: vector<4xf64>) -> vector<4xf32> {
|
||||
// CHECK: llvm.fptrunc %{{.*}} : vector<4xf64> to vector<4xf32>
|
||||
// CHECK-NOT: nvvm
|
||||
%out = nvgpu.truncf %arg0 : vector<4xf64> to vector<4xf32>
|
||||
return %out : vector<4xf32>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f64_to_f16
|
||||
func.func @fptrunc_f64_to_f16(%arg0: vector<2xf64>) -> vector<2xf16> {
|
||||
// CHECK: vector.insert_strided_slice
|
||||
// CHECK: llvm.fptrunc %{{.*}} : vector<4xf64> to vector<4xf16>
|
||||
// CHECK-NOT: nvvm.convert
|
||||
// CHECK: vector.extract_strided_slice
|
||||
%out = nvgpu.truncf %arg0 : vector<2xf64> to vector<2xf16>
|
||||
return %out : vector<2xf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f64_to_bf16
|
||||
func.func @fptrunc_f64_to_bf16(%arg0: vector<2xf64>) -> vector<2xbf16> {
|
||||
// CHECK: vector.insert_strided_slice
|
||||
// CHECK: llvm.fptrunc %{{.*}} : vector<4xf64> to vector<4xbf16>
|
||||
// CHECK-NOT: nvvm.convert
|
||||
// CHECK: vector.extract_strided_slice
|
||||
%out = nvgpu.truncf %arg0 : vector<2xf64> to vector<2xbf16>
|
||||
return %out : vector<2xbf16>
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f64_to_f8
|
||||
func.func @fptrunc_f64_to_f8(%arg0: vector<4xf64>) -> vector<4xf8E4M3FN> {
|
||||
// CHECK: vector.insert_strided_slice
|
||||
// CHECK: llvm.fptrunc %{{.*}} : vector<8xf64> to vector<8xf32>
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK: vector.extract_strided_slice
|
||||
%out = nvgpu.truncf %arg0 : vector<4xf64> to vector<4xf8E4M3FN>
|
||||
return %out : vector<4xf8E4M3FN>
|
||||
}
|
||||
|
||||
// Saturation and relu attributes.
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f32_to_f8_satfinite
|
||||
func.func @fptrunc_f32_to_f8_satfinite(%in : vector<8xf32>) {
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK-SAME: sat = #nvvm.sat_mode<satfinite>
|
||||
%out = nvgpu.truncf %in {sat = #nvvm.sat_mode<satfinite>}
|
||||
: vector<8xf32> to vector<8xf8E4M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f32_to_f16_relu
|
||||
func.func @fptrunc_f32_to_f16_relu(%in : vector<4xf32>) {
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-SAME: relu = true
|
||||
%out = nvgpu.truncf %in {relu = true}
|
||||
: vector<4xf32> to vector<4xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// Default SATFINITE behavior.
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f32_to_f8_default_sat
|
||||
func.func @fptrunc_f32_to_f8_default_sat(%in : vector<8xf32>) {
|
||||
// CHECK: nvvm.convert.f32x2.to.f8x2
|
||||
// CHECK-SAME: sat = #nvvm.sat_mode<satfinite>
|
||||
%out = nvgpu.truncf %in : vector<8xf32> to vector<8xf8E4M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f32_to_f16_default_sat
|
||||
func.func @fptrunc_f32_to_f16_default_sat(%in : vector<4xf32>) {
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-SAME: sat = #nvvm.sat_mode<satfinite>
|
||||
%out = nvgpu.truncf %in : vector<4xf32> to vector<4xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f32_to_f16_explicit_none
|
||||
func.func @fptrunc_f32_to_f16_explicit_none(%in : vector<4xf32>) {
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-NOT: satfinite
|
||||
%out = nvgpu.truncf %in {sat = #nvvm.sat_mode<none>}
|
||||
: vector<4xf32> to vector<4xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// Stochastic rounding (RS + random_bits).
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f32_to_f16_rs
|
||||
func.func @fptrunc_f32_to_f16_rs(%in : vector<4xf32>, %rbits : i32) {
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rs>
|
||||
%out = nvgpu.truncf %in, %rbits {rnd = #nvvm.fp_rnd_mode<rs>}
|
||||
: vector<4xf32> to vector<4xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f32_to_bf16_rs
|
||||
func.func @fptrunc_f32_to_bf16_rs(%in : vector<4xf32>, %rbits : i32) {
|
||||
// CHECK: nvvm.convert.f32x2.to.bf16x2
|
||||
// CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rs>
|
||||
%out = nvgpu.truncf %in, %rbits {rnd = #nvvm.fp_rnd_mode<rs>}
|
||||
: vector<4xf32> to vector<4xbf16>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f32_to_f16_rz
|
||||
func.func @fptrunc_f32_to_f16_rz(%in : vector<4xf32>) {
|
||||
// CHECK: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rz>
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rz>}
|
||||
: vector<4xf32> to vector<4xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// CHECK-LABEL: @fptrunc_f32_to_bf16_rz
|
||||
func.func @fptrunc_f32_to_bf16_rz(%in : vector<4xf32>) {
|
||||
// CHECK: nvvm.convert.f32x2.to.bf16x2
|
||||
// CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rz>
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rz>}
|
||||
: vector<4xf32> to vector<4xbf16>
|
||||
return
|
||||
}
|
||||
|
||||
// End-to-end: no residual vector ops after full lowering.
|
||||
|
||||
// CHECK-E2E-LABEL: @e2e_scalar_f32_to_f16
|
||||
// CHECK-E2E-NOT: vector.broadcast
|
||||
// CHECK-E2E-NOT: vector.insert_strided_slice
|
||||
// CHECK-E2E-NOT: vector.extract_strided_slice
|
||||
// CHECK-E2E-NOT: vector.extract
|
||||
// CHECK-E2E-NOT: vector.shape_cast
|
||||
// CHECK-E2E: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-E2E: return
|
||||
func.func @e2e_scalar_f32_to_f16(%in : f32) -> f16 {
|
||||
%out = nvgpu.truncf %in : f32 to f16
|
||||
return %out : f16
|
||||
}
|
||||
|
||||
// CHECK-E2E-LABEL: @e2e_v2x4_f32_to_f16
|
||||
// CHECK-E2E-NOT: vector.shape_cast
|
||||
// CHECK-E2E: nvvm.convert.f32x2.to.f16x2
|
||||
// CHECK-E2E: return
|
||||
func.func @e2e_v2x4_f32_to_f16(%in : vector<2x4xf32>) -> vector<2x4xf16> {
|
||||
%out = nvgpu.truncf %in : vector<2x4xf32> to vector<2x4xf16>
|
||||
return %out : vector<2x4xf16>
|
||||
}
|
||||
|
||||
// CHECK-E2E-LABEL: @e2e_v3f16_to_v3f8
|
||||
// CHECK-E2E-NOT: vector.insert_strided_slice
|
||||
// CHECK-E2E-NOT: vector.extract_strided_slice
|
||||
// CHECK-E2E: nvvm.convert.f16x2.to.f8x2
|
||||
// CHECK-E2E: return
|
||||
func.func @e2e_v3f16_to_v3f8(%in : vector<3xf16>) -> vector<3xf8E4M3FN> {
|
||||
%out = nvgpu.truncf %in : vector<3xf16> to vector<3xf8E4M3FN>
|
||||
return %out : vector<3xf8E4M3FN>
|
||||
}
|
||||
66
mlir/test/Dialect/NVGPU/nvgpu-extf-invalid.mlir
Normal file
66
mlir/test/Dialect/NVGPU/nvgpu-extf-invalid.mlir
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
// RUN: mlir-opt -split-input-file -verify-diagnostics %s
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fpext_wider(%in : vector<16xf8E5M2>) {
|
||||
// expected-error @+1 {{'nvgpu.extf' op result type 'f4E2M1FN' must be wider than operand type 'f8E5M2'}}
|
||||
%out = nvgpu.extf %in : vector<16xf8E5M2> to vector<16xf4E2M1FN>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fpext_dst_bitwidth(%in : vector<16xf4E2M1FN>) {
|
||||
// expected-error @+1 {{'nvgpu.extf' op result type must be 16, 32, or 64 bitwidth, but got 8}}
|
||||
%out = nvgpu.extf %in : vector<16xf4E2M1FN> to vector<16xf8E4M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fpext_e8m0_to_f16(%in : vector<16xf8E8M0FNU>) {
|
||||
// expected-error @+1 {{'nvgpu.extf' op expects bf16 or f32 output type when input type is e8m0.}}
|
||||
%out = nvgpu.extf %in : vector<16xf8E8M0FNU> to vector<16xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fpext_bad_rounding(%in : vector<16xf8E5M2>) {
|
||||
// expected-error @+1 {{'nvgpu.extf' op expects RN rounding mode, but got #nvvm.fp_rnd_mode<rz>}}
|
||||
%out = nvgpu.extf %in {rnd = #nvvm.fp_rnd_mode<rz>}
|
||||
: vector<16xf8E5M2> to vector<16xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fpext_relu_bf16(%in : vector<8xf8E5M2>) {
|
||||
// expected-error @+1 {{'nvgpu.extf' op relu is not supported for bf16 destination}}
|
||||
%out = nvgpu.extf %in {relu = true} : vector<8xf8E5M2> to vector<8xbf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fpext_shape_mismatch(%in : vector<16xf8E5M2>) {
|
||||
// expected-error @+1 {{'nvgpu.extf' op input and output shapes must match}}
|
||||
%out = nvgpu.extf %in : vector<16xf8E5M2> to vector<8xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fpext_scalar_vector_mismatch(%in : f8E4M3FN) {
|
||||
// expected-error @+1 {{'nvgpu.extf' op input and output must both be scalars or both be vectors}}
|
||||
%out = nvgpu.extf %in : f8E4M3FN to vector<1xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fpext_rank0_vector(%in : vector<f8E4M3FN>) {
|
||||
// expected-error @+1 {{'nvgpu.extf' op operand #0 must be floating-point or vector of floating-point values, but got 'vector<f8E4M3FN>'}}
|
||||
%out = nvgpu.extf %in : vector<f8E4M3FN> to vector<f16>
|
||||
return
|
||||
}
|
||||
113
mlir/test/Dialect/NVGPU/nvgpu-truncf-invalid.mlir
Normal file
113
mlir/test/Dialect/NVGPU/nvgpu-truncf-invalid.mlir
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
// RUN: mlir-opt -split-input-file -verify-diagnostics %s
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_narrower(%in : vector<16xf16>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op result type 'f32' must be narrower than operand type 'f16'}}
|
||||
%out = nvgpu.truncf %in : vector<16xf16> to vector<16xf32>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_src_bitwidth(%in : vector<16xf8E5M2>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op input type must be 64/32/16 bitwidth, but got 8}}
|
||||
%out = nvgpu.truncf %in : vector<16xf8E5M2> to vector<16xf4E2M1FN>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_e8m0_bad_rounding(%in : vector<16xf32>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op expects RZ or RP rounding mode when result type is e8m0, but got #nvvm.fp_rnd_mode<rn>}}
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rn>}
|
||||
: vector<16xf32> to vector<16xf8E8M0FNU>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_unsupported_sat_mode(%in : vector<8xf32>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op attribute 'sat' failed to satisfy constraint: Describes the saturation mode whose value is one of {none, satfinite}}}
|
||||
%out = nvgpu.truncf %in {sat = #nvvm.sat_mode<sat>}
|
||||
: vector<8xf32> to vector<8xf8E4M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_f32_to_f8_rz(%in : vector<16xf32>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op expects RN rounding mode, but got #nvvm.fp_rnd_mode<rz>}}
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rz>}
|
||||
: vector<16xf32> to vector<16xf8E4M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_f64_to_f16_rz(%in : vector<4xf64>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op expects RN rounding mode for f64 input, but got #nvvm.fp_rnd_mode<rz>}}
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rz>}
|
||||
: vector<4xf64> to vector<4xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_rs_unsupported_types(%in : vector<16xf32>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op RS (stochastic) rounding is only supported for f32->f16/bf16, got 'f32' -> 'f8E4M3FN'}}
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rs>}
|
||||
: vector<16xf32> to vector<16xf8E4M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_rs_no_random_bits(%in : vector<4xf32>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op random_bits operand is required with RS rounding}}
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rs>}
|
||||
: vector<4xf32> to vector<4xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_bad_rounding(%in : vector<16xf32>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op expects RN rounding mode, but got #nvvm.fp_rnd_mode<rp>}}
|
||||
%out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode<rp>}
|
||||
: vector<16xf32> to vector<16xf8E4M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_random_bits_no_rs(%in : vector<4xf32>, %rbits : i32) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op random_bits can only be used with RS rounding mode}}
|
||||
%out = nvgpu.truncf %in, %rbits
|
||||
: vector<4xf32> to vector<4xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_shape_mismatch(%in : vector<16xf32>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op input and output shapes must match}}
|
||||
%out = nvgpu.truncf %in : vector<16xf32> to vector<8xf8E4M3FN>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_scalar_vector_mismatch(%in : f32) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op input and output must both be scalars or both be vectors}}
|
||||
%out = nvgpu.truncf %in : f32 to vector<1xf16>
|
||||
return
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
func.func @fptrunc_rank0_vector(%in : vector<f32>) {
|
||||
// expected-error @+1 {{'nvgpu.truncf' op operand #0 must be floating-point or vector of floating-point values, but got 'vector<f32>'}}
|
||||
%out = nvgpu.truncf %in : vector<f32> to vector<f16>
|
||||
return
|
||||
}
|
||||
Loading…
Reference in a new issue