mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
TensorGlue-IEIT
e84a0ee699
* [Refactor] Extract shared-access analysis into a reusable header - Move AccessEntry/AccessType definitions from thread_storage_sync.cc to src/transform/common/shared_access_analysis.h so other passes (notably MergeSharedMemoryAllocations) can reuse the same precise access representation. - Add tl.debug_thread_storage_sync_summary pass-config option for optional ThreadSync planner summary logging. - Add regression tests covering explicit shared sync deduplication and shared.dyn alias-aware dependency planning. * [Feature] Rewrite MergeSharedMemoryAllocations with precise access analysis - Build per-access live ranges using shared_access_analysis utilities so branch-alternative dataflow and barrier-driven epoch boundaries are modeled directly instead of through whole-buffer approximations. - Plan arena layout starting from arena_size_const so allocations whose starting offsets are non-zero no longer collide. - Preserve single-slot semantics for mutually exclusive branch writes followed by a post-dominating read (fixes the K_norm_pairs_shared style miscompile where stage-1 alternatives were split across epochs). - Cover BufferRegion block reads/writes, tvm_access_ptr, tl.access_ptr, cp_async branches, loop-carried buffers and partial-subregion liveness cases. - Keep aggressive-merge disabled under warp specialization and retain whole-buffer liveness as the conservative default. - Add 18 transform regressions in test_tilelang_transform_merge_shared_memory_allocations.py covering basic merge, sync, branch, cp_async, BufferRegion and opaque-pointer paths. * [Cleanup] Remove dead FreeList/FreeBlock/ActiveInterval helpers These helpers were superseded by the segment-aware LinearScanPack implementation that performs interval-overlap walking directly. They have no remaining callers after the rewrite; remove them to reduce maintenance burden. * [Test] Cover tail-allocation arena growth scenario Add a regression case mirroring upstreamb4c913be(#2106): when the only free slot reusable by an incoming buffer sits at the arena tail and is smaller than the request, the arena must grow by the deficit only, not push the new buffer past the freed slot. The legacy free-list-based packer would emit a 384 B arena for this pattern (X=128, Y=64-then-freed, Z=192), bumping Z to offset 192. The interval-based LinearScanPack reuses Y's offset and grows the tail to fit Z, producing a 320 B arena -- a 16.7%% reduction that matches the peak-live lower bound. * [Feature] Add EpochGraph and use per-epoch liveness to relax aliasing Introduce src/transform/common/epoch_graph.h, a control-flow-aware builder that partitions a flattened StmtEntry sequence into hard-sync and cp_async-wait epochs across nested For/IfThenElse scopes, exposing loop-carry, branch-join and parent-edge structure required for sound liveness analysis. Wire MergeSharedMemoryAllocations to compute per-buffer live epochs as the intersection of: * forward def-reach: epochs reachable from any write of the buffer along epoch-graph successors; * backward use-reach: epochs that can reach any read of the buffer along predecessors. A buffer is live in epoch e iff e is in both sets, modeling 'data written upstream is consumed downstream' rather than the legacy [start_stmt, end_stmt) bounding box. The result feeds two relaxers: * LinearScanPack: two intervals whose live-epoch sets are disjoint may share an offset even if their stmt ranges overlap. * PlanMemory: the symmetric overlap probe consults the same set before forcing a non-aliased layout. Hard syncs (kSync) and cp_async waits are treated as epoch boundaries; ptx_cp_async_wait is conservatively classified as a hard barrier. IfThenElse joins take the union of branch live sets so a buffer written in only one branch remains live at the join. The whole-buffer EpochGraph + liveness pipeline always runs; verbose [MSMA-EPOCH-GRAPH] / [MSMA-EPOCH-LIVE] / [MSMA-EPOCH-ACCESS] / [MSMA-EPOCHSEG] / [MSMA-STAGE0-RAW] dumps remain gated behind the existing tl.debug_merge_shared_memory_allocations pass-config flag. Also remove SplitLocalSegmentsAtEpochSyncs, a Phase-1 helper that was defined but never wired into the planner; the new liveness pipeline subsumes its role. Validated on a real lowbit MHA decode kernel (seq_lens=[(3,96),(2,64)], heads=4, dim=64, page_size=32, bits=4): arena drops from 35,232 B (upstream baseline) to 31,136 B, an 11.6%% reduction driven by K_shared and V_shared sharing one offset thanks to disjoint live-epoch sets. Add testing/python/transform/test_tilelang_transform_epoch_graph.py covering the five S1.4 epoch-graph topology cases (sequential __syncthreads, ptx_wait_group, For body/back/exit edges, IfThenElse branch_join, nested For). Existing test_tilelang_transform_merge_shared_memory_allocations.py and test_tilelang_transform_thread_sync.py suites continue to pass. * [Feature] Gate per-epoch alias relaxer on annotated layout equality Two shared-memory buffers carrying distinct layouts under `block.annotations[attr::kLayoutMap]` (e.g. via `T.annotate_layout(...)`) must not share storage even when per-epoch dataflow proves their liveness sets disjoint: downstream layout-driven access lowering can emit indexing arithmetic that is only sound under each buffer's original layout, so naively aliasing them risks silent miscompiles. Add a `LayoutSignatureCollector` that harvests the per-Var `Layout` entries from `Block` annotations and a `LayoutAliasIncompatible` helper that returns true iff both candidates have non-trivial layouts and the layouts are not structurally equal. The two relaxer call sites in `SharedMemoryRewriter` (`LinearScanPack` and the post-pack overlap audit) consult the helper before dropping a conflict, so the legacy decision is preserved exactly when the gate fires. Buffers without a layout entry remain layout-agnostic and continue to participate in arena packing, keeping the existing reuse for the common un-annotated case. Verified that the lowbit reference kernel arena footprints stay at 31136/61856/61808/61792 B and the merge / thread-sync / epoch-graph trio passes 49/49. * [Test] Add transform regression that fails on upstream MergeSharedMemoryAllocations Add test_general_cond_branch_alternatives_share_offset_with_postdom_buffer to lock the per-epoch / branch-alternative analysis behavior on a mutually-exclusive if/else with a runtime cond plus a post-dominator buffer after a sync. Upstream's linear-scan plan keeps the two branch buffers at distinct offsets (128-byte arena); the per-epoch relaxer collapses them onto the same offset and reuses it for the post-dominator buffer (64-byte arena, 50% smaller). * [Example] Demo MergeSharedMemoryAllocations branch-alternative arena reuse (64 vs 128 B) Add runnable demo under examples/merge_shared_memory_allocations/ that constructs a three-buffer PrimFunc with mutually-exclusive if/else branches, lowers through the merge pass, and reads the merged arena size from the dyn_shared_memory_buf attr. The demo also compiles the kernel end-to-end and validates output correctness with torch.allclose. On this branch: arena = 64 B (saved 128 B / 66.7% vs naive 192). On upstream main (1fbd994c): arena = 128 B (-50% worse). * [Test] Add end-to-end torch correctness regression for merge pass Add two tests (test_merged_branch_kernel_generates_correct_results, test_merged_sync_phased_buffers_generate_correct_results) that compile the merged kernel and verify bit-exact output via torch.allclose, ensuring the arena packing does not alter numerical results. * [Test] Add torch correctness and dynamic-buffer fallback regression - test_merged_branch_kernel_generates_correct_results: compiles a merged branch-alternative kernel with torch.allclose validation. - test_merged_sync_phased_buffers_generate_correct_results: sync-phased two-buffer kernel end-to-end check. - test_dynamic_size_buffer_fallback_does_not_incorrectly_alias: verifies that a symbolic-extent shared.dyn buffer correctly aliases with a sync-separated constant-size buffer without unsafe overlap. * [Fix] Disable per-epoch relaxer for loop-live buffers to avoid pipelined iteration overlap * [Fix] Fallback to structural_segments_by_var for buffers missing planner segments to avoid buffer_byte_offsets crash and shared memory bloat * Merge branch 'main' into merge_pass and fix merge comflicts * [Fix] Resolve post-merge artifacts in CMakeLists, hipcc, and deepseek example - Remove duplicate TILELANG_OUTPUT_TARGETS in CMakeLists.txt - Fix duplicate _resolve_artifact_paths function in hipcc.py - Add missing G=kv_group and fix assert in sparse_mla_fwd_pipelined.py - Add ruff noqa pragma to suppress lint warnings * [Style] Apply clang-format to modified source files