Skip to content

Commit 77ac61c

Browse files
stefankoncarevicstefankoncarevic
committed
[mlir][AMDGPU] Implement AMDGPU DPP operation in MLIR.
Defined AMDGPU DPP operation in mlir to represent semantics. Introduced a new enumeration attribute for different permutations and allowed for different types of arguments.Implemented constant attribute handling for ROCDL::DPPUpdateOp operation. The operation now correctly accepts constant attributes for dppCtrl, rowMask, bankMask, boundCtrl, and passes them to the corresponding LLVM intrinsic.
1 parent 220756f commit 77ac61c

File tree

6 files changed

+438
-2
lines changed

6 files changed

+438
-2
lines changed

‎llvm/include/llvm/IR/IntrinsicsAMDGPU.td

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2411,6 +2411,8 @@ def int_amdgcn_buffer_wbinvl1_vol :
24112411
// VI Intrinsics
24122412
//===----------------------------------------------------------------------===//
24132413

2414+
// The llvm.amdgcn.mov.dpp.i32 intrinsic represents the mov.dpp operation in AMDGPU.
2415+
// This operation is being deprecated and can be replaced with llvm.amdgcn.update.dpp.i32.
24142416
// llvm.amdgcn.mov.dpp.i32 <src> <dpp_ctrl> <row_mask> <bank_mask> <bound_ctrl>
24152417
def int_amdgcn_mov_dpp :
24162418
Intrinsic<[llvm_anyint_ty],
@@ -2420,6 +2422,10 @@ def int_amdgcn_mov_dpp :
24202422
ImmArg<ArgIndex<1>>, ImmArg<ArgIndex<2>>,
24212423
ImmArg<ArgIndex<3>>, ImmArg<ArgIndex<4>>, IntrNoCallback, IntrNoFree]>;
24222424

2425+
// The llvm.amdgcn.update.dpp.i32 intrinsic represents the update.dpp operation in AMDGPU.
2426+
// It takes an old value, a source operand, a DPP control operand, a row mask, a bank mask, and a bound control.
2427+
// This operation is equivalent to a sequence of v_mov_b32 operations.
2428+
// It is preferred over llvm.amdgcn.mov.dpp.i32 for future use.
24232429
// llvm.amdgcn.update.dpp.i32 <old> <src> <dpp_ctrl> <row_mask> <bank_mask> <bound_ctrl>
24242430
// Should be equivalent to:
24252431
// v_mov_b32 <dest> <old>

‎mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -410,6 +410,61 @@ def AMDGPU_RawBufferAtomicUminOp :
410410
let hasVerifier = 1;
411411
}
412412

413+
def AMDGPU_DPPPerm : I32EnumAttr<"DPPPerm",
414+
"The possible permutations for a DPP operation",
415+
[
416+
I32EnumAttrCase<"quad_perm", 0>,
417+
I32EnumAttrCase<"row_shl", 1>,
418+
I32EnumAttrCase<"row_shr", 2>,
419+
I32EnumAttrCase<"row_ror", 3>,
420+
I32EnumAttrCase<"wave_shl", 4>,
421+
I32EnumAttrCase<"wave_shr", 5>,
422+
I32EnumAttrCase<"wave_ror", 6>,
423+
I32EnumAttrCase<"wave_rol", 7>,
424+
I32EnumAttrCase<"row_mirror", 8>,
425+
I32EnumAttrCase<"row_half_mirror", 9>,
426+
I32EnumAttrCase<"row_bcast_15", 10>,
427+
I32EnumAttrCase<"row_bcast_31", 11>
428+
]> {
429+
let genSpecializedAttr = 0;
430+
let cppNamespace = "::mlir::amdgpu";
431+
}
432+
433+
def AMDGPU_DPPPermAttr : EnumAttr<AMDGPU_Dialect, AMDGPU_DPPPerm,
434+
"dpp_perm">;
435+
436+
def AMDGPU_DPPOp : AMDGPU_Op<"dpp", [SameTypeOperands, AllTypesMatch<["result", "src", "old"]>]>,
437+
Arguments<(ins AnyType:$src,
438+
AnyType:$old,
439+
AMDGPU_DPPPermAttr:$kind,
440+
OptionalAttr<AnyAttrOf<[I32Attr, ArrayAttr, UnitAttr]>>:$permArgument,
441+
DefaultValuedAttr<I32Attr, "0xf">:$row_mask,
442+
DefaultValuedAttr<I32Attr, "0xf">:$bank_mask,
443+
DefaultValuedAttr<BoolAttr, "false">:$bound_ctrl)> {
444+
let summary = "AMDGPU DPP operation";
445+
let description = [{
446+
This operation represents DPP functionality in a GPU program.
447+
DPP provides the following operations:
448+
- Full crossbar in a group of four (`quad_perm`)
449+
- Wavefront shift left by one lane (`wave_shl`)
450+
- Wavefront shift right by one lane (`wave_shr`)
451+
- Wavefront rotate right by one lane (`wave_ror`)
452+
- Wavefront rotate left by one lane (`wave_rol`)
453+
- Row shift left by 1–15 lanes (`row_shl`)
454+
- Row shift right by 1–15 lanes (`row_shr`)
455+
- Row rotate right by 1–15 lanes (`row_ror`)
456+
- Reverse within a row (`row_mirror`)
457+
- Reverse within a half-row (`row_half_mirror`)
458+
- Broadcast the 15th lane of each row to the next row (`row_bcast`)
459+
- Broadcast lane 31 to rows 2 and 3 (`row_bcast`)
460+
}];
461+
let results = (outs AnyType:$result);
462+
let assemblyFormat = [{
463+
$src $old $kind (`(` $permArgument^ `)`)? attr-dict `:` type($result)
464+
}];
465+
let hasVerifier = 1;
466+
}
467+
413468
def AMDGPU_LDSBarrierOp : AMDGPU_Op<"lds_barrier"> {
414469
let summary = "Barrier that includes a wait for LDS memory operations.";
415470
let description = [{

‎mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -591,6 +591,30 @@ def ROCDL_RawBufferAtomicUMinOp :
591591
let hasCustomAssemblyFormat = 1;
592592
}
593593

594+
// DPP Update intrinsic
595+
def ROCDL_DPPUpdateOp : ROCDL_IntrOp<"update.dpp", [], [0],
596+
[AllTypesMatch<["res", "src", "old"]>], 1>,
597+
Arguments<(ins LLVM_Type:$old, LLVM_Type:$src, I32Attr:$dppCtrl, I32Attr:$rowMask,
598+
I32Attr:$bankMask, I1Attr:$boundCtrl)> {
599+
let results = (outs LLVM_Type:$res);
600+
let assemblyFormat = [{
601+
attr-dict $old `,` $src `with` $dppCtrl `,` $rowMask `,` $bankMask `,` $boundCtrl `:` type($src)
602+
}];
603+
string llvmBuilder = [{
604+
auto vdataType = moduleTranslation.convertType(op.getSrc().getType());
605+
llvm::Value *args[] = {
606+
moduleTranslation.lookupValue(op.getOld()),
607+
moduleTranslation.lookupValue(op.getSrc()),
608+
builder.getInt32(op.getDppCtrl()),
609+
builder.getInt32(op.getRowMask()),
610+
builder.getInt32(op.getBankMask()),
611+
builder.getInt1(op.getBoundCtrl())
612+
};
613+
$res = createIntrinsicCall(builder,
614+
llvm::Intrinsic::amdgcn_update_dpp, args, {vdataType});
615+
}];
616+
}
617+
594618
//===---------------------------------------------------------------------===//
595619
// 8-bit float intrinsics
596620
//===---------------------------------------------------------------------===//

‎mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

Lines changed: 143 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -828,6 +828,147 @@ LogicalResult PackedStochRoundFp8OpLowering::matchAndRewrite(
828828
return success();
829829
}
830830

831+
// Implement the AMDGPU_DPPLowering class that will convert the amdgpu.dpp
832+
// operation into the corresponding ROCDL instructions.
833+
struct AMDGPUDPPLowering : public ConvertOpToLLVMPattern<DPPOp> {
834+
AMDGPUDPPLowering(LLVMTypeConverter &converter, Chipset chipset)
835+
: ConvertOpToLLVMPattern<DPPOp>(converter), chipset(chipset) {}
836+
Chipset chipset;
837+
838+
LogicalResult
839+
matchAndRewrite(DPPOp DppOp, DPPOp::Adaptor adaptor,
840+
ConversionPatternRewriter &rewriter) const override {
841+
842+
// Convert the source operand to the corresponding LLVM type
843+
Location loc = DppOp.getLoc();
844+
Value src = adaptor.getSrc();
845+
Value old = adaptor.getOld();
846+
Type srcType = src.getType();
847+
Type oldType = old.getType();
848+
auto llvmI32Type = typeConverter->convertType(rewriter.getI32Type());
849+
auto llvmSrcIntType = typeConverter->convertType(
850+
rewriter.getIntegerType(srcType.getIntOrFloatBitWidth()));
851+
852+
// If the source type is less or equal to i32 or f32, use bitcast to convert
853+
// it to i32.
854+
auto convertOperand = [&](Value operand, Type operandType) {
855+
if (llvm::isa<FloatType>(operandType)) {
856+
operand =
857+
rewriter.create<LLVM::BitcastOp>(loc, llvmSrcIntType, operand);
858+
}
859+
860+
if (operandType.getIntOrFloatBitWidth() < 32) {
861+
auto llvmVecType = typeConverter->convertType(mlir::VectorType::get(
862+
32 / operandType.getIntOrFloatBitWidth(), llvmSrcIntType));
863+
Value undefVec = rewriter.create<LLVM::UndefOp>(loc, llvmVecType);
864+
operand = rewriter.create<LLVM::InsertElementOp>(
865+
loc, undefVec, operand, createI32Constant(rewriter, loc, 0));
866+
operand = rewriter.create<LLVM::BitcastOp>(loc, llvmI32Type, operand);
867+
}
868+
return operand;
869+
};
870+
871+
src = convertOperand(src, srcType);
872+
old = convertOperand(old, oldType);
873+
874+
// This is taken from the following file llvm/lib/Target/AMDGPU/SIDefines.h
875+
enum DppCtrl : unsigned {
876+
ROW_SHL0 = 0x100,
877+
ROW_SHR0 = 0x110,
878+
ROW_ROR0 = 0x120,
879+
WAVE_SHL1 = 0x130,
880+
WAVE_ROL1 = 0x134,
881+
WAVE_SHR1 = 0x138,
882+
WAVE_ROR1 = 0x13C,
883+
ROW_MIRROR = 0x140,
884+
ROW_HALF_MIRROR = 0x141,
885+
BCAST15 = 0x142,
886+
BCAST31 = 0x143,
887+
};
888+
889+
auto kind = DppOp.getKind();
890+
auto permArgument = DppOp.getPermArgument();
891+
uint32_t DppCtrl = 0;
892+
893+
switch (kind) {
894+
895+
case DPPPerm::quad_perm:
896+
if (auto quadPermAttr = cast<ArrayAttr>(*permArgument)) {
897+
int32_t i = 0;
898+
for (auto elem : quadPermAttr.getAsRange<IntegerAttr>()) {
899+
uint32_t num = elem.getInt();
900+
DppCtrl |= num << (i * 2);
901+
i++;
902+
}
903+
}
904+
break;
905+
case DPPPerm::row_shl:
906+
if (auto intAttr = cast<IntegerAttr>(*permArgument)) {
907+
DppCtrl = intAttr.getInt() + DppCtrl::ROW_SHL0;
908+
}
909+
break;
910+
case DPPPerm::row_shr:
911+
if (auto intAttr = cast<IntegerAttr>(*permArgument)) {
912+
DppCtrl = intAttr.getInt() + DppCtrl::ROW_SHR0;
913+
}
914+
break;
915+
case DPPPerm::row_ror:
916+
if (auto intAttr = cast<IntegerAttr>(*permArgument)) {
917+
DppCtrl = intAttr.getInt() + DppCtrl::ROW_ROR0;
918+
}
919+
break;
920+
case DPPPerm::wave_shl:
921+
DppCtrl = DppCtrl::WAVE_SHL1;
922+
break;
923+
case DPPPerm::wave_shr:
924+
DppCtrl = DppCtrl::WAVE_SHR1;
925+
break;
926+
case DPPPerm::wave_rol:
927+
DppCtrl = DppCtrl::WAVE_ROL1;
928+
break;
929+
case DPPPerm::wave_ror:
930+
DppCtrl = DppCtrl::WAVE_ROR1;
931+
break;
932+
case DPPPerm::row_mirror:
933+
DppCtrl = DppCtrl::ROW_MIRROR;
934+
break;
935+
case DPPPerm::row_half_mirror:
936+
DppCtrl = DppCtrl::ROW_HALF_MIRROR;
937+
break;
938+
case DPPPerm::row_bcast_15:
939+
DppCtrl = DppCtrl::BCAST15;
940+
break;
941+
case DPPPerm::row_bcast_31:
942+
DppCtrl = DppCtrl::BCAST31;
943+
break;
944+
}
945+
946+
// Check for row_mask, bank_mask, bound_ctrl if they exist and create
947+
// constants
948+
auto rowMask = DppOp->getAttrOfType<IntegerAttr>("row_mask").getInt();
949+
auto bankMask = DppOp->getAttrOfType<IntegerAttr>("bank_mask").getInt();
950+
bool boundCtrl = DppOp->getAttrOfType<BoolAttr>("bound_ctrl").getValue();
951+
952+
// create a ROCDL_DPPMovOp instruction with the appropriate attributes
953+
auto dppMovOp = rewriter.create<ROCDL::DPPUpdateOp>(
954+
loc, llvmI32Type, old, src, DppCtrl, rowMask, bankMask, boundCtrl);
955+
956+
Value result = dppMovOp.getRes();
957+
if (srcType.getIntOrFloatBitWidth() < 32) {
958+
result = rewriter.create<LLVM::TruncOp>(loc, llvmSrcIntType, result);
959+
}
960+
961+
if (!llvm::isa<IntegerType>(srcType)) {
962+
result = rewriter.create<LLVM::BitcastOp>(loc, srcType, result);
963+
}
964+
965+
// We are replacing the AMDGPU_DPPOp instruction with the new
966+
// ROCDL_DPPMovOp instruction
967+
rewriter.replaceOp(DppOp, ValueRange(result));
968+
return success();
969+
}
970+
};
971+
831972
struct ConvertAMDGPUToROCDLPass
832973
: public impl::ConvertAMDGPUToROCDLBase<ConvertAMDGPUToROCDLPass> {
833974
ConvertAMDGPUToROCDLPass() = default;
@@ -879,8 +1020,8 @@ void mlir::populateAMDGPUToROCDLConversionPatterns(LLVMTypeConverter &converter,
8791020
ROCDL::RawPtrBufferAtomicUminOp>,
8801021
RawBufferOpLowering<RawBufferAtomicCmpswapOp,
8811022
ROCDL::RawPtrBufferAtomicCmpSwap>,
882-
LDSBarrierOpLowering, MFMAOpLowering, WMMAOpLowering,
883-
ExtPackedFp8OpLowering, PackedTrunc2xFp8OpLowering,
1023+
AMDGPUDPPLowering, LDSBarrierOpLowering, MFMAOpLowering,
1024+
WMMAOpLowering, ExtPackedFp8OpLowering, PackedTrunc2xFp8OpLowering,
8841025
PackedStochRoundFp8OpLowering>(converter, chipset);
8851026
}
8861027

‎mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -326,6 +326,68 @@ LogicalResult MFMAOp::verify() {
326326
return success();
327327
}
328328

329+
//===----------------------------------------------------------------------===//
330+
// DPPOp
331+
//===----------------------------------------------------------------------===//
332+
LogicalResult DPPOp::verify() {
333+
Type srcType = getSrc().getType();
334+
if (srcType.getIntOrFloatBitWidth() > 32) {
335+
return emitOpError("integer and floating point types larger than 32 bits "
336+
"are not supported");
337+
}
338+
339+
DPPPerm kind = getKind();
340+
Attribute permArgument = getPermArgument().value_or(Attribute{});
341+
342+
switch (kind) {
343+
344+
case DPPPerm::quad_perm: {
345+
auto quadPermAttr = dyn_cast_or_null<ArrayAttr>(permArgument);
346+
if (!quadPermAttr || quadPermAttr.size() != 4) {
347+
return emitOpError("quad_perm attribute must have exactly 4 elements");
348+
}
349+
for (auto elem : quadPermAttr.getAsRange<IntegerAttr>()) {
350+
uint32_t num = elem.getInt();
351+
if (num < 0 || num > 3) {
352+
return emitOpError(
353+
"Each element of quad_perm must be in the range [0, 3]");
354+
}
355+
}
356+
} break;
357+
358+
case DPPPerm::row_shl:
359+
case DPPPerm::row_shr:
360+
case DPPPerm::row_ror: {
361+
if (!permArgument) {
362+
return emitOpError("Attribute '" + Twine(stringifyDPPPerm(kind)) +
363+
"' value not specified");
364+
}
365+
if (auto intAttr = dyn_cast<IntegerAttr>(permArgument)) {
366+
uint32_t attrValue = intAttr.getInt();
367+
if (attrValue < 1 || attrValue > 15) {
368+
return emitOpError("Attribute value must be between 1 and 15");
369+
}
370+
}
371+
} break;
372+
373+
case DPPPerm::wave_shl:
374+
case DPPPerm::wave_shr:
375+
case DPPPerm::wave_rol:
376+
case DPPPerm::wave_ror:
377+
case DPPPerm::row_mirror:
378+
case DPPPerm::row_half_mirror:
379+
case DPPPerm::row_bcast_15:
380+
case DPPPerm::row_bcast_31: {
381+
if (permArgument && !isa<UnitAttr>(permArgument)) {
382+
return emitOpError("Expected unit attribute for permArgument, but found "
383+
"non-trivial argument");
384+
}
385+
break;
386+
}
387+
}
388+
return success();
389+
}
390+
329391
#include "mlir/Dialect/AMDGPU/IR/AMDGPUEnums.cpp.inc"
330392

331393
#define GET_ATTRDEF_CLASSES

0 commit comments

Comments
 (0)