Skip to content

[mlir][gpu] Add extra value types for gpu::ShuffleOp #104605

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Aug 20, 2024
Merged

[mlir][gpu] Add extra value types for gpu::ShuffleOp #104605

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
482035c
[mlir] Add extra value types for gpu.shuffle op
FMarno Aug 16, 2024
7041ef7
fixup! [mlir] Add extra value types for gpu.shuffle op
FMarno Aug 19, 2024
7623fea
fixup! [mlir] Add extra value types for gpu.shuffle op
FMarno Aug 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1250,17 +1250,11 @@ def GPU_ShuffleMode : I32EnumAttr<"ShuffleMode",
def GPU_ShuffleModeAttr : EnumAttr<GPU_Dialect, GPU_ShuffleMode,
"shuffle_mode">;

def I32I64F32OrF64 : TypeConstraint<Or<[I32.predicate,
I64.predicate,
F32.predicate,
F64.predicate]>,
"i32, i64, f32 or f64">;

def GPU_ShuffleOp : GPU_Op<
"shuffle", [Pure, AllTypesMatch<["value", "shuffleResult"]>]>,
Arguments<(ins I32I64F32OrF64:$value, I32:$offset, I32:$width,
Arguments<(ins AnyIntegerOrFloatOr1DVector:$value, I32:$offset, I32:$width,
GPU_ShuffleModeAttr:$mode)>,
Results<(outs I32I64F32OrF64:$shuffleResult, I1:$valid)> {
Results<(outs AnyIntegerOrFloatOr1DVector:$shuffleResult, I1:$valid)> {
let summary = "Shuffles values within a subgroup.";
let description = [{
The "shuffle" op moves values to a across lanes (a.k.a., invocations,
Expand Down
30 changes: 20 additions & 10 deletions mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,26 +241,34 @@ struct GPUShuffleConversion final : ConvertOpToLLVMPattern<gpu::ShuffleOp> {
llvm_unreachable("Unhandled shuffle mode");
}

static StringRef getTypeMangling(Type type) {
return TypeSwitch<Type, StringRef>(type)
static std::optional<StringRef> getTypeMangling(Type type) {
return TypeSwitch<Type, std::optional<StringRef>>(type)
.Case<Float16Type>([](auto) { return "Dhj"; })
.Case<Float32Type>([](auto) { return "fj"; })
.Case<Float64Type>([](auto) { return "dj"; })
.Case<IntegerType>([](auto intTy) {
.Case<IntegerType>([](auto intTy) -> std::optional<StringRef> {
switch (intTy.getWidth()) {
case 8:
return "cj";
case 16:
return "sj";
case 32:
return "ij";
case 64:
return "lj";
}
llvm_unreachable("Invalid integer width");
});
return std::nullopt;
})
.Default([](auto) { return std::nullopt; });
}

static std::string getFuncName(gpu::ShuffleOp op) {
static std::optional<std::string> getFuncName(gpu::ShuffleOp op) {
StringRef baseName = getBaseName(op.getMode());
StringRef typeMangling = getTypeMangling(op.getType(0));
std::optional<StringRef> typeMangling = getTypeMangling(op.getType(0));
if (!typeMangling)
return std::nullopt;
return llvm::formatv("_Z{0}{1}{2}", baseName.size(), baseName,
typeMangling);
typeMangling.value());
}

/// Get the subgroup size from the target or return a default.
Expand All @@ -284,15 +292,17 @@ struct GPUShuffleConversion final : ConvertOpToLLVMPattern<gpu::ShuffleOp> {
return rewriter.notifyMatchFailure(
op, "shuffle width and subgroup size mismatch");

std::string funcName = getFuncName(op);
std::optional<std::string> funcName = getFuncName(op);
if (!funcName)
return rewriter.notifyMatchFailure(op, "unsupported value type");

Operation *moduleOp = op->getParentWithTrait<OpTrait::SymbolTable>();
assert(moduleOp && "Expecting module");
Type valueType = adaptor.getValue().getType();
Type offsetType = adaptor.getOffset().getType();
Type resultType = valueType;
LLVM::LLVMFuncOp func = lookupOrCreateSPIRVFn(
moduleOp, funcName, {valueType, offsetType}, resultType,
moduleOp, funcName.value(), {valueType, offsetType}, resultType,
/*isMemNone=*/false, /*isConvergent=*/true);

Location loc = op->getLoc();
Expand Down
74 changes: 61 additions & 13 deletions mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,6 +317,12 @@ gpu.module @shuffles attributes {
// CHECK-SAME-DAG: will_return
// CHECK-NOT: memory_effects = #llvm.memory_effects
// CHECK-SAME: }
// CHECK: llvm.func spir_funccc @_Z20sub_group_shuffle_upDhj(f16, i32) -> f16 attributes {
// CHECK-SAME-DAG: no_unwind
// CHECK-SAME-DAG: convergent
// CHECK-SAME-DAG: will_return
// CHECK-NOT: memory_effects = #llvm.memory_effects
// CHECK-SAME: }
// CHECK: llvm.func spir_funccc @_Z21sub_group_shuffle_xorlj(i64, i32) -> i64 attributes {
// CHECK-SAME-DAG: no_unwind
// CHECK-SAME-DAG: convergent
Expand All @@ -329,26 +335,54 @@ gpu.module @shuffles attributes {
// CHECK-SAME-DAG: will_return
// CHECK-NOT: memory_effects = #llvm.memory_effects
// CHECK-SAME: }
// CHECK: llvm.func spir_funccc @_Z21sub_group_shuffle_xorsj(i16, i32) -> i16 attributes {
// CHECK-SAME-DAG: no_unwind
// CHECK-SAME-DAG: convergent
// CHECK-SAME-DAG: will_return
// CHECK-NOT: memory_effects = #llvm.memory_effects
// CHECK-SAME: }
// CHECK: llvm.func spir_funccc @_Z17sub_group_shufflecj(i8, i32) -> i8 attributes {
// CHECK-SAME-DAG: no_unwind
// CHECK-SAME-DAG: convergent
// CHECK-SAME-DAG: will_return
// CHECK-NOT: memory_effects = #llvm.memory_effects
// CHECK-SAME: }

// CHECK-LABEL: gpu_shuffles
// CHECK-SAME: (%[[VAL_0:.*]]: i32, %[[VAL_1:.*]]: i32, %[[VAL_2:.*]]: i64, %[[VAL_3:.*]]: i32, %[[VAL_4:.*]]: f32, %[[VAL_5:.*]]: i32, %[[VAL_6:.*]]: f64, %[[VAL_7:.*]]: i32)
func.func @gpu_shuffles(%val0: i32, %id: i32,
%val1: i64, %mask: i32,
%val2: f32, %delta_up: i32,
%val3: f64, %delta_down: i32) {
// CHECK-SAME: (%[[I8_VAL:.*]]: i8, %[[I16_VAL:.*]]: i16,
// CHECK-SAME: %[[I32_VAL:.*]]: i32, %[[I64_VAL:.*]]: i64,
// CHECK-SAME: %[[F16_VAL:.*]]: f16, %[[F32_VAL:.*]]: f32,
// CHECK-SAME: %[[F64_VAL:.*]]: f64, %[[OFFSET:.*]]: i32) {
func.func @gpu_shuffles(%i8_val: i8,
%i16_val: i16,
%i32_val: i32,
%i64_val: i64,
%f16_val: f16,
%f32_val: f32,
%f64_val: f64,
%offset: i32) {
%width = arith.constant 16 : i32
// CHECK: llvm.call spir_funccc @_Z17sub_group_shuffleij(%[[VAL_0]], %[[VAL_1]])
// CHECK: llvm.call spir_funccc @_Z17sub_group_shufflecj(%[[I8_VAL]], %[[OFFSET]])
// CHECK: llvm.mlir.constant(true) : i1
// CHECK: llvm.call spir_funccc @_Z21sub_group_shuffle_xorlj(%[[VAL_2]], %[[VAL_3]])
// CHECK: llvm.call spir_funccc @_Z21sub_group_shuffle_xorsj(%[[I16_VAL]], %[[OFFSET]])
// CHECK: llvm.mlir.constant(true) : i1
// CHECK: llvm.call spir_funccc @_Z20sub_group_shuffle_upfj(%[[VAL_4]], %[[VAL_5]])
// CHECK: llvm.call spir_funccc @_Z17sub_group_shuffleij(%[[I32_VAL]], %[[OFFSET]])
// CHECK: llvm.mlir.constant(true) : i1
// CHECK: llvm.call spir_funccc @_Z22sub_group_shuffle_downdj(%[[VAL_6]], %[[VAL_7]])
// CHECK: llvm.call spir_funccc @_Z21sub_group_shuffle_xorlj(%[[I64_VAL]], %[[OFFSET]])
// CHECK: llvm.mlir.constant(true) : i1
%shuffleResult0, %valid0 = gpu.shuffle idx %val0, %id, %width : i32
%shuffleResult1, %valid1 = gpu.shuffle xor %val1, %mask, %width : i64
%shuffleResult2, %valid2 = gpu.shuffle up %val2, %delta_up, %width : f32
%shuffleResult3, %valid3 = gpu.shuffle down %val3, %delta_down, %width : f64
// CHECK: llvm.call spir_funccc @_Z20sub_group_shuffle_upDhj(%[[F16_VAL]], %[[OFFSET]])
// CHECK: llvm.mlir.constant(true) : i1
// CHECK: llvm.call spir_funccc @_Z20sub_group_shuffle_upfj(%[[F32_VAL]], %[[OFFSET]])
// CHECK: llvm.mlir.constant(true) : i1
// CHECK: llvm.call spir_funccc @_Z22sub_group_shuffle_downdj(%[[F64_VAL]], %[[OFFSET]])
// CHECK: llvm.mlir.constant(true) : i1
%shuffleResult0, %valid0 = gpu.shuffle idx %i8_val, %offset, %width : i8
%shuffleResult1, %valid1 = gpu.shuffle xor %i16_val, %offset, %width : i16
%shuffleResult2, %valid2 = gpu.shuffle idx %i32_val, %offset, %width : i32
%shuffleResult3, %valid3 = gpu.shuffle xor %i64_val, %offset, %width : i64
%shuffleResult4, %valid4 = gpu.shuffle up %f16_val, %offset, %width : f16
%shuffleResult5, %valid5 = gpu.shuffle up %f32_val, %offset, %width : f32
%shuffleResult6, %valid6 = gpu.shuffle down %f64_val, %offset, %width : f64
return
}
}
Expand Down Expand Up @@ -378,6 +412,20 @@ gpu.module @shuffles_mismatch {
}
}

// -----

// Cannot convert due to value type not being supported by the conversion

gpu.module @not_supported_lowering {
func.func @gpu_shuffles(%val: i1, %id: i32) {
%width = arith.constant 32 : i32
// expected-error@below {{failed to legalize operation 'gpu.shuffle' that was explicitly marked illegal}}
%shuffleResult, %valid = gpu.shuffle xor %val, %id, %width : i1
return
}
}


// -----

gpu.module @kernels {
Expand Down
2 changes: 1 addition & 1 deletion mlir/test/Dialect/GPU/invalid.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -430,7 +430,7 @@ func.func @shuffle_mismatching_type(%arg0 : f32, %arg1 : i32, %arg2 : i32) {
// -----

func.func @shuffle_unsupported_type(%arg0 : index, %arg1 : i32, %arg2 : i32) {
// expected-error@+1 {{operand #0 must be i32, i64, f32 or f64}}
// expected-error@+1 {{op operand #0 must be Integer or Float or vector of Integer or Float values of ranks 1, but got 'index'}}
%shfl, %pred = gpu.shuffle xor %arg0, %arg1, %arg2 : index
return
}
Expand Down
Loading