include!("../build_support/msl_postpass.rs");
const GENERATED: &str = include_str!("../src/backend/shaders/slang/gemm_q8_0.metal");
fn patched() -> String {
postpass_gemm_msl(GENERATED).expect("post-pass should apply to committed slangc output")
}
#[test]
fn applies_all_three_conditions() {
let out = patched();
assert!(
out.contains("threadgroup char* shmem_p_0 [[threadgroup(0)]]"),
"kernel should take a threadgroup scratch parameter"
);
assert!(
!out.contains("threadgroup array<half, int("),
"static half scratch should be gone"
);
assert!(
!out.contains("threadgroup array<float, int("),
"static float scratch should be gone"
);
assert!(
out.contains("threadgroup const half* pa_1 = pa_0 + "),
"half scratch should advance by pointer"
);
assert!(
out.contains("threadgroup const float* pb_1 = pb_0 + "),
"float scratch should advance by pointer"
);
assert!(
!out.contains("(_sa)[0]) + (lsma"),
"no matrix load should still address through the uint index"
);
assert!(
out.contains("#pragma unroll(4)"),
"k-loop should carry its trip count as an unroll factor"
);
}
#[test]
fn derives_the_scratch_split_from_the_declared_extent() {
assert!(
patched().contains("(threadgroup float*)(shmem_p_0 + 4096)"),
"2048 halves should put the float scratch at byte 4096"
);
let retiled = GENERATED
.replace("array<half, int(2048)>", "array<half, int(1024)>")
.replace("array<float, int(1024)>", "array<float, int(1536)>");
assert_ne!(retiled, GENERATED, "fixture should contain both extents");
let out = postpass_gemm_msl(&retiled).expect("retiled fixture should still apply");
assert!(
out.contains("(threadgroup float*)(shmem_p_0 + 2048)"),
"1024 halves should move the float scratch to byte 2048"
);
}
#[test]
fn declines_when_scratch_exceeds_what_callers_bind() {
let big = GENERATED.replace("array<half, int(2048)>", "array<half, int(4096)>");
assert_ne!(big, GENERATED, "fixture should contain the extent");
let err = postpass_gemm_msl(&big).expect_err("should decline on oversized scratch");
assert!(err.contains("callers bind"), "unexpected error: {err}");
}
#[test]
fn declines_on_a_misaligned_scratch_split() {
let odd = GENERATED
.replace("array<half, int(2048)>", "array<half, int(2050)>")
.replace("array<float, int(1024)>", "array<float, int(1023)>");
assert_ne!(odd, GENERATED, "fixture should contain both extents");
let err = postpass_gemm_msl(&odd).expect_err("should decline on a misaligned split");
assert!(err.contains("multiple of 16"), "unexpected error: {err}");
}
#[test]
fn declines_when_the_k_loop_index_is_written_again() {
let step = " uint lsma_1 = lsma_0 + 512U;";
let after_step = " lsmb_0 = lsmb_1;";
let header = " for(;;)";
let cases = [
(step, format!(" lsma_0 = lsma_0 ^ 64U;\n{step}")),
(
after_step,
format!("{after_step}\n lsma_0 = lsma_0 + 8U;"),
),
(
header,
format!(" lsma_0 = lsma_0 + 8U;\n{header}"),
),
(step, format!(" lsma_0 += 8U;\n{step}")),
(step, format!(" lsma_0++;\n{step}")),
];
for (anchor, replacement) in cases {
let mutated = GENERATED.replacen(anchor, &replacement, 1);
assert_ne!(mutated, GENERATED, "fixture should contain {anchor:?}");
let err =
postpass_gemm_msl(&mutated).expect_err(&format!("should decline for {replacement:?}"));
assert!(
err.contains("to be written 2 time(s)"),
"unexpected error for {replacement:?}: {err}"
);
}
}
#[test]
fn seeds_the_pointer_from_the_index_not_a_second_copy() {
let out = patched();
assert!(
out.contains("pa_0 = _sa + lsma_0;") && out.contains("pb_0 = _sb + lsmb_0;"),
"pointers should be seeded from the index variables"
);
let seed = " lsma_0 = _S26;";
let odd = GENERATED.replace(seed, " lsma_0 = _S26++;");
assert_ne!(odd, GENERATED, "fixture should contain the seed");
let out = postpass_gemm_msl(&odd).expect("a side-effecting seed should still apply");
assert_eq!(
out.matches("_S26++").count(),
1,
"the seed expression must not be duplicated onto the pointer"
);
let odd = GENERATED.replace(seed, " lsma_0 = _S26, zz_0 = 1U;");
assert_ne!(odd, GENERATED, "fixture should contain the seed");
postpass_gemm_msl(&odd).expect("a comma seed should apply, since it is not re-spliced");
}
#[test]
fn declines_on_a_non_literal_index_step() {
let step = " uint lsma_1 = lsma_0 + 512U;";
for spelling in [
" uint lsma_1 = lsma_0 + 256U << 1U;",
" uint lsma_1 = lsma_0 + 512U | 1U;",
" uint lsma_1 = lsma_0 + kStride;",
] {
let odd = GENERATED.replacen(step, spelling, 1);
assert_ne!(odd, GENERATED, "fixture should contain the step");
let err = postpass_gemm_msl(&odd).expect_err(&format!("should decline for {spelling:?}"));
assert!(
err.contains("not a plain integer literal"),
"unexpected error for {spelling:?}: {err}"
);
}
let out = patched();
assert!(
out.contains("pa_1 = pa_0 + (512U);"),
"the plain literal step should still be mirrored onto the pointer"
);
}
#[test]
fn declines_when_a_load_respelling_defeats_the_rewrite() {
let spaced = "&((*(((&kernelContext_0)->sa_0)))[0]) + (lsma_0 + 64U)";
for respelled in [
"&((*(((&kernelContext_0)->sa_0)))[0]) + (lsma_0+64U)",
"&((*(((&kernelContext_0)->sa_0)))[int(0)]) + (lsma_0 + 64U)",
"&((*(((&kernelContext_0)->sa_0)))[0])+(lsma_0 + 64U)",
] {
let odd = GENERATED.replacen(spaced, respelled, 1);
assert_ne!(odd, GENERATED, "fixture should contain the load");
let Err(err) = postpass_gemm_msl(&odd) else {
panic!("should decline on the respelled load {respelled}");
};
assert!(
err.contains("not the pointer walk"),
"unexpected error for {respelled}: {err}"
);
}
}
#[test]
fn declines_when_the_load_base_is_hoisted_out_of_the_loop() {
let seed = " uint lsmb_0 = ";
let hoist = format!(
" threadgroup half* base_a_0 = &((*(((&kernelContext_0)->sa_0)))[0]);\n{seed}"
);
let hoisted = GENERATED.replacen(seed, &hoist, 1);
assert_ne!(hoisted, GENERATED, "fixture should contain the index seed");
let all = hoisted.replace(
"&((*(((&kernelContext_0)->sa_0)))[0]) + (lsma_0",
"base_a_0 + (lsma_0",
);
let one = hoisted.replacen(
"&((*(((&kernelContext_0)->sa_0)))[0]) + (lsma_0 + 64U)",
"base_a_0 + (lsma_0 + 64U)",
1,
);
for (label, odd) in [("every load", all), ("one load", one)] {
assert_ne!(odd, hoisted, "fixture should contain the loads");
let Err(err) = postpass_gemm_msl(&odd) else {
panic!("should decline when {label} reads a hoisted base");
};
assert!(
err.contains("not the pointer walk"),
"unexpected error for {label}: {err}"
);
}
}
#[test]
fn reads_the_load_address_up_to_the_call_boundary() {
let load = "&((*(((&kernelContext_0)->sa_0)))[0]) + (lsma_0), (ulong)(8U)";
for (spelling, expected) in [
("zz_0", "\"zz_0\""),
("(zz_0), (ulong)(8U)", "\"(zz_0)\""),
] {
let odd = GENERATED.replacen(load, spelling, 1);
assert_ne!(odd, GENERATED, "fixture should contain the load");
let Err(err) = postpass_gemm_msl(&odd) else {
panic!("should decline on a load reading {spelling}");
};
assert!(
err.contains(&format!("a matrix load reads {expected},")),
"address read for {spelling} is not bounded by the call: {err}"
);
}
}
#[test]
fn declines_when_no_matrix_load_is_recognized() {
let renamed = GENERATED.replace("_slang_simdgroup_load<", "_slang_coopmat_read<");
assert_ne!(
renamed, GENERATED,
"fixture should contain the load wrapper"
);
let err = postpass_gemm_msl(&renamed).expect_err("should decline when no load is recognized");
assert!(
err.contains("nothing proves the loads were repointed"),
"unexpected error: {err}"
);
}
#[test]
fn declines_on_an_index_form_load_outside_the_k_loop() {
let epilogue = " lsma_0 = sv_groupindex_0;";
let injected = format!(
"{epilogue}\n (void)(&((*(((&kernelContext_0)->sa_0)))[0]) + (lsma_0 + 32U));"
);
let odd = GENERATED.replacen(epilogue, &injected, 1);
assert_ne!(odd, GENERATED, "fixture should contain the epilogue anchor");
let err = postpass_gemm_msl(&odd).expect_err("should decline on an epilogue index load");
assert!(
err.contains("still uses the index form"),
"unexpected error: {err}"
);
}
#[test]
fn declines_when_the_shader_already_has_a_directive() {
let header = " for(;;)\n";
for injected in [
format!(" #pragma unroll(2)\n{header}"),
format!(" #pragma unroll(2)\n#line 278\n{header}"),
format!(" #pragma unroll(2)\n // slang: k loop\n{header}"),
format!(" #pragma unroll(2)\n /* k */\n{header}"),
format!(" _Pragma(\"unroll(2)\")\n{header}"),
format!(" # pragma unroll(2)\n{header}"),
] {
let odd = GENERATED.replacen(header, &injected, 1);
assert_ne!(odd, GENERATED, "fixture should contain the loop header");
let err = postpass_gemm_msl(&odd).expect_err("should decline on an existing directive");
assert!(
err.contains("which this pass would collide with"),
"unexpected error for {injected:?}: {err}"
);
}
let far = GENERATED.replacen("[[kernel]]", "#pragma clang diagnostic push\n[[kernel]]", 1);
assert_ne!(far, GENERATED, "fixture should contain the entry point");
postpass_gemm_msl(&far).expect_err("should decline on a directive far from the k-loop");
}
#[test]
fn declines_when_the_k_loop_header_is_spelled_differently() {
let header = " for(;;)\n";
let odd = GENERATED.replacen(header, " while(true)\n", 1);
assert_ne!(odd, GENERATED, "fixture should contain the loop header");
let err = postpass_gemm_msl(&odd).expect_err("should decline on an unrecognized header");
assert!(
err.contains("not spelled for(;;)"),
"unexpected error: {err}"
);
}
#[test]
fn declines_when_an_introduced_name_is_already_taken() {
for taken in ["pa_0", "_sa", "shmem_p_0"] {
let odd = GENERATED.replacen(
" uint lsma_0;\n",
&format!(" uint {taken};\n uint lsma_0;\n"),
1,
);
assert_ne!(odd, GENERATED, "fixture should contain the anchor");
let err = postpass_gemm_msl(&odd).expect_err(&format!("should decline for {taken}"));
assert!(
err.contains(&format!("already defines {taken}")),
"unexpected error for {taken}: {err}"
);
}
}
#[test]
fn declines_when_a_step_temporary_is_written_again() {
let step = " uint lsmb_1 = lsmb_0 + 256U;";
let odd = GENERATED.replacen(
step,
&format!("{step}\n lsma_1 = lsma_1 + 8U;"),
1,
);
assert_ne!(odd, GENERATED, "fixture should contain the step");
let err = postpass_gemm_msl(&odd).expect_err("should decline on a second step write");
assert!(
err.contains("lsma_1 to be written 1 time(s)"),
"unexpected error: {err}"
);
}
#[test]
fn declines_when_a_deleted_scratch_array_is_still_referenced() {
let bind = " (&kernelContext_0)->sa_0 = &sa_1;";
let odd = GENERATED.replacen(bind, &format!("{bind}\n sa_1[0] = (half)(0.0);"), 1);
assert_ne!(odd, GENERATED, "fixture should contain the binding");
let err = postpass_gemm_msl(&odd).expect_err("should decline on a dangling array reference");
assert!(
err.contains("sa_1 is still referenced"),
"unexpected error: {err}"
);
}
#[test]
fn declines_when_a_threadgroup_slot_is_already_bound() {
let sig = "float device* dst_1 [[buffer(2)]])";
let odd = GENERATED.replacen(
sig,
"float device* dst_1 [[buffer(2)]], threadgroup half* scratch_0 [[threadgroup(0)]])",
1,
);
assert_ne!(odd, GENERATED, "fixture should contain the signature");
let err = postpass_gemm_msl(&odd).expect_err("should decline on an occupied threadgroup slot");
assert!(
err.contains("which this pass would collide with"),
"unexpected error: {err}"
);
let spaced = GENERATED.replacen(
sig,
"float device* dst_1 [[buffer(2)]], threadgroup half* scratch_0 [[ threadgroup(0) ]])",
1,
);
assert_ne!(spaced, GENERATED, "fixture should contain the signature");
postpass_gemm_msl(&spaced).expect_err("should decline on the spaced attribute spelling");
let inner = GENERATED.replacen(
sig,
"float device* dst_1 [[buffer(2)]], threadgroup half* s_0 [[ threadgroup (0) ]])",
1,
);
assert_ne!(inner, GENERATED, "fixture should contain the signature");
postpass_gemm_msl(&inner).expect_err("should decline on space before the attribute paren");
}
#[test]
fn tolerates_an_unrelated_threadgroup_attribute() {
let odd = GENERATED.replacen(
"[[kernel]] void gemm_q8_0(",
"[[max_total_threads_per_threadgroup(128)]] [[kernel]] void gemm_q8_0(",
1,
);
assert_ne!(odd, GENERATED, "fixture should contain the entry point");
let out =
postpass_gemm_msl(&odd).expect("an unrelated threadgroup attribute should not decline");
assert!(
out.contains("threadgroup char* shmem_p_0 [[threadgroup(0)]]"),
"the scratch parameter should still have been spliced in"
);
}
#[test]
fn declines_on_an_unrecognized_access_spelling() {
let extra = GENERATED.replace(
" (&kernelContext_0)->sa_0 = &sa_1;\n",
" (&kernelContext_0)->sa_0 = &sa_1;\n (void)((&kernelContext_0)->sa_0);\n",
);
assert_ne!(extra, GENERATED, "fixture should contain the binding");
let err = postpass_gemm_msl(&extra).expect_err("should decline on an unknown access form");
assert!(err.contains("surviving sa_0"), "unexpected error: {err}");
}
#[test]
fn is_idempotent() {
let once = patched();
let twice = postpass_gemm_msl(&once).expect("re-running should succeed");
assert_eq!(once, twice, "post-pass should be idempotent");
assert_eq!(once.matches("#pragma unroll(4)").count(), 1);
assert_eq!(once.matches("pa_0 = pa_1;").count(), 1);
assert_eq!(once.matches(POSTPASS_MARKER).count(), 1);
}
#[test]
fn declines_rather_than_half_applying() {
let moved = GENERATED.replace("lsma_1", "lsmaPrime_1");
assert_ne!(moved, GENERATED, "fixture should contain the anchor");
let err = postpass_gemm_msl(&moved).expect_err("should decline on a moved anchor");
assert!(
err.contains("lsma_1"),
"error should name the anchor: {err}"
);
}
#[test]
fn declines_on_a_duplicated_anchor() {
let dup = GENERATED.replace(" uint lsma_0;\n", " uint lsma_0;\n uint lsma_0;\n");
assert_ne!(dup, GENERATED, "fixture should contain the anchor");
let err = postpass_gemm_msl(&dup).expect_err("should decline when an anchor is ambiguous");
assert!(
err.contains("uint lsma_0;") && err.contains("found 2"),
"should decline on the duplicated anchor itself: {err}"
);
}
#[test]
fn declines_on_an_implausible_trip_count() {
let wide = GENERATED.replace("if(ik_0 < 4U)", "if(ik_0 < 4096U)");
assert_ne!(wide, GENERATED, "fixture should contain the trip test");
let err = postpass_gemm_msl(&wide).expect_err("should decline on a huge trip count");
assert!(
err.contains("refusing to unroll"),
"unexpected error: {err}"
);
}