Function alloc_layout

tests/python/tirx/test_layout.py:1424–1429

(shape, dtype, swizzle_mode="auto")

def test_pool_allocator_alloc_mma():
    def alloc_layout(shape, dtype, swizzle_mode="auto"):
        with IRBuilder():
            with Tx_builder.prim_func():
                pool = T.SMEMPool(Var("smem_ptr", PointerType(PrimType("uint8"))))
                buf = pool.alloc_mma(shape, dtype, swizzle_mode=swizzle_mode)
        return buf.layout

    cases = [
        ("uint8", (3, 64, 256)),