ksgr5566 commented on code in PR #18823:
URL: https://github.com/apache/tvm/pull/18823#discussion_r2870109063
##########
src/target/target_kind.cc:
##########
@@ -424,8 +424,27 @@ TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan)
// Tags
.set_default_keys({"vulkan", "gpu"});
+/*!
+ * \brief Update WebGPU target attributes based on subgroup support.
+ * When supports_subgroups is true, set thread_warp_size to 32 so that
+ * TIR lowering uses warp-level shuffle reductions instead of shared memory.
+ */
+TargetJSON UpdateWebGPUAttrs(TargetJSON target) {
+ if (target.count("supports_subgroups")) {
+ bool subgroups = Downcast<Bool>(target.at("supports_subgroups"));
+ if (subgroups) {
+ target.Set("thread_warp_size", int64_t(32));
+ }
+ }
+ return target;
+}
+
TVM_REGISTER_TARGET_KIND("webgpu", kDLWebGPU)
.add_attr_option<int64_t>("max_num_threads", 256)
+ .add_attr_option<bool>("supports_subgroups", false)
+ // thread_warp_size=1: is_subwarp_reduction and is_multiwarp_reduction
returns false, so no subgroup ops are emitted.
+ .add_attr_option<int64_t>("thread_warp_size", 1)
+ .set_target_parser(UpdateWebGPUAttrs)
.set_default_keys({"webgpu", "gpu"});
Review Comment:
added
##########
src/target/target_kind.cc:
##########
@@ -424,8 +424,27 @@ TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan)
// Tags
.set_default_keys({"vulkan", "gpu"});
+/*!
+ * \brief Update WebGPU target attributes based on subgroup support.
+ * When supports_subgroups is true, set thread_warp_size to 32 so that
+ * TIR lowering uses warp-level shuffle reductions instead of shared memory.
+ */
+TargetJSON UpdateWebGPUAttrs(TargetJSON target) {
+ if (target.count("supports_subgroups")) {
+ bool subgroups = Downcast<Bool>(target.at("supports_subgroups"));
+ if (subgroups) {
+ target.Set("thread_warp_size", int64_t(32));
+ }
Review Comment:
supports_subgroups flag is the single intentional entry point, and
UpdateWebGPUAttrs already sets thread_warp_size=32 when it's true. Manually
setting thread_warp_size without supports_subgroups isn't a supported
configuration.
##########
src/target/source/codegen_webgpu.cc:
##########
@@ -118,7 +121,9 @@ void CodeGenWebGPU::InitFuncState(const PrimFunc& f) {
}
}
-CodeGenWebGPU::CodeGenWebGPU(Target target) : target_(target) {}
+CodeGenWebGPU::CodeGenWebGPU(Target target) : target_(target) {
+ enable_subgroups_ =
target_->GetAttr<Bool>("supports_subgroups").value_or(Bool(false));
Review Comment:
supports_subgroups flag is the single intentional entry point, and
UpdateWebGPUAttrs already sets thread_warp_size=32 when it's true. Manually
setting thread_warp_size without supports_subgroups isn't a supported
configuration.
##########
src/s_tir/transform/lower_thread_allreduce.cc:
##########
Review Comment:
CodeGenWebGPU::PrintType already handles this. I think adding one more check
here is redundant.
##########
src/s_tir/transform/lower_thread_allreduce.cc:
##########
@@ -510,7 +510,9 @@ class ThreadAllreduceBuilder final : public StmtExprMutator
{
//
// The former may cause dead lock as there is a divergent
// branch with a warp sync call inside.
- PrimExpr other = WarpShuffle(builtin::tvm_warp_shuffle_down(),
mask_buffer, val, offset);
+ bool cast_offset_to_uint = target_->kind->name == "webgpu";
Review Comment:
added
##########
src/target/source/codegen_webgpu.cc:
##########
@@ -120,7 +123,9 @@ void CodeGenWebGPU::InitFuncState(const PrimFunc& f) {
}
}
-CodeGenWebGPU::CodeGenWebGPU(Target target) : target_(target) {}
+CodeGenWebGPU::CodeGenWebGPU(Target target) : target_(target) {
+ enable_subgroups_ =
target_->GetAttr<Bool>("supports_subgroups").value_or(Bool(false));
Review Comment:
I added validation in UpdateWebGPUAttrs so that combination is rejected up
front
##########
src/s_tir/transform/lower_thread_allreduce.cc:
##########
@@ -719,11 +723,11 @@ class ThreadAllreduceBuilder final : public
StmtExprMutator {
bool IsWarpReduction(const std::vector<DataType>& types, int group_extent,
int reduce_extent,
int contiguous_reduce_extent) {
if ((target_->kind->name != "cuda") && (target_->kind->name != "rocm") &&
- (target_->kind->name != "metal")) {
+ (target_->kind->name != "metal") && (target_->kind->name != "webgpu"))
{
return false;
}
Review Comment:
Instead of adding another WebGPU-specific branch in here, I fixed the
invariant at target construction in target_kind.cc. Now that target shape is
rejected up front, so this scenario never comes up in lowering/codegen. I also
added a test for that in test_target_target.py.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
