Skip to content

Commit 47552dd

Browse files
maksleventalmakslevental
committed
address comments
1 parent c26179d commit 47552dd

File tree

8 files changed

+26
-43
lines changed

8 files changed

+26
-43
lines changed

‎include/triton/Conversion/MLIRTypes.h‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,19 +26,15 @@ inline Type f32Ty(MLIRContext *ctx) { return Float32Type::get(ctx); }
2626
inline Type f64Ty(MLIRContext *ctx) { return Float64Type::get(ctx); }
2727
inline Type bf16Ty(MLIRContext *ctx) { return BFloat16Type::get(ctx); }
2828

29+
inline bool isFloat8(Type type) {
30+
return isa<Float8E4M3B11FNUZType, Float8E4M3FNType, Float8E4M3FNUZType,
31+
Float8E5M2Type, Float8E5M2FNUZType>(type);
32+
}
33+
2934
inline bool isFloat(Type type) {
3035
return type.isF32() || type.isF64() || type.isF16() || type.isF128() ||
3136
type.isBF16() || llvm::isa<Float8E4M3B11FNUZType>(type) ||
32-
llvm::isa<Float8E4M3FNType>(type) ||
33-
llvm::isa<Float8E4M3FNUZType>(type) ||
34-
llvm::isa<Float8E5M2Type>(type) || llvm::isa<Float8E5M2FNUZType>(type);
35-
}
36-
37-
inline bool isFloat8(Type type) {
38-
return llvm::isa<Float8E4M3B11FNUZType>(type) ||
39-
llvm::isa<Float8E4M3FNType>(type) ||
40-
llvm::isa<Float8E4M3FNUZType>(type) ||
41-
llvm::isa<Float8E5M2Type>(type) || llvm::isa<Float8E5M2FNUZType>(type);
37+
isFloat8(type);
4238
}
4339

4440
inline bool isInt(Type type) { return type.isIntOrFloat() && !isFloat(type); }

‎lib/Analysis/Utility.cpp‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -750,15 +750,14 @@ bool supportMMA(triton::DotOp op, int version) {
750750
return false;
751751
if (!(numWarps % 4 == 0 && retShapePerCTA[rank - 2] % 64 == 0 &&
752752
retShapePerCTA[rank - 1] % 8 == 0 &&
753-
(llvm::isa<Float8E5M2Type>(aElemTy) ||
754-
llvm::isa<Float8E4M3FNType>(aElemTy) || aElemTy.isInteger(8) ||
755-
aElemTy.isF16() || aElemTy.isBF16() || aElemTy.isF32()))) {
753+
(llvm::isa<Float8E5M2Type, Float8E4M3FNType>(aElemTy) ||
754+
aElemTy.isInteger(8) || aElemTy.isF16() || aElemTy.isBF16() ||
755+
aElemTy.isF32()))) {
756756
return false;
757757
}
758758
// We cannot use MMA_V3 if we need to accumulate in F32 within the MMA op.
759759
if (op.getMaxNumImpreciseAcc() < 32 &&
760-
(llvm::isa<Float8E5M2Type>(aElemTy) ||
761-
llvm::isa<Float8E4M3FNType>(aElemTy)) &&
760+
(llvm::isa<Float8E5M2Type, Float8E4M3FNType>(aElemTy)) &&
762761
cast<RankedTensorType>(op.getType()).getElementType().isF32()) {
763762
return false;
764763
}
@@ -779,10 +778,8 @@ bool supportMMA(Value value, int version) {
779778
cast<triton::gpu::TensorOrMemDesc>(value.getType()).getElementType();
780779
// FP8 is not natively supported on all mma versions but it can always be
781780
// promoted to fp16 therefore we can always support it.
782-
bool isFP8 = llvm::isa<Float8E5M2Type>(elemTy) ||
783-
llvm::isa<Float8E4M3FNType>(elemTy) ||
784-
llvm::isa<Float8E5M2FNUZType>(elemTy) ||
785-
llvm::isa<Float8E4M3FNUZType>(elemTy);
781+
bool isFP8 = llvm::isa<Float8E5M2Type, Float8E4M3FNType, Float8E5M2FNUZType,
782+
Float8E4M3FNUZType>(elemTy);
786783
return isFP8 || elemTy.isF16() || elemTy.isBF16() ||
787784
(elemTy.isF32() && version >= 2) ||
788785
(elemTy.isInteger(8) && version >= 2);

‎lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -632,8 +632,7 @@ static void decomposeMixedModeDotOp(ModuleOp mod, int computeCapability) {
632632
NvidiaMmaEncodingAttr mmaLayout =
633633
dyn_cast<NvidiaMmaEncodingAttr>(D.getType().getEncoding());
634634
if (mmaLayout) {
635-
bool isNativeFP8 = llvm::isa<Float8E5M2Type>(AElType) ||
636-
llvm::isa<Float8E4M3FNType>(AElType);
635+
bool isNativeFP8 = llvm::isa<Float8E5M2Type, Float8E4M3FNType>(AElType);
637636
// promote operands for sm < 89 since fp8 mma is not natively supported
638637
// promote operands for sm >= 90 when mma is not v3
639638
if (!isNativeFP8 ||

‎lib/Dialect/TritonGPU/Transforms/Utility.cpp‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,10 +45,9 @@ SmallVector<unsigned, 3> mmaVersionToInstrShape(int version,
4545
SmallVector<unsigned> validN;
4646

4747
// MMAv3 with larger instruction shape is preferred.
48-
if (llvm::isa<Float8E5M2Type>(eltType) ||
49-
llvm::isa<Float8E4M3FNType>(eltType) ||
50-
llvm::isa<Float8E4M3FNUZType>(eltType) || eltType.isF16() ||
51-
eltType.isBF16() || eltType.isF32()) {
48+
if (llvm::isa<Float8E5M2Type, Float8E4M3FNType, Float8E4M3FNUZType>(
49+
eltType) ||
50+
eltType.isF16() || eltType.isBF16() || eltType.isF32()) {
5251
validN.assign({256, 248, 240, 232, 224, 216, 208, 200, 192, 184, 176,
5352
168, 160, 152, 144, 136, 128, 120, 112, 104, 96, 88,
5453
80, 72, 64, 56, 48, 40, 32, 24, 16, 8});

‎lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -77,10 +77,8 @@ bool WarpGroupDotOp::needsPartialAccumulator() {
7777
const auto &d = getD();
7878
auto aTensorTy = cast<triton::gpu::TensorOrMemDesc>(a.getType());
7979
auto aElTy = cast<triton::gpu::TensorOrMemDesc>(a.getType()).getElementType();
80-
bool isFP8 = llvm::isa<Float8E5M2Type>(aElTy) ||
81-
llvm::isa<Float8E4M3FNType>(aElTy) ||
82-
llvm::isa<Float8E5M2FNUZType>(aElTy) ||
83-
llvm::isa<Float8E4M3FNUZType>(aElTy);
80+
bool isFP8 = llvm::isa<Float8E5M2Type, Float8E4M3FNType, Float8E5M2FNUZType,
81+
Float8E4M3FNUZType>(aElTy);
8482
bool accFP32 =
8583
cast<triton::gpu::TensorOrMemDesc>(d.getType()).getElementType().isF32();
8684
uint32_t maxNumImpreciseAcc = getMaxNumImpreciseAcc();

‎third_party/amd/lib/TritonAMDGPUToLLVM/ElementwiseOpToLLVM.cpp‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1043,19 +1043,16 @@ struct FpToFpOpConversion
10431043
return outVals;
10441044
}
10451045
size_t numElements = 4;
1046-
if (llvm::isa<Float8E4M3FNType>(srcElementType) ||
1047-
llvm::isa<Float8E4M3FNType>(dstElementType) ||
1048-
llvm::isa<Float8E4M3FNUZType>(srcElementType) ||
1049-
llvm::isa<Float8E4M3FNUZType>(dstElementType) ||
1050-
llvm::isa<Float8E5M2FNUZType>(srcElementType) ||
1051-
llvm::isa<Float8E5M2FNUZType>(dstElementType)) {
1046+
if (llvm::isa<Float8E4M3FNType, Float8E4M3FNUZType, Float8E5M2FNUZType>(
1047+
srcElementType) ||
1048+
llvm::isa<Float8E4M3FNType, Float8E4M3FNUZType, Float8E5M2FNUZType>(
1049+
dstElementType)) {
10521050
numElements = 2;
10531051
}
10541052
bool useFP16IntermediateSrc =
10551053
srcElementType.isF32() &&
10561054
!(isaFamily == AMD::ISAFamily::CDNA3 &&
1057-
(llvm::isa<Float8E4M3FNUZType>(dstElementType) ||
1058-
llvm::isa<Float8E5M2FNUZType>(dstElementType)));
1055+
(llvm::isa<Float8E4M3FNUZType, Float8E5M2FNUZType>(dstElementType)));
10591056
bool isDstFP32 = dstElementType.isF32();
10601057
Type srcType = useFP16IntermediateSrc ? f16_ty : srcElementType;
10611058
Type dstType = isDstFP32 ? f16_ty : dstElementType;

‎third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -416,8 +416,7 @@ class BlockedToMFMA : public OpRewritePattern<tt::DotOp> {
416416
// store instructions, except for fp8 matmul kernels due to regression
417417
// TODO (lixun): investigate the regression and enable this feature again
418418
auto aElemTy = mfmaInstr.getElementTypeA();
419-
bool isFP8 = llvm::isa<Float8E5M2FNUZType>(aElemTy) ||
420-
llvm::isa<Float8E4M3FNUZType>(aElemTy);
419+
bool isFP8 = llvm::isa<Float8E5M2FNUZType, Float8E4M3FNUZType>(aElemTy);
421420
bool isTransposed = isChainDot(dotOp) || !isFP8;
422421
mfmaEnc = ttg::AMDMfmaEncodingAttr::get(
423422
oldRetType.getContext(),

‎third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -475,8 +475,7 @@ struct FpToFpOpConversion
475475
auto dstElementType = getElementType(op.getResult());
476476
auto roundingMode = op.getRounding();
477477

478-
if (llvm::isa<Float8E5M2Type>(dstElementType) ||
479-
llvm::isa<Float8E4M3FNType>(dstElementType)) {
478+
if (llvm::isa<Float8E5M2Type, Float8E4M3FNType>(dstElementType)) {
480479
assert(roundingMode.has_value() &&
481480
"Rounding mode must be specified for convertsions to fp8");
482481

@@ -514,8 +513,7 @@ struct FpToFpOpConversion
514513
bool useFP16IntermediateSrc =
515514
srcElementType.isF32() &&
516515
(!(computeCapability >= 90 &&
517-
(llvm::isa<Float8E4M3FNType>(dstElementType) ||
518-
llvm::isa<Float8E5M2Type>(dstElementType))) ||
516+
(llvm::isa<Float8E4M3FNType, Float8E5M2Type>(dstElementType))) ||
519517
roundingMode.value() == RoundingMode::RTZ);
520518
bool isDstFP32 = dstElementType.isF32();
521519
Type srcType = useFP16IntermediateSrc ? f16_ty : srcElementType;

0 commit comments

Comments
 (0)