Create a thread-local view of this buffer. When called with no shape arguments, auto-infers a 1D shape from the layout's non-thread component (i.e. ``layout.storage().shard``). Parameters ---------- shape : tuple of Expr The shape of the local vi
(self, *shape, layout=None)
| 381 | ) |
| 382 | |
| 383 | def local(self, *shape, layout=None) -> "Buffer": |
| 384 | """Create a thread-local view of this buffer. |
| 385 | |
| 386 | When called with no shape arguments, auto-infers a 1D shape from |
| 387 | the layout's non-thread component (i.e. ``layout.storage().shard``). |
| 388 | |
| 389 | Parameters |
| 390 | ---------- |
| 391 | shape : tuple of Expr |
| 392 | The shape of the local view for indexing. If omitted, a 1D |
| 393 | shape is computed automatically. |
| 394 | |
| 395 | layout : optional |
| 396 | Override layout. If None, uses the storage layout |
| 397 | (parent layout with thread axes removed). |
| 398 | |
| 399 | Returns |
| 400 | ------- |
| 401 | local : DeclBufferFrame |
| 402 | The corresponding local buffer. |
| 403 | """ |
| 404 | if not shape: |
| 405 | local_layout = self.layout.storage() |
| 406 | total = functools.reduce( |
| 407 | lambda x, y: x * y, [it.extent for it in local_layout.shard], 1 |
| 408 | ) |
| 409 | shape = (total,) |
| 410 | return tvm.tirx.script.builder.decl_buffer( |
| 411 | shape, |
| 412 | self.dtype, |
| 413 | self.data, |
| 414 | self.strides, |
| 415 | self.elem_offset, |
| 416 | None, |
| 417 | self.scope(), |
| 418 | self.data_alignment, |
| 419 | self.offset_factor, |
| 420 | "", |
| 421 | self.axis_separators, |
| 422 | self.layout.storage() if layout is None else layout, |
| 423 | ) |
| 424 | |
| 425 | def permute(self, *dims) -> "Buffer": |
| 426 | """Permute the dimensions of the buffer. |