Expression for allocating a tensor (returns a tensor object). Mimics PyTorch's `torch.empty`, `torch.zeros`, etc. Supports multi-dimensional shapes, device placement, and initialization modes.
| 518 | |
| 519 | |
| 520 | class AllocateTensorExpr(Expr): |
| 521 | """Expression for allocating a tensor (returns a tensor object). |
| 522 | |
| 523 | Mimics PyTorch's `torch.empty`, `torch.zeros`, etc. Supports multi-dimensional shapes, |
| 524 | device placement, and initialization modes. |
| 525 | """ |
| 526 | |
| 527 | def __init__(self, shape: List[Expr], dtype: LLType, mode: AllocateMode = AllocateMode.EMPTY): |
| 528 | """ |
| 529 | Args: |
| 530 | shape: List of expressions defining tensor dimensions (e.g., [Literal(32), threadIdx.x + 1]) |
| 531 | dtype: Data type of tensor elements (e.g., float32, int32) |
| 532 | mode: Initialization mode (empty/zeros/ones/rand) |
| 533 | |
| 534 | Raises: |
| 535 | TypeError: If shape elements are not integers or dtype is invalid |
| 536 | """ |
| 537 | # Tensor type is represented as a specialized type (assume LLType has a tensor constructor) |
| 538 | tensor_type = Tensor[dtype] |
| 539 | super().__init__(expr_type=tensor_type) |
| 540 | |
| 541 | # Validate shape: all dimensions must be integer expressions |
| 542 | for dim in shape: |
| 543 | if not isinstance(dim.type, IntType): |
| 544 | raise TypeError(f"Tensor shape dimension must be integer, got {dim.type}") |
| 545 | |
| 546 | self.shape = shape |
| 547 | self.dtype = dtype |
| 548 | self.mode = mode |
| 549 | |
| 550 | def __repr__(self) -> str: |
| 551 | shape_str = ", ".join(repr(dim) for dim in self.shape) |
| 552 | return (f"AllocateTensorExpr(shape=({shape_str}), dtype={self.dtype}, " |
| 553 | f"mode={self.mode.name.lower()})") |
| 554 | |
| 555 | def __str__(self) -> str: |
| 556 | shape_str = ", ".join(repr(dim) for dim in self.shape) |
| 557 | return (f"AllocateTensorExpr(shape=({shape_str}), dtype={self.dtype}, " |
| 558 | f"mode={self.mode.name.lower()})") |
no outgoing calls