Creates a new view of the buffer. (used by parser) Supported signatures are ``view(*shape, layout=None)``, where shape can contain ``-1`` to indicate that the dimension size is auto-inferred, and ``view(dtype: Union[str, tvm.DataType])``. Returns -------
(self, *args, **kwargs)
| 279 | return tvm.tirx.address_of(self[tuple(indices)]) |
| 280 | |
| 281 | def view(self, *args, **kwargs) -> "Buffer": |
| 282 | """Creates a new view of the buffer. (used by parser) |
| 283 | |
| 284 | Supported signatures are ``view(*shape, layout=None)``, where shape can contain |
| 285 | ``-1`` to indicate that the dimension size is auto-inferred, and |
| 286 | ``view(dtype: Union[str, tvm.DataType])``. |
| 287 | |
| 288 | Returns |
| 289 | ------- |
| 290 | view : DeclBufferFrame |
| 291 | The corresponding view buffer. |
| 292 | """ |
| 293 | |
| 294 | def _infer_shape(shape): |
| 295 | shape = list(shape) |
| 296 | if -1 in shape and shape.count(-1) == 1: |
| 297 | size = functools.reduce(lambda x, y: x * y, self.shape) |
| 298 | n_size = functools.reduce(lambda x, y: x * y, [s for s in shape if s != -1], 1) |
| 299 | shape[shape.index(-1)] = size // n_size |
| 300 | else: |
| 301 | # Only validate the shape product when both old and new shapes |
| 302 | # are fully concrete: a PrimExpr `==` returns an `EQ` node, not |
| 303 | # a Python bool, and `assert <PrimExpr>` raises (no __bool__). |
| 304 | if all(isinstance(s, int) for s in shape) and all( |
| 305 | isinstance(s, int) for s in self.shape |
| 306 | ): |
| 307 | assert functools.reduce(lambda x, y: x * y, shape) == functools.reduce( |
| 308 | lambda x, y: x * y, self.shape |
| 309 | ), ( |
| 310 | "The shape of the buffer " |
| 311 | + str(self.shape) |
| 312 | + " and the new shape " |
| 313 | + str(shape) |
| 314 | + " are not compatible" |
| 315 | ) |
| 316 | return shape |
| 317 | |
| 318 | if len(args) == 1 and isinstance(args[0], str | tvm.DataType) and not kwargs: |
| 319 | cast_dtype = tvm.DataType(args[0]) |
| 320 | cur_dtype = tvm.DataType(self.dtype) |
| 321 | if cast_dtype.bits > cur_dtype.bits: |
| 322 | # cast up |
| 323 | assert cast_dtype.bits % cur_dtype.bits == 0 |
| 324 | ratio = cast_dtype.bits // cur_dtype.bits |
| 325 | layout = self.layout.pack(ratio) |
| 326 | shape = [s for s in self.shape[:-1]] + [self.shape[-1] // ratio] |
| 327 | new_elem_offset = self.elem_offset // ratio |
| 328 | else: |
| 329 | # cast down |
| 330 | assert cur_dtype.bits % cast_dtype.bits == 0 |
| 331 | ratio = cur_dtype.bits // cast_dtype.bits |
| 332 | layout = self.layout.unpack(ratio) |
| 333 | shape = [s for s in self.shape[:-1]] + [self.shape[-1] * ratio] |
| 334 | new_elem_offset = self.elem_offset * ratio |
| 335 | return tvm.tirx.script.builder.decl_buffer( |
| 336 | shape, |
| 337 | cast_dtype, |
| 338 | self.data, |
no test coverage detected