Update

[ghstack-poisoned]

Author: JCNTH

backends/webgpu/test/native/test_dynamic_shape.cpp

@@ -102,35 +102,43 @@ void check_s(Module& m, const std::string& prefix, int s) {
                       << " golden.size=" << golden.size() << ")";
 }
 
-// Dynamic quantized linear: input [M, kLinK] -> output [M, kLinN].
+// Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the
+// register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM.
 constexpr int kLinK = 64;
 constexpr int kLinN = 128;
-// Run dyn_linear at [m_rows, kLinK] on an already-loaded module (so it can be
+constexpr int kLinNShmem = 2048;
+// Run <prefix> at [m_rows, kLinK] on an already-loaded module (so it can be
 // reused across M without a fresh load), and compare to the golden.
-void run_linear(Module& m, int m_rows) {
-  const std::string base = g_dir + "/dyn_linear.S" + std::to_string(m_rows);
+void run_linear(Module& m, int m_rows, const char* prefix, int n) {
+  const std::string base = g_dir + "/" + prefix + ".S" + std::to_string(m_rows);
   auto input = read_bin(base + ".input.bin");
   auto golden = read_bin(base + ".golden.bin");
-  ASSERT_FALSE(input.empty()) << "missing dyn_linear.S" << m_rows;
+  ASSERT_FALSE(input.empty()) << "missing " << prefix << ".S" << m_rows;
   auto t = make_tensor_ptr({m_rows, kLinK}, std::move(input));
   auto r = m.forward({EValue(t)});
   ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor())
-      << "dyn_linear M=" << m_rows << " forward failed";
+      << prefix << " M=" << m_rows << " forward failed";
   const auto& out = r.get()[0].toTensor();
-  const size_t numel = static_cast<size_t>(m_rows) * kLinN;
+  const size_t numel = static_cast<size_t>(m_rows) * n;
   ASSERT_EQ(static_cast<size_t>(out.numel()), numel)
-      << "dyn_linear M=" << m_rows << " output numel mismatch";
+      << prefix << " M=" << m_rows << " output numel mismatch";
   std::vector<float> got(
       out.const_data_ptr<float>(), out.const_data_ptr<float>() + numel);
   const float e = max_err(got, golden);
   // 4-bit quant: looser tol (the kernel mirrors the dequant-matmul reference).
-  EXPECT_LT(e, 5e-3f) << "dyn_linear M=" << m_rows << " max_err=" << e;
+  EXPECT_LT(e, 5e-3f) << prefix << " M=" << m_rows << " max_err=" << e;
 }
 
 void check_linear(int m_rows) {
   Module m(g_dir + "/dyn_linear.pte");
   ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear.pte";
-  run_linear(m, m_rows);
+  run_linear(m, m_rows, "dyn_linear", kLinN);
+}
+
+void check_linear_shmem(int m_rows) {
+  Module m(g_dir + "/dyn_linear_shmem.pte");
+  ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte";
+  run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem);
 }
 
 // Dynamic SDPA (GQA prefill, input_pos=0): q[1,s,hq,d] k/v[1,s,hkv,d]
@@ -395,7 +403,25 @@ TEST(DynamicShape, QuantizedLinearReusedGraph) {
   Module m(g_dir + "/dyn_linear.pte");
   ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear.pte";
   for (int m_rows : {128, 32, 1, 128}) {
-    run_linear(m, m_rows);
+    run_linear(m, m_rows, "dyn_linear", kLinN);
+  }
+}
+
+// I3: dynamic linear at N=2048 -> the shmem-GEMM route (K>=4096||N>=2048); the
+// resize hook recomputes the shmem tile count for the live M on the fixed shmem
+// pipeline (M=1 exercises a partial row-tile).
+TEST(DynamicShape, QuantizedLinearShmem) {
+  for (int m_rows : {128, 32, 1}) {
+    check_linear_shmem(m_rows);
+  }
+}
+
+// I4: shmem-routed linear reusing ONE loaded graph across M.
+TEST(DynamicShape, QuantizedLinearShmemReusedGraph) {
+  Module m(g_dir + "/dyn_linear_shmem.pte");
+  ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte";
+  for (int m_rows : {128, 32, 1, 128}) {
+    run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem);
   }
 }
 

backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py

@@ -183,8 +183,10 @@ def export_dynamic_shape_cases(out_dir: str) -> None:
     )
     _write_goldens(rmsmul, "dyn_rmsmul", out_dir, [MAXS, 32, 1])
 
-    # 2d) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM.
+    # 2d) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM
+    # (register-tiled N=128) + a shmem-GEMM-routed variant (N=2048).
     _export_dynamic_linear(out_dir)
+    _export_dynamic_linear(out_dir, n=LIN_SHMEM_N, prefix="dyn_linear_shmem")
 
     # 2e) Fused SDPA with a DYNAMIC seq-len S (prefill, input_pos=0).
     _export_dynamic_sdpa(out_dir)
@@ -222,17 +224,20 @@ def export_dynamic_shape_cases(out_dir: str) -> None:
 # Quantized linear: K x N weight, dynamic rows M; input [M, K], output [M, N].
 LIN_K = 64
 LIN_N = 128
+LIN_SHMEM_N = 2048  # N>=2048 routes linear_q4gsw to the shmem-GEMM path
 LIN_GROUP = 32
 LIN_MAXM = 128
 
 
-def _export_dynamic_linear(out_dir: str) -> None:
+def _export_dynamic_linear(
+    out_dir: str, n: int = LIN_N, prefix: str = "dyn_linear"
+) -> None:
     from executorch.backends.webgpu.test.ops.quantized_linear.test_quantized_linear import (
         _fp64_golden,
         _make_quantized_model,
     )
 
-    model = _make_quantized_model(LIN_K, LIN_N, LIN_GROUP)
+    model = _make_quantized_model(LIN_K, n, LIN_GROUP)
     x = _ramp((LIN_MAXM, LIN_K))
     m_dim = torch.export.Dim("m", min=1, max=LIN_MAXM)
     ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},))
@@ -244,17 +249,17 @@ def _export_dynamic_linear(out_dir: str) -> None:
         for plan in et.executorch_program.execution_plan
         for d in plan.delegates
     ), "linear_q4gsw not delegated"
-    with open(os.path.join(out_dir, "dyn_linear.pte"), "wb") as f:
+    with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f:
         f.write(et.buffer)
-    print("Exported dyn_linear.pte")
+    print(f"Exported {prefix}.pte")
     for m in [LIN_MAXM, 32, 1]:
         xm = _ramp((m, LIN_K))
         g = _fp64_golden(model, xm).astype("<f4")  # [m, N]
         xm.detach().numpy().astype("<f4").tofile(
-            os.path.join(out_dir, f"dyn_linear.S{m}.input.bin")
+            os.path.join(out_dir, f"{prefix}.S{m}.input.bin")
         )
-        g.tofile(os.path.join(out_dir, f"dyn_linear.S{m}.golden.bin"))
-        print(f"  golden dyn_linear M={m}")
+        g.tofile(os.path.join(out_dir, f"{prefix}.S{m}.golden.bin"))
+        print(f"  golden {prefix} M={m}")
 
 
 # Dynamic SDPA: GQA prefill (input_pos=0), q/k/v seq-len dynamic.