[mlir][amdgpu] lowerings for ScaledExtPacked816 #168123
Conversation
| if (is_fp8 && firstScaleLane == 1 && firstScaleByte == 2 && is_block_16) { | ||
| return 0b1101; | ||
| } | ||
| if (is_fp8 && firstScaleLane == 1 && firstScaleByte == 3 && !is_block_16) { |
There was a problem hiding this comment.
Is there any way to make this a bit less of an exhaustive enumeration of the cases?
| ScaledExtPacked816Op op, ScaledExtPacked816OpAdaptor adaptor, | ||
| ConversionPatternRewriter &rewriter) const { | ||
| | ||
| int firstScaleLane = op.getFirstScaleLane(); |
There was a problem hiding this comment.
Is int the right type here? Probably int64_t for index and int32_t otherwise.
| } else if (isa<Float6E3M2FNType>(srcElemType) && destElemType.isF16()) { | ||
| rewriter.replaceOpWithNewOp<ROCDL::CvtPkScalePk16F16Bf6Op>( | ||
| op, op.getResult().getType(), castedSource, castedScale, scaleSel); | ||
| } else if (isa<Float6E2M3FNType>(srcElemType) && destElemType.isBF16()) { |
There was a problem hiding this comment.
I think we could pull a page out of the MFMA book and have a general create that is relying on an OperationName, so that we don't have to repeat the ::create part N different times.
| // When lowering amdgpu.scaled_ext_packed816 to | ||
| // rocdl.cvt.scale.pk*.f*.f* operations, the | ||
| // attributes blockSize, sourceType, firstScaleLane and firstScaleByte | ||
| // are merged into a single attribute scaleSel. |
| | ||
| llvm_unreachable("invalid combination of firstScaleLane, firstScaleByte, " | ||
| "blockSize and type."); | ||
| return 0; |
| ConversionPatternRewriter &rewriter) const { | ||
| | ||
| int firstScaleLane = op.getFirstScaleLane(); |
There was a problem hiding this comment.
| ConversionPatternRewriter &rewriter) const { | |
| int firstScaleLane = op.getFirstScaleLane(); | |
| ConversionPatternRewriter &rewriter) const { | |
| int firstScaleLane = op.getFirstScaleLane(); |
nit: superfluous empty line
| Type packedType; | ||
| if (isa<Float4E2M1FNType>(srcElemType)) { | ||
| packedType = i32; | ||
| packedType = getTypeConverter()->convertType(packedType); |
There was a problem hiding this comment.
In the most general case, type conversion can fail -- are we checking this anywhere?
| Value castedSource = | ||
| LLVM::BitcastOp::create(rewriter, loc, packedType, source); | ||
| | ||
| if (isa<Float4E2M1FNType>(srcElemType) && destElemType.isF16()) { |
There was a problem hiding this comment.
nit: you can define type aliases for these, e.g.: using fp4 = Float4E2M1FNType to make the code more concise.
| return failure(); | ||
| } | ||
| | ||
| return success(); |
3 participants
Depends on #167369