Permute the dimensions of the buffer. Parameters ---------- dims : tuple of int The permutation of dimensions. Returns ------- permuted : DeclBufferFrame The buffer with permuted dimensions.
(self, *dims)
| 423 | ) |
| 424 | |
| 425 | def permute(self, *dims) -> "Buffer": |
| 426 | """Permute the dimensions of the buffer. |
| 427 | |
| 428 | Parameters |
| 429 | ---------- |
| 430 | dims : tuple of int |
| 431 | The permutation of dimensions. |
| 432 | |
| 433 | Returns |
| 434 | ------- |
| 435 | permuted : DeclBufferFrame |
| 436 | The buffer with permuted dimensions. |
| 437 | """ |
| 438 | new_shape = [self.shape[d] for d in dims] |
| 439 | # Permute *logical* dims, not the layout's fine-grained shard iters: a |
| 440 | # tcgen05/atom layout maps several shard iters to each logical axis, so |
| 441 | # group by the current shape first and permute whole groups. ``group`` |
| 442 | # returns a regrouped layout (degenerate extent-1 iters folded away) |
| 443 | # plus seps over *that* layout — permute the regrouped one, not |
| 444 | # ``self.layout``. For a simple layout (one shard iter per axis) this |
| 445 | # reduces to ``permute_dims(dims)``. |
| 446 | grouped, seps = self.layout.group(list(self.shape)) |
| 447 | new_layout = grouped.permute_by_groups(seps, list(dims)) |
| 448 | return tvm.tirx.script.builder.decl_buffer( |
| 449 | new_shape, |
| 450 | self.dtype, |
| 451 | self.data, |
| 452 | self.strides, |
| 453 | self.elem_offset, |
| 454 | None, |
| 455 | self.scope(), |
| 456 | self.data_alignment, |
| 457 | self.offset_factor, |
| 458 | "", |
| 459 | self.axis_separators, |
| 460 | new_layout, |
| 461 | ) |
| 462 | |
| 463 | def __getitem__(self, indices): |
| 464 | from ..arith import Analyzer # pylint: disable=import-outside-toplevel |