Support register-backed tensor caching across `hl.static_range` iterations (equivalent to Triton's `VAR_ARGS_ARRAY`)
**Is your feature request related to a problem? Please describe.**
When writing two-pass kernels (e.g., layernorm: pass 1 gathers data + accumulates stats, pass 2 normalizes + writes output), there's no way to cache tensors in registers across two `hl.static_range` loops. This forces a re-gather from global memory in pass 2, adding ~67% more memory traffic compared to the equivalent Triton kernel.
In Triton, this is solved with `VAR_ARGS_ARRAY`:
```python
out_values: "VAR_ARGS_ARRAY"
for i in range(len(input_group_ptrs)):
out_values[i] = tl.load(...) # store in registers
# Later, zero-cost register read:
for i in range(len(input_group_ptrs)):
value = out_values[i] # no memory op
```
In Helion, all attempts to cache tensors across `hl.static_range` iterations fail:
- `list.append(tensor)` → not supported by tracer
- `pre_allocated_list[i] = tensor` → `TypeInferenceError: Subscript assignment not supported` (SymInt key into SequenceType)
- Temporary global memory buffer → works but adds write+read round-trip, negating the benefit
**Describe the solution you'd like**
Support storing and retrieving tensors from a list-like structure inside `hl.static_range` loops. Since `static_range` is unrolled at trace time, each `gathered_values[i]` resolves to a distinct tensor reference. The type propagation system should track this.
Possible API options:
```python
# Option A: Allow list subscript assignment with static_range index
gathered_values = [None] * G
for i in hl.static_range(G):
gathered_values[i] = input_group[i][idx, tile_d]
# Option B: Dedicated API
gathered_values = hl.register_tensor_list(G, shape=[tile_o, tile_d], dtype=dtype)
for i in hl.static_range(G):
gathered_values[i] = input_group[i][idx, tile_d]
```
**Describe alternatives you've considered**
1. **Re-gather from global memory** (current approach): Works but 1.8x slower than Triton for the fused group_index_select + layernorm kernel (G=10, D=48, N=4096, O=4096 on H100).
2. **Temporary buffer in global memory**: Allocate `torch.empty([O, FULL_D])`, write in pass 1, read in pass 2. Compiles but adds 2G extra memory ops (G stores + G loads) — same total traffic as re-gather, with worse L2 sharing.
**Additional context**
- Benchmark: Triton fwd 0.0068ms (1204 GB/s) vs Helion fwd 0.0123ms (666 GB/s) for `G=10, D=48, N=4096, O=4096` on H100 bf16
- The gap is entirely due to the extra memory traffic in pass 2 (Triton: 3 mem ops/group, Helion: 5 mem ops/group)
- This pattern is common in any two-pass reduction kernel (layernorm, softmax, etc.) where intermediate values need reuse
1 条评论