ec3171ab7a
## Summary
These three passes are logically a single host/device split step;
having intermediaries between them obscures the model and blocks
folding them into one pass. This PR moves each intermediary to the
position its actual ordering constraint allows, so that
`AnnotateDeviceRegions`, `SplitHostDevice`, and
`LowerDeviceKernelLaunch` run consecutively in every pipeline.
## Rationale
- `MergeSharedMemoryAllocations` moves **before**
`AnnotateDeviceRegions`
(the only legal position: `LowerDeviceKernelLaunch` requires at most
one dyn-shmem allocation per kernel, so Merge cannot move past Lower).
- `MakePackedAPI` moves **after** `LowerDeviceKernelLaunch` (Lower's
`kCallingConv = kDeviceKernelLaunch` flag causes `MakePackedAPI` to
correctly skip device kernels; the host body's lowered
`tvm_call_packed` is transparent to `MakePackedAPI`'s subroutine
rewriter).
- `FP8StorageLegalize` / `BF16StorageLegalize` move **after**
`MakePackedAPI` (their `buffer_map.size()==0` ICHECK requires
`MakePackedAPI` to have cleared the map).
Prereq for Phase 2: collapsing the three consecutive passes into a
single `tirx.transform.SplitHostDevice` with three commented regions.
## Test plan
- [x] tests/python/tirx-transform/ target-pass unit tests (25 pass)
- [x]
tests/python/s_tir/transform/test_merge_dynamic_shared_memory_allocations.py
(5 pass)
- [x] tests/python/tirx-transform/test_tir_transform_fp8_legalize.py /
test_tir_transform_bf16_legalize.py (13 pass)
- [x] tests/python/codegen/test_target_codegen_c_host.py /
test_target_codegen_device.py (6 pass including
test_subroutine_call — verifies Risk #2)
- [x] pre-commit run --all-files clean
- [ ] CI: lint / Windows / MacOS