Allocate MMA-compatible shared memory with an inferred swizzle layout.
(self, shape, dtype="float16", swizzle_mode="auto", align=1024)
| 476 | return res |
| 477 | |
| 478 | def alloc_mma(self, shape, dtype="float16", swizzle_mode="auto", align=1024): |
| 479 | """Allocate MMA-compatible shared memory with an inferred swizzle layout.""" |
| 480 | from tvm.backend.cuda.operator.tile_primitive.tma_utils import ( |
| 481 | SwizzleMode, |
| 482 | mma_shared_layout, |
| 483 | ) |
| 484 | |
| 485 | if isinstance(swizzle_mode, str): |
| 486 | if swizzle_mode == "auto": |
| 487 | swizzle_mode = _auto_swizzle_mode(dtype) |
| 488 | elif swizzle_mode == "none": |
| 489 | swizzle_mode = SwizzleMode.SWIZZLE_NONE |
| 490 | else: |
| 491 | raise ValueError( |
| 492 | f"Unsupported swizzle_mode={swizzle_mode!r}; expected 'auto', 'none', " |
| 493 | "or SwizzleMode" |
| 494 | ) |
| 495 | _validate_mma_alloc_shape(shape, dtype, swizzle_mode) |
| 496 | layout = mma_shared_layout(dtype, swizzle_mode, shape) |
| 497 | return self.alloc(shape, dtype, align=align, layout=layout) |
| 498 | |
| 499 | def move_base_to(self, offset): |
| 500 | self.offset = offset |