[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup

backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl

@@ -1,6 +1,8 @@
-@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
-@group(0) @binding(1) var<storage, read> t_in: array<f32>;
-@group(0) @binding(2) var<storage, read> t_weight: array<f32>;
+$if DTYPE == "half":
+  enable f16;
+@group(0) @binding(0) var<storage, read_write> t_out: array<${buffer_gvec_type(DTYPE, VEC)}>;
+@group(0) @binding(1) var<storage, read> t_in: array<${buffer_gvec_type(DTYPE, VEC)}>;
+@group(0) @binding(2) var<storage, read> t_weight: array<${buffer_gvec_type(DTYPE, VEC)}>;
 
 struct Params {
   num_rows: u32,
@@ -12,7 +14,7 @@ struct Params {
 
 const WG_SIZE: u32 = 64u;
 
-var<workgroup> shared_sum: array<f32, WG_SIZE>;
+var<workgroup> shared_sum: array<${accum_scalar_type(DTYPE)}, WG_SIZE>;
 
 fn reduce_shared(worker_id: u32) {
   workgroupBarrier();
@@ -29,6 +31,10 @@ fn reduce_shared(worker_id: u32) {
   }
 }
 
+$if VEC == 4:
+  // vec4 variant of rms_norm: each lane strides by WG_SIZE over rw4 = row_width/4
+  // texels and accumulates dot(v, v). row_width is the ELEMENT count, so mean_sq
+  // divides by it (not rw4). The host selects this only when row_width % 4 == 0.
 @compute @workgroup_size(64, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
@@ -40,33 +46,70 @@ fn main(
     return;
   }
 
-  let base = row_idx * params.row_width;
+  $if VEC == 4:
+    let rw4 = params.row_width / 4u;
+    let base4 = row_idx * rw4;
+  $else:
+    let base = row_idx * params.row_width;
 
-  var local_sq_sum: f32 = 0.0;
-  var x: u32 = worker_id;
-  loop {
-    if (x >= params.row_width) {
-      break;
+  var local_sq_sum: ${accum_scalar_type(DTYPE)} = 0.0;
+  $if VEC == 4:
+    var x4: u32 = worker_id;
+    loop {
+      if (x4 >= rw4) {
+        break;
+      }
+      let v = t_in[base4 + x4];
+      $if DTYPE == "half":
+        local_sq_sum = local_sq_sum + dot(vec4<f32>(v), vec4<f32>(v));
+      $else:
+        local_sq_sum = local_sq_sum + dot(v, v);
+      x4 = x4 + WG_SIZE;
+    }
+  $else:
+    var x: u32 = worker_id;
+    loop {
+      if (x >= params.row_width) {
+        break;
+      }
+      let v = t_in[base + x];
+      $if DTYPE == "half":
+        local_sq_sum = local_sq_sum + f32(v) * f32(v);
+      $else:
+        local_sq_sum = local_sq_sum + v * v;
+      x = x + WG_SIZE;
     }
-    let v = t_in[base + x];
-    local_sq_sum = local_sq_sum + v * v;
-    x = x + WG_SIZE;
-  }
 
   shared_sum[worker_id] = local_sq_sum;
   reduce_shared(worker_id);
 
   let mean_sq = shared_sum[0] / f32(params.row_width);
   let rstd = inverseSqrt(mean_sq + params.epsilon);
 
-  x = worker_id;
-  loop {
-    if (x >= params.row_width) {
-      break;
+  $if VEC == 4:
+    x4 = worker_id;
+    loop {
+      if (x4 >= rw4) {
+        break;
+      }
+      $if DTYPE == "half":
+        t_out[base4 + x4] = vec4<f16>(vec4<f32>(t_in[base4 + x4]) * rstd * vec4<f32>(t_weight[x4]));
+      $else:
+        t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4];
+      x4 = x4 + WG_SIZE;
+    }
+  $else:
+    x = worker_id;
+    loop {
+      if (x >= params.row_width) {
+        break;
+      }
+      let v = t_in[base + x];
+      let w = t_weight[x];
+      $if DTYPE == "half":
+        t_out[base + x] = f16(f32(v) * rstd * f32(w));
+      $else:
+        t_out[base + x] = v * rstd * w;
+      x = x + WG_SIZE;
     }
-    let v = t_in[base + x];
-    let w = t_weight[x];
-    t_out[base + x] = v * rstd * w;
-    x = x + WG_SIZE;
-  }
 }

backends/webgpu/runtime/ops/rms_norm/rms_norm.yaml

@@ -0,0 +1,12 @@
+rms_norm:
+  parameter_names_with_default_values:
+    DTYPE: float
+    VEC: 1
+  generate_variant_forall:
+    VEC:
+      - VALUE: 1
+        SUFFIX: ""
+      - VALUE: 4
+        SUFFIX: vec4
+  shader_variants:
+    - NAME: rms_norm

backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h

@@ -12,7 +12,7 @@
 
 namespace executorch::backends::webgpu {
 
-// @generated from rms_norm_vec4.wgsl - DO NOT EDIT.
+// @generated from rms_norm.wgsl - DO NOT EDIT.
 // wgsl-sha256: 4c0ba56708bf125a7ec6ea3c51d1288e05ac00a8e2cfa10e38e9a208e230b8df
 inline constexpr const char* kRmsNormVec4WGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<vec4<f32>>;