From 16094a86a6cfe79fa50dcca39b539ab6465b8ac0 Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Tue, 14 Jul 2026 16:05:05 +0530 Subject: [PATCH] [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. --- mlir/include/mlir/Conversion/Passes.td | 3 +- .../include/mlir/Dialect/NVGPU/IR/NVGPUOps.td | 75 ++ .../lib/Conversion/NVGPUToNVVM/CMakeLists.txt | 1 + .../Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp | 669 ++++++++++++++++++ mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp | 116 +++ .../NVGPUToNVVM/nvgpu-extf-large.mlir | 78 ++ .../Conversion/NVGPUToNVVM/nvgpu-extf.mlir | 506 +++++++++++++ .../NVGPUToNVVM/nvgpu-truncf-large.mlir | 76 ++ .../Conversion/NVGPUToNVVM/nvgpu-truncf.mlir | 401 +++++++++++ .../Dialect/NVGPU/nvgpu-extf-invalid.mlir | 66 ++ .../Dialect/NVGPU/nvgpu-truncf-invalid.mlir | 113 +++ 11 files changed, 2103 insertions(+), 1 deletion(-) create mode 100644 mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf-large.mlir create mode 100644 mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf.mlir create mode 100644 mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf-large.mlir create mode 100644 mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf.mlir create mode 100644 mlir/test/Dialect/NVGPU/nvgpu-extf-invalid.mlir create mode 100644 mlir/test/Dialect/NVGPU/nvgpu-truncf-invalid.mlir diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td index b92bc1cd3172..a675ce431f56 100644 --- a/mlir/include/mlir/Conversion/Passes.td +++ b/mlir/include/mlir/Conversion/Passes.td @@ -1106,7 +1106,8 @@ def ConvertNVGPUToNVVMPass : Pass<"convert-nvgpu-to-nvvm"> { "arith::ArithDialect", "LLVM::LLVMDialect", "memref::MemRefDialect", - "NVVM::NVVMDialect" + "NVVM::NVVMDialect", + "vector::VectorDialect" ]; } diff --git a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPUOps.td b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPUOps.td index 679dfdb9a474..69b6b57de54b 100644 --- a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPUOps.td +++ b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPUOps.td @@ -671,4 +671,79 @@ def NVGPU_RcpOp : NVGPU_Op<"rcp", [Pure, let hasVerifier = 1; } +//===----------------------------------------------------------------------===// +// NVGPU Conversion Ops +//===----------------------------------------------------------------------===// + +def AnyI32Like : TypeOrValueSemanticsContainer; + +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:$rnd, + DefaultValuedAttr]>, + "NVVM::SaturationMode::SATFINITE">:$sat, + DefaultValuedAttr:$relu, + Optional:$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:$rnd, + DefaultValuedAttr:$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 diff --git a/mlir/lib/Conversion/NVGPUToNVVM/CMakeLists.txt b/mlir/lib/Conversion/NVGPUToNVVM/CMakeLists.txt index a050749eb7da..2e0a417bd823 100644 --- a/mlir/lib/Conversion/NVGPUToNVVM/CMakeLists.txt +++ b/mlir/lib/Conversion/NVGPUToNVVM/CMakeLists.txt @@ -18,6 +18,7 @@ add_mlir_conversion_library(MLIRNVGPUToNVVM MLIRNVGPUDialect MLIRNVVMDialect MLIRArithDialect + MLIRVectorDialect MLIRPass MLIRSCFTransforms MLIRTransforms diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp index e566449ffadf..948229bb5332 100644 --- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp +++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp @@ -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 { 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 classifyFPType(Type t) { + static constexpr auto isConvertibleF8Type = [](Type t) { + return isa(t); + }; + static constexpr auto isConvertibleF6Type = [](Type t) { + return isa(t); + }; + static constexpr auto isConvertibleF4Type = [](Type t) { + return isa(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 +static std::optional +lookupConvOp(const TableEntry (&table)[N], Type srcElemType, Type dstElemType) { + std::optional srcKind = classifyFPType(srcElemType); + std::optional 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(srcVec.getType()); + assert(idx >= 0 && idx < vecTy.getNumElements() && + "extractElement: index out of bounds"); + IntegerType i64Ty = b.getI64Type(); + return b.create( + srcVec, b.create(i64Ty, b.getI64IntegerAttr(idx))); +} + +/// Extract a pair of f32 values from an i32 vector at the given base index. +static std::pair 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(f32Ty, elem0), + b.create(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(vecTy, elem); +} + +/// Create a sub-byte conversion from an f32 pair source and return the native +/// result. +template +static Value convertFromF32Pair(ImplicitLocOpBuilder &b, Value srcI32Vec, + int srcBaseIdx, Type resultTy, Args &&...args) { + auto [lo, hi] = extractF32Pair(b, srcI32Vec, srcBaseIdx); + return b.create(resultTy, hi, lo, std::forward(args)...); +} + +/// Create a sub-byte conversion from a packed f16x2/bf16x2 source and return +/// the native result. +template +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(resultTy, src, std::forward(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( + VectorType::get(2, b.getF16Type()), hi, lo, randomBits, rndAttr, + satAttr, reluAttr); + return b.create(i32Ty, r); + } + case FPTruncConvOp::F32x2_TO_BF16x2: { + auto [lo, hi] = extractF32Pair(b, srcI32Vec, srcBaseIdx); + Value r = b.create( + VectorType::get(2, b.getBF16Type()), hi, lo, randomBits, rndAttr, + satAttr, reluAttr); + return b.create(i32Ty, r); + } + case FPTruncConvOp::F32x2_TO_F8x2: + return convertFromF32Pair( + b, srcI32Vec, srcBaseIdx, i16Ty, rndAttr, satAttr, reluAttr, dstTyAttr); + case FPTruncConvOp::F32x2_TO_F6x2: + return convertFromF32Pair( + b, srcI32Vec, srcBaseIdx, i16Ty, reluAttr, actualDstTyAttr); + case FPTruncConvOp::F32x2_TO_F4x2: + return convertFromF32Pair( + b, srcI32Vec, srcBaseIdx, i8Ty, reluAttr, dstTyAttr); + case FPTruncConvOp::F16x2_TO_F8x2: + return convertFromPacked( + b, srcI32Vec, srcBaseIdx, b.getF16Type(), i16Ty, reluAttr, dstTyAttr); + case FPTruncConvOp::F16x2_TO_F6x2: + return convertFromPacked( + b, srcI32Vec, srcBaseIdx, b.getF16Type(), i16Ty, reluAttr, + actualDstTyAttr); + case FPTruncConvOp::F16x2_TO_F4x2: + return convertFromPacked( + b, srcI32Vec, srcBaseIdx, b.getF16Type(), i8Ty, reluAttr, + actualDstTyAttr); + case FPTruncConvOp::BF16x2_TO_F8x2: + return convertFromPacked( + b, srcI32Vec, srcBaseIdx, b.getBF16Type(), i16Ty, rndAttr, satAttr, + reluAttr, dstTyAttr); + case FPTruncConvOp::BF16x2_TO_F6x2: + return convertFromPacked( + b, srcI32Vec, srcBaseIdx, b.getBF16Type(), i16Ty, reluAttr, + actualDstTyAttr); + case FPTruncConvOp::BF16x2_TO_F4x2: + return convertFromPacked( + 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(op.getIn().getType()); + auto dstType = llvm::dyn_cast(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(convertedType, input); + rewriter.replaceOp(op, result); + return success(); + } + auto f32VecTy = VectorType::get(srcType.getShape(), b.getF32Type()); + input = b.create(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(VectorType::get(srcI32Elems, i32Ty), input); + Value dstI32Vec = + b.create(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(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(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( + subVec, subResult, + b.create(i64Ty, b.getI64IntegerAttr(insertIdx))); + insertIdx--; + } + + dstValue = b.create(i32Ty, subVec); + } + + dstI32Vec = + b.create(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(i8VecTy, dstI32Vec); + Value truncVec = b.create(convertedType, i8Vec); + rewriter.replaceOp(op, truncVec); + } else { + auto dstVec = b.create(convertedType, dstI32Vec); + rewriter.replaceOp(op, dstVec); + } + return success(); +} + +struct NVGPUTruncfOpLowering : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::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(i32Ty, r); + } + case FPExtConvOp::F8x2_TO_BF16x2: { + Value r = NVVM::ConvertF8x2ToBF16x2Op::create( + b, VectorType::get(2, b.getBF16Type()), src, extScaleFactor, srcTyAttr); + return b.create(i32Ty, r); + } + case FPExtConvOp::F6x2_TO_F16x2: { + Value r = NVVM::ConvertF6x2ToF16x2Op::create( + b, VectorType::get(2, b.getF16Type()), src, srcTyAttr, reluAttr); + return b.create(i32Ty, r); + } + case FPExtConvOp::F6x2_TO_BF16x2: { + Value r = NVVM::ConvertF6x2ToBF16x2Op::create( + b, VectorType::get(2, b.getBF16Type()), src, extScaleFactor, srcTyAttr); + return b.create(i32Ty, r); + } + case FPExtConvOp::F4x2_TO_F16x2: { + Value r = NVVM::ConvertF4x2ToF16x2Op::create( + b, VectorType::get(2, b.getF16Type()), src, srcTyAttr, reluAttr); + return b.create(i32Ty, r); + } + case FPExtConvOp::F4x2_TO_BF16x2: { + Value r = NVVM::ConvertF4x2ToBF16x2Op::create( + b, VectorType::get(2, b.getBF16Type()), src, extScaleFactor, srcTyAttr); + return b.create(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(op.getIn().getType()); + auto dstType = llvm::dyn_cast(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(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(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 to vector, then bitcast to i32s. + Value inputVec = adaptor.getIn(); + if (srcBW == 6) { + auto i8VecTy = VectorType::get(numElems, i8Ty); + inputVec = b.create(i8VecTy, inputVec); + } + + int srcI32Elems = numElems * effectiveSrcBW / regBits; + int dstI32Elems = numElems * intermediateDstBW / regBits; + Value srcI32Vec = + b.create(VectorType::get(srcI32Elems, i32Ty), inputVec); + Value dstI32Vec = + b.create(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( + srcI32Vec, + b.create(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(VectorType::get(2, i16Ty), srcI32); + for (int half = 0; half < 2; half++) { + Value halfI16 = b.create( + i16Vec, + b.create(i64Ty, b.getI64IntegerAttr(half))); + Value src = + b.create(VectorType::get(2, i8Ty), halfI16); + Value dstValue = + createExtConversion(b, ctx, convOp, src, reluBoolAttr, + actualSrcFloatType, extScaleFactor); + Value dstIdxConst = + b.create(i64Ty, b.getI64IntegerAttr(dstIdx)); + dstI32Vec = + b.create(dstI32Vec, dstValue, dstIdxConst); + dstIdx++; + } + } else { + // f4: one i32 holds 4 bytes -> each byte is one conversion input. + Value i8Vec = b.create(VectorType::get(4, i8Ty), srcI32); + for (int byteIdx = 0; byteIdx < 4; byteIdx++) { + Value src = b.create( + i8Vec, + b.create(i64Ty, b.getI64IntegerAttr(byteIdx))); + Value dstValue = + createExtConversion(b, ctx, convOp, src, reluBoolAttr, + actualSrcFloatType, extScaleFactor); + Value dstIdxConst = + b.create(i64Ty, b.getI64IntegerAttr(dstIdx)); + dstI32Vec = + b.create(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(intermediateVecTy, dstI32Vec); + result = b.create(convertedType, intermediateVec); + } else { + result = b.create(convertedType, dstI32Vec); + } + rewriter.replaceOp(op, result); + return success(); +} + +struct NVGPUExtfOpLowering : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::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 +struct NVGPUFPCanonicalizePattern : public OpRewritePattern { + using OpRewritePattern::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(inType); + auto srcVecTy = dyn_cast(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{0}, SmallVector{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{0}, SmallVector{numElems}, + SmallVector{1}); + } + + if (isMultiRank) { + result = + vector::ShapeCastOp::create(b, cast(outType), result); + } + + if (isScalar) { + result = vector::ExtractOp::create(b, result, SmallVector{0}); + } + + rewriter.replaceOp(op, result); + return success(); + } +}; + +using NVGPUTruncfCanonicalizePattern = + NVGPUFPCanonicalizePattern; +using NVGPUExtfCanonicalizePattern = + NVGPUFPCanonicalizePattern; } // 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( + patterns.getContext()); } diff --git a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp index 2b2686a9163b..c77ecb8431ad 100644 --- a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp +++ b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp @@ -700,6 +700,122 @@ LogicalResult RcpOp::verify() { return success(); } +//===----------------------------------------------------------------------===// +// NVGPU_TruncfOp +//===----------------------------------------------------------------------===// + +static LogicalResult verifyConversionShapes(Operation *op, Type inType, + Type outType) { + bool srcIsVector = llvm::isa(inType); + bool dstIsVector = llvm::isa(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(inType); + auto dstVector = llvm::cast(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(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(srcType) && + !llvm::isa(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(dstType)) + return emitOpError("relu is not supported for bf16 destination"); + + return success(); +} + //===----------------------------------------------------------------------===// // TableGen'd dialect, type, and op definitions //===----------------------------------------------------------------------===// diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf-large.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf-large.mlir new file mode 100644 index 000000000000..9f2ddc5270c5 --- /dev/null +++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf-large.mlir @@ -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> +} diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf.mlir new file mode 100644 index 000000000000..5924f984f870 --- /dev/null +++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-extf.mlir @@ -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> +} diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf-large.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf-large.mlir new file mode 100644 index 000000000000..c5a66ccc3751 --- /dev/null +++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf-large.mlir @@ -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> +} diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf.mlir new file mode 100644 index 000000000000..d34402fcf1ce --- /dev/null +++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-truncf.mlir @@ -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 + // CHECK-SAME: : vector<2xf16> + // CHECK: nvvm.convert.f32x2.to.f16x2 + // CHECK-SAME: rnd = #nvvm.fp_rnd_mode + // 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 + // 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 + // 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 + // CHECK-SAME: : i16(f8E8M0FNU) + %out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode} + : 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 + // CHECK-SAME: : i16(f8E8M0FNU) + %out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode} + : 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 + %out = nvgpu.truncf %in {sat = #nvvm.sat_mode} + : 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 + %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 + %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} + : 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 + %out = nvgpu.truncf %in, %rbits {rnd = #nvvm.fp_rnd_mode} + : 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 + %out = nvgpu.truncf %in, %rbits {rnd = #nvvm.fp_rnd_mode} + : 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 + %out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode} + : 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 + %out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode} + : 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> +} diff --git a/mlir/test/Dialect/NVGPU/nvgpu-extf-invalid.mlir b/mlir/test/Dialect/NVGPU/nvgpu-extf-invalid.mlir new file mode 100644 index 000000000000..ef659efb96f8 --- /dev/null +++ b/mlir/test/Dialect/NVGPU/nvgpu-extf-invalid.mlir @@ -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}} + %out = nvgpu.extf %in {rnd = #nvvm.fp_rnd_mode} + : 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) { + // expected-error @+1 {{'nvgpu.extf' op operand #0 must be floating-point or vector of floating-point values, but got 'vector'}} + %out = nvgpu.extf %in : vector to vector + return +} diff --git a/mlir/test/Dialect/NVGPU/nvgpu-truncf-invalid.mlir b/mlir/test/Dialect/NVGPU/nvgpu-truncf-invalid.mlir new file mode 100644 index 000000000000..1b2fb3db5cec --- /dev/null +++ b/mlir/test/Dialect/NVGPU/nvgpu-truncf-invalid.mlir @@ -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}} + %out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode} + : 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} + : 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}} + %out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode} + : 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}} + %out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode} + : 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} + : 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} + : 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}} + %out = nvgpu.truncf %in {rnd = #nvvm.fp_rnd_mode} + : 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) { + // expected-error @+1 {{'nvgpu.truncf' op operand #0 must be floating-point or vector of floating-point values, but got 'vector'}} + %out = nvgpu.truncf %in : vector to vector + return +}