()
| 269 | |
| 270 | |
| 271 | def test_ex_8_8_fa4_pattern(): |
| 272 | root = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA) |
| 273 | |
| 274 | # Branch 1: warp 12 (single warp, tcgen05 MMA elected) |
| 275 | b1 = root.with_filter(WARP_FLAT, 12, 13) |
| 276 | assert b1.A.warpid == AxisRange(1, 12) |
| 277 | assert b1.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(1, 12)} |
| 278 | |
| 279 | # Branch 2: softmax warpgroups (warps 0-7) |
| 280 | b2 = root.with_filter(WARP_FLAT, 0, 8).with_scope_switch(WARPGROUP) |
| 281 | assert b2.inter == {"wgid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)} |
| 282 | assert b2.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)} |
| 283 | |
| 284 | # Branch 3: correction warpgroup (warps 8-11 = wg2) |
| 285 | b3 = root.with_filter(WARP_FLAT, 8, 12) |
| 286 | assert b3.A.warpid == AxisRange(4, 8) |
| 287 | assert b3.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(4, 8)} |
| 288 | # And should factor cleanly when entering warpgroup |
| 289 | b3wg = b3.with_scope_switch(WARPGROUP) |
| 290 | assert b3wg.inter == {"wgid": AxisRange(1, 2), "cta_id": AxisRange(1, 0)} |
| 291 | assert b3wg.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)} |
| 292 | |
| 293 | |
| 294 | # --------------------------------------------------------------------------- |
nothing calls this directly
no test coverage detected
searching dependent graphs…