mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
056ccfc267f03498d8dcc79f6efd2aba2bae8abd
63
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Lei Wang
|
66ef211815 |
[TIR][Transform] Partition parallel loops with fragment access (#2340)
Fix fragment parallel loop partitioning |
||
|
Tong WUandLeiWang1999
|
c2998920cc |
[BugFix] Fix eager JIT sub-btye shape binding (#2334)
* Fix eager JIT packed FP4 shape binding * Add logical shape and stride calculations for buffers in eager builder This update introduces helper functions to compute logical shape and stride values for buffers, enhancing the handling of packed storage factors. The changes streamline the process of determining dimensions and strides in the context of tensor operations, improving overall efficiency and clarity in the eager execution flow. --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei Wang
|
d988dd0491 |
[Transform] Place auto WS producers in first warp group (#2315)
* Place auto WS producers in first warp group * [Transform] Update layout inference to bind thread range for fully replicated fragments; add tests for TMA fragment mask copy and producer-consumer warp specialization * [Refactor] Update thread variable calculation to use zero-based indexing across GEMM operations * [Transform] Enhance simplification process with ContextSingletonFloorDivFolder and add tests for floor division simplification * [Refactor] Simplify tensor conversion in dequantization functions and improve performance in GEMM examples |
||
|
Chaofan Lin
|
8d315be580 |
[Backend][Transform] Move backend-specific transforms into separate namespaces (#2297)
* [Backend][Transform] Move backend-specific transforms into separate namespaces * cleanup * fix * Restore transform pass docstrings |
||
|
Chaofan Lin
|
4f0e127de2 |
[Backend] Promote PassPipeline to backend sub-folder and cleanup Metal Leftover (#2291)
* [Backend] Promote PassPipeline to backend sub-folder and fix metal transform position * cleanup metal gemm |
||
|
Chaofan Lin
|
0a9b651244 |
[Backend] Refactor Transform Pipeline to support different backends (#2189)
* [Backend] Refactor Pipeline to support different backends Introduce Pipeline abstraction in backend/pipeline.py with per-backend registration. Each backend (cuda, hip, c, llvm) now registers its own compilation pass pipeline. engine/lower.py resolves the pipeline from the target instead of hardcoding phase imports. * enhance the structure * fuse two stage * fix after rebase * update template pass pipeline * pipeline * fix test * fix test * fix cuda pipeline prologue |
||
|
|
c66cadf87f |
[Feature] Support named barrier arrive (#2194)
* [Fix] Suppress spurious __syncthreads() between T.ws blocks (#2192) and add T.named_barrier_arrive (#2188) Issue #2192 - BugFix: Unexpected CTA-wide __syncthreads() between T.ws producer/consumer - Root cause: FindConflict() in thread_storage_sync.cc did not check for mutually exclusive thread constraints. When T.ws(0) (tx<128) reads and T.ws(1) (tx>=128) writes the same shared buffer, no thread satisfies both constraints, so no __syncthreads() is needed — mbarriers already handle the synchronization. - Fix: add mutual-exclusivity check (prove NOT(P AND C)) early in FindConflict(). If the access constraints are mutually exclusive (disjoint thread sets), return false (no conflict) immediately, before any index-range analysis. - Measured improvement: barrierpipe_stage2 gemm latency 44ms → 25ms on H200. Issue #2188 - Feature: expose native T.named_barrier_arrive(barrier_id, thread_count) - Adds tl.named_barrier_arrive intrinsic op (bar.arrive PTX, one-sided arrive). - Registers op in builtin.cc, declares in builtin.h. - Adds tl::__named_barrier_arrive<id, cnt>() template in common.h. - Adds codegen case in codegen_cuda.cc. - Exposes T.named_barrier_arrive() in tilelang/language/builtin.py. - Adds regression test testing/python/issue/test_tilelang_issue_2188_2192.py. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * lint: clang-format + ruff fixes after #2188/#2192 changes Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * test: simplify issue 2188/2192 test, keep only #2192, add requires_cuda Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [Fix] Make named_barrier_arrive accept variable barrier_id/thread_count * Update documentation for named barriers and remove obsolete test file - Clarified documentation for `named_barrier_arrive` in `builtin.h` and `builtin.py` to specify the barrier index range. - Removed obsolete regression test for issues #2188 and #2192 as it is no longer needed. * lint fix --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
VitalyRandLeiWang1999
|
b1083d72b1 |
Fix atomic_load access_ptr lowering for dynamic indices (#2157)
* [Bugfix] Fix issue 2123 access_ptr lowering Add the issue 2123 regression test and keep the access_ptr base BufferLoad intact through safe-memory legalization and lowering. Closes #2123 * Deduplicate access_ptr base load handling Move shared tl.access_ptr BufferLoad base visitation into a common transform helper so LowerAccessPtr and safe-memory legalization stay aligned. Extend the issue 2123 regression with a direct LowerAccessPtr pass check while keeping the LowerAndLegalize pipeline coverage. * Format access_ptr helper header * Narrow access_ptr helper scope * Share access_ptr base helper * Address access_ptr helper review comments * [Transform] Preserve access_ptr buffer loads in safe memory legalization * Refactor AccessPtrInfo return structure in legalize_safe_memory_access.cc Updated the return structure of AccessPtrInfo to improve readability and maintainability. The logic for handling the tvm_access_ptr operation has been streamlined, ensuring clearer checks and consistent formatting of returned values. * fix --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
|
b939fa0137 |
[TIR][IR] Update to use tirx (#2216)
* [Feature] Implement Hoist Loop-Invariant If Pass - Added a new pass to hoist loop-invariant if statements out of loops, improving optimization opportunities. - Introduced classes for collecting and checking conditions, as well as for rewriting statements. - Integrated the new pass into the optimization pipeline and provided a corresponding API for usage. * Add tilelang semantics guide to programming guides (#1740) * Add CallNodeChecker to Loop Unswitching Pass and integrate tilelang cache disabling in benchmark script - Introduced CallNodeChecker class to identify CallNode expressions in loop conditions, enhancing loop-invariant checks. - Updated IsLoopInvariant function to reject conditions containing CallNodes, preventing potential side effects. - Added tilelang.disable_cache() in benchmark_mha_sink_fwd.py to optimize performance during benchmarking. * lint fix * lint fix * Enhance Loop Unswitching Pass with Let Binding Support - Added support for Let-bound variables in the WrittenBufferReadChecker to improve buffer read checks. - Introduced UsesLoopVarThroughLetBindings function to check if conditions depend on loop variables through Let bindings. - Updated IsLoopInvariant function to account for Let bindings when determining loop invariance. - Enhanced HoistableIfFinder to track Let bindings for variables bound to BufferLoad expressions. - Added debug print statements in the OptimizeForTarget function to visualize the module state before and after loop unswitching. * fix * lint fix * Enhance Loop Unswitching with CallNode Exclusion Check - Added CallCheckerExcludingIf class to ensure function calls outside of hoisted if statements are identified, preventing potential synchronization issues during loop unswitching. - Updated loop unswitching logic to incorporate the new call checker, enhancing safety and correctness. - Integrated debug print statements in OptimizeForTarget to visualize module state before and after loop unswitching. - Disabled tilelang cache in the benchmark script for improved performance. * Update TVM subproject to latest commit and remove debug print statements in OptimizeForTarget function * Add support for disabling loop unswitching optimization - Introduced a new configuration option `tl.disable_loop_unswitching` to allow users to disable the loop unswitching optimization. - Updated the Loop Unswitching pass to check this configuration and return the original function if the option is enabled. - Added relevant documentation in the PassConfigKey enumeration for clarity. * update tvm * Refactor: Update TileLang to use tirx instead of tir - Removed references to `tvm.tir` and replaced them with `tvm.tirx` across various files, including examples and backend operations. - Updated target configuration in documentation to reflect the new usage of target config dictionaries instead of CLI-style strings. - Cleaned up `pyproject.toml` and `load_tvm.cmake` by removing obsolete paths. - Enhanced examples for dequantization and GEMM operations to utilize the new `tirx` constructs. - Adjusted various backend operations to ensure compatibility with the new `tirx` namespace. * pass fix * Update submodule 'tvm' to latest commit f24c7b19 * Apply formatting fixes * Update CMakeLists.txt to set HIDE_PRIVATE_SYMBOLS for TileLang and modify target_host initialization in lower.py for compatibility with TVM's target handling. * Refactor CMakeLists.txt and kernel_cache.py to update target names and library dependencies for TVM integration. Adjusted output targets and library names for better compatibility across platforms. * Update repair-wheel-command in pyproject.toml to exclude libz3.so with wildcard for better compatibility with dependency packages. * Fix macOS Z3 dylib lookup * Move macOS Z3 install-name fix to install step * Update TVM submodule for runtime import fix * Update ROCm backend for TVM module factory changes * Fix build-tree Z3 runtime lookup * Fix macOS build-tree Z3 lookup * Support legacy target option strings * Normalize CUDA target in transform test * Qualify Metal stream FFI lookup * Fix PyPI Z3 packaging links * Update TVM for Windows DLL lookup * Register Windows dependency DLL directories * Bump tvm-ffi to match tirx * tvm update * Apply lint fixes * Enhance reduction operations with support for nan propagation and vectorization. Introduce new functions for preferred vectorized size and update existing reduction logic to handle packed operations for bfloat16 and float16 types. Add nan-aware min and max operations in CUDA and ROCm backends, and update related tests to validate functionality. * Refactor ReduceLowerer: Remove unnecessary buffer allocation logic for packed operations in reduce.h * Normalize PyPI Z3 dylib before linking * Drop legacy target option string support * Fix type reference for barrier annotations in CUDA copy operations. Updated `BufferLoadNode` to `tirx::BufferLoadNode` in `GetBarrier` and `LowerCluster` methods to ensure compatibility with recent changes in the TIR API. --------- Co-authored-by: Yu Cheng <yu.cheng@pku.edu.cn> Co-authored-by: Yichen Yan <oraluben@outlook.com> |
||
|
Lei Wang
|
256fbc6e60 |
[Pipeline] Fix guarded TMA pipeline handling (#2224)
* Fix guarded TMA pipeline handling * Apply ruff fixes for latest main |
||
|
|
f11954cb36 |
[Backend] Refactor gemm_sp (#2048)
* [Refactor] refactor gemm_sp following
|
||
|
Chenhao Xu
|
fa76eaf014 |
Allow variable barrier id in T.sync_threads() (#2197)
* fix: allow variable barrier id in T.sync_threads() * ruff format |
||
|
Kasper Nielsen
|
d135bd1cd2 |
fix: TMA alignment to 1024 bytes on Blackwell (#2134)
* fix: TMA alignment to 1024 bytes on Blackwell Previously 1024 byte alignemnt was guarded only by `IsHopper`. This commit changes it to be guarded by `TargetHasBulkCopy` with the same threshold as Hopper is using. The kernel added in the test was failing on my RTX5090 GPU. * Fix lint |
||
|
|
74fc9806fb |
[Refactor] Refactor DecoupleTypeCast Pass (#2026)
* draft
* fix lint
* fix lint
* Improve pass comments
* treat loop-invariant memory access as scalar in vectorization
* Add unit test
* fix lint
* fix b[i] = a[i] + a[i+32] case
* Minor fix
* Tighten DecoupleTypeCast cast detection and add RMW regression test
Skip BufferLoad/BufferStore indices when searching for Cast nodes so an
index-type conversion does not spuriously trigger the decoupling
transformation. Clarify the load-replacement table in visit_for_ — store
entries must feed into it so RMW loads map to the store-side cast buffer —
and cover the a[i] = a[i] + a[i+32] case with a regression test.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Keep loop-invariant stores in memory bucket during vectorize planning
Commit
|
||
|
Lei WangandClaude Opus 4.6
|
3ee0988c86 |
[Refactor][Pipeline] Run pipeline rewriting before layout inference and stabilize tiled WS (#2002)
* Refactor optimization pipeline in phase.py to streamline barrier handling
- Consolidated the handling of shared barriers and pipeline planning by removing redundant conditional checks.
- Ensured that `LowerSharedBarrier`, `PipelinePlanning`, and `InjectSoftwarePipeline` are consistently applied, enhancing the clarity and efficiency of the optimization process.
This change improves the maintainability of the code while preserving existing functionality.
* Enhance tile operation handling with instruction annotation and refactor gemm.h
- Added a new `InstructionAnnotation` pass to annotate tile operations with their instruction kind before layout inference, improving the optimization pipeline's ability to reason about instruction mixes.
- Refactored `gemm.h` to move the `allowTcgen5Mma` and `allowWgmma` methods under the private section, enhancing code organization and encapsulation.
These changes improve the clarity and maintainability of the code while preserving existing functionality.
* Enhance multi-version buffer handling and introduce tile-level warp specialization
- Updated `multi_version_buffer_rewriter.cc` to improve read/write access detection for tile operations by analyzing `tl.tileop.region` calls, ensuring accurate buffer access tracking.
- Modified `phase.py` to integrate `ProducerConsumerWarpSpecializedTiled` before layout inference, allowing for high-level tile-op IR transformations that enhance producer/consumer splits.
- Added a new `ProducerConsumerWarpSpecializedTiled` function in `__init__.py` to facilitate tile-level warp specialization, improving the optimization pipeline's efficiency.
These changes enhance the handling of multi-version buffers and optimize the transformation process for tiled operations.
* Enable tiled WS for stage-1 pipelines
* Move Hopper pipeline planning before layout inference
* Run pipeline rewriting before layout inference
* Add pipeline refactor WIP: subtree_modified_ guard, debug prints, and plan
- inject_pipeline.cc: guard reads/writes recalculation with subtree_modified_
flag to prevent local.var buffer promotion to kernel parameters
- phase.py: add temporary debug prints after PipelinePlanning and
InjectSoftwarePipeline (to be removed during CI fix work)
- example_group_per_split_token_cast_to_fp8.py: add disable_cache for debugging
- docs/plan.md: implementation plan for fixing CI test failures
- draft.md: original design draft
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Remove debug artifacts: prints in phase.py, disable_cache in example
- Remove print("After PipelinePlanning"), print(mod),
print("After InjectSoftwarePipeline"), print(mod) from
LowerAndLegalize in tilelang/engine/phase.py
- Remove tilelang.disable_cache() from
examples/cast/example_group_per_split_token_cast_to_fp8.py
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix LayoutInference crash on multi-versioned pipeline buffers
After InjectSoftwarePipeline, multi-versioned buffers share the same
data Var as the original but have an extra leading dimension (num_stages).
LayoutInference's alias propagation and annotation handling tried to
Reshape layouts between these buffers, which failed because the total
element counts differ.
Guard three Reshape call sites in layout_inference.cc to skip sibling
buffers whose total storage size is incompatible with the source layout.
This lets multi-versioned buffers get their own layout inference instead
of inheriting an incompatible layout from the original buffer.
Fixes compilation failures in dequantize_gemm, GDN, and other kernels
that use software pipelining with shared memory buffers.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Strip pipeline annotations on WS fallback for TMA kernels
When ProducerConsumerWarpSpecializedTiled identifies a TMA kernel as a
warp-specialization candidate but the tiled rewriter cannot handle it
(e.g., conditional loop bodies like sparse block masks), the fallback
previously returned the original function with num_stages annotations
intact. PipelinePlanning and InjectSoftwarePipeline would then generate
non-WS TMA pipeline code with broken barrier phase tracking for
conditional pipeline bodies (barrier waits outside conditionals cause
deadlocks when the condition is false).
Fix: on WS fallback, strip num_stages annotations from pipeline loops
so that the pipeline passes skip the function. The kernel runs
unpipelined but correctly.
Fixes CUDA_ERROR_LAUNCH_FAILED in blocksparse_gemm and related TMA
kernels with conditional loop bodies.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add guarded phase-counter support to tiled WS pass
Port PhaseCounter and StageExprReplacer from the legacy
ProducerConsumerWarpSpecialized pass into the tiled WS pass to handle
conditional loop bodies (e.g., sparse block masks).
When the pipeline loop body is wrapped in an IfThenElse without else:
1. Unwrap the condition before classifying statements
2. Create separate producer/consumer PhaseCounters (local int32 buffers)
3. Use counter-based stage/parity expressions instead of loop-variable
4. Wrap producer and consumer bodies in the original condition
5. Increment counters at end of each guarded iteration
6. Rewrite shared-buffer stage indices via StageExprReplacer
This ensures barrier parity stays correct when iterations are
conditionally skipped, fixing CUDA_ERROR_LAUNCH_FAILED in
blocksparse_gemm and related TMA kernels with conditional execution.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add GemmSPNode support to pipeline planning and injection
Add GemmSPNode handling to:
- inject_pipeline.cc AddReadsWritesForTileOp: model A, E, B as reads
and C as write (E is the sparse metadata buffer)
- pipeline_planning.cc: same access model for dependency analysis
This makes sparse GEMM visible to the pipeline machinery for correct
stage assignment and buffer multi-versioning. However, the tile-op's
consumer-side buffer accesses still don't get stage-indexed because
the pipeline body rewriter can't rewrite high-level tile-op Call
arguments (an architectural limitation of running InjectSoftwarePipeline
before LowerTileOp).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix GemmSPNode::Lower to use MakeAccessPtrFromRegion
Replace whole-buffer access_ptr(1) calls with MakeAccessPtrFromRegion
for A, B, C, and E buffers in sparse GEMM lowering. This preserves
stage-specific region offsets from pipeline multi-versioning, matching
the dense GemmNode::Lower pattern.
CUDA output now shows correct stage-indexed consumer accesses:
gemm_sp_ss(..., (k%3)*8192, (k%3)*8192+27648, C_local, (k%3)*2048+49152)
instead of always using stage-0 offsets.
Note: gemm_sp still produces incorrect results because the kernel needs
warp specialization but TiledWSCandidate::Check doesn't recognize it
as a TMA candidate. The non-WS pipeline path generates structurally
different code from the reference WS path.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix gemm_sp and seer_attention regressions
Three changes in producer_consumer_ws_tiled.cc:
1. Require num_stages >= 2 for WS candidacy (was >= 1). Single-stage
kernels like seer_attention don't need WS and the transformation
produces incorrect results for them.
2. Add HasTmaPipeline() check to detect TMA kernels with pipeline
annotations that are rejected by the full WS candidate check (e.g.,
kernels with manual layout annotations like gemm_sp).
3. Strip num_stages annotations for rejected TMA pipeline kernels to
prevent InjectSoftwarePipeline from generating broken non-WS TMA
pipeline code.
One change in gemm_sp.cc:
- Use MakeAccessPtrFromRegion for A, B, C, E buffer access pointers
instead of whole-buffer access_ptr. This preserves stage-specific
region offsets from pipeline multi-versioning.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix WS conditional body unwrapping for complex loop bodies
Only unwrap IfThenElse wrapper when the then-branch is a simple flat
sequence of tile-op Evaluate calls. Skip unwrapping for complex bodies
with LetStmt, For, or other control flow that could break variable
scoping when split into producer/consumer for WS.
Fixes variable-used-before-definition error in
blocksparse_attention sparse_gqa_decode_varlen_indice, which has a
conditional loop body containing LetStmt bindings.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Revert num_stages >= 2 to >= 1 for WS candidacy
The >= 2 threshold broke test_num_stages_one_pure_tma_keeps_auto_warp_specialize.
Pure TMA kernels with num_stages=1 should still be WS candidates.
The seer_attention issue (num_stages=1 with manual layout) is handled
by the has_manual_layout_ check, not the num_stages threshold.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Propagate manual layouts onto MVB-expanded buffers via Layout::Expand
When LayoutInference encounters an MVB-expanded buffer (with leading
stage dimensions) whose trailing dimensions match the original layout,
use Layout::Expand to propagate the manual layout instead of rejecting
or skipping the buffer. Applied to all 3 layout propagation paths:
annotated layout map, alias propagation, and finalization.
Also remove the blanket !has_manual_layout_ WS candidate rejection
since manual layouts now survive onto versioned shared buffers via
Layout::Expand.
Fixes test_sparse_ws_regular_metadata_copy_stays_in_producer.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Extend conditional body unwrapping for LetStmt chains and restore
has_manual_layout_ guard
1. LetStmt chain peeling: when IfThenElse then_case starts with
LetStmt bindings, peel them and append to let_bindings before
checking the simple-body guard. This allows WS for conditional
bodies with variable definitions (e.g., sparse attention patterns).
2. Restore !has_manual_layout_ in WS candidacy check: removing it
caused dequant_groupedgemm_bf16_mxfp4_hopper to fail because
MXFP4 layouts don't survive MVB expansion. The Layout::Expand fix
handles sparse metadata layouts but not all manual layout types.
3. Layout::Expand propagation (from Round 6) remains in place for
future use when MVB learns to handle all manual layout types.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Remove is_simple_body guard from conditional WS unwrapping
The simple-body guard only accepted flat Evaluate sequences inside
IfThenElse, blocking legitimate WS for complex conditional bodies
like sparse flash attention (T.clear, T.reduce_max, T.gemm inside
the guard). The LetStmt peeling already handles variable scoping.
Fixes test_pure_tma_consumer_local_init_does_not_leak_into_producer
and test_sparse_ws_regular_metadata_copy_stays_in_producer.
The remaining test_mixed_tma_cp_async_shared_stage_barriers failure
is a pre-existing issue on the original branch: the tiled WS pass
produces SIMT copies instead of cp.async because LowerPTXAsyncCopy
is not in the pass pipeline (the comment at phase.py:308 says it
runs earlier but no actual call exists on either branch or reference).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Preserve pipeline context for WS producer cp.async injection
Wrap the WS producer loop body in a kPipelineContextNumStages AttrStmt
so that LowerTileOp's pipelined_depth_ is > 0 when processing SIMT
producer copies. This enables InjectPTXAsyncCopy to generate cp.async
for global-to-shared copies in the WS producer branch.
Without this, the WS rewriter strips all pipeline annotations from
the rewritten loops, causing LowerTileOp to skip cp.async injection
for SIMT producers. The consumer loop stays annotation-free since it
doesn't need cp.async.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix mixed barrier protocol, consumer-init sinking, restore manual-layout guard
1. Mixed TMA+cp.async barrier: use ptx_cp_async_barrier_noinc for
forward barrier arrival in mixed producer groups, matching the
reference producer_consumer_ws.cc protocol.
2. Consumer-only pre-loop init sinking: in ReplacePipelineLoopInStmt,
guard pre-loop siblings as consumer-only when they're not classified
as producer (TMA/SIMT/cp.async). Fragment init (T.fill, T.clear)
and local buffer init are placed in the consumer branch instead of
the shared prelude.
3. Restored blanket has_manual_layout_ guard: the dtype-based heuristic
to distinguish sparse metadata from MXFP4 layouts doesn't work
because both use uint8. Dequant_groupedgemm_bf16_mxfp4_hopper
requires the guard to prevent broken WS.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Narrow consumer-init sinking, targeted manual-layout check, per-group barrier
1. Consumer-init sinking: only sink pre-loop stmts that are FillNode
writing to fragment/local buffers. Keep block_mask setup and shared
state in the shared prelude.
2. Manual-layout: attempt targeted check that only rejects when TMA
copy destinations match layout_map entries. Collect layout_map vars
and compare against TMA copy destinations.
3. Per-group cp.async barrier: use group-level cp.async flag (single
group for now) instead of function-wide boolean.
The 3 WS issue tests still fail because the layout_map annotation
parsing falls through to the conservative rejection path.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix layout_map parsing with Var→Buffer mapping for manual-layout check
The layout_map annotation uses Map<Var, Layout> before LayoutInference
(not Map<Buffer, Layout>). Parse as Map<ObjectRef, ObjectRef> and
handle both key types: Buffer (post-inference) and Var (pre-inference).
For Var keys, look up the corresponding alloc_buffer by data Var match.
Compare collected manual-layout buffers against pipeline copy
destinations: reject only when a manually-laid-out buffer is also
a producer copy target (TMA/SIMT/cp.async) inside the pipeline.
This allows sparse metadata (E_shared, SIMT-copied) onto the WS path
while rejecting dequant MXFP4 (B_shared, SIMT-copied with swizzle).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix SEGFAULT in manual-layout check, use DetectSwizzleMode
1. Fixed SEGFAULT: removed dangling layout_map_layouts_ vector that
was never populated, causing OOB access. Now stores Buffer+Layout
pairs in layout_map_entries_.
2. Use DetectSwizzleMode to distinguish swizzled layouts (MXFP4,
incompatible with MVB) from non-swizzled (sparse metadata, safe).
Swizzled layouts reject WS candidacy; non-swizzled layouts allow it.
3. Removed debug LOG(WARNING) from hot path.
4. Parse layout_map annotation keys as both Buffer and Var (via
Map<ObjectRef, ObjectRef>), resolving Var keys to alloc_buffers.
Results: sparse metadata WS test PASSES, dequant PASSES (protected).
Down to 2 failures: mixed barrier pattern + consumer init sinking.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Simplify consumer-init sinking to Evaluate-only classification
Sink pre-loop Evaluate nodes classified as kConsumer into consumer
branch. Keep For loops (block_mask setup), producer copies, and other
control flow in the shared prelude. This is simpler and safer than
the FillNode scope check approach.
The test_pure_tma_consumer_local_init test still fails because the
T.fill statements are at a different structural level than the
SeqStmt where the pipeline loop lives. Fixing this requires deeper
IR structure analysis.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Sink consumer-only pre-loop init into WS consumer branch
Extract consumer-only pre-loop Evaluate statements (T.fill on
fragments) from the shared prelude and prepend them to the consumer
branch inside the WS if/else structure. This ensures fragment init
like acc_o, logsum, scores_max appears only in the consumer branch,
not in the shared prelude or producer branch.
Uses a two-pass approach: first ReplacePipelineLoopInStmt extracts
consumer-only stmts into extracted_consumer_init_, then the WS body
is rebuilt with the extracted stmts prepended to the consumer branch.
Fixes test_pure_tma_consumer_local_init_does_not_leak_into_producer.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Finalize pipeline refactor fixes
* Refactor CopyNode stride checks for TMA bulk load/store
- Extracted the global stride validation logic into a new static method `CheckGlobalStrides` for better readability and reusability.
- Updated `CheckBulkLoad` and `CheckBulkStore` methods to utilize the new stride checking function, improving code clarity and maintainability.
- Enhanced documentation for the new method to clarify its purpose and requirements.
* Remove redundant copy direction notes from documentation in `copy.h` for stride checks. This simplifies the comments while maintaining clarity on TMA requirements.
* refactor
* layout related fix
* refactor mbarrier with software pipeline.
* remove legacy pass
* Enhance TMA barrier handling and merging logic
- Updated `CopyToTmaCopyRewriter` to conditionally emit arrive barriers based on the last TMA copy.
- Refactored barrier creation to allow merging of TMA barriers when conditions are met, reducing the number of barriers needed.
- Adjusted consumer and producer logic to accommodate merged barriers, ensuring correct synchronization across TMA copies.
- Improved documentation and comments for clarity on the new barrier merging behavior.
* Update example_mhc_post.py to disable main execution and add test function call
- Commented out the main function call in `example_mhc_post.py`.
- Added a call to `tilelang.disable_cache()` and a new `test(n=4096, h=2560)` function for testing purposes.
* Enhance pipeline barrier handling and TMA copy detection
- Updated `RewritePipelineTmaBarriers` to accept additional parameters for loop variable and stage count, improving barrier synchronization.
- Modified `HandleTileOp` to ensure only user-defined T.copy patterns are marked for pipeline barriers, preventing interference with T.tma_copy.
- Adjusted test cases to enforce specific CUDA compute version requirements for better compatibility and consistency.
* lint fix
* Enhance convolution example and TMA handling
- Added kernel source printing in `example_convolution.py` for debugging.
- Disabled the main function call and invoked `run_regression_perf()` to streamline execution.
- Introduced `MakeTmaLeaderCondition` function in `copy.cc` to improve TMA leader-thread condition handling.
- Updated `LowerBulkCopy` and `LowerBulkCopy1D` methods to utilize the new TMA leader condition for better thread management.
- Enhanced `Conv2DIm2ColOpNode` to support barrier annotations, improving synchronization in TMA operations.
* enhance
* Remove example_dequant_groupedgemm_bf16_mxfp4_hopper.py and clean up phase.py by removing unnecessary whitespace. Enhance TMA handling in producer_consumer_ws_tiled.cc to improve barrier synchronization and streamline TMA copy operations.
* fix
* fix
* Eliminate MVB(barrier_only=true) late fixup from OptimizeForTarget
Move pipeline barrier ownership into InjectSoftwarePipeline: create
pipeline_mbar[num_stages] at final expanded size instead of
pipeline_mbar[1] that required late MVB expansion.
Key changes:
- RewritePipelineTmaBarriers creates barriers at num_stages size
- Barrier indices use FloorMod(loop_var - loop_min, num_stages)
- barrier_init has num_stages entries (one per slot, arrive_count=1)
- CopyToTmaCopyRewriter accepts PrimExpr barrier_id (was int)
- phase.py: remove MVB(barrier_only=True), unify both paths under
PlanAndUpdateBufferAllocationLocation
- Fix pre-existing test expectation for tma_copies annotation
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix Codex review issues: num_stages=1 regression, im2col pipeline barriers, test fixes
1. Fix num_stages=1 regression: use tl_pipelined_num_stages annotation
for barrier sizing instead of max_stage+1. Gate barrier creation on
pipeline_depth > 1 so num_stages=1 kernels don't get multi-versioned
pipeline barriers.
2. Extend RewritePipelineTmaBarriers to handle c2d_im2col: annotate
im2col calls with pipeline barrier in CopyToTmaCopyRewriter. Add
im2col TMA recognition to PipelinePlanning. Fix im2col Lower() to
respect emit_arrive annotation from pipeline barriers.
3. Fix test_simple_pipeline: use annotation-checking approach instead of
structural equality with hardcoded tma_copies annotation that fails
on non-Hopper targets.
4. Remove dead kPipelineMVBStageExpr/kPipelineMVBParityExpr/
kPipelineMVBContextNumStages emission from inject_pipeline.cc since
pipeline barriers are now created at final size.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix depth-1 barrier ownership, add regression tests, remove dead kPipelineMVB*
1. Fix depth-1 ownership gap: always create shared pipeline barrier for
TMA copies even when pipeline_depth=1. This prevents LowerTileOp from
allocating separate per-copy internal barriers, keeping num_stages=1
kernels at pipeline_mbar[1] (single slot).
2. Add checked-in lowering regression tests:
- non-WS num_stages=3 TMA GEMM → asserts pipeline_mbar[3], no fallback
- non-WS num_stages=1 TMA GEMM → asserts pipeline_mbar[1], no multi-slot
- non-WS num_stages=3 im2col → asserts pipeline_mbar[3] feeds tma_load_im2col
3. Remove dead kPipelineMVB* definitions from pipeline_utils.h and
consumption code from multi_version_buffer_rewriter.cc (stacks,
explicit parity/version index logic, attr stripping).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix GemmWMMA.lower() signature and autotuner cache backwards compatibility
[P1] Add missing mbar_phase_expr parameter to GemmWMMA.lower() to match
the interface expected by GemmPy.lower() dispatch. Without this, any
RDNA kernel using the WMMA path would fail with TypeError.
[P2] Make out_idx.json loading optional in autotuner load_from_disk().
Older cache directories don't have this file; fallback to
compile_args.out_idx when the file is absent.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add .humanize to gitignore
* Fix HIP swizzle codegen and CPU scalar GEMM region handling
[P2] Update codegen_hip.cc to parse tvm_tuple(device_func, panel_size)
format for threadblock_swizzle_pattern, matching the annotation format
now emitted by T.use_swizzle(). Previously expected StringImmNode which
would ICHECK-fail.
[P2] Fix gemm_scalar.py: clear only the output tile region (not the
whole buffer), and use the last two dimensions from regions to handle
rank>2 buffers with leading singleton dimensions.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix HIP codegen to parse tvm_tuple swizzle annotation format
Update codegen_hip.cc threadblock_swizzle_pattern handler to parse
tvm_tuple(device_func, panel_size) format, matching CUDA and cutedsl
codegens. Previously expected StringImmNode which ICHECK-fails with
the current annotation format.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Restore kPipelineMVB* annotations — still needed by WS tiled full MVB
The kPipelineMVBStageExpr/ParityExpr/ContextNumStages annotations and
their MVB consumption code cannot be removed: they are emitted by
EmitImpl for ALL pipelines and consumed by the full MVB call inside
ProducerConsumerWarpSpecializedTiled. Removing them broke non-TMA
pipeline kernels like mhc_post that go through the WS path.
Restores: annotation emission in inject_pipeline.cc, constant
definitions in pipeline_utils.h, and consumption code in
multi_version_buffer_rewriter.cc.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Unify barrier multi-versioning in InjectSoftwarePipeline (Direction A)
Move all pipeline barrier multi-versioning into InjectSoftwarePipeline
via ExpandPipelineBarriers, eliminating the need for late
MVB(barrier_only=true) in OptimizeForTarget.
ExpandPipelineBarriers runs before BuildPipeline and handles:
- ISP-created pipeline_mbar (for non-WS TMA pipelines)
- User-written T.alloc_barrier (for manual WS pipelines like softpipe)
Key design: only barriers with explicit ptx_arrive_barrier calls OR
ISP-created local barriers are expanded. Barriers whose arrival is
managed internally by tile-ops (e.g., tcgen05 MMA arrive) are left
unchanged. This distinguishes pipeline sync barriers from hardware-
managed barriers.
Buffer expansion: barrier[N] -> barrier[N * num_stages]
Index rewriting: barrier[idx] -> barrier[stage_expr * N + idx]
Parity rewriting: user_parity -> (iteration_block + offset) % 2
barrier_init replication: [c0,c1] -> [c0,c1,c0,c1] for num_stages=2
Expanded buffers propagate to outer blocks via pending_buffer_remap_.
barrier_init annotations are replicated in VisitStmt_(BlockNode).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix Python 3.9 compat: replace star expression in index with tuple concat
* fix
* fix
* revert changes
* fix
* enhance pipeline
* enhance pipeline
* implement descriptor reuse pass
* Refactor pipeline management and enhance async copy handling
- Updated `example_mhc_pre.py` to replace the main function call with a test function for better testing flexibility.
- Modified `example_gqa_decode.py` to disable argument parsing and added latency measurement for regression performance.
- Enhanced `copy.cc` and `copy.h` with new checks for pipeline-managed cp.async synchronization and improved async copy handling.
- Updated `inject_pipeline.cc` and `pipeline_planning.cc` to refine the handling of tile operations and global/shared buffer checks.
- Added tests to ensure correct behavior of async pipeline and descriptor allocation reuse.
This commit improves the overall pipeline management and async copy capabilities, ensuring better performance and flexibility in the codebase.
* enhance pipeline
* enhance pipeline
* enhance pipeline
* enhance pipeline
* enhance pipeline
* Revert "enhance pipeline"
This reverts commit
|
||
|
KurisandClaude Opus 4.6
|
5f70374c48 |
[Test] Add 1D TMA regression test for issue #1842 (#2005)
Add a regression test covering 1D single-dimension tensor TMA copy (global -> shared -> global) with warp specialization disabled. The underlying bug was fixed in #1840, but the test suite only covered 2D descriptor-based TMA paths. This test ensures the 1D bulk copy path (cp.async.bulk) also works correctly with proper mbarrier allocation. Closes #1842 Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
||
|
Lei Wang
|
1c561f6d46 |
[Bugfix] Raise error on zero grid dimension instead of silent clamp (#1994)
* [Bugfix] Raise error on zero grid dimension instead of silent clamp (#1993) Fix ThreadWorkLoad::Extract() silently clamping zero grid dims to 1, which caused either CUDA_ERROR_ILLEGAL_ADDRESS crashes (dynamic case) or silent wrong results (static case). Closes #1993 * lint fix |
||
|
Lei WangandClaude Opus 4.6
|
ded6a99221 |
[Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API (#1909)
* [Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API
This PR introduces Producer-Consumer Warp Specialization for sm90+ TMA
pipelines and a new T.tma_copy() API for explicit mbarrier management.
Key changes:
1. **ProducerConsumerWarpSpecialized pass** (new):
Splits pipelined TMA loops into producer (TMA loads) and consumer
(compute) warp groups with back-pressure barriers for buffer reuse.
Works with num_stages >= 1. Producer warp (128 threads) handles
arrive_expect_tx + tma_load; consumer warps handle compute + arrive
on back-pressure barriers.
2. **T.tma_copy() API** (new):
Fire-and-forget TMA copy with a required `barrier` parameter. Unlike
T.copy() which emits producer+wait pairs, T.tma_copy() emits only
arrive_expect_tx + tma_load. User manages synchronization via
T.mbarrier_wait_parity().
3. **MultiVersionBuffer barrier expansion**:
Extended to handle `shared.barrier` scope buffers, expanding them for
pipelining (1D size multiplication instead of prepending a dimension).
Auto-computes mbarrier parity as `(k // num_stages) % 2`.
4. **Pass ordering** (phase.py):
LowerSharedBarrier now runs after MultiVersionBuffer in the TMA path
so barrier buffers retain their `shared.barrier` scope during expansion.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* fix for annotate reg alloc
* Enhance debugging and performance measurement in examples
- Added print statements to output kernel source for both forward and backward sparse MLA examples, aiding in debugging.
- Updated `sparse_mla_fwd_pipelined.py` to run performance regression tests and print average latency, improving performance tracking.
- Disabled auto cache in `example_mha_sink_fwd_bhsd.py` for better control during testing.
* fix
* Enhance debugging and performance tracking in examples
- Added print statements to output kernel source in various example scripts, including `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py`, aiding in debugging.
- Updated default `window_size` in `example_gqa_sink_bwd_bhsd.py` for improved configuration.
- Disabled auto cache in `example_gemm_schedule.py` to enhance control during testing.
- Introduced a new test for preserving protected auto-injected wait groups in the async copy optimization process.
* clean
* Remove unnecessary print statements for kernel source in example scripts
- Eliminated print statements from `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to clean up output and improve readability.
- These changes streamline the examples while maintaining functionality.
* clean
* lint fix
* Refactor T.tma_copy() and barrier synchronization in examples and tests
- Updated `example_warp_specialize_gemm_copy_1_gemm_0.py` to use `T.tma_copy()` for optimized memory copying with explicit barrier management.
- Enhanced test cases in `test_tilelang_issue_tma_no_ws.py` and `test_tilelang_language_tma_copy.py` to reflect changes in barrier synchronization, ensuring proper functionality with the new T.tma_copy() API.
- Removed unnecessary print statements and adjusted barrier allocations for clarity and performance.
- Cleaned up the `test_tilelang_transform_inject_tma_barrier.py` file by deleting it, as it was no longer needed.
* Add buffer data mapping for remapped buffers in multi_version_buffer_rewriter
- Introduced a check to ensure that the data variable of a buffer is discoverable, allowing the barrier_init annotation update to find the remapped buffer correctly.
- This change enhances the functionality of the buffer remapping process, ensuring proper synchronization in the transformation pipeline.
* Implement TMA im2col load synchronization with mbarrier allocation
- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier creation and integration into the TMA copy process, ensuring proper synchronization and data transfer.
- Updated the handling of the mbarrier in the TMA copy statement to improve performance and maintain functionality within the thread-gated block.
* Implement TMA im2col load synchronization with mbarrier allocation
- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier initialization and integration within the TMA copy process, ensuring proper synchronization during memory operations.
- Updated the handling of the mbarrier in the im2col operation to improve performance and maintain consistency in multi-threaded environments.
* disable tma lower for test_cutedsl_barrier
* Refactor code formatting in test_tilelang_jit_cutedsl.py
- Improved readability by adjusting the formatting of the `tilelang.compile` function call.
- Enabled the main testing function in the script to run without being commented out.
* Enhance sparse_mla_fwd.py and thread_storage_sync.cc with code cleanup and new functionality
- Added a call to disable cache in the main execution block of `sparse_mla_fwd.py` for improved performance.
- Commented out unused benchmarking code in `test_sparse_mla_fwd` to streamline testing.
- Introduced new utility functions in `thread_storage_sync.cc` to analyze and optimize shared memory access patterns, enhancing synchronization management.
- Improved readability and organization of the code by restructuring and commenting on key sections.
* Add thread storage synchronization utilities and optimizations
- Introduced functions to identify shared-like scopes, buffers, and pointer variables, enhancing memory management.
- Implemented a `BarrierTransparentStmtAnalyzer` to analyze statement transparency concerning shared memory.
- Developed a `ThreadStorageSyncOptimizer` to optimize sequences of statements by removing unnecessary synchronization calls, improving performance in multi-threaded environments.
- Enhanced handling of shared memory access in various call types, ensuring proper synchronization during memory operations.
* remove optimize thread_storage_sync
* Refactor pipeline planning to enhance cp.async group scheduling
- Introduced a last-use index for cp.async groups to improve scheduling accuracy.
- Removed unnecessary tracking of first consumers, simplifying the scheduling logic.
- Updated the handling of cp.async and commit statements to ensure proper group ordering based on last-use.
- Added a new test to verify that cp.async groups are ordered correctly by their last-use index.
* lint fix
* Add CUDA requirements to TMA copy tests and update main execution call
- Added decorators to require CUDA and a minimum compute version for the `test_tma_copy_pipeline_2_stages` and `test_tma_copy_pipeline_3_stages` functions.
- Updated the main execution block to call `tilelang.testing.main()` instead of directly invoking the test functions, improving test execution management.
* Refactor attention sink examples by removing pipelined implementations
- Updated imports in `benchmark_gqa_sink_fwd.py` and `benchmark_mha_sink_fwd.py` to use the non-pipelined versions of the attention sink functions.
- Deleted the pipelined implementations of `example_gqa_sink_fwd_bhsd_wgmma_pipelined.py` and `example_mha_sink_fwd_bhsd_wgmma_pipelined.py`.
- Cleaned up regression and test files to remove references to the deleted pipelined examples, ensuring consistency across the codebase.
* performance fix
* Enhance example_blocksparse_gemm.py and inject_pipeline.cc with debugging and utility improvements
- Added a print statement to output the kernel source in `example_blocksparse_gemm.py` for debugging purposes.
- Updated `inject_pipeline.cc` to include a new utility function for stripping TMA copy write buffer attributes, improving code clarity and functionality.
- Modified `producer_consumer_ws.cc` to utilize the new utility function, enhancing the handling of TMA barriers.
- Expanded the `_compile_tvm_ffi` function in the test file to accept additional keyword arguments, improving flexibility for test configurations.
- Introduced a new test to verify that pure TMA warp specialization correctly emits barriers, ensuring compliance with expected behavior.
* lint fix
* Refactor guard variable declarations in producer_consumer_ws.cc for clarity
- Changed guard variable declarations from non-const to const references to improve performance and readability.
- This minor adjustment enhances the handling of producer guards in the context of TMA barriers.
* lint fix
* fix
* lint fix
* perf fix
* Add print statement for kernel source and update main execution in example_dequant_gemm_w4a8.py
- Introduced a print statement to output the kernel source for debugging purposes.
- Updated the main execution block to call `run_regression_perf()` and print the latency, enhancing performance measurement capabilities.
* fix
* Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py
- Added a print statement to display the kernel source for better debugging.
- Updated the main execution block to disable cache and call `run_regression_perf()`, printing the resulting latency for performance measurement.
- Enhanced the producer-consumer logic in producer_consumer_ws.cc with new utility functions to manage local variable access and improve code clarity.
- Introduced a new test case in test_tilelang_issue_tma_no_ws.py to ensure consumer-only local initializations do not leak into the producer, verifying correct behavior in TMA contexts.
* Revert "Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py"
This reverts commit
|
||
|
Henry
|
911385fdd1 |
[BugFix] Fix Hopper TMA lowering without warp specialization (#1840)
* [BugFix] Fix Hopper TMA lowering without warp specialization Decouple the Hopper pipeline entry condition from warp specialization so TMA lowering works correctly when `tl.disable_warp_specialized=True`. Also fix barrier id promotion for multi-dim TMA loads in inject_tma_barrier.cc. Fixes #737, #1648 * Unify is_1d_tma_load logic and improve test diagnostics Address CodeRabbit review feedback: - Extract shared Is1DTmaLoad() helper to ensure consistent barrier argument index selection across TmaExpectTxRewriter and TmaBarrierRewriter - Add failure message to regex assertion in regression test * Fix ruff-format: merge f-string into single line |
||
|
Tong WUandLeiWang1999
|
ffa0cdd1e1 |
[BugFix] Fix eager mode where there is no tensor args (#1851)
* [BugFix] Fix eager mode where there is no tensor args * fix --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei Wang
|
4ff81c7d40 |
[Refactor] Introduce T.access_of to combine T.address_of and access_ptr (#1827)
* [Refactor] Update atomic operations to use access_ptr for buffer access * Replaced instances of address_of with access_ptr in atomic operations to improve clarity and maintainability. * Enhanced atomic operation implementations in various files, including atomic_add, atomic_max, and atomic_min. * Updated related logic in loop vectorization and tile language to support new access patterns. * Added a utility function to extract destination scalar dtype for atomic operations, supporting both BufferLoad and tvm_access_ptr. * Improved kernel cache logic to include a build-stamp for the TileLang runtime library, ensuring cache validity during development. * lint fix * Fix access_ptr layout rewrite and atomic vectorization * lint fix * Refactor atomic operations to remove const qualifiers from pointer parameters in atomicAddx2 and ToFloat functions. This change enhances type flexibility and consistency across atomic operations. * Refactor buffer access logic in lower_tile_op and vectorize_loop to enhance clarity and maintainability. Removed unused code and improved offset calculations for tile access. Updated atomic operation handling to streamline dtype extraction and ensure correct vectorization behavior. * Add logging for loop layout and lowered loop in atomic_add; implement tvm_access_ptr handling in loop_vectorize and lower_tile_op This commit introduces logging statements to track the loop layout and the result of the lowered loop in the AtomicAddNode. Additionally, it implements a new method to handle tvm_access_ptr calls in the loop_vectorize and lower_tile_op files, enhancing the handling of memory access patterns and ensuring correct vectorization behavior. * fix * lint fix * fix * Update ToFloat2 and ToFloat4 functions to use const qualifiers for pointer parameters in atomic.h; modify test script to call test_atomic_addx2_float instead of main. * Refactor tvm_access_ptr handling in loop_vectorize to improve argument validation and buffer resolution logic. Ensure at least three arguments are provided and streamline the search for matching buffers in layout_map_. Update vector size calculations to enhance clarity and maintainability. * Add handling for tvm_access_ptr in vectorize_loop to support offset substitution and buffer creation This commit introduces logic to handle tvm_access_ptr calls, allowing for the substitution of loop variables in offsets and the creation of dummy buffers. The new method ensures correct base address calculation and enhances the overall handling of memory access patterns in vectorization. * Add handling for tvm_access_ptr offset validation in loop_vectorize This commit enhances the loop_vectorize functionality by adding checks for the element offset in tvm_access_ptr calls. It ensures that offsets are invariant within the vector boundary, improving the robustness of vectorization logic and memory access patterns. * fix * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update improves the robustness of vectorization logic and ensures offsets remain invariant within vector boundaries. * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update ensures offsets remain invariant within vector boundaries, improving memory access robustness and vectorization logic. * fix * Add tl.access_ptr support in TileLang This commit introduces the `tl.access_ptr` operation, which carries pointer access metadata in the frontend. It allows for better handling of buffer loads by retaining necessary information for downstream analysis. The new operation is integrated into various components, including the loop vectorization and memory access legalizations, ensuring compatibility with existing functionality while enhancing the robustness of memory access patterns. * lint fix * fix * Add support for tl.access_ptr in atomic operations This commit updates the AtomicAddNode and AtomicOpBaseNode implementations to utilize the new tl.access_ptr operation for building access pointers to destination elements. It enhances the handling of buffer loads by incorporating type checks for index data types, improving the robustness of memory access patterns. Additionally, it removes unnecessary logging statements in the lowering process, streamlining the codebase. * lint fix * Refactor buffer load extraction in vectorize_loop to check for non-empty arguments in access_ptr calls. This change improves the robustness of buffer load handling by ensuring valid argument presence before processing. * Remove unused fallback logic in KernelCache and clean up import statements in test_tilelang_language_access_ptr.py for improved code clarity. * Refactor memory access in BitNet and dequantization examples to use tl.access_ptr This update replaces T.address_of calls with T.access_ptr in multiple files, enhancing memory access handling by specifying read and write permissions. This change improves code clarity and aligns with recent enhancements in pointer access management. |
||
|
Lei Wang
|
6722690697 |
[BugFix] Fix fp16 annotate_l2_hit_ratio host stub compilation (issue #1810) (#1818)
* Fix L2 persistent lowering for fp16 host stubs Avoid typed access_ptr casts in stream access policy window setup and add issue #1810 regression test. * test and lint fix * Add regression test for issue #1810 Add a regression test for issue #1810 related to CUDA stream access policy. |
||
|
Jiawei Lin
|
af30ac2152 |
[Bugfix] Support loop-dependent conditions in IfThenElse within T.Pipelined (#1799)
* [BugFix] Fix loop-dependent conditions in IfThenElse within T.Pipelined This commit applies the same strategy used for LetStmt to IfThenElse conditions: 1.Introduced IfWrapper struct to track if conditions that depend on the loop variable 2.Added dependency detection that checks if an if condition uses: - The pipeline loop variable directly, OR - Any variable transitively dependent on the loop variable 3.Loop-dependent conditions are pushed inside each pipeline stage with the loop variable properly substituted for that iteration * Add test for loop-dependent condtions within T.Pipelined * Fix code format |
||
|
Yu ChengandLeiWang1999
|
8c45417584 |
[Feature] Support Pass LoopUnswitching (#1747)
* [Feature] Implement Hoist Loop-Invariant If Pass - Added a new pass to hoist loop-invariant if statements out of loops, improving optimization opportunities. - Introduced classes for collecting and checking conditions, as well as for rewriting statements. - Integrated the new pass into the optimization pipeline and provided a corresponding API for usage. * Add tilelang semantics guide to programming guides (#1740) * Add CallNodeChecker to Loop Unswitching Pass and integrate tilelang cache disabling in benchmark script - Introduced CallNodeChecker class to identify CallNode expressions in loop conditions, enhancing loop-invariant checks. - Updated IsLoopInvariant function to reject conditions containing CallNodes, preventing potential side effects. - Added tilelang.disable_cache() in benchmark_mha_sink_fwd.py to optimize performance during benchmarking. * lint fix * lint fix * Enhance Loop Unswitching Pass with Let Binding Support - Added support for Let-bound variables in the WrittenBufferReadChecker to improve buffer read checks. - Introduced UsesLoopVarThroughLetBindings function to check if conditions depend on loop variables through Let bindings. - Updated IsLoopInvariant function to account for Let bindings when determining loop invariance. - Enhanced HoistableIfFinder to track Let bindings for variables bound to BufferLoad expressions. - Added debug print statements in the OptimizeForTarget function to visualize the module state before and after loop unswitching. * fix * lint fix * Enhance Loop Unswitching with CallNode Exclusion Check - Added CallCheckerExcludingIf class to ensure function calls outside of hoisted if statements are identified, preventing potential synchronization issues during loop unswitching. - Updated loop unswitching logic to incorporate the new call checker, enhancing safety and correctness. - Integrated debug print statements in OptimizeForTarget to visualize module state before and after loop unswitching. - Disabled tilelang cache in the benchmark script for improved performance. * Update TVM subproject to latest commit and remove debug print statements in OptimizeForTarget function --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei Wang
|
fbf28a143f |
[ThreadSync] Use Z3 for constraint equivalence checking (#1760)
* [ThreadSync] Use Z3 for constraint equivalence checking Replace const_int_bound based range comparison with Z3 SMT solver for checking thread constraint equivalence in FindConflict. * [ThreadSync] Clarify conflict detection for constraint equivalence Added comments to improve understanding of the logic determining whether constraints are equivalent or in conflict in the thread storage synchronization process. * add test * Update test_issue_1106 to check syncthreads is outside for loop * test fix |
||
|
Lei Wang
|
4fa2d24d27 |
[BugFix] Always build guard in loop partitioning to prevent out-of-bounds access (#1756)
* [Refactor] Simplify guard construction in loop partitioning - Removed the need for a separate guard check based on bijectivity, allowing the analyzer to determine the necessity of the guard. - Streamlined the guard construction process by consolidating index bounds checks into a single loop, improving code clarity and maintainability. - Enhanced handling of non-bijective layouts and cases where loop extents differ from layout input shapes. * add a reproducable test |
||
|
Zhengju Tang
|
9ddf577042 |
[BugFix] Stride check and fix for tensors with zero-stride argument (#1749)
* [BugFix] Stride check and fix for empty tensor in tvm_ffi * Lint * Lint * Revert * Lint * Fix zero stride check in C host code * Add warning for zero stride * Modify test scripts * Lint * Rename test * Refactor * simplify * Lint * Fix |
||
|
Kuris
|
fa9660bb5b |
[Layout] Fix Layout Bugs in Parallel and Reduce (#1713)
* [Layout] Add read write check in layout inference * remove debug stmts * fix typo in validater * [Layout] Prevent widening of layout in ReduceOpNode::InferLayout * [Layout] Enhance layout conflict handling in ReduceOpNode::InferLayout * [Layout] Fix layout issue in parallel and reduce * Fix lint error * Fix inconsistent predicate bug * Fix bug in test * Add cuda restriction on test * Add cuda restriction on test * Fix comments * Add assertion * fix bug in tests |
||
|
Kuris
|
7ec5542db9 |
[BugFix] Remove unnecessary binding in loop variable analysis and add test for issue 1728 (#1735)
* [Fix] Remove unnecessary binding in loop variable analysis and add test for issue 1728 * [Test] Add test |
||
|
Kuris
|
5358f5a272 |
[EagerJIT] Lazy Evaluation of Kernel Body in Eager JIT (#1690) (#1694)
* [Refactor] Support lazy evaluate in eager jit * add test script * fix lint error * canonicalize signature and error of `T.empty` * Fix comment from genai * [EagerJIT] Add comment for eager jit stage |
||
|
Lei Wang
|
c9cec62557 |
[Bugfix][Refactor] Always disable light storage reuse (#1691)
* fix * lint fix * Remove print statement in test_tilelang_issue_1678.py Remove print statement for kernel source in test. |
||
|
silentCoder-devandLeiWang1999
|
1d71bed2db |
[Feature] Reimplement Threadsync with ConstrVisitor (#1631)
* reimplement ThreadSync pass * fix bugs * remove old storage_access * remove deadcode for debugging * format * add WhileOp & expose MakeGuard in constr_visitor * bugfix * bugfix * bugfix * bugfix * trigger ci * add try for checking * bugfix * update tvm with better analyzer * update tvm * Log for debugging in thread_storage_sync * format * consider let node * typo * remove debugging log * consider while op in thread_sync & fix bug sbout letnode * add testing for thread_sync issues * update testing for thread_sync * format * consider T.Ramp expr * remove useless header * format --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei Wang
|
4fb96a0f56 |
[Bugfix] Fix variable scoping issue in InjectSoftwarePipeline for transitive LetStmt dependencies (#1657)
* [Enhancement] Update global load/store functions for CUDA compatibility (#1652) Refactor the `ld_global_256` and `st_global_256` functions to support both CUDA versions above 12.9 and earlier versions. This change ensures that 256-bit loads and stores are handled correctly across different CUDA versions, improving performance and compatibility. The implementation now uses two 128-bit loads/stores for older versions, enhancing the robustness of the codebase. * Update comments in global load/store functions for CUDA compatibility Clarified comments in `ld_global_256` and `st_global_256` functions to indicate that the fallback for CUDA versions below 12.9 may have performance regressions. This change enhances code readability and provides better context for developers working with different CUDA versions. * Update submodule and enhance LetStmt handling in inject_pipeline.cc - Updated the TVM submodule to the latest commit. - Improved the handling of LetStmt in the inject_pipeline.cc file to account for transitive dependencies on loop variables, ensuring correct variable substitution in rewritten blocks. - Adjusted test_tilelang_issue_1263.py to remove unnecessary jit decorator and updated the kernel compilation process with specific pass configurations. * lint fix * revert tvm * remove unused test * test fix |
||
|
Kuris
|
326be4dd3f |
[Feat] Allow dangling producer in wasp pipeline planning (#1263) (#1647)
* [Fix] Handle cases where producers lack corresponding consumers in synchronization logic * add test for pipeline without consumer * fix lint error |
||
|
Lei Wang
|
6e439530b5 |
[Refactor] Unify @jit and @lazy_jit into a single @jit decorator (#1632)
* Refactor: Replace @tilelang.lazy_jit with @tilelang.jit in examples and tests Updated all instances of the @tilelang.lazy_jit decorator to @tilelang.jit in the lazyjit example notebooks and related test files. This change aligns the code with the new JIT compilation approach, enhancing consistency across the codebase. Additionally, removed the lazy_jit import from the module initialization to streamline the API. * Enhancement: Improve layout annotation handling in T.copy and parallel loops Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify layout hints for the outermost parallel loop. Enhanced layout annotation handling in CopyNode and AtomicAddNode classes to ensure compatibility with SIMT operations. Added tests to validate the functionality of loop layout annotations, improving robustness and clarity in layout management for parallel loops. * Add out_idx parameter to @jit and validate it's only used in lazy mode * lintfix * Refactor: centralize mode inference in JITImpl.__call__ before parse_args * test fix * Enhancement: Introduce JITNoBuilderError for eager mode checks in T.Kernel and T.const functions. Ensure proper error handling when Builder is unavailable, improving robustness in JIT context management. * lint fix * lint fix * Enhancement: Update example_triton_sparse_gqa_decode_varlen_indice.py to improve performance metrics output by adding average FLOPS calculations. Modify layout_inference.cc to refine buffer layout inference checks by ensuring fragment buffers are considered. * lint fix * Enhancement: Refactor JIT implementation to centralize mode initialization and improve error handling for eager mode. Update matmul tests to utilize both Cython and TVM FFI JIT kernels, ensuring consistency in output validation. * Refactor: Remove unused simplify_prim_func decorator from example_gemm_intrinsics.py and related test files. Enhance layout.cc to simplify integer set analysis for multi-threaded compilation, ensuring constant integer outputs. Update lower_tile_op.cc to enforce constant integer checks in layout output shapes. Improve error handling in builder.py for JIT context management. * remove claude.md * Refactor: Clean up code formatting in layout.cc, lower_tile_op.cc, and proxy.py for improved readability. Add missing import in test_tilelang_kernel_int4_gemm_mma.py and ensure consistent spacing in builder.py. * Refactor: Replace logical conditions with tir_all in complex expressions within test_arith_hard.py for improved clarity and consistency. Update main execution flow to call test_hard_prove instead of tilelang.testing.main. * Enhancement: Add @tilelang.testing.requires_cuda decorator to test functions in test_tilelang_issue_1549.py and test_tilelang_issue_1601.py to ensure CUDA requirements are explicitly stated for relevant tests. * lint fix * Refactor: Simplify output shape calculation in layout.cc by directly setting max value from ist.max(). Update test_valid_loop in test_tilelang_fragment_loop_checker.py to disable certain tests and ensure proper execution flow. Enhance error handling in builder.py to correctly manage PrimFunc return types. * Refactor: Update boolean condition in test_boolop to use Or and Not for improved clarity and consistency in logical expressions. * Refactor: Update test execution flow in test_tilelang_fragment_loop_checker.py to call tilelang.testing.main. Modify JIT implementation to handle callable functions more generically. Enhance TirTemplate class in builder.py to return an empty tuple if matcher is None, improving robustness. * Refactor JIT function handling: Replace LazyJITFunc with JITFunc across the codebase, updating related function signatures and logic to support eager JIT compilation mode. This change enhances the clarity and functionality of JIT operations. * Update JITImpl documentation: Clarify output tensor index attribute description and remove outdated comment regarding implicit TIR construction. * Refactor TileLang to transition from v2 to eager language module. Update imports across multiple files to replace references to the deprecated v2 module with the new eager module, ensuring compatibility with the latest JIT compilation features. Remove obsolete v2 module files to streamline the codebase. * Enhance JIT compilation support: Update JIT function handling to improve compatibility with eager execution. Refactor related function signatures and logic, ensuring seamless integration with the latest compilation features. This change aims to streamline JIT operations and enhance overall code clarity. |
||
|
silentCoder-devandkurisu6912
|
b5be1a16ab |
[Feature] Add hoist_broadcast_values pass (#1606)
* add hoist_broadcast_values pass * only hoist intImm or floatImm * refactor * add test for hoist_broadcast_values * add a test for issue 1601 * add transform test for hoist_broadcast_value --------- Co-authored-by: kurisu6912 <227995639+kurisu6912@users.noreply.github.com> |
||
|
Lei Wang
|
358f899951 |
[Bugfix] Do not consider local.var as local buffer during LowerTileOP (#1628)
* [Refactor] Simplify IsLocalBuffer checks in lower_tile_op.cc by removing redundant boolean argument * lint fix |
||
|
Kuris
|
5d806298de |
[Feat] Add strong checker to detect data racing in T.Parallel (#1615)
* add verifier for parallel loop data racing * fix lint error * fix typo * fix lint error * fix bug in metal test |
||
|
Chaofan Lin
|
1b0022015f |
[BugFix] Fix some bugs in lowering ParallelOp and VectorizeLoop (#1607)
* fix * [BugFix] Fix some bugs in lowering ParallelOp and VectorizeLoop * remove debug output |
||
|
Gongen-Ali
|
4e53b50750 |
[Bugfix] Open Rocm ci test and fix some bugs. (#1443)
* Fix some bugs on ci and open rocm ci test * fix pre-commit * fix build issue on hip * refactor on hip * pre-commit fix * fix test with pdl * fix pdl test |
||
|
Lei Wang
|
0fa16b4c41 |
[Enhancement] Improve equality checks in layout nodes and fragment validation (#1573)
* [Enhancement] Improve equality checks in layout nodes and fragment validation * Enhanced the IsEqual method in LayoutNode and FragmentNode to include detailed comparisons of forward mappings, ensuring accurate equality checks. * Introduced a new parameter in ProveFragmentContains to validate physical indices when checking fragment containment, improving correctness in layout validation. * Removed obsolete test file related to layout inference. This update aims to strengthen the integrity of layout comparisons and fragment validations in the system. * lint fix |
||
|
Lei Wang
|
0f9bbd759c |
[Parallel][Infer] Free-mode chooses minimal replication between buffer-based and PlanLoopPartition (#1559)
* [Enhancement] Improve layout inference in ParallelOpNode * Enhanced the layout inference mechanism in ParallelOpNode to utilize two strategies: compute_loop_layout_from_buffer and PlanLoopPartition, selecting the one that minimizes replication while ensuring compatibility. * Updated the logic to choose the best candidate layout based on replication size and containment checks. * Refactored the HasKnownLayoutAnchor function to clarify its purpose in prioritizing buffer layouts. * Added a new test case to validate the layout inference behavior, ensuring the correct fragments are generated in the output. This update aims to optimize layout inference for parallel operations, improving performance and resource utilization. * lint fix * bug fix and refactor * lint fix |
||
|
KurisandLeiWang1999
|
d5d959ef7e |
[Fix] Add support for non-var complement arithmetic computation (#1374) (#1533)
* [Fix] Add complement for non-var IterExpr * fix lint error * lint fix --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei Wang
|
c750fb8a98 |
[Enhancement] Update examples and tests for improved type handling functionality (#1448)
* [Enhancement] Update examples and tests for improved type handling and functionality
- Enhanced various example scripts to support new data types and improve compatibility with PyTorch.
- Updated tests across multiple modules to ensure correct functionality with the latest changes in type handling.
- Refactored code in examples to streamline operations and improve clarity, particularly in tensor operations and memory management.
- Added comprehensive tests for new features and fixed existing issues related to type conversions and buffer handling.
* [Refactor] Update accumulation data type to float32 across examples
- Changed accumulation data type from "float" to T.float32 in multiple example scripts to ensure consistency and improve numerical stability.
- This update affects various modules including flash attention, GEMM analysis, convolution, and deepseek MLA examples, enhancing type handling across the board.
* [Refactor] Standardize data type usage across benchmark scripts
- Updated data type definitions in benchmark scripts to use T.float16 and T.float32 consistently, enhancing clarity and type handling.
- Adjusted dtype assignments in matmul functions and configuration setups to align with the new standard.
- Improved overall code consistency and maintainability by ensuring uniform data type usage across various modules.
* [Refactor] Standardize data type usage in templates and scripts
- Updated data type definitions in various templates and scripts to use string representations (e.g., "float16", "int32") instead of T.float16 and T.int32 for improved consistency and clarity.
- Enhanced overall code maintainability by ensuring uniform data type usage across multiple modules, including convolution, elementwise operations, and matrix multiplication templates.
- This change aims to streamline type handling and improve compatibility with existing workflows.
* [Refactor] Standardize data type usage in examples and benchmarks
- Updated data type definitions in various example and benchmark scripts to use T.float16 and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in kernel functions and configuration setups to align with the new standard.
- Improved overall code consistency by ensuring uniform data type usage across multiple modules, including attention mechanisms, matrix multiplication, and GEMM examples.
* [Refactor] Import dtypes from language.v2 module
- Added import statement for dtypes from the language.v2 module to enhance type handling and maintain consistency across the codebase.
- This change aims to streamline data type management and improve overall code clarity.
* fix
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use string representations (e.g., "float16", "int8") instead of T.float16 and T.int8 for improved consistency and clarity.
- Adjusted dtype assignments in functions and configuration setups to align with the new standard, enhancing overall code maintainability.
- This change affects multiple modules, including benchmark and attention mechanisms, ensuring uniform data type usage throughout the codebase.
* [Refactor] Update data type handling for consistency and clarity
- Changed string representations of data types in the Hint class to use T.float32 and T.int32 for improved consistency.
- Added new data types "int4" and "int16" to the dtypes module, enhancing type support across the codebase.
- Updated function signatures and assertions in the lop3 and mxfp modules to utilize the new data types, ensuring uniformity in type handling.
- This refactor aims to streamline data type management and improve overall code clarity and maintainability.
* [Enhancement] Improve data type handling and error messaging
- Introduced a mapping for canonical data types to their display strings, enhancing clarity in type representation.
- Updated the dtype creation logic to utilize the new mapping, ensuring more intuitive handling of string inputs.
- Refined error messages in the lop3 module to provide clearer feedback on invalid source formats, improving debugging and user experience.
* [Fix] Correct boolean flag in GEMM SP test case
- Updated the boolean flag in the test_gemm_sp_sm90 function to ensure proper functionality in the test case.
- This change enhances the accuracy of the test and aligns it with expected behavior for the GEMM SP implementation.
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use T.float16 and T.bfloat16 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in function signatures and argument parsing to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change affects multiple modules, including benchmarks and examples, improving overall code consistency and readability.
* [Refactor] Standardize data type usage in various modules
- Updated data type assignments in multiple scripts to utilize T.float32, T.int8, and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted function signatures and parameter types across benchmarks, examples, and tests to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change improves overall code consistency and readability, impacting modules related to matrix multiplication, GEMM, and tensor operations.
* [Refactor] Update argument parsing for data types in benchmarks
- Changed argument parsing for data types in benchmark_matmul_intrinsic.py and benchmark_matmul_sp.py to use string representations ("float16", "int8", "float") instead of T.float16 and T.float.
- This update enhances consistency in data type handling across benchmark scripts, improving clarity and maintainability.
* [Refactor] Update data type handling in benchmark and example scripts
- Changed data type arguments in benchmark and example scripts to use string representations ("float16") instead of T.float16 for improved consistency.
- Updated function signatures and argument parsing to align with the new standard, enhancing clarity and maintainability across the codebase.
- This change affects multiple modules related to attention mechanisms and tensor operations, ensuring uniform data type usage throughout the examples.
* [Refactor] Fix data type conversion in multiple scripts
- Corrected the usage of the data type conversion method from dtype..as_torch() to dtype.as_torch() across various benchmark and example scripts.
- This change enhances consistency in data type handling and improves code readability, impacting modules related to attention mechanisms and tensor operations.
* [Refactor] Update float8 data type usage across multiple scripts
- Changed instances of T.float8_e4m3 to T.float8_e4m3fn in various benchmark, example, and test scripts to ensure consistency in data type handling.
- This update enhances clarity and maintainability across the codebase, particularly in modules related to matrix multiplication and tensor operations.
* [Refactor] Enhance float8 data type handling in CUDA code generation
- Updated the handling of float8 data types in the CUDA code generation to include additional float8 variants, improving type conversion logic.
- Adjusted conditions to ensure proper type checks for float8 conversions, enhancing clarity and maintainability in the codebase.
- Modified layout inference to streamline float8 type checks, ensuring consistency across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Streamline float8 data type handling in CUDA and related modules
- Enhanced float8 data type handling in CUDA code generation by refining type conversion logic and ensuring consistent type checks.
- Updated layout inference for float8 types to improve clarity and maintainability across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Remove unnecessary cache disabling in float8 example script
- Eliminated the call to tilelang.disable_cache() in example_group_per_split_token_cast_to_fp8.py to streamline the code.
- This change enhances clarity and maintainability of the example script without affecting its functionality.
* [Refactor] Update data type usage in debug print tests
- Changed the argument for dtype in the test_debug_print_buffer function from a string representation to the corresponding T.bool type.
- This update enhances consistency in data type handling within the test suite, improving clarity and maintainability.
* lint fix
* Update function parameter types from `str` to `T.dtype` for improved type safety in attention sink and related examples
* Refactor `gemv_alloc_reducer` function signature for improved readability by formatting parameters across multiple lines.
|
||
|
Lei Wang
|
29051439db | [Lint] Phaseout Yapf format and embrace ruff format (#1417) | ||
|
Lei WangandZhiwen Mo
|
f5d9da4678 |
[Refactor] Phaseout vmap for Tile Operators (#1334)
* Refactor GEMM and Reduce operations by moving NormalizeToBufferRegion and MakeAccessPtrFromRegion to utils.{h,cc} for better code organization and reuse.
* lint fix
* Refactor region handling by removing the RegionOp and updating NormalizeToBufferRegion to only accept BufferLoad and BufferRegion. This change improves code organization and simplifies the handling of memory regions across various operations.
* fix
* Refactor memory region handling by introducing `tl.region` calls across various operations, including GEMM and fill functions. This change enhances the consistency of region management and improves code organization by utilizing utility functions for buffer region conversions.
* fix
* fix
* test fix
* lint fix
* Refactor GEMM operations to improve memory region handling by replacing `mbarPtr_` with `mbarRegion_` and updating related logic in both C++ and Python implementations. This change enhances the clarity and consistency of buffer region management.
* fix
* lint fix
* fix
* fix
* test fix
* lint fix
* lint fix
* minor fix
* fix
---------
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
|
||
|
Chaofan LinandLei Wang
|
9dda774aff |
[BugFix] Use BufferRegion in tl.cumsum to infer buffer shape (#1321)
* [BugFix] Use BufferRegion in tl.cumsum to infer buffer shape * remove debug lines * remove rubbish * Fix decorator syntax for atomic_different_memory_orders_program --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com> |
||
|
Lei Wang
|
2c0072a888 |
[Refactor] Update buffer handling in copy and atomic operations (#1247)
* [Refactor] Update buffer handling in copy and atomic operations
* Refactored the `copy` and `atomic_add` functions to use element-wise minimum for defining copy extents, ensuring correct handling of overlapping regions.
* Updated utility functions to create `BufferLoad` instances with explicit extents, improving memory management and clarity.
* Removed unused imports from `atomic.py` and `copy.py` to streamline the codebase.
* Adjusted logging in `copy.cc` to provide clearer warnings for fallback scenarios in bulk copy operations.
* Remove obsolete .git_commit.txt file
* Add unit test for dynamic copy extent handling in TileLang
* Introduced a new test file `test_tilelang_issue_1237.py` to verify that the `T.copy` function correctly manages dynamic extents during primitive function building.
* The test reproduces a specific issue related to dynamic slice lengths and static buffer sizes, ensuring robustness in the handling of such scenarios.
* The test does not require execution of the kernel, as building the primitive function is sufficient to validate the fix.
* lint fix
* fix
* Revert "fix"
This reverts commit
|
||
|
Kuris
|
d5fda276e0 |
[Fix] Fix buffer re-import typo in tilelang.languge (#1214)
* Fix Buffer re-import typo in tilelang.langugage * fix lint error |
||
|
Lei Wang
|
85218bd973 |
[Bugfix] Enhane LetStmt Handling in Pipeline Transform (#1212)
* [Enhancement] Introduce LetWrapper for handling loop variable substitutions in pipeline rewriting * Added LetWrapper struct to encapsulate variable and value pairs for loop variable substitutions. * Updated PipelineRewriter to accept a vector of LetWrapper instances, allowing for proper handling of Let statements that depend on the pipeline loop variable. * Enhanced the BuildPipeline method to incorporate LetWrapper instances into rewritten blocks, ensuring correct substitutions during pipeline execution. * Refactored logic for processing Let statements to differentiate between those that use the loop variable and those that do not, improving the flexibility of the pipeline transformation. * Refactor lambda expression for clarity in loop variable usage check in inject_pipeline.cc * [Test] Add regression test for loop variable handling in kernel compilation * Introduced a new test case to verify correct handling of loop variables in the kernel compilation process, addressing a regression issue with InjectSoftwarePipeline. * The test ensures that the loop variable is not left as a free variable, which previously caused failures in MakePackedAPI. * Configurations are set to disable warp specialization and TMA lowering to align with the original issue reproduction. * Remove unused import in regression test for loop variable handling in kernel compilation |