feat(igla): hslm_bwd_ptx kernel 3 -- split-K dW, 4.2-4.9x faster than kernel 2 at K=1782 (Closes #10895)

.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json

@@ -1,29 +1,29 @@
 {
-  "built_by": "t27c-bootstrap@0.5.2+4e6abf0ec",
+  "built_by": "t27c-bootstrap@0.5.2+ce14e2b",
   "closure": [
     {
       "spec": "specs/igla/coder/ptx_emit.t27",
       "spec_hash": "sha256:83ef298864db7814120b22b0f0621ee2ed80e484a9f4bbfb4bec59caa2883262"
     }
   ],
   "config": "gen=c,rust,verilog,zig;tests=zig",
-  "gen_hash_c": "sha256:5680fdfb166ad3d0f26bb2e3caae8e1e1896df8271ee20fe162c8c7c5ef52aa7",
-  "gen_hash_rust": "sha256:c6e4550b86cf46546df827e1e1a6427242b9bdd1755e64c73e0600f04d3923b4",
-  "gen_hash_verilog": "sha256:e5234465d9e51e5a2e244a964ef6aabd16ab80f3c09d193da971a30cc1469fc1",
-  "gen_hash_zig": "sha256:0217dc3f12608f12adaa5b79d9d655e7d750ff2d6f1652de4b39f95bac52f30e",
+  "gen_hash_c": "sha256:22aa0e3bed6a5dc83976ed5094e7f8e434e85fb0ae7ea2ee3ab28ec96cb9ebba",
+  "gen_hash_rust": "sha256:bd4040475901d80ad7316c3c29303e10de70924ec6046fb9dc2d2958681fa109",
+  "gen_hash_verilog": "sha256:634e8992f0c93c2c59316a11d1b66beccb433b150692b5289cccbacc01c3e917",
+  "gen_hash_zig": "sha256:cb73ef03b2a6b6fce7b2c9d4889dc1b53921a326627cd0251e54b14a3b556664",
   "module": "igla-coder-hslm-bwd-ptx",
   "ring": 12,
   "seal_schema": 2,
-  "sealed_at": "2026-10-19T08:01:23Z",
+  "sealed_at": "2026-10-19T08:35:21Z",
   "sealed_by": "t27c-bootstrap@0.5.2",
-  "spec_hash": "sha256:adc270a216694ead6050bcd7d1df1f64f0f99e5fc3447556399d50b0ce3e7f19",
+  "spec_hash": "sha256:c6e574a517e847a24751ecc3b85956b6fc6d0ba09516bafbb3fa61fe450eac44",
   "spec_path": "specs/igla/coder/hslm_bwd_ptx.t27",
   "tests": {
     "failed": 0,
     "failing": [],
     "forced": false,
-    "passed": 19,
-    "total": 19,
+    "passed": 26,
+    "total": 26,
     "vacuous": 0
   },
   "toolchain": {

bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs

@@ -1,6 +1,6 @@
 // Generated from .t27 spec
 // DO NOT EDIT — generated by t27c
-// NOT LOWERED BY THIS BACKEND: 19 test(s), 7 invariant(s), 0 function(s) with no body.
+// NOT LOWERED BY THIS BACKEND: 26 test(s), 7 invariant(s), 0 function(s) with no body.
 // This backend emits declarations only. The spec's checks live in
 // the Zig and Verilog outputs; do not read this file as verified.
 
@@ -10,6 +10,10 @@ pub const KERNEL_DX: u32 = 1;
 
 pub const KERNEL_DW: u32 = 2;
 
+pub const KERNEL_DW_SPLITK: u32 = 3;
+
+pub const SPLIT_TARGET_BLOCKS: usize = 64;
+
 pub const BM: u32 = 128;
 
 pub const BN: u32 = 128;
@@ -56,6 +60,9 @@ pub fn put_name_at(out: &mut [u8], cap: usize, w: usize, kernel: u32) -> usize {
     if (kernel == KERNEL_DW) {
         return put_s(out, cap, w, "hslm_dw_s8");
     }
+    if (kernel == KERNEL_DW_SPLITK) {
+        return put_s(out, cap, w, "hslm_dw_splitk_s8");
+    }
     return w;
 }
 
@@ -129,9 +136,15 @@ pub fn emit_setup_dx(out: &mut [u8], cap: usize, w: usize) -> usize {
     return ins(out, cap, x, "mov.b32 %r45, 0;");
 }
 
-pub fn emit_setup_dw(out: &mut [u8], cap: usize, w: usize) -> usize {
-    let mut x: usize = ins(out, cap, w, "shr.u32 %r12, %r4, 3;");
-    x = ins(out, cap, x, "and.b32 %r13, %r4, 7;");
+pub fn emit_setup_dw(out: &mut [u8], cap: usize, w: usize, split: bool) -> usize {
+    let mut x: usize = w;
+    if split {
+        x = ins(out, cap, x, "and.b32 %r12, %r4, 31;");
+        x = ins(out, cap, x, "shr.u32 %r13, %r4, 5;");
+    } else {
+        x = ins(out, cap, x, "shr.u32 %r12, %r4, 3;");
+        x = ins(out, cap, x, "and.b32 %r13, %r4, 7;");
+    }
     x = ins(out, cap, x, "shl.b32 %r13, %r13, 4;");
     x = ins(out, cap, x, "shl.b32 %r9, %r6, 7;");
     x = ins(out, cap, x, "add.u32 %r9, %r9, %r13;");
@@ -280,18 +293,87 @@ pub fn emit_loads_dw(out: &mut [u8], cap: usize, w: usize) -> usize {
     return emit_transposed(out, cap, x, 50, "%r16", "%r37");
 }
 
+pub fn emit_slice(out: &mut [u8], cap: usize, w: usize) -> usize {
+    let mut x: usize = ins(out, cap, w, "mov.u32 %r46, %ctaid.z;");
+    x = ins(out, cap, x, "mov.u32 %r47, %nctaid.z;");
+    x = ins(out, cap, x, "add.u32 %r48, %r3, 31;");
+    x = ins(out, cap, x, "shr.u32 %r48, %r48, 5;");
+    x = ins(out, cap, x, "add.u32 %r48, %r48, %r47;");
+    x = ins(out, cap, x, "sub.u32 %r48, %r48, 1;");
+    x = ins(out, cap, x, "div.u32 %r48, %r48, %r47;");
+    x = ins(out, cap, x, "shl.b32 %r48, %r48, 5;");
+    x = ins(out, cap, x, "mul.lo.u32 %r49, %r46, %r48;");
+    x = ins(out, cap, x, "add.u32 %r46, %r46, 1;");
+    x = ins(out, cap, x, "mul.lo.u32 %r46, %r46, %r48;");
+    x = ins(out, cap, x, "min.u32 %r3, %r3, %r46;");
+    x = ins(out, cap, x, "setp.ge.u32 %p9, %r49, %r3;");
+    return ins(out, cap, x, "@%p9 ret;");
+}
+
+pub fn emit_red_store(out: &mut [u8], cap: usize, w: usize, dr: u32, dc: u32, b: u32) -> usize {
+    let mut x: usize = put_s(out, cap, w, "    add.u32 %r41, %r39, ");
+    x = put_u(out, cap, x, dr);
+    x = put_s(out, cap, x, ";\n");
+    x = put_s(out, cap, x, "    add.u32 %r42, %r40, ");
+    x = put_u(out, cap, x, dc);
+    x = put_s(out, cap, x, ";\n");
+    x = ins(out, cap, x, "setp.lt.u32 %p6, %r41, %r1;");
+    x = ins(out, cap, x, "setp.lt.u32 %p7, %r42, %r2;");
+    x = ins(out, cap, x, "and.pred %p6, %p6, %p7;");
+    x = ins(out, cap, x, "cvt.u64.u32 %rd10, %r42;");
+    x = ins(out, cap, x, "mad.wide.u32 %rd10, %r41, %r2, %rd10;");
+    x = ins(out, cap, x, "shl.b64 %rd10, %rd10, 2;");
+    x = ins(out, cap, x, "add.u64 %rd10, %rd10, %rd3;");
+    x = put_s(out, cap, x, "    @%p6 red.global.add.s32 [%rd10], %acc");
+    x = put_u(out, cap, x, b);
+    x = put_s(out, cap, x, ";\n    @%p6 red.global.add.s32 [%rd10+4], %acc");
+    x = put_u(out, cap, x, (b + 1));
+    return put_s(out, cap, x, ";\n");
+}
+
+pub fn emit_red_epilogue(out: &mut [u8], cap: usize, w: usize) -> usize {
+    let mut x: usize = ins(out, cap, w, "shl.b32 %r39, %r6, 7;");
+    x = ins(out, cap, x, "shl.b32 %r41, %r23, 6;");
+    x = ins(out, cap, x, "add.u32 %r39, %r39, %r41;");
+    x = ins(out, cap, x, "add.u32 %r39, %r39, %r21;");
+    x = ins(out, cap, x, "shl.b32 %r40, %r5, 7;");
+    x = ins(out, cap, x, "shl.b32 %r41, %r24, 5;");
+    x = ins(out, cap, x, "add.u32 %r40, %r40, %r41;");
+    x = ins(out, cap, x, "shl.b32 %r41, %r22, 1;");
+    x = ins(out, cap, x, "add.u32 %r40, %r40, %r41;");
+    let mut mi: u32 = 0;
+    while (mi < 4) {
+        let mut ni: u32 = 0;
+        while (ni < 4) {
+            let b: u32 = (((mi * 4) + ni) * 4);
+            x = emit_red_store(out, cap, x, (mi * 16), (ni * 8), b);
+            x = emit_red_store(out, cap, x, ((mi * 16) + 8), (ni * 8), (b + 2));
+            ni += 1;
+        }
+        mi += 1;
+    }
+    return x;
+}
+
 pub fn emit_ptx(kernel: u32, out: &mut [u8], cap: usize) -> usize {
-    if ((kernel != KERNEL_DX) && (kernel != KERNEL_DW)) {
+    if (((kernel != KERNEL_DX) && (kernel != KERNEL_DW)) && (kernel != KERNEL_DW_SPLITK)) {
         return 0;
     }
+    let split: bool = (kernel == KERNEL_DW_SPLITK);
     let mut x: usize = emit_header(out, cap, 0, kernel);
     x = emit_params(out, cap, x);
     if (kernel == KERNEL_DX) {
         x = emit_setup_dx(out, cap, x);
     } else {
-        x = emit_setup_dw(out, cap, x);
+        x = emit_setup_dw(out, cap, x, split);
+    }
+    if split {
+        x = emit_slice(out, cap, x);
     }
     x = emit_warp_setup(out, cap, x);
+    if split {
+        x = ins(out, cap, x, "mov.u32 %r28, %r49;");
+    }
     x = put_s(out, cap, x, "LOOP_K:\n");
     x = ins(out, cap, x, "setp.ge.u32 %p3, %r28, %r3;");
     x = ins(out, cap, x, "@%p3 bra DONE_K;");
@@ -306,7 +388,11 @@ pub fn emit_ptx(kernel: u32, out: &mut [u8], cap: usize) -> usize {
     x = ins(out, cap, x, "add.u32 %r28, %r28, 32;");
     x = ins(out, cap, x, "bra LOOP_K;");
     x = put_s(out, cap, x, "DONE_K:\n");
-    x = emit_epilogue(out, cap, x);
+    if split {
+        x = emit_red_epilogue(out, cap, x);
+    } else {
+        x = emit_epilogue(out, cap, x);
+    }
     x = ins(out, cap, x, "ret;");
     return put_s(out, cap, x, "}\n");
 }
@@ -315,6 +401,41 @@ pub fn grid_dim(x: usize) -> u32 {
     return (((x + 127) / 128) as u32);
 }
 
+pub fn grid_z(kernel: u32, m: usize, n: usize, k: usize) -> u32 {
+    if (kernel != KERNEL_DW_SPLITK) {
+        return 1;
+    }
+    let steps: usize = ((k + 31) / 32);
+    let blocks: usize = (((m + 127) / 128) * ((n + 127) / 128));
+    let mut s: usize = (((SPLIT_TARGET_BLOCKS + blocks) - 1) / blocks);
+    if (s > steps) {
+        s = steps;
+    }
+    if (s < 1) {
+        s = 1;
+    }
+    let per: usize = (((steps + s) - 1) / s);
+    return ((((steps + per) - 1) / per) as u32);
+}
+
+pub fn zero_first(kernel: u32) -> bool {
+    return (kernel == KERNEL_DW_SPLITK);
+}
+
+pub fn slice_lo(z: usize, nz: usize, k: usize) -> usize {
+    let per: usize = (((((k + 31) / 32) + nz) - 1) / nz);
+    return ((z * per) * 32);
+}
+
+pub fn slice_hi(z: usize, nz: usize, k: usize) -> usize {
+    let per: usize = (((((k + 31) / 32) + nz) - 1) / nz);
+    let hi: usize = (((z + 1) * per) * 32);
+    if (hi < k) {
+        return hi;
+    }
+    return k;
+}
+
 pub fn shape_ok(kernel: u32, m: usize, n: usize, k: usize) -> bool {
     if (((m < 1) || (n < 16)) || ((n % 16) != 0)) {
         return false;
@@ -325,14 +446,14 @@ pub fn shape_ok(kernel: u32, m: usize, n: usize, k: usize) -> bool {
     if (kernel == KERNEL_DX) {
         return (((k >= 16) && ((k % 16) == 0)) && (k <= K_MAX_DX));
     }
-    if (kernel == KERNEL_DW) {
+    if ((kernel == KERNEL_DW) || (kernel == KERNEL_DW_SPLITK)) {
         return ((((m % 16) == 0) && (k >= 1)) && (k <= K_MAX_DW));
     }
     return false;
 }
 
 pub fn b_len(kernel: u32, n: usize, k: usize) -> usize {
-    if (kernel == KERNEL_DW) {
+    if ((kernel == KERNEL_DW) || (kernel == KERNEL_DW_SPLITK)) {
         return (k * n);
     }
     return ((k * n) / 4);

bootstrap/src/host/driver.rs

@@ -299,14 +299,14 @@ mod tests {
     }
 }
 
-// `t27c ptx-run` (#10404, #10847): load the PTX that ternary_gemm_ptx.t27 (kernel 0) or hslm_bwd_ptx.t27 (1, 2) wrote, run it, count mismatches with its reference.
+// `t27c ptx-run` (#10404, #10847): load the PTX that ternary_gemm_ptx.t27 (kernel 0) or hslm_bwd_ptx.t27 (1, 2, 3) wrote, run it, count mismatches with its reference.
 #[cfg(feature = "cuda")] #[link(name = "cuda")]
 extern "C" {
     fn cuInit(flags: u32) -> i32; fn cuDeviceGet(dev: *mut i32, ordinal: i32) -> i32;
     fn cuCtxCreate_v2(ctx: *mut usize, flags: u32, dev: i32) -> i32; fn cuCtxSynchronize() -> i32;
     fn cuModuleLoadData(module: *mut usize, image: *const u8) -> i32; fn cuModuleGetFunction(func: *mut usize, module: usize, name: *const u8) -> i32;
     fn cuMemAlloc_v2(ptr: *mut u64, bytes: usize) -> i32; fn cuMemcpyHtoD_v2(dst: u64, src: *const u8, bytes: usize) -> i32;
-    fn cuMemcpyDtoH_v2(dst: *mut u8, src: u64, bytes: usize) -> i32;
+    fn cuMemcpyDtoH_v2(dst: *mut u8, src: u64, bytes: usize) -> i32; fn cuMemsetD32_v2(dst: u64, v: u32, n: usize) -> i32;
     fn cuLaunchKernel(f: usize, gx: u32, gy: u32, gz: u32, bx: u32, by: u32, bz: u32, smem: u32, stream: usize, params: *mut *mut u8, extra: usize) -> i32;
 }
 #[cfg(feature = "cuda")]
@@ -328,7 +328,7 @@ pub fn ptx_run(kernel: u32, m: usize, n: usize, k: usize, iters: u32, stride: us
         ck(cuMemcpyHtoD_v2(da, a.as_ptr() as *const u8, a.len()))?; ck(cuMemcpyHtoD_v2(dw, wp.as_ptr(), wp.len()))?;
         let mut args: [*mut u8; 6] = [&mut da as *mut u64 as *mut u8, &mut dw as *mut u64 as *mut u8, &mut dc as *mut u64 as *mut u8,
             &mut pm as *mut u32 as *mut u8, &mut pn as *mut u32 as *mut u8, &mut pk as *mut u32 as *mut u8];
-        let mut launch = || ck(cuLaunchKernel(f, t::grid_dim(n), t::grid_dim(m), 1, t::THREADS, 1, 1, 0, 0, args.as_mut_ptr(), 0));
+        let mut launch = || { if u::zero_first(kernel) { ck(cuMemsetD32_v2(dc, 0, m * n))?; } ck(cuLaunchKernel(f, t::grid_dim(n), t::grid_dim(m), u::grid_z(kernel, m, n, k), t::THREADS, 1, 1, 0, 0, args.as_mut_ptr(), 0)) };
         launch()?; ck(cuCtxSynchronize())?; ck(cuMemcpyDtoH_v2(c.as_mut_ptr() as *mut u8, dc, c.len() * 4))?;
         let t0 = std::time::Instant::now();
         for _ in 0..iters { launch()?; } ck(cuCtxSynchronize())?;