Repository navigation
[Feature] Tensor::view overload that reduces rank by squeezing dropped axes #785
Description
Summary
Add a Tensor::view overload that produces a lower-rank view by squeezing
out size-1 axes (the result of a rank-reducing slice). The current
view(shapes[], offsets[]) inherits ndims from the parent, so it cannot
express "slice + drop axis" in a single step, and view(...).reshape(...)
does not compose for sliced views.
Motivation / Use Case
PyPTO RFC #1338 / PR #1343 added a drop_dims operand to tensor.slice so
that numpy-style indexing (C[i], C[i, j], C[i, j, :, :], ...) can
produce a lower-rank Tensor. The orchestration codegen
(src/codegen/tensor_op_codegen.cpp, REGISTER_ORCHESTRATION_OP(tensor_slice))
needs to emit a runtime Tensor at the reduced rank so that downstream
kernel-call bindings see the correct ndims. Without this, a kernel that
takes a sub-tensor via numpy-style indexing generates wrong-rank code. See
PyPTO follow-up issue: hw-native-sys/pypto#1349.
The composition view(...).reshape(...) does not work because:
view(view_shapes[], view_offsets[])(tensor.h:331) inherits
ndims = other.ndims(init_with_viewattensor.h:233). The view
stays at the parent's rank with size-1 entries in dropped positions.reshape(new_shapes[], new_ndims)(tensor.h:358) hard-asserts
is_contiguous()(tensor.h:360). A rank-reducing slice produces
shapes[i]=1, raw_shapes[i]=Bfor any non-leading dropped axis, which
failsis_contiguous()(tensor.h:341only allows divergence in dim 0
viais_raw_eq_shapes).- Even when the contiguity check passes,
reshapeclobbers the per-dim
offsets thatviewjust installed:result.is_all_offset_zero = true; result.is_raw_eq_shapes = true;(tensor.h:364-365). The data offset
from the slice is lost.
Proposed API / Behavior
Add an overload (under runtime/src/{a2a3,a5}/runtime/tensormap_and_ringbuffer/runtime/tensor.h):
```cpp
// Slice-and-squeeze: produce a view at a lower rank by collapsing size-1
// axes listed in drop_dims. view_shapes / view_offsets are at the
// parent's rank; view_shapes[d] for d in drop_dims must equal 1.
Tensor view(
const uint32_t view_shapes[],
const uint32_t view_offsets[],
const uint32_t drop_dims[],
uint32_t num_drop_dims,
bool manual_dep = false) const;
```
Semantics:
- The result's
ndimsisparent.ndims - num_drop_dims. - The result's
shapes[]andraw_shapes[]are the parent's
view_shapes[]andparent_raw[]with the entries at indices in
drop_dimsremoved. - The element offset contributed by the dropped axes
(sum(view_offsets[d] * stride(d))ford in drop_dims) is folded into
start_offset(or equivalently intooffsets[]of the surviving
leading axis), so element addressing in the lower-rank view points at
the same memory as the parent slice. - Surviving axes keep their
view_offsets[]as-is.
Constraint: every d in drop_dims must have view_shapes[d] == 1.
This matches what a rank-reducing slice produces and lets us collapse the
dimension without needing a real reshape (no contiguity requirement).
Alternatives Considered
view().reshape()composition — Rejected, see Motivation Refactor AICPU core assignment to support dynamic block distribution #2/efactor: Rename Graph to Runtime to Better Reflect Its Responsibility #3
above. Breaks for any non-leading drop and loses per-dim offsets.- Standalone
squeeze(keep_mask[])operation — Composable but
requires two C++ calls per slice in generated code and a separate
intermediateTensor. The overload above is one call and one result. - Keep full-rank in runtime; only drop dims in IR / kernel-call binding
metadata — Rejected, this diverges runtime tensor rank from IR rank
and breaks the orchestration → kernel-call ABI inference path.
Additional Context
- Same change is needed in both
src/a2a3/.../tensor.hand
src/a5/.../tensor.hso the orchestration codegen path is uniform
across architectures. - PyPTO codegen call site that will adopt the new API:
src/codegen/tensor_op_codegen.cpp,REGISTER_ORCHESTRATION_OP(tensor_slice)
(around line 209). It already constructs_shapesand_offsetsarrays;
it would emit an additional_drop_dimsarray and call the new overload. - The
tile.slicerank-reducing path in PyPTO is purely codegen-side and
does not need runtime support.
Activity
ChaoZheng109 commentedon Jun 18, 2026
Leaning toward making the rank reduction (squeezing dropped axes) a PyPTO-side feature, and keeping the simpler runtime at full rank.
Rationale:
The runtime has no functional need for rank reduction. A simpler
Tensoris fundamentally a memory-addressing descriptor; addressing goes throughstart_offset + Σ indices[d]·strides[d](compute_flat_offset). The axes hit bydrop_dimshaveshape == 1and an index that is always 0, so they contribute 0 to the address — keeping those size-1 axes is addressing-equivalent to removing them. In other words, rank reduction is IR/logical-level semantics, not a runtime addressing requirement.The interface cost should live in PyPTO, not in two runtimes. Otherwise we add and maintain a
view(..., drop_dims[], num_drop_dims)overload in botha2a3anda5copies oftensor.h. The smaller the simpler runtime surface, the better.On the issue's reason for rejecting "Keep full-rank in runtime" (runtime rank diverging from IR rank, breaking the orchestration → kernel-call ABI inference): that actually argues that the source of truth for
ndimsbelongs in PyPTO. The reducedndimsthat kernel-call binding needs can be emitted explicitly by PyPTO when it generates the binding metadata — there's no need to force the runtime tensor to match the IR rank instead.Proposal: land
drop_dims/ squeeze as a PyPTO feature (cf. pypto#1349), and keep using full dimensions when it reaches simpler.