[WebGPU] Add gating logic for subgroup shuffle primitives

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";
+        PrimExpr other = WarpShuffle(builtin::tvm_warp_shuffle_down(), mask_buffer, val, offset,
+                                     cast_offset_to_uint);
         Buffer local_buf = local_bufs[i];
         Stmt s = BufferStore(local_buf, other, zero_indices);
         seq->push_back(s);
@@ -699,7 +701,10 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
 
   // Emit warp shuffle  calls.
   PrimExpr WarpShuffle(const Op& op, ffi::Optional<Buffer> mask_buffer, PrimExpr val,
-                       PrimExpr delta_or_lane) {
+                       PrimExpr delta_or_lane, bool cast_delta_to_uint = false) {
+    if (cast_delta_to_uint) {
+      delta_or_lane = cast(DataType::UInt(32, delta_or_lane.dtype().lanes()), delta_or_lane);
+    }
     ffi::Array<PrimExpr> indices = {0};
     PrimExpr mask;
     if (mask_buffer.defined()) {
@@ -719,11 +724,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;
     }
 
-    need_warp_shuffle_mask_ = target_->kind->name != "metal";
+    need_warp_shuffle_mask_ = target_->kind->name != "metal" && target_->kind->name != "webgpu";
 
     // rocm only supports 32 bit operands for shuffling at the moment
     if ((target_->kind->name == "rocm") &&

src/target/source/codegen_webgpu.cc

@@ -107,6 +107,9 @@ std::string CodeGenWebGPU::Finish() {
   if (enable_fp16_) {
     header_stream << "enable f16;\n\n";
   }
+  if (enable_subgroups_) {
+    header_stream << "enable subgroups;\n\n";
+  }
   return header_stream.str() + decl_stream.str() + this->fwd_decl_stream.str() + stream.str();
 }
 
@@ -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));
+}
 
 runtime::FunctionInfo CodeGenWebGPU::AddFunction(const PrimFunc& f, bool skip_readonly_decl) {
   // clear previous generated state.

src/target/source/codegen_webgpu.h

@@ -92,6 +92,8 @@ class CodeGenWebGPU final : public CodeGenC {
 
   // whether enable fp16
   bool enable_fp16_{false};
+  // whether enable subgroups
+  bool enable_subgroups_{false};
 
   /*! \brief the header stream for function label and enable directive if any, goes before any other
    * declaration */

src/target/source/intrin_rule_webgpu.cc

@@ -32,6 +32,29 @@ namespace intrin {
 
 using tir::FLowerIntrinsic;
 
+// warp-level primitives. Follows implementation in intrin_rule_metal.cc
+struct WebGPUWarpIntrinsic {
+    const Op operator()(DataType t, const Op& orig_op) const {
+        if (orig_op.same_as(builtin::tvm_warp_shuffle())) {
+            return Op::Get("tir.webgpu.subgroup_shuffle");
+        } else if (orig_op.same_as(builtin::tvm_warp_shuffle_up())) {
+            return Op::Get("tir.webgpu.subgroup_shuffle_up");
+        } else {
+            ICHECK(orig_op.same_as(builtin::tvm_warp_shuffle_down()));
+            return Op::Get("tir.webgpu.subgroup_shuffle_down");
+        }
+    }
+};
+
+template <typename T>
+static PrimExpr DispatchWebGPUShuffle(const PrimExpr& e) {
+    const CallNode* call = e.as<CallNode>();
+    ICHECK(call != nullptr);
+    ICHECK_EQ(call->args.size(), 5);  // mask, value, warp_id, width, warp_size
+    ffi::Array<PrimExpr> webgpu_args{{call->args[1], call->args[2]}};
+    return Call(call->dtype, T()(call->dtype, Downcast<Op>(call->op)), webgpu_args);
+}
+
 // See full list of builtin: https://www.w3.org/TR/WGSL/#builtin-functions
 
 struct ReturnAbs {
@@ -113,6 +136,39 @@ TVM_REGISTER_OP("tir.trunc")
 // extra dispatch
 TVM_REGISTER_OP("tir.erf").set_attr<FLowerIntrinsic>("webgpu.FLowerIntrinsic", DispatchFastErf);
 
+// warp-level primitives. Follows implementation in intrin_rule_metal.cc
+TVM_REGISTER_OP("tir.tvm_warp_shuffle")
+    .set_attr<FLowerIntrinsic>("webgpu.FLowerIntrinsic", DispatchWebGPUShuffle<WebGPUWarpIntrinsic>);
+
+TVM_REGISTER_OP("tir.tvm_warp_shuffle_up")
+    .set_attr<FLowerIntrinsic>("webgpu.FLowerIntrinsic", DispatchWebGPUShuffle<WebGPUWarpIntrinsic>);
+
+TVM_REGISTER_OP("tir.tvm_warp_shuffle_down")
+    .set_attr<FLowerIntrinsic>("webgpu.FLowerIntrinsic", DispatchWebGPUShuffle<WebGPUWarpIntrinsic>);
+
+// Register low-level builtin ops.
+TVM_REGISTER_OP("tir.webgpu.subgroup_shuffle")
+    .set_num_inputs(2)
+    .add_argument("var", "Expr", "The variable to sync.")
+    .add_argument("lane", "Expr", "The source thread id.")
+    .set_attr<TGlobalSymbol>("TGlobalSymbol", "subgroupShuffle")
+    .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
+
+TVM_REGISTER_OP("tir.webgpu.subgroup_shuffle_up")
+    .set_num_inputs(2)
+    .add_argument("var", "Expr", "The variable to sync.")
+    .add_argument("delta", "Expr", "The source lane id offset to be added.")
+    .set_attr<TGlobalSymbol>("TGlobalSymbol", "subgroupShuffleUp")
+    .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
+
+TVM_REGISTER_OP("tir.webgpu.subgroup_shuffle_down")
+    .set_num_inputs(2)
+    .add_argument("var", "Expr", "The variable to sync.")
+    .add_argument("delta", "Expr", "The source lane id offset to be subtracted.")
+    .set_attr<TGlobalSymbol>("TGlobalSymbol", "subgroupShuffleDown")
+    .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
+
+
 }  // namespace intrin
 }  // namespace codegen
 }  // namespace tvm

src/target/target_kind.cc

@@ -427,8 +427,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", refl::DefaultValue(32));
+    }
+  }
+  return target;
+}
+
 TVM_REGISTER_TARGET_KIND("webgpu", kDLWebGPU)
     .add_attr_option<int64_t>("max_num_threads", refl::DefaultValue(256))
+    .add_attr_option<bool>("supports_subgroups", refl::DefaultValue(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", refl::DefaultValue(1))
+    .set_target_parser(UpdateWebGPUAttrs)
     .set_default_keys({"webgpu", "gpu"});
 
 TVM_REGISTER_TARGET_KIND("hexagon", kDLHexagon)

web/src/webgpu.ts

@@ -127,6 +127,9 @@ export async function detectGPUDevice(powerPreference: "low-power" | "high-perfo
     if (adapter.features.has("shader-f16")) {
       requiredFeatures.push("shader-f16");
     }
+    if (adapter.features.has("subgroups")) {
+      requiredFeatures.push("subgroups");
+    }
 
     // requestAdapterInfo() is deprecated, causing requestAdapterInfo to raise
     // issue when building. However, it is still needed for older browsers, hence `as any`.