diff --git a/.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json b/.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json index 20ffd993cf..38325ab3ac 100644 --- a/.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json +++ b/.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json @@ -1,17 +1,22 @@ { - "built_by": "t27c-bootstrap@0.5.2+ae19b93cd", - "closure": [], + "built_by": "t27c-bootstrap@0.5.2+4e6abf0ec", + "closure": [ + { + "spec": "specs/igla/coder/ptx_emit.t27", + "spec_hash": "sha256:83ef298864db7814120b22b0f0621ee2ed80e484a9f4bbfb4bec59caa2883262" + } + ], "config": "gen=c,rust,verilog,zig;tests=zig", - "gen_hash_c": "sha256:101422e82dd5bea1e45f3b8285e4466c73efbeb4f6acd467551519905ed7708d", - "gen_hash_rust": "sha256:14974c23c38d513812ba4d466f1875d8e6362727b3f47bba837cee79e032e25a", - "gen_hash_verilog": "sha256:40bdc27019e915b9ec734639e75c7f4e9d179b97b4465dc69830e494217107b7", - "gen_hash_zig": "sha256:adbca1945a80534c19f3bcb3f3084a8e513e7d42747508f032e0dd69580d2724", + "gen_hash_c": "sha256:5680fdfb166ad3d0f26bb2e3caae8e1e1896df8271ee20fe162c8c7c5ef52aa7", + "gen_hash_rust": "sha256:c6e4550b86cf46546df827e1e1a6427242b9bdd1755e64c73e0600f04d3923b4", + "gen_hash_verilog": "sha256:e5234465d9e51e5a2e244a964ef6aabd16ab80f3c09d193da971a30cc1469fc1", + "gen_hash_zig": "sha256:0217dc3f12608f12adaa5b79d9d655e7d750ff2d6f1652de4b39f95bac52f30e", "module": "igla-coder-hslm-bwd-ptx", "ring": 12, "seal_schema": 2, - "sealed_at": "2026-10-19T05:03:29Z", + "sealed_at": "2026-10-19T08:01:23Z", "sealed_by": "t27c-bootstrap@0.5.2", - "spec_hash": "sha256:08849f6ede4e95b030c40d819aa95c727a2c5ee83ec5dec8d7ce7fac0cdbe4c2", + "spec_hash": "sha256:adc270a216694ead6050bcd7d1df1f64f0f99e5fc3447556399d50b0ce3e7f19", "spec_path": "specs/igla/coder/hslm_bwd_ptx.t27", "tests": { "failed": 0, diff --git a/.trinity/seals/coder_igla-coder-ptx-emit.json b/.trinity/seals/coder_igla-coder-ptx-emit.json new file mode 100644 index 0000000000..5af23a70a5 --- /dev/null +++ b/.trinity/seals/coder_igla-coder-ptx-emit.json @@ -0,0 +1,28 @@ +{ + "built_by": "t27c-bootstrap@0.5.2+4e6abf0ec", + "closure": [], + "config": "gen=c,rust,verilog,zig;tests=zig", + "gen_hash_c": "sha256:26ae27227a3743696ab4c10e9e433b05299fa4ea60993d9e0e30d60dbfd8b612", + "gen_hash_rust": "sha256:5f1d8f73dc4b664698d57255127ed3b097771b4848eeb0e763c6898f097acee3", + "gen_hash_verilog": "sha256:65e15b1cafc340410fcb913f1368e12ff7189fe068d305d3af8c12f5457995d5", + "gen_hash_zig": "sha256:5527415604d9e59cb7d83309697b486a9d1751c40b6bfd8f6248d04ba9ffca97", + "module": "igla-coder-ptx-emit", + "ring": 12, + "seal_schema": 2, + "sealed_at": "2026-10-19T08:01:19Z", + "sealed_by": "t27c-bootstrap@0.5.2", + "spec_hash": "sha256:83ef298864db7814120b22b0f0621ee2ed80e484a9f4bbfb4bec59caa2883262", + "spec_path": "specs/igla/coder/ptx_emit.t27", + "tests": { + "failed": 0, + "failing": [], + "forced": false, + "passed": 10, + "total": 10, + "vacuous": 0 + }, + "toolchain": { + "test_runner": "sha256:8512b80da533b79289110fca4a65f3336abd7a059ca4b72ea595d5e33ff8ad10", + "zig": "0.16.0" + } +} \ No newline at end of file diff --git a/.trinity/seals/coder_igla-coder-ternary-gemm-ptx.json b/.trinity/seals/coder_igla-coder-ternary-gemm-ptx.json index d6c0212a91..08c9e450b2 100644 --- a/.trinity/seals/coder_igla-coder-ternary-gemm-ptx.json +++ b/.trinity/seals/coder_igla-coder-ternary-gemm-ptx.json @@ -1,17 +1,22 @@ { - "built_by": "t27c-bootstrap@0.5.2+2406a4804", - "closure": [], + "built_by": "t27c-bootstrap@0.5.2+4e6abf0ec", + "closure": [ + { + "spec": "specs/igla/coder/ptx_emit.t27", + "spec_hash": "sha256:83ef298864db7814120b22b0f0621ee2ed80e484a9f4bbfb4bec59caa2883262" + } + ], "config": "gen=c,rust,verilog,zig;tests=zig", - "gen_hash_c": "sha256:7ac537a19563bff12ac28e1a1669de534e956d202b1adf4dc30f2e94d9f96b47", - "gen_hash_rust": "sha256:df047da747040c5a33657f13557fee1d3469efc588c9771ab969b612ba14531e", - "gen_hash_verilog": "sha256:8f0b6dd58e6c0094639dbb97e1d308dd21c19936c66c6db3c9043ed17c03c3aa", - "gen_hash_zig": "sha256:fd96535350bba53d44225aa26dba60a9dcf420d3eb3eaf63565a07864cf688fb", + "gen_hash_c": "sha256:ed03842467ddc0a5387129e23993d484aa2986bc55e64a7db7f6004e7bfd7fcb", + "gen_hash_rust": "sha256:8ed3eaebdee6f3f8cfc9d88ccdceb0f4def7943c9f19ae84339e62d8d5c3ad02", + "gen_hash_verilog": "sha256:2e5252e7a61417c2be10fdc4ab88d545ce089d0f1cd073f4ae4fd702ddb8e241", + "gen_hash_zig": "sha256:c8514b781c2365fd0f892673cd2f04f827e0c2745dbdc71f83815fa8899c4bda", "module": "igla-coder-ternary-gemm-ptx", "ring": 12, "seal_schema": 2, - "sealed_at": "2026-10-18T11:38:54Z", + "sealed_at": "2026-10-19T08:01:22Z", "sealed_by": "t27c-bootstrap@0.5.2", - "spec_hash": "sha256:2ead1690d3f7fdbff7c9ac62cc6f98a89bb210bf359dccc225d63f236ba03c98", + "spec_hash": "sha256:11b3b4abd2381d5b06d52431dc1469f3e4611cf426b94ebb6443080840fa66e3", "spec_path": "specs/igla/coder/ternary_gemm_ptx.t27", "tests": { "failed": 0, @@ -22,7 +27,7 @@ "vacuous": 0 }, "toolchain": { - "test_runner": "sha256:1982ae948e24d00c5d221ccec12dfce1eb01204b53e37b0f857f93b16d64629d", + "test_runner": "sha256:8512b80da533b79289110fca4a65f3336abd7a059ca4b72ea595d5e33ff8ad10", "zig": "0.16.0" } } \ No newline at end of file diff --git a/bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs b/bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs index 921e73b1f4..01d024dee5 100644 --- a/bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs +++ b/bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs @@ -46,54 +46,6 @@ pub const HSLM_EMBED_PAD: usize = 256; pub const HSLM_FF_PAD: usize = 736; -pub fn put_s(out: &mut [u8], cap: usize, w: usize, s: &'static str) -> usize { - let mut i: usize = 0; - while (i < s.len()) { - if ((w + i) < cap) { - out[((w + i)) as usize] = s.as_bytes()[(i) as usize]; - } - i += 1; - } - return (w + s.len()); -} - -pub fn put_u(out: &mut [u8], cap: usize, w: usize, v: u32) -> usize { - let mut d: u32 = 1; - while ((v / d) >= 10) { - d = (d * 10); - } - let mut x: usize = w; - while (d > 0) { - let c: u8 = (b'0' + (((v / d) % 10) as u8)); - if (x < cap) { - out[(x) as usize] = c; - } - x += 1; - d = (d / 10); - } - return x; -} - -pub fn ins(out: &mut [u8], cap: usize, w: usize, s: &'static str) -> usize { - let mut x: usize = put_s(out, cap, w, " "); - x = put_s(out, cap, x, s); - return put_s(out, cap, x, "\n"); -} - -pub fn put_quad(out: &mut [u8], cap: usize, w: usize, name: &'static str, first: u32) -> usize { - let mut x: usize = put_s(out, cap, w, "{"); - let mut j: u32 = 0; - while (j < 4) { - if (j > 0) { - x = put_s(out, cap, x, ", "); - } - x = put_s(out, cap, x, name); - x = put_u(out, cap, x, (first + j)); - j += 1; - } - return put_s(out, cap, x, "}"); -} - pub fn put_name_at(out: &mut [u8], cap: usize, w: usize, kernel: u32) -> usize { if (kernel == KERNEL_FWD) { return put_s(out, cap, w, "ternary_gemm_s8"); @@ -328,102 +280,6 @@ 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_lds(out: &mut [u8], cap: usize, w: usize, reg: &'static str, i: u32, base: &'static str, off: u32) -> usize { - let mut x: usize = put_s(out, cap, w, " ld.shared.u32 "); - x = put_s(out, cap, x, reg); - x = put_u(out, cap, x, i); - x = put_s(out, cap, x, ", ["); - x = put_s(out, cap, x, base); - x = put_s(out, cap, x, "+"); - x = put_u(out, cap, x, off); - return put_s(out, cap, x, "];\n"); -} - -pub fn emit_mma_step(out: &mut [u8], cap: usize, w: usize) -> usize { - let mut x: usize = w; - let mut mi: u32 = 0; - while (mi < 4) { - let off: u32 = ((mi * 16) * SMEM_STRIDE); - x = emit_lds(out, cap, x, "%fa", (mi * 4), "%r25", off); - x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 1), "%r25", (off + (8 * SMEM_STRIDE))); - x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 2), "%r25", (off + 16)); - x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 3), "%r25", ((off + (8 * SMEM_STRIDE)) + 16)); - mi += 1; - } - let mut ni: u32 = 0; - while (ni < 4) { - x = emit_lds(out, cap, x, "%fb", (ni * 2), "%r27", ((ni * 8) * SMEM_STRIDE)); - x = emit_lds(out, cap, x, "%fb", ((ni * 2) + 1), "%r27", (((ni * 8) * SMEM_STRIDE) + 16)); - ni += 1; - } - mi = 0; - while (mi < 4) { - ni = 0; - while (ni < 4) { - let b: u32 = (((mi * 4) + ni) * 4); - x = put_s(out, cap, x, " mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "); - x = put_quad(out, cap, x, "%acc", b); - x = put_s(out, cap, x, ", "); - x = put_quad(out, cap, x, "%fa", (mi * 4)); - x = put_s(out, cap, x, ", {%fb"); - x = put_u(out, cap, x, (ni * 2)); - x = put_s(out, cap, x, ", %fb"); - x = put_u(out, cap, x, ((ni * 2) + 1)); - x = put_s(out, cap, x, "}, "); - x = put_quad(out, cap, x, "%acc", b); - x = put_s(out, cap, x, ";\n"); - ni += 1; - } - mi += 1; - } - return x; -} - -pub fn emit_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 st.global.v2.u32 [%rd10], {%acc"); - x = put_u(out, cap, x, b); - x = put_s(out, cap, x, ", %acc"); - x = put_u(out, cap, x, (b + 1)); - return put_s(out, cap, x, "};\n"); -} - -pub fn emit_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_store(out, cap, x, (mi * 16), (ni * 8), b); - x = emit_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)) { return 0; @@ -482,16 +338,6 @@ pub fn b_len(kernel: u32, n: usize, k: usize) -> usize { return ((k * n) / 4); } -pub fn code_at(wp: &[u8], i: usize) -> u32 { - let b: u32 = (wp[((i / 4)) as usize] as u32); - let sh: u32 = (((i % 4) * 2) as u32); - return ((b >> sh) & 3); -} - -pub fn weight(code: u32) -> i32 { - return (((code & 1) as i32) - (((code >> 1) & 1) as i32)); -} - pub fn sbyte(v: u8) -> i32 { let u: i32 = (v as i32); if (u >= 128) { @@ -620,6 +466,14 @@ pub fn decoded_word(word: u32, wq: u32) -> u32 { return prmt_table((sp >> 16)); } +pub fn lcg(s: u64) -> u64 { + return (((s * 1103515245) + 12345) % 2147483648); +} + +pub fn unit(s: u64) -> f64 { + return (((((s >> 16) % 2001) as f64) / 1000.0) - 1.0); +} + pub fn at_text(t: &[u8], n: usize, at: usize, s: &'static str) -> bool { if ((at + s.len()) > n) { return false; @@ -634,6 +488,12 @@ pub fn at_text(t: &[u8], n: usize, at: usize, s: &'static str) -> bool { return true; } +pub fn code_at(wp: &[u8], i: usize) -> u32 { + let b: u32 = (wp[((i / 4)) as usize] as u32); + let sh: u32 = (((i % 4) * 2) as u32); + return ((b >> sh) & 3); +} + pub fn count_text(t: &[u8], n: usize, s: &'static str) -> usize { let mut c: usize = 0; let mut i: usize = 0; @@ -646,11 +506,151 @@ pub fn count_text(t: &[u8], n: usize, s: &'static str) -> usize { return c; } -pub fn lcg(s: u64) -> u64 { - return (((s * 1103515245) + 12345) % 2147483648); +pub fn emit_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_store(out, cap, x, (mi * 16), (ni * 8), b); + x = emit_store(out, cap, x, ((mi * 16) + 8), (ni * 8), (b + 2)); + ni += 1; + } + mi += 1; + } + return x; } -pub fn unit(s: u64) -> f64 { - return (((((s >> 16) % 2001) as f64) / 1000.0) - 1.0); +pub fn emit_lds(out: &mut [u8], cap: usize, w: usize, reg: &'static str, i: u32, base: &'static str, off: u32) -> usize { + let mut x: usize = put_s(out, cap, w, " ld.shared.u32 "); + x = put_s(out, cap, x, reg); + x = put_u(out, cap, x, i); + x = put_s(out, cap, x, ", ["); + x = put_s(out, cap, x, base); + x = put_s(out, cap, x, "+"); + x = put_u(out, cap, x, off); + return put_s(out, cap, x, "];\n"); +} + +pub fn emit_mma_step(out: &mut [u8], cap: usize, w: usize) -> usize { + let mut x: usize = w; + let mut mi: u32 = 0; + while (mi < 4) { + let off: u32 = ((mi * 16) * SMEM_STRIDE); + x = emit_lds(out, cap, x, "%fa", (mi * 4), "%r25", off); + x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 1), "%r25", (off + (8 * SMEM_STRIDE))); + x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 2), "%r25", (off + 16)); + x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 3), "%r25", ((off + (8 * SMEM_STRIDE)) + 16)); + mi += 1; + } + let mut ni: u32 = 0; + while (ni < 4) { + x = emit_lds(out, cap, x, "%fb", (ni * 2), "%r27", ((ni * 8) * SMEM_STRIDE)); + x = emit_lds(out, cap, x, "%fb", ((ni * 2) + 1), "%r27", (((ni * 8) * SMEM_STRIDE) + 16)); + ni += 1; + } + mi = 0; + while (mi < 4) { + ni = 0; + while (ni < 4) { + let b: u32 = (((mi * 4) + ni) * 4); + x = put_s(out, cap, x, " mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "); + x = put_quad(out, cap, x, "%acc", b); + x = put_s(out, cap, x, ", "); + x = put_quad(out, cap, x, "%fa", (mi * 4)); + x = put_s(out, cap, x, ", {%fb"); + x = put_u(out, cap, x, (ni * 2)); + x = put_s(out, cap, x, ", %fb"); + x = put_u(out, cap, x, ((ni * 2) + 1)); + x = put_s(out, cap, x, "}, "); + x = put_quad(out, cap, x, "%acc", b); + x = put_s(out, cap, x, ";\n"); + ni += 1; + } + mi += 1; + } + return x; +} + +pub fn emit_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 st.global.v2.u32 [%rd10], {%acc"); + x = put_u(out, cap, x, b); + x = put_s(out, cap, x, ", %acc"); + x = put_u(out, cap, x, (b + 1)); + return put_s(out, cap, x, "};\n"); +} + +pub fn ins(out: &mut [u8], cap: usize, w: usize, s: &'static str) -> usize { + let mut x: usize = put_s(out, cap, w, " "); + x = put_s(out, cap, x, s); + return put_s(out, cap, x, "\n"); +} + +pub fn put_quad(out: &mut [u8], cap: usize, w: usize, name: &'static str, first: u32) -> usize { + let mut x: usize = put_s(out, cap, w, "{"); + let mut j: u32 = 0; + while (j < 4) { + if (j > 0) { + x = put_s(out, cap, x, ", "); + } + x = put_s(out, cap, x, name); + x = put_u(out, cap, x, (first + j)); + j += 1; + } + return put_s(out, cap, x, "}"); +} + +pub fn put_s(out: &mut [u8], cap: usize, w: usize, s: &'static str) -> usize { + let mut i: usize = 0; + while (i < s.len()) { + if ((w + i) < cap) { + out[((w + i)) as usize] = s.as_bytes()[(i) as usize]; + } + i += 1; + } + return (w + s.len()); +} + +pub fn put_u(out: &mut [u8], cap: usize, w: usize, v: u32) -> usize { + let mut d: u32 = 1; + while ((v / d) >= 10) { + d = (d * 10); + } + let mut x: usize = w; + while (d > 0) { + let c: u8 = (b'0' + (((v / d) % 10) as u8)); + if (x < cap) { + out[(x) as usize] = c; + } + x += 1; + d = (d / 10); + } + return x; +} + +pub fn weight(code: u32) -> i32 { + return (((code & 1) as i32) - (((code >> 1) & 1) as i32)); } diff --git a/bootstrap/gen/rust/igla/coder/ternary_gemm_ptx.rs b/bootstrap/gen/rust/igla/coder/ternary_gemm_ptx.rs index b57d529427..c51085eaec 100644 --- a/bootstrap/gen/rust/igla/coder/ternary_gemm_ptx.rs +++ b/bootstrap/gen/rust/igla/coder/ternary_gemm_ptx.rs @@ -28,54 +28,6 @@ pub const D_MODEL_500M: usize = 1280; pub const D_FF_500M: usize = 4096; -pub fn put_s(out: &mut [u8], cap: usize, w: usize, s: &'static str) -> usize { - let mut i: usize = 0; - while (i < s.len()) { - if ((w + i) < cap) { - out[((w + i)) as usize] = s.as_bytes()[(i) as usize]; - } - i += 1; - } - return (w + s.len()); -} - -pub fn put_u(out: &mut [u8], cap: usize, w: usize, v: u32) -> usize { - let mut d: u32 = 1; - while ((v / d) >= 10) { - d = (d * 10); - } - let mut x: usize = w; - while (d > 0) { - let c: u8 = (b'0' + (((v / d) % 10) as u8)); - if (x < cap) { - out[(x) as usize] = c; - } - x += 1; - d = (d / 10); - } - return x; -} - -pub fn ins(out: &mut [u8], cap: usize, w: usize, s: &'static str) -> usize { - let mut x: usize = put_s(out, cap, w, " "); - x = put_s(out, cap, x, s); - return put_s(out, cap, x, "\n"); -} - -pub fn put_quad(out: &mut [u8], cap: usize, w: usize, name: &'static str, first: u32) -> usize { - let mut x: usize = put_s(out, cap, w, "{"); - let mut j: u32 = 0; - while (j < 4) { - if (j > 0) { - x = put_s(out, cap, x, ", "); - } - x = put_s(out, cap, x, name); - x = put_u(out, cap, x, (first + j)); - j += 1; - } - return put_s(out, cap, x, "}"); -} - pub fn emit_header(out: &mut [u8], cap: usize, w: usize) -> usize { let mut x: usize = put_s(out, cap, w, "// ternary_gemm_s8 -- emitted by t27c from specs/igla/coder/ternary_gemm_ptx.t27 (#10404). Do not edit.\n"); x = put_s(out, cap, x, ".version 7.0\n.target sm_80\n.address_size 64\n\n"); @@ -197,102 +149,6 @@ pub fn emit_tile_loads(out: &mut [u8], cap: usize, w: usize) -> usize { return x; } -pub fn emit_lds(out: &mut [u8], cap: usize, w: usize, reg: &'static str, i: u32, base: &'static str, off: u32) -> usize { - let mut x: usize = put_s(out, cap, w, " ld.shared.u32 "); - x = put_s(out, cap, x, reg); - x = put_u(out, cap, x, i); - x = put_s(out, cap, x, ", ["); - x = put_s(out, cap, x, base); - x = put_s(out, cap, x, "+"); - x = put_u(out, cap, x, off); - return put_s(out, cap, x, "];\n"); -} - -pub fn emit_mma_step(out: &mut [u8], cap: usize, w: usize) -> usize { - let mut x: usize = w; - let mut mi: u32 = 0; - while (mi < 4) { - let off: u32 = ((mi * 16) * SMEM_STRIDE); - x = emit_lds(out, cap, x, "%fa", (mi * 4), "%r25", off); - x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 1), "%r25", (off + (8 * SMEM_STRIDE))); - x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 2), "%r25", (off + 16)); - x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 3), "%r25", ((off + (8 * SMEM_STRIDE)) + 16)); - mi += 1; - } - let mut ni: u32 = 0; - while (ni < 4) { - x = emit_lds(out, cap, x, "%fb", (ni * 2), "%r27", ((ni * 8) * SMEM_STRIDE)); - x = emit_lds(out, cap, x, "%fb", ((ni * 2) + 1), "%r27", (((ni * 8) * SMEM_STRIDE) + 16)); - ni += 1; - } - mi = 0; - while (mi < 4) { - ni = 0; - while (ni < 4) { - let b: u32 = (((mi * 4) + ni) * 4); - x = put_s(out, cap, x, " mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "); - x = put_quad(out, cap, x, "%acc", b); - x = put_s(out, cap, x, ", "); - x = put_quad(out, cap, x, "%fa", (mi * 4)); - x = put_s(out, cap, x, ", {%fb"); - x = put_u(out, cap, x, (ni * 2)); - x = put_s(out, cap, x, ", %fb"); - x = put_u(out, cap, x, ((ni * 2) + 1)); - x = put_s(out, cap, x, "}, "); - x = put_quad(out, cap, x, "%acc", b); - x = put_s(out, cap, x, ";\n"); - ni += 1; - } - mi += 1; - } - return x; -} - -pub fn emit_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 st.global.v2.u32 [%rd10], {%acc"); - x = put_u(out, cap, x, b); - x = put_s(out, cap, x, ", %acc"); - x = put_u(out, cap, x, (b + 1)); - return put_s(out, cap, x, "};\n"); -} - -pub fn emit_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_store(out, cap, x, (mi * 16), (ni * 8), b); - x = emit_store(out, cap, x, ((mi * 16) + 8), (ni * 8), (b + 2)); - ni += 1; - } - mi += 1; - } - return x; -} - pub fn emit_ptx(out: &mut [u8], cap: usize) -> usize { let mut x: usize = emit_header(out, cap, 0); x = emit_prologue(out, cap, x); @@ -361,16 +217,6 @@ pub fn gen_wp(wp: &mut [u8], count: usize, seed: u64) -> u64 { return s; } -pub fn code_at(wp: &[u8], i: usize) -> u32 { - let b: u32 = (wp[((i / 4)) as usize] as u32); - let sh: u32 = (((i % 4) * 2) as u32); - return ((b >> sh) & 3); -} - -pub fn weight(code: u32) -> i32 { - return (((code & 1) as i32) - (((code >> 1) & 1) as i32)); -} - pub fn ref_row(a: impl AsRef<[i8]>, wp: &[u8], n: usize, k: usize, row: usize, acc: &mut [i32]) -> () { let a: &[i8] = a.as_ref(); let mut j: usize = 0; @@ -434,6 +280,12 @@ pub fn at_text(t: &[u8], n: usize, at: usize, s: &'static str) -> bool { return true; } +pub fn code_at(wp: &[u8], i: usize) -> u32 { + let b: u32 = (wp[((i / 4)) as usize] as u32); + let sh: u32 = (((i % 4) * 2) as u32); + return ((b >> sh) & 3); +} + pub fn count_text(t: &[u8], n: usize, s: &'static str) -> usize { let mut c: usize = 0; let mut i: usize = 0; @@ -446,3 +298,151 @@ pub fn count_text(t: &[u8], n: usize, s: &'static str) -> usize { return c; } +pub fn emit_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_store(out, cap, x, (mi * 16), (ni * 8), b); + x = emit_store(out, cap, x, ((mi * 16) + 8), (ni * 8), (b + 2)); + ni += 1; + } + mi += 1; + } + return x; +} + +pub fn emit_lds(out: &mut [u8], cap: usize, w: usize, reg: &'static str, i: u32, base: &'static str, off: u32) -> usize { + let mut x: usize = put_s(out, cap, w, " ld.shared.u32 "); + x = put_s(out, cap, x, reg); + x = put_u(out, cap, x, i); + x = put_s(out, cap, x, ", ["); + x = put_s(out, cap, x, base); + x = put_s(out, cap, x, "+"); + x = put_u(out, cap, x, off); + return put_s(out, cap, x, "];\n"); +} + +pub fn emit_mma_step(out: &mut [u8], cap: usize, w: usize) -> usize { + let mut x: usize = w; + let mut mi: u32 = 0; + while (mi < 4) { + let off: u32 = ((mi * 16) * SMEM_STRIDE); + x = emit_lds(out, cap, x, "%fa", (mi * 4), "%r25", off); + x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 1), "%r25", (off + (8 * SMEM_STRIDE))); + x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 2), "%r25", (off + 16)); + x = emit_lds(out, cap, x, "%fa", ((mi * 4) + 3), "%r25", ((off + (8 * SMEM_STRIDE)) + 16)); + mi += 1; + } + let mut ni: u32 = 0; + while (ni < 4) { + x = emit_lds(out, cap, x, "%fb", (ni * 2), "%r27", ((ni * 8) * SMEM_STRIDE)); + x = emit_lds(out, cap, x, "%fb", ((ni * 2) + 1), "%r27", (((ni * 8) * SMEM_STRIDE) + 16)); + ni += 1; + } + mi = 0; + while (mi < 4) { + ni = 0; + while (ni < 4) { + let b: u32 = (((mi * 4) + ni) * 4); + x = put_s(out, cap, x, " mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "); + x = put_quad(out, cap, x, "%acc", b); + x = put_s(out, cap, x, ", "); + x = put_quad(out, cap, x, "%fa", (mi * 4)); + x = put_s(out, cap, x, ", {%fb"); + x = put_u(out, cap, x, (ni * 2)); + x = put_s(out, cap, x, ", %fb"); + x = put_u(out, cap, x, ((ni * 2) + 1)); + x = put_s(out, cap, x, "}, "); + x = put_quad(out, cap, x, "%acc", b); + x = put_s(out, cap, x, ";\n"); + ni += 1; + } + mi += 1; + } + return x; +} + +pub fn emit_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 st.global.v2.u32 [%rd10], {%acc"); + x = put_u(out, cap, x, b); + x = put_s(out, cap, x, ", %acc"); + x = put_u(out, cap, x, (b + 1)); + return put_s(out, cap, x, "};\n"); +} + +pub fn ins(out: &mut [u8], cap: usize, w: usize, s: &'static str) -> usize { + let mut x: usize = put_s(out, cap, w, " "); + x = put_s(out, cap, x, s); + return put_s(out, cap, x, "\n"); +} + +pub fn put_quad(out: &mut [u8], cap: usize, w: usize, name: &'static str, first: u32) -> usize { + let mut x: usize = put_s(out, cap, w, "{"); + let mut j: u32 = 0; + while (j < 4) { + if (j > 0) { + x = put_s(out, cap, x, ", "); + } + x = put_s(out, cap, x, name); + x = put_u(out, cap, x, (first + j)); + j += 1; + } + return put_s(out, cap, x, "}"); +} + +pub fn put_s(out: &mut [u8], cap: usize, w: usize, s: &'static str) -> usize { + let mut i: usize = 0; + while (i < s.len()) { + if ((w + i) < cap) { + out[((w + i)) as usize] = s.as_bytes()[(i) as usize]; + } + i += 1; + } + return (w + s.len()); +} + +pub fn put_u(out: &mut [u8], cap: usize, w: usize, v: u32) -> usize { + let mut d: u32 = 1; + while ((v / d) >= 10) { + d = (d * 10); + } + let mut x: usize = w; + while (d > 0) { + let c: u8 = (b'0' + (((v / d) % 10) as u8)); + if (x < cap) { + out[(x) as usize] = c; + } + x += 1; + d = (d / 10); + } + return x; +} + +pub fn weight(code: u32) -> i32 { + return (((code & 1) as i32) - (((code >> 1) & 1) as i32)); +} + diff --git a/specs/igla/coder/hslm_bwd_ptx.t27 b/specs/igla/coder/hslm_bwd_ptx.t27 index 96f8771af2..d922270533 100644 --- a/specs/igla/coder/hslm_bwd_ptx.t27 +++ b/specs/igla/coder/hslm_bwd_ptx.t27 @@ -41,6 +41,7 @@ ; ASCII only (L3). phi^2 + 1/phi^2 = 3 | TRINITY module igla-coder-hslm-bwd-ptx; +use igla::coder::ptx_emit; /// Kernel ids as `t27c ptx-run --kernel` takes them; 0 is the forward kernel of ternary_gemm_ptx.t27. pub const KERNEL_FWD : u32 = 0; @@ -70,51 +71,6 @@ pub const HSLM_ROWS : usize = 1782; pub const HSLM_EMBED_PAD : usize = 256; pub const HSLM_FF_PAD : usize = 736; -// ============================================================================ -// 1. A bounded text writer (as in ternary_gemm_ptx.t27) -// ============================================================================ - -fn put_s(out : []u8, cap : usize, w : usize, s : str) -> usize { - var i : usize = 0; - while (i < s.len) { - if (w + i < cap) { out[w + i] = s[i]; } - i += 1; - } - return w + s.len; -} - -fn put_u(out : []u8, cap : usize, w : usize, v : u32) -> usize { - var d : u32 = 1; - while (v / d >= 10) { d = d * 10; } - var x : usize = w; - while (d > 0) { - const c : u8 = '0' + (((v / d) % 10) as u8); - if (x < cap) { out[x] = c; } - x += 1; - d = d / 10; - } - return x; -} - -fn ins(out : []u8, cap : usize, w : usize, s : str) -> usize { - // One writer call per statement: a nested call would borrow `out` twice in the Rust that gen-rust writes. - var x : usize = put_s(out, cap, w, " "); - x = put_s(out, cap, x, s); - return put_s(out, cap, x, "\n"); -} - -fn put_quad(out : []u8, cap : usize, w : usize, name : str, first : u32) -> usize { - var x : usize = put_s(out, cap, w, "{"); - var j : u32 = 0; - while (j < 4) { - if (j > 0) { x = put_s(out, cap, x, ", "); } - x = put_s(out, cap, x, name); - x = put_u(out, cap, x, first + j); - j += 1; - } - return put_s(out, cap, x, "}"); -} - fn put_name_at(out : []u8, cap : usize, w : usize, kernel : u32) -> usize { if (kernel == KERNEL_FWD) { return put_s(out, cap, w, "ternary_gemm_s8"); } if (kernel == KERNEL_DX) { return put_s(out, cap, w, "hslm_dx_s8"); } @@ -359,104 +315,6 @@ fn emit_loads_dw(out : []u8, cap : usize, w : usize) -> usize { return emit_transposed(out, cap, x, 50, "%r16", "%r37"); } -fn emit_lds(out : []u8, cap : usize, w : usize, reg : str, i : u32, base : str, off : u32) -> usize { - var x : usize = put_s(out, cap, w, " ld.shared.u32 "); - x = put_s(out, cap, x, reg); - x = put_u(out, cap, x, i); - x = put_s(out, cap, x, ", ["); - x = put_s(out, cap, x, base); - x = put_s(out, cap, x, "+"); - x = put_u(out, cap, x, off); - return put_s(out, cap, x, "];\n"); -} - -/// The fragment loads and the 16 mma.sync of one k-step, the forward kernel's. -fn emit_mma_step(out : []u8, cap : usize, w : usize) -> usize { - var x : usize = w; - var mi : u32 = 0; - while (mi < 4) { - const off : u32 = mi * 16 * SMEM_STRIDE; - x = emit_lds(out, cap, x, "%fa", mi * 4, "%r25", off); - x = emit_lds(out, cap, x, "%fa", mi * 4 + 1, "%r25", off + 8 * SMEM_STRIDE); - x = emit_lds(out, cap, x, "%fa", mi * 4 + 2, "%r25", off + 16); - x = emit_lds(out, cap, x, "%fa", mi * 4 + 3, "%r25", off + 8 * SMEM_STRIDE + 16); - mi += 1; - } - var ni : u32 = 0; - while (ni < 4) { - x = emit_lds(out, cap, x, "%fb", ni * 2, "%r27", ni * 8 * SMEM_STRIDE); - x = emit_lds(out, cap, x, "%fb", ni * 2 + 1, "%r27", ni * 8 * SMEM_STRIDE + 16); - ni += 1; - } - mi = 0; - while (mi < 4) { - ni = 0; - while (ni < 4) { - const b : u32 = (mi * 4 + ni) * 4; - x = put_s(out, cap, x, " mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "); - x = put_quad(out, cap, x, "%acc", b); - x = put_s(out, cap, x, ", "); - x = put_quad(out, cap, x, "%fa", mi * 4); - x = put_s(out, cap, x, ", {%fb"); - x = put_u(out, cap, x, ni * 2); - x = put_s(out, cap, x, ", %fb"); - x = put_u(out, cap, x, ni * 2 + 1); - x = put_s(out, cap, x, "}, "); - x = put_quad(out, cap, x, "%acc", b); - x = put_s(out, cap, x, ";\n"); - ni += 1; - } - mi += 1; - } - return x; -} - -fn emit_store(out : []u8, cap : usize, w : usize, dr : u32, dc : u32, b : u32) -> usize { - var 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 st.global.v2.u32 [%rd10], {%acc"); - x = put_u(out, cap, x, b); - x = put_s(out, cap, x, ", %acc"); - x = put_u(out, cap, x, b + 1); - return put_s(out, cap, x, "};\n"); -} - -/// The 32 guarded pair stores of the forward kernel's epilogue. -fn emit_epilogue(out : []u8, cap : usize, w : usize) -> usize { - var 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;"); - var mi : u32 = 0; - while (mi < 4) { - var ni : u32 = 0; - while (ni < 4) { - const b : u32 = (mi * 4 + ni) * 4; - x = emit_store(out, cap, x, mi * 16, ni * 8, b); - x = emit_store(out, cap, x, mi * 16 + 8, ni * 8, b + 2); - ni += 1; - } - mi += 1; - } - return x; -} - /// The whole text of kernel 1 or 2 into out[0..cap). Returns the bytes it needs (more than cap means /// a clipped text); an unknown id writes nothing and returns 0. pub fn emit_ptx(kernel : u32, out : []u8, cap : usize) -> usize { @@ -517,18 +375,6 @@ pub fn b_len(kernel : u32, n : usize, k : usize) -> usize { // 4. The integer references and the mismatch count // ============================================================================ -/// Code i of a packed code matrix, lane 0 in the low bits. -pub fn code_at(wp : []const u8, i : usize) -> u32 { - const b : u32 = wp[i / 4] as u32; - const sh : u32 = ((i % 4) * 2) as u32; - return (b >> sh) & 3; -} - -/// The decode of ternary_gemm_gpu.t27: 0 -> 0, 1 -> +1, 2 -> -1, 3 -> 0. -pub fn weight(code : u32) -> i32 { - return ((code & 1) as i32) - (((code >> 1) & 1) as i32); -} - /// A byte of B read as s8. pub fn sbyte(v : u8) -> i32 { const u : i32 = v as i32; @@ -648,30 +494,6 @@ pub fn decoded_word(word : u32, wq : u32) -> u32 { return prmt_table(sp >> 16); } -// ============================================================================ -// 6. Test helpers -// ============================================================================ - -fn at_text(t : []const u8, n : usize, at : usize, s : str) -> bool { - if (at + s.len > n) { return false; } - var i : usize = 0; - while (i < s.len) { - if (t[at + i] != s[i]) { return false; } - i += 1; - } - return true; -} - -fn count_text(t : []const u8, n : usize, s : str) -> usize { - var c : usize = 0; - var i : usize = 0; - while (i < n) { - if (at_text(t, n, i, s)) { c += 1; } - i += 1; - } - return c; -} - /// 31-bit LCG (glibc constants), the generator of ternary_gemm_ptx.t27. fn lcg(s : u64) -> u64 { return (s * 1103515245 + 12345) % 2147483648; diff --git a/specs/igla/coder/ptx_emit.t27 b/specs/igla/coder/ptx_emit.t27 new file mode 100644 index 0000000000..bfea1af7e4 --- /dev/null +++ b/specs/igla/coder/ptx_emit.t27 @@ -0,0 +1,295 @@ +// SPDX-License-Identifier: Apache-2.0 +; specs/igla/coder/ptx_emit.t27 -- the PTX text helpers the two HSLM kernel specs share. +; Issue #10930 (duplicate-bodies, 12 helpers copied between ternary_gemm_ptx.t27 and hslm_bwd_ptx.t27). +; -- +; What is here. A bounded text writer (put_s, put_u, ins, put_quad), the forward kernel's fragment +; loads, mma step and epilogue (emit_lds, emit_mma_step, emit_store, emit_epilogue), the 2-bit code +; decode (code_at, weight) and two searches over emitted text (at_text, count_text). Both +; ternary_gemm_ptx.t27 (#10404) and hslm_bwd_ptx.t27 (#10847) `use` this module; neither keeps a copy. +; The emitters read SMEM_STRIDE below, which is the value both kernels declare. +; ASCII only (L3). phi^2 + 1/phi^2 = 3 | TRINITY + +module igla-coder-ptx-emit; + +/// sA and sB row stride in bytes: 32 data + 16 pad. +pub const SMEM_STRIDE : u32 = 48; + +// ============================================================================ +// 1. A bounded text writer +// ============================================================================ + +/// Write s at out[w..]; bytes past cap are counted but not written. Returns the new end. +pub fn put_s(out : []u8, cap : usize, w : usize, s : str) -> usize { + var i : usize = 0; + while (i < s.len) { + if (w + i < cap) { out[w + i] = s[i]; } + i += 1; + } + return w + s.len; +} + +/// Write v in decimal at out[w..], bounded by cap like put_s. +pub fn put_u(out : []u8, cap : usize, w : usize, v : u32) -> usize { + var d : u32 = 1; + while (v / d >= 10) { d = d * 10; } + var x : usize = w; + while (d > 0) { + const c : u8 = '0' + (((v / d) % 10) as u8); + if (x < cap) { out[x] = c; } + x += 1; + d = d / 10; + } + return x; +} + +/// One indented instruction line that has no operand to fill in. +pub fn ins(out : []u8, cap : usize, w : usize, s : str) -> usize { + // One writer call per statement: a nested call would borrow `out` twice in the Rust that gen-rust writes. + var x : usize = put_s(out, cap, w, " "); + x = put_s(out, cap, x, s); + return put_s(out, cap, x, "\n"); +} + +/// "{%name, %name, %name, %name}" +pub fn put_quad(out : []u8, cap : usize, w : usize, name : str, first : u32) -> usize { + var x : usize = put_s(out, cap, w, "{"); + var j : u32 = 0; + while (j < 4) { + if (j > 0) { x = put_s(out, cap, x, ", "); } + x = put_s(out, cap, x, name); + x = put_u(out, cap, x, first + j); + j += 1; + } + return put_s(out, cap, x, "}"); +} + +// ============================================================================ +// 2. Fragment loads, mma step and epilogue of the forward kernel +// ============================================================================ + +/// " ld.shared.u32 %, [%+];" +pub fn emit_lds(out : []u8, cap : usize, w : usize, reg : str, i : u32, base : str, off : u32) -> usize { + var x : usize = put_s(out, cap, w, " ld.shared.u32 "); + x = put_s(out, cap, x, reg); + x = put_u(out, cap, x, i); + x = put_s(out, cap, x, ", ["); + x = put_s(out, cap, x, base); + x = put_s(out, cap, x, "+"); + x = put_u(out, cap, x, off); + return put_s(out, cap, x, "];\n"); +} + +/// The fragment loads and the 16 mma.sync of one k-step. +pub fn emit_mma_step(out : []u8, cap : usize, w : usize) -> usize { + var x : usize = w; + var mi : u32 = 0; + while (mi < 4) { + const off : u32 = mi * 16 * SMEM_STRIDE; + x = emit_lds(out, cap, x, "%fa", mi * 4, "%r25", off); + x = emit_lds(out, cap, x, "%fa", mi * 4 + 1, "%r25", off + 8 * SMEM_STRIDE); + x = emit_lds(out, cap, x, "%fa", mi * 4 + 2, "%r25", off + 16); + x = emit_lds(out, cap, x, "%fa", mi * 4 + 3, "%r25", off + 8 * SMEM_STRIDE + 16); + mi += 1; + } + var ni : u32 = 0; + while (ni < 4) { + x = emit_lds(out, cap, x, "%fb", ni * 2, "%r27", ni * 8 * SMEM_STRIDE); + x = emit_lds(out, cap, x, "%fb", ni * 2 + 1, "%r27", ni * 8 * SMEM_STRIDE + 16); + ni += 1; + } + mi = 0; + while (mi < 4) { + ni = 0; + while (ni < 4) { + const b : u32 = (mi * 4 + ni) * 4; + x = put_s(out, cap, x, " mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "); + x = put_quad(out, cap, x, "%acc", b); + x = put_s(out, cap, x, ", "); + x = put_quad(out, cap, x, "%fa", mi * 4); + x = put_s(out, cap, x, ", {%fb"); + x = put_u(out, cap, x, ni * 2); + x = put_s(out, cap, x, ", %fb"); + x = put_u(out, cap, x, ni * 2 + 1); + x = put_s(out, cap, x, "}, "); + x = put_quad(out, cap, x, "%acc", b); + x = put_s(out, cap, x, ";\n"); + ni += 1; + } + mi += 1; + } + return x; +} + +/// One guarded pair store: accumulators b, b+1 to C[row0 + dr][col0 + dc .. +1]. +pub fn emit_store(out : []u8, cap : usize, w : usize, dr : u32, dc : u32, b : u32) -> usize { + var 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 st.global.v2.u32 [%rd10], {%acc"); + x = put_u(out, cap, x, b); + x = put_s(out, cap, x, ", %acc"); + x = put_u(out, cap, x, b + 1); + return put_s(out, cap, x, "};\n"); +} + +/// Row base r39 = 128 by + 64 wm + g and column base r40 = 128 bx + 32 wn + 2t, then 32 stores. +pub fn emit_epilogue(out : []u8, cap : usize, w : usize) -> usize { + var 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;"); + var mi : u32 = 0; + while (mi < 4) { + var ni : u32 = 0; + while (ni < 4) { + const b : u32 = (mi * 4 + ni) * 4; + x = emit_store(out, cap, x, mi * 16, ni * 8, b); + x = emit_store(out, cap, x, mi * 16 + 8, ni * 8, b + 2); + ni += 1; + } + mi += 1; + } + return x; +} + +// ============================================================================ +// 3. The 2-bit code decode +// ============================================================================ + +/// Code i of the packed matrix. +pub fn code_at(wp : []const u8, i : usize) -> u32 { + const b : u32 = wp[i / 4] as u32; + const sh : u32 = ((i % 4) * 2) as u32; + return (b >> sh) & 3; +} + +/// The decode the kernel runs: (c & 1) - ((c >> 1) & 1); 0 -> 0, 1 -> +1, 2 -> -1, 3 -> 0. +pub fn weight(code : u32) -> i32 { + return ((code & 1) as i32) - (((code >> 1) & 1) as i32); +} + +// ============================================================================ +// 4. Search the emitted text +// ============================================================================ + +/// Whether t[at..] starts with s. +pub fn at_text(t : []const u8, n : usize, at : usize, s : str) -> bool { + if (at + s.len > n) { return false; } + var i : usize = 0; + while (i < s.len) { + if (t[at + i] != s[i]) { return false; } + i += 1; + } + return true; +} + +/// How many times s occurs in t[0..n). +pub fn count_text(t : []const u8, n : usize, s : str) -> usize { + var c : usize = 0; + var i : usize = 0; + while (i < n) { + if (at_text(t, n, i, s)) { c += 1; } + i += 1; + } + return c; +} + + +// ============================================================================ +// 5. Tests +// ============================================================================ + +test put_s_counts_every_byte_but_writes_only_below_cap { + var buf : [8]u8 = [0; 8]; + const e : usize = put_s(&buf, 4, 0, "abcdef"); + assert e == 6; + assert buf[0] == 'a'; + assert buf[3] == 'd'; + assert buf[4] == 0; +} + +test put_u_writes_decimal { + var buf : [8]u8 = [0; 8]; + assert put_u(&buf, 8, 0, 0) == 1; + assert buf[0] == '0'; + assert put_u(&buf, 8, 0, 4096) == 4; + assert buf[0] == '4'; + assert buf[1] == '0'; + assert buf[2] == '9'; + assert buf[3] == '6'; +} + +test ins_indents_by_four_and_ends_the_line { + var buf : [16]u8 = [0; 16]; + const e : usize = ins(&buf, 16, 0, "nop;"); + assert e == 9; + assert at_text(&buf, e, 0, " nop;\n"); +} + +test put_quad_names_four_consecutive_registers { + var buf : [32]u8 = [0; 32]; + const e : usize = put_quad(&buf, 32, 0, "%acc", 4); + assert e == 28; + assert at_text(&buf, e, 0, "{%acc4, %acc5, %acc6, %acc7}"); +} + +test emit_lds_writes_one_shared_load { + var buf : [64]u8 = [0; 64]; + const e : usize = emit_lds(&buf, 64, 0, "%fa", 3, "%r25", 48); + assert at_text(&buf, e, 0, " ld.shared.u32 %fa3, [%r25+48];\n"); +} + +test at_text_and_count_text_search_the_text { + var buf : [16]u8 = [0; 16]; + const e : usize = put_s(&buf, 16, 0, "abab ab"); + assert at_text(&buf, e, 2, "ab"); + assert at_text(&buf, e, 6, "abc") == false; + assert count_text(&buf, e, "ab") == 3; + assert count_text(&buf, e, "ba") == 1; +} + +test one_k_step_has_sixteen_mma_and_the_fragment_loads { + var buf : [8192]u8 = undefined; + const e : usize = emit_mma_step(&buf, 8192, 0); + assert e < 8192; + assert count_text(&buf, e, "mma.sync") == 16; + assert count_text(&buf, e, "ld.shared.u32") == 24; +} + +test the_epilogue_has_thirty_two_guarded_stores { + var buf : [16384]u8 = undefined; + const e : usize = emit_epilogue(&buf, 16384, 0); + assert e < 16384; + assert count_text(&buf, e, "st.global.v2.u32") == 32; +} + +test the_decode_maps_codes_to_zero_plus_one_minus_one_zero { + assert weight(0) == 0; + assert weight(1) == 1; + assert weight(2) == -1; + assert weight(3) == 0; +} + +test code_at_reads_lane_zero_from_the_low_bits { + var wp : [2]u8 = [0; 2]; + wp[0] = 0x1B; + assert code_at(&wp, 0) == 3; + assert code_at(&wp, 1) == 2; + assert code_at(&wp, 2) == 1; + assert code_at(&wp, 3) == 0; +} diff --git a/specs/igla/coder/ternary_gemm_ptx.t27 b/specs/igla/coder/ternary_gemm_ptx.t27 index 9df0bb900a..1485c3f4c3 100644 --- a/specs/igla/coder/ternary_gemm_ptx.t27 +++ b/specs/igla/coder/ternary_gemm_ptx.t27 @@ -46,6 +46,7 @@ ; ASCII only (L3). phi^2 + 1/phi^2 = 3 | TRINITY module igla-coder-ternary-gemm-ptx; +use igla::coder::ptx_emit; /// The kernel's entry name, the one cuModuleGetFunction looks up. pub const KERNEL_NAME : str = "ternary_gemm_s8"; @@ -68,55 +69,6 @@ pub const PTX_CAP : usize = 32768; pub const D_MODEL_500M : usize = 1280; pub const D_FF_500M : usize = 4096; -// ============================================================================ -// 1. A bounded text writer -// ============================================================================ - -/// Write s at out[w..]; bytes past cap are counted but not written. Returns the new end. -fn put_s(out : []u8, cap : usize, w : usize, s : str) -> usize { - var i : usize = 0; - while (i < s.len) { - if (w + i < cap) { out[w + i] = s[i]; } - i += 1; - } - return w + s.len; -} - -/// Write v in decimal at out[w..], bounded by cap like put_s. -fn put_u(out : []u8, cap : usize, w : usize, v : u32) -> usize { - var d : u32 = 1; - while (v / d >= 10) { d = d * 10; } - var x : usize = w; - while (d > 0) { - const c : u8 = '0' + (((v / d) % 10) as u8); - if (x < cap) { out[x] = c; } - x += 1; - d = d / 10; - } - return x; -} - -/// One indented instruction line that has no operand to fill in. -fn ins(out : []u8, cap : usize, w : usize, s : str) -> usize { - // One writer call per statement: a nested call would borrow `out` twice in the Rust that gen-rust writes. - var x : usize = put_s(out, cap, w, " "); - x = put_s(out, cap, x, s); - return put_s(out, cap, x, "\n"); -} - -/// "{%name, %name, %name, %name}" -fn put_quad(out : []u8, cap : usize, w : usize, name : str, first : u32) -> usize { - var x : usize = put_s(out, cap, w, "{"); - var j : u32 = 0; - while (j < 4) { - if (j > 0) { x = put_s(out, cap, x, ", "); } - x = put_s(out, cap, x, name); - x = put_u(out, cap, x, first + j); - j += 1; - } - return put_s(out, cap, x, "}"); -} - // ============================================================================ // 2. The kernel, section by section // ============================================================================ @@ -248,106 +200,6 @@ fn emit_tile_loads(out : []u8, cap : usize, w : usize) -> usize { return x; } -/// " ld.shared.u32 %, [%+];" -fn emit_lds(out : []u8, cap : usize, w : usize, reg : str, i : u32, base : str, off : u32) -> usize { - var x : usize = put_s(out, cap, w, " ld.shared.u32 "); - x = put_s(out, cap, x, reg); - x = put_u(out, cap, x, i); - x = put_s(out, cap, x, ", ["); - x = put_s(out, cap, x, base); - x = put_s(out, cap, x, "+"); - x = put_u(out, cap, x, off); - return put_s(out, cap, x, "];\n"); -} - -/// The fragment loads and the 16 mma.sync of one k-step. -fn emit_mma_step(out : []u8, cap : usize, w : usize) -> usize { - var x : usize = w; - var mi : u32 = 0; - while (mi < 4) { - const off : u32 = mi * 16 * SMEM_STRIDE; - x = emit_lds(out, cap, x, "%fa", mi * 4, "%r25", off); - x = emit_lds(out, cap, x, "%fa", mi * 4 + 1, "%r25", off + 8 * SMEM_STRIDE); - x = emit_lds(out, cap, x, "%fa", mi * 4 + 2, "%r25", off + 16); - x = emit_lds(out, cap, x, "%fa", mi * 4 + 3, "%r25", off + 8 * SMEM_STRIDE + 16); - mi += 1; - } - var ni : u32 = 0; - while (ni < 4) { - x = emit_lds(out, cap, x, "%fb", ni * 2, "%r27", ni * 8 * SMEM_STRIDE); - x = emit_lds(out, cap, x, "%fb", ni * 2 + 1, "%r27", ni * 8 * SMEM_STRIDE + 16); - ni += 1; - } - mi = 0; - while (mi < 4) { - ni = 0; - while (ni < 4) { - const b : u32 = (mi * 4 + ni) * 4; - x = put_s(out, cap, x, " mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "); - x = put_quad(out, cap, x, "%acc", b); - x = put_s(out, cap, x, ", "); - x = put_quad(out, cap, x, "%fa", mi * 4); - x = put_s(out, cap, x, ", {%fb"); - x = put_u(out, cap, x, ni * 2); - x = put_s(out, cap, x, ", %fb"); - x = put_u(out, cap, x, ni * 2 + 1); - x = put_s(out, cap, x, "}, "); - x = put_quad(out, cap, x, "%acc", b); - x = put_s(out, cap, x, ";\n"); - ni += 1; - } - mi += 1; - } - return x; -} - -/// One guarded pair store: accumulators b, b+1 to C[row0 + dr][col0 + dc .. +1]. -fn emit_store(out : []u8, cap : usize, w : usize, dr : u32, dc : u32, b : u32) -> usize { - var 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 st.global.v2.u32 [%rd10], {%acc"); - x = put_u(out, cap, x, b); - x = put_s(out, cap, x, ", %acc"); - x = put_u(out, cap, x, b + 1); - return put_s(out, cap, x, "};\n"); -} - -/// Row base r39 = 128 by + 64 wm + g and column base r40 = 128 bx + 32 wn + 2t, then 32 stores. -fn emit_epilogue(out : []u8, cap : usize, w : usize) -> usize { - var 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;"); - var mi : u32 = 0; - while (mi < 4) { - var ni : u32 = 0; - while (ni < 4) { - const b : u32 = (mi * 4 + ni) * 4; - x = emit_store(out, cap, x, mi * 16, ni * 8, b); - x = emit_store(out, cap, x, mi * 16 + 8, ni * 8, b + 2); - ni += 1; - } - mi += 1; - } - return x; -} - /// The whole kernel into out[0..cap). Returns the bytes it needs; when that exceeds cap the buffer /// holds a clipped text the caller must not use. pub fn emit_ptx(out : []u8, cap : usize) -> usize { @@ -426,18 +278,6 @@ pub fn gen_wp(wp : []u8, count : usize, seed : u64) -> u64 { return s; } -/// Code i of the packed matrix. -pub fn code_at(wp : []const u8, i : usize) -> u32 { - const b : u32 = wp[i / 4] as u32; - const sh : u32 = ((i % 4) * 2) as u32; - return (b >> sh) & 3; -} - -/// The decode the kernel runs: (c & 1) - ((c >> 1) & 1); 0 -> 0, 1 -> +1, 2 -> -1, 3 -> 0. -pub fn weight(code : u32) -> i32 { - return ((code & 1) as i32) - (((code >> 1) & 1) as i32); -} - /// acc[0..n) = row `row` of A x W, k-outer so the codes are read in memory order. pub fn ref_row(a : []const i8, wp : []const u8, n : usize, k : usize, row : usize, acc : []i32) -> void { var j : usize = 0; @@ -483,32 +323,6 @@ pub fn mismatches(a : []const i8, wp : []const u8, c : []const i32, m : usize, n return bad; } -// ============================================================================ -// 5. Test helpers: search the emitted text -// ============================================================================ - -/// Whether t[at..] starts with s. -fn at_text(t : []const u8, n : usize, at : usize, s : str) -> bool { - if (at + s.len > n) { return false; } - var i : usize = 0; - while (i < s.len) { - if (t[at + i] != s[i]) { return false; } - i += 1; - } - return true; -} - -/// How many times s occurs in t[0..n). -fn count_text(t : []const u8, n : usize, s : str) -> usize { - var c : usize = 0; - var i : usize = 0; - while (i < n) { - if (at_text(t, n, i, s)) { c += 1; } - i += 1; - } - return c; -} - // ============================================================================ // 6. Tests // ============================================================================