Method
alloc
(
self,
shape,
dtype="float32",
strides=None,
scope="shared.dyn",
align=0,
buffer_type="",
axis_separators=None,
layout="default",
)
Source from the content-addressed store, hash-verified
| 442 | self.max_offset = 0 |
| 443 | |
| 444 | def alloc( |
| 445 | self, |
| 446 | shape, |
| 447 | dtype="float32", |
| 448 | strides=None, |
| 449 | scope="shared.dyn", |
| 450 | align=0, |
| 451 | buffer_type="", |
| 452 | axis_separators=None, |
| 453 | layout="default", |
| 454 | ): |
| 455 | ir = _get_ir() |
| 456 | if align > 0: |
| 457 | self.offset = (self.offset + align - 1) // align * align |
| 458 | res = ir.decl_buffer( |
| 459 | shape, |
| 460 | dtype, |
| 461 | data=self.ptr, |
| 462 | strides=strides, |
| 463 | byte_offset=self.offset, |
| 464 | scope=scope, |
| 465 | align=align, |
| 466 | buffer_type=buffer_type, |
| 467 | axis_separators=axis_separators, |
| 468 | layout=layout, |
| 469 | ) |
| 470 | # Advance in bits then round up to bytes so sub-byte dtypes (e.g. |
| 471 | # float4_e2m1fn = 4 bits) still bump the cursor instead of leaving it |
| 472 | # at 0 (bits // 8) and silently overlapping the next allocation. |
| 473 | self.offset += (_shape_product(shape) * DataType(dtype).bits + 7) // 8 |
| 474 | if self._owns_buffer: |
| 475 | self.max_offset = max(self.max_offset, self.offset) |
| 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.""" |