()
| 1421 | |
| 1422 | |
| 1423 | def test_pool_allocator_alloc_mma(): |
| 1424 | def alloc_layout(shape, dtype, swizzle_mode="auto"): |
| 1425 | with IRBuilder(): |
| 1426 | with Tx_builder.prim_func(): |
| 1427 | pool = T.SMEMPool(Var("smem_ptr", PointerType(PrimType("uint8")))) |
| 1428 | buf = pool.alloc_mma(shape, dtype, swizzle_mode=swizzle_mode) |
| 1429 | return buf.layout |
| 1430 | |
| 1431 | cases = [ |
| 1432 | ("uint8", (3, 64, 256)), |
| 1433 | ("float16", (3, 64, 256)), |
| 1434 | ("bfloat16", (3, 64, 256)), |
| 1435 | ("float32", (3, 64, 256)), |
| 1436 | ("float4_e2m1fn", (3, 64, 256)), |
| 1437 | ] |
| 1438 | for dtype, shape in cases: |
| 1439 | layout = alloc_layout(shape, dtype) |
| 1440 | expected = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, shape) |
| 1441 | assert_structural_equal(layout, expected) |
| 1442 | |
| 1443 | shape = (3, 64, 256) |
| 1444 | layout_64b = alloc_layout(shape, "float32", SwizzleMode.SWIZZLE_64B_ATOM) |
| 1445 | expected_64b = mma_shared_layout("float32", SwizzleMode.SWIZZLE_64B_ATOM, shape) |
| 1446 | assert_structural_equal(layout_64b, expected_64b) |
| 1447 | |
| 1448 | layout_none = alloc_layout(shape, "float16", "none") |
| 1449 | expected_none = mma_shared_layout("float16", SwizzleMode.SWIZZLE_NONE, shape) |
| 1450 | assert_structural_equal(layout_none, expected_none) |
| 1451 | |
| 1452 | |
| 1453 | def test_storage(): |
nothing calls this directly
no test coverage detected
searching dependent graphs…