forked from ccf-ai-infra/TileOPs-Metax
2de15909e4
## Add payload support and implicit thread blocks to trace API
### Summary
This PR adds two enhancements to the trace API:
1. **Payload support for duration ranges** - Brings `trace.range()` in
line with `trace.record()` for payload capability. The `AnnoToken` now
carries payload through both `RANGE_BEGIN` and `RANGE_END` markers.
2. **Implicit thread blocks support** - Enables tracing in kernels using
simple `T.Kernel(..., threads=N)` syntax without explicit threadIdx.x
binding. Uses a new `__tl_thread_idx_x()` CUDA helper when explicit
binding is not present.
### Changes
#### 1. Payload Support
**API Changes**:
- **`trace.range(name, lane, payload=None)`** - Added optional `payload`
parameter
- **`trace.range_start(name, lane, payload=None)`** - Added optional
`payload` parameter
- Payload is recorded in both `RANGE_BEGIN` and `RANGE_END` events
- `AnnoToken` now stores payload value for the entire range scope
**Implementation**:
- Updated `AnnoToken` to store payload value
- Modified `open_range()` and `close_range()` to accept and emit payload
- Updated `RangeScope` to pass payload through context manager
- Simplified payload handling by letting `_emit()` handle None
conversion
**Usage Example**:
```python
# Tag iterations with their index
for i in range(N):
with trace.range("iteration", payload=i):
work(i)
# Payload appears in timeline UI (displayed separately and in hover tooltip)
```
#### 2. Implicit Thread Blocks Support
**Problem**: Kernels using simple `T.Kernel(1, threads=16)` syntax have
no explicit threadIdx.x binding, causing trace lowering to fail with
`KeyError: 'threadIdx.x'`.
**Solution**:
- Added `__tl_thread_idx_x()` CUDA device helper in
`tileops/trace/device.py`
- Modified `tileops/trace/passes.py` to detect missing threadIdx.x
binding
- Falls back to `T.call_extern("int32", "__tl_thread_idx_x")` when
binding not found
- Emits a warning to stderr when fallback is used
**Before** (required explicit binding):
```python
with T.Kernel(1, threads=(1, 16)):
ty = T.thread_binding(0, 1, thread="threadIdx.y")
tx = T.thread_binding(0, 16, thread="threadIdx.x") # Required!
with trace.range("work"):
work(tx)
```
**After** (simple syntax works):
```python
with T.Kernel(1, threads=16):
tx = T.get_thread_binding() # Works now!
with trace.range("work"):
work(tx)
```
### Backward Compatibility
✅ Fully backward compatible:
- `payload` parameter is optional and defaults to `None` (recorded as 0)
- Explicit threadIdx.x binding still works as before
- Existing code continues to work without changes
```python
# Both still work
with trace.range("compute"): # No payload
work()
with T.Kernel(1, threads=(1, 16)): # Explicit binding
tx = T.thread_binding(0, 16, thread="threadIdx.x")
with trace.range("work"):
work(tx)
```
### Testing
**Test coverage** (`tests/trace/test_payload.py`):
- ✅ API signature verification (compilation tests)
- ✅ Constant payload with decode (payload=42)
- ✅ Dynamic payload with runtime PrimExpr (loop index → [0, 1, 2, 3])
- ✅ range_start/range_end with payload decode
- ✅ Backward compatibility without payload
- ✅ Implicit thread blocks with trace enabled/disabled
- ✅ Process-level trace state preservation (no test pollution)
- ✅ All tests marked with `pytest.mark.full` tier
**Test strategy**:
- GPU end-to-end verification with trace.decode()
- Tests verify payload is written to slots and can be decoded
- Tests verify implicit thread blocks fallback compiles and runs
- Precise assertions (e.g., "exactly one slice with payload=42")
- Coverage of both constant and runtime expression payloads
### Known Limitations
⚠️ **Unverified behavior**: Payload behavior inside `T.Pipelined()`
loops has not been validated. Compiler optimizations (such as
loop-invariant code motion or loop unrolling) may affect whether payload
values from individual iterations are captured.
**Recommendation**: Use payload for:
- ✅ Non-pipelined loops
- ✅ Conditional tracing
- ✅ High-level event tagging
- ✅ Simple kernels with implicit thread blocks
**Use with caution**:
- ⚠️ `T.Pipelined()` loop iterations (behavior not yet verified)
For precise pipeline profiling, consider using NVIDIA NSight Compute.
### Future Work
Potential improvements (out of scope for this PR):
- Add tests that decode trace slots and assert payload values
- Validate payload behavior in pipelined loops
- Compiler pass to preserve trace markers through optimizations
- Pipeline-aware trace API
- Eliminate warning for implicit thread blocks (make it silent)
### Checklist
- [x] API changes are backward compatible
- [x] Unit tests added (compilation and API verification)
- [x] Tests preserve trace state (no pollution)
- [x] Implicit thread blocks support working
- [x] Code follows project style
- [x] Known limitations documented
- [ ] Tests that decode and assert payload values (follow-up)
- [ ] Documentation updated (follow-up)
---
**Note**: This PR combines two related trace API enhancements. Both
features improve the usability of the trace system for common kernel
patterns. Feedback welcome on:
1. API design and payload propagation through token
2. Implicit thread blocks fallback strategy
3. Test coverage approach
|
||
|---|---|---|
| .. | ||
| test_payload.py | ||