(self, indices)
| 461 | ) |
| 462 | |
| 463 | def __getitem__(self, indices): |
| 464 | from ..arith import Analyzer # pylint: disable=import-outside-toplevel |
| 465 | from .expr import BufferLoad, Ramp, const # pylint: disable=import-outside-toplevel |
| 466 | from .stmt import BufferRegion # pylint: disable=import-outside-toplevel |
| 467 | |
| 468 | if not isinstance(indices, tuple | list): |
| 469 | indices = [indices] |
| 470 | has_slice = any(isinstance(i, slice) for i in indices) |
| 471 | has_step = any( |
| 472 | isinstance(i, slice) and (i.step is not None and i.step != 1) for i in indices |
| 473 | ) |
| 474 | has_implicit_slice = len(indices) < len(self.shape) |
| 475 | analyzer = Analyzer() |
| 476 | if (has_slice and not has_step) or has_implicit_slice: |
| 477 | region = [] |
| 478 | for i, index in enumerate(indices): |
| 479 | if isinstance(index, slice): |
| 480 | start = 0 if index.start is None else index.start |
| 481 | stop = self.shape[i] if index.stop is None else index.stop |
| 482 | region.append(Range.from_min_extent(start, analyzer.simplify(stop - start))) |
| 483 | else: |
| 484 | region.append( |
| 485 | Range.from_min_extent( |
| 486 | index, const(1, index.dtype) if isinstance(index, PrimExpr) else 1 |
| 487 | ) |
| 488 | ) |
| 489 | if has_implicit_slice: |
| 490 | for i in range(len(indices), len(self.shape)): |
| 491 | region.append(Range.from_min_extent(0, self.shape[i])) |
| 492 | return BufferRegion(self, region) |
| 493 | else: |
| 494 | expr_indices = [] |
| 495 | for i, index in enumerate(indices): |
| 496 | if isinstance(index, slice): |
| 497 | start = 0 if index.start is None else index.start |
| 498 | stop = self.shape[i] if index.stop is None else index.stop |
| 499 | step = 1 if index.step is None else index.step |
| 500 | # We should ensure the dtype of start is the same with that of step. |
| 501 | if isinstance(start, tvm.tirx.expr.PrimExpr) and isinstance(step, int): |
| 502 | step = tvm.tirx.expr.IntImm(start.dtype, step) |
| 503 | lanes = analyzer.simplify((stop - start + step - 1) // step) |
| 504 | if lanes == 1: |
| 505 | expr_indices.append(start) |
| 506 | else: |
| 507 | expr_indices.append(Ramp(start, step, int(lanes))) |
| 508 | else: |
| 509 | expr_indices.append(index) |
| 510 | return BufferLoad(self, expr_indices) |
| 511 | |
| 512 | |
| 513 | def decl_buffer( |
nothing calls this directly
no test coverage detected