mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
6873e19da2b267d3cb7d0305561ba0bed63feae3
15
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 andyluo7
|
dd9c20dbf7 |
[Testing] Gate CUDA-only tests so non-CUDA backends can run the suite (#2863)
A large number of tests import CUDA-only internals directly (for example
`tilelang.cuda.pipeline.CUDAPassPipelineBodyPrologue`) or exercise
NVIDIA-specific features such as TMA, LDGSTG, mbarrier, warp specialization,
tcgen05 and tfloat32. `src/cuda/CMakeLists.txt` only compiles a small
"ALWAYS" subset when `USE_CUDA=OFF`, so those FFI symbols are never
registered and the tests fail with:
AttributeError: module 'tilelang.cuda._ffi_api' has no attribute
'AnnotateDeviceBoundTmaCopies'
This is not ROCm-specific -- the same tests fail on any build without the
CUDA backend, including Metal and CPU-only.
Add `@tilelang.testing.requires_cuda` to those tests, plus three
portability fixes that keep coverage on non-CUDA backends instead of
skipping it:
* `test_tilelang_transform_layout_inference.py`: CUDA emits
`__launch_bounds__(N, 1)` while HIP emits `__launch_bounds__(N)`.
Match both.
* `test_tilelang_transform_pipeline_planning.py`: PipelinePlanning
deliberately strips pipeline annotations on ROCm targets other than
gfx950 (see the guard in src/transform/pipeline_planning.cc), so assert
the documented sequential fallback there.
* `test_tilelang_kernel_gemm.py`: the two f32f32f32 cases use `T.tfloat32`,
which is NVIDIA-only, so narrow `requires_cuda_or_cdna` to
`requires_cuda`.
Also refresh three stale RDNA tests: gfx12 is now a supported generation
(`get_arch` accepts 11 and 12), so they can no longer assert that gfx1200
is rejected. Force an unsupported generation instead, which is the only
way to reach the guard now that `target_is_rdna` reports only gfx11/gfx12.
Verified on 8x MI300X (gfx942, ROCm 7.2 container, torch nightly rocm7.2):
the ROCm suite goes from 105 failures to 0 with no loss of passing tests.
|
||
|
Chaofan Lin
|
e01c498b64 | [JIT] Remove legacy DLPack execution backend (#2816) | ||
|
KellyFrogandchenkf
|
0bc1913daf |
[BugFix] Resolve partial scalar reduce barrier participation (#2777)
* [CUDA][Reduce] Resolve partial AllReduce participants
* refactor: inline ResolveAllReduceThreadRange into reduce.h and remove allreduce_barriers doc
- Inline ResolveAllReduceThreadRange body from reduce.cc into reduce.h,
add #include <tvm/ir/with_context.h> for With<ConstraintContext> usage.
- Delete src/backend/common/op/reduce.cc (function now in header).
- Remove reduce.cc from CMakeLists.txt TILE_LANG_SRCS.
- Delete docs/compiler_internals/allreduce_barriers.md.
- Remove toctree entry from docs/index.md.
* [CUDA][Reduce] Resolve batched AllReduce participants
* Revert "[CUDA][Reduce] Resolve batched AllReduce participants"
This reverts commit
|
||
|
Lei Wang
|
0d0e96861a |
[Env] Require JSON for default target config (#2491)
Require JSON for default target config |
||
|
Lei Wang
|
23d2f25673 |
[CUDA] Add target code attribute support (#2454)
* Add CUDA target code attr support * Add CUDA target normalization support Introduced a new function to normalize CUDA targets, ensuring that the target is correctly identified and transformed based on the detected architecture. Added a test to verify that the bare CUDA target uses the detected architecture accurately. |
||
|
Jayce Suanddingsg
|
e7da9025c4 |
[BugFix][CuTeDSL] Fix TileKernels scan, optional-shape, and e5m6 paths (#2369)
* [BugFix] Complete CuTeDSL scan lowering for TileKernels TileKernels exercises CuTeDSL scan and routing patterns that were not fully covered by the existing Python backend helpers. The MoE get_fused_mapping kernel computes expert base offsets with T.cumsum over a 256-entry shared buffer. CuTeDSL CumSum1D only performed a single warp-local prefix scan, while the CUDA scan template uses the first warp to scan all 32-element segments and carry the result between segments. As a result, expert prefixes reset at later warp-sized chunks, positions for experts such as 32 and 64 overlapped earlier expert ranges, and pos_to_expert validation failed on H100. Bring the CuTeDSL scan helpers in line with the CUDA InclusiveScanLine behavior: add reusable one-warp line scans with carry across 32-element segments, route CumSum1D and CumMax1D through them, and use the same line scan for CumSum2D/CumMax2D on both row-wise and column-wise axes. This removes the previous dim=0 H<=32 limitation and fixes multi-segment cummax as well as cumsum. The same TileKernels path also relies on CuTeDSL codegen accepting non-flat buffer accesses after TIR lowering. Add a shared index linearization helper for loads, stores, address_of/reinterpret, and atomic pointer generation, and lower tl.match_any_sync through the existing CuTeDSL warp intrinsic. Thread binding now falls back to the IterVar name when thread_tag has already been normalized away. Verified on H100 CUDA/CuTeDSL with CUDA_VISIBLE_DEVICES=0 and TILELANG_TARGET=cutedsl: - cmake --build build -j$(nproc) - ruff format --check tilelang/contrib/cutedsl/reduce.py - ruff check tilelang/contrib/cutedsl/reduce.py - git diff --check - pytest testing/python/target/test_tilelang_codegen_cutedsl_scan.py testing/python/language/test_tilelang_language_scan.py -q -s --tb=short: 13 passed - pytest testing/python/language/test_tilelang_language_warp_vote.py::test_match_any_sync -q -s --tb=short: 1 passed - pytest tests/moe/test_get_fused_mapping.py::test_get_fused_mapping[num_send_tokens=4001-num_topk=2-num_experts=72-num_ep_ranks=1-alignment=64] -q -s --tb=short: 1 passed Co-authored-by: dingsg <shengge.ding@enflame-tech.com> * [BugFix] Resolve CuTeDSL dynamic shapes from live tensors CuTeDSL records dynamic symbolic dimensions in the same ordering as the CUDA wrapper, but it previously kept only the first buffer shape or stride location for each symbol. That is not sufficient for kernels with optional tensor arguments. TileKernels reduce_fused declares num_tokens on topk_weights, token_topk_to_pos, and out; when with_weights is false the topk_weights argument is None, so the first recorded source is not a live tensor. The previous adapter fallback substituted 0 when that first source was None. On H100 with TILELANG_TARGET=cutedsl this passed num_tokens=0 to the generated launcher even though token_topk_to_pos had shape (4001, 2), producing grid=[0, 1, 1] and a CUDA error 1 launch failure in reduce_fused. Keep all shape/stride candidates for each dynamic symbol and resolve the runtime value from the first candidate backed by a real torch.Tensor. Reuse the same resolution helper for allocated output shapes and for dynamic arguments passed to the generated CuTeDSL module, and fail clearly if no live tensor source exists. Verified with a CuTeDSL adapter regression test for repeated dynamic shape symbols with a None first candidate, and with the TileKernels reduce_fused H100 failure case that previously launched with num_tokens=0. Co-authored-by: dingsg <shengge.ding@enflame-tech.com> * [BugFix] Permit dead optional CuTeDSL stride symbols TileKernels no-scale MoE expand paths can keep a dynamic stride symbol in the generated CuTeDSL host wrapper ABI even when the optional scale tensor is absent from the lowered device kernel variant. In that case the runtime argument is intentionally None, so the stride-only symbol has no live tensor from which the adapter can read torch stride metadata. The previous dynamic-symbol fix made missing live tensor sources strict to avoid unsafe shape fallbacks such as launching with num_tokens=0 when the first candidate tensor was optional. That strict behavior is still required for shape symbols because they control output allocation and launch dimensions. Limit the fallback to symbols whose candidates are stride-only. Such values are ABI placeholders for optional strided tensor parameters in pruned variants; returning 0 preserves the generated call signature without weakening dynamic shape resolution. Add adapter regressions covering both the optional stride fallback and the strict missing-shape error. Verified on H100 with TILELANG_TARGET=cutedsl: - ruff format/check on the touched files - pytest testing/python/jit/test_tilelang_jit_cutedsl_host_codegen.py::test_cutedsl_adapter_resolves_dynamic_symbol_from_live_tensor_candidate testing/python/jit/test_tilelang_jit_cutedsl_host_codegen.py::test_cutedsl_adapter_allows_optional_stride_symbol_without_live_tensor -q - pytest tests/moe/test_expand_to_fused.py::test_expand_to_fused[num_send_tokens=4001-num_topk=2-num_experts=9-num_ep_ranks=8-hidden=576] -vv -s --tb=short Co-authored-by: dingsg <shengge.ding@enflame-tech.com> * [BugFix] Fix CuTeDSL optional ABI and early-return codegen TileKernels full functional coverage with TILELANG_TARGET=cutedsl exposed additional CuTeDSL gaps in top2_sum_gate after the broader TileKernels fixes. First, the host adapter still appended dynamic shape symbols for optional tensors that had been compiled out by static flags. In top2_sum_gate, to_physical_map/logical_count can be None when logical expert mapping is disabled, but their shape-only symbol remains in the wrapper ABI. Keep output allocation strict, but allow dead host ABI dynamic shape arguments to use a zero placeholder when no live tensor source exists. Second, CuTeDSL cannot emit dynamic thread returns directly, so codegen rewrites if cond: thread_return(); rest into if not cond: rest. The existing recognizer only handled a bare thread_return then body, and missed the common if cond: stores; thread_return() pattern. Preserve the pre-return stores and guard the following statements so masked tokens cannot be overwritten by the fallthrough path. Third, vector binary codegen reused SSA aliases for operands containing BufferLoad nodes. That is not valid for mutable local tensors after an intervening store. The softmax top2_sum_gate path restored raw logits to scores_local, but the subsequent scores_local + bias expression reused an older softmax load and ranked by softmax(logits)+bias instead of logits+bias. Avoid SSA alias reuse for vector operands containing BufferLoad so the generated code reloads the current tensor value. Verified on H100 with the local TileLang build and TileKernels CuTeDSL backend: - cmake --build build -j 32 - ruff format/check on touched Python files - pytest testing/python/jit/test_tilelang_jit_cutedsl_host_codegen.py::{dynamic symbol candidate tests} -q: 3 passed - tests/moe/test_top2_sum_gate.py::test_top2_sum_gate[num_groups=0-num_topk_groups=0-num_routed_experts=72-num_shared_experts=1-num_topk=6] with TILELANG_TARGET=cutedsl TK_FULL_TEST=1: 1 passed Co-authored-by: dingsg <shengge.ding@enflame-tech.com> * [BugFix] Promote narrow CuTeDSL shifts before wide casts TileKernels e5m6 quantization kernels pack 8 truncated fp16 values into three uint32 words. The TileLang source writes expressions such as T.cast(half_u16[i] << 20, T.uint32), which the CUDA C++ backend evaluates with normal integer promotion before the final uint32 cast. CuTeDSL uses strongly typed cutlass integer wrappers instead of CUDA C integer promotion. The previous CuTeDSL lowering cast the shift result back to the narrow TIR result type first, producing code like cutlass.Uint32(cutlass.Uint16((half_u16[0] << cutlass.Uint16(20)))). For shifts by more than the uint16 width this drops all high bits before the packed uint32 store, corrupting e5m6 payload bytes. Cast-back then faithfully decoded the corrupted payload into nan/inf values, so TileKernels full correctness failed in cast_back_e5m6 and per_token_cast_to_e5m6 cases. When a narrow integer shift is immediately cast to a wider integer type, emit the CuTeDSL shift with the left operand promoted to the target integer type before shifting. This preserves the intended packed-bit construction while leaving narrow shift expressions that are not widened on the existing path. Add a CuTeDSL codegen regression test for uint16 << 20 widened to uint32 so future changes do not reintroduce a narrow in-place shift. Verified on H100 CUDA/CuTeDSL with TILELANG_TARGET=cutedsl and TILELANG_DISABLE_CACHE=1: - cmake --build build -j 32 - python -m pytest testing/python/target/test_tilelang_codegen_cutedsl_integer_shift.py -q - TileKernels focused e5m6 quant tests: tests/quant/test_per_token_cast_to_e5m6.py tests/quant/test_cast_back_e5m6.py -n 2: 168 passed, 112 skipped Co-authored-by: dingsg <shengge.ding@enflame-tech.com> --------- Co-authored-by: dingsg <shengge.ding@enflame-tech.com> |
||
|
Chaofan Lin
|
76571366e7 |
[Backend] Add target detector/normalizer registeration and merge ExecutionBackend into backend module (#2323)
* [Target] Add TileLang target registry * extend * fix gitignore * [Backend] Register target and execution backend resolution * [Backend] Remove JIT execution backend shim * [Backend] Move target helpers under backend * [Backend] Use registration order for target detection * [Backend] Use registration order for execution backend * [Backend] Split target helpers by backend * [Backend] Remove supported target registry * remove lazy detector * fix debug print |
||
|
Xuyang LiandClaude Opus 4.8
|
3953e4b5fc |
[BugFix] Reject T.alloc_barrier() on pre-Hopper targets with a clear error (#2345)
LowerSharedBarrier lowers T.alloc_barrier() to hardware mbarrier code (the cutlass Barrier type, tl::tl_shuffle_elect, tl::fence_barrier_init), all of which only exist on sm_90+. On pre-Hopper targets (e.g. sm_86) the frontend accepted the call silently and nvcc later failed with cryptic "identifier 'Barrier' is undefined" / "namespace 'tl' has no member ..." errors. Add a static compute-capability guard in CUDAPassPipelineBody, before LowerSharedBarrier: if the target is < sm_90 and the module uses a shared barrier, raise a clear ValueError naming the offending primitives and pointing to __syncthreads()/named barriers as the pre-Hopper alternative. Fixes the sm_86 alloc_barrier compile crash. Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> |
||
|
Zhang JasonandLeiWang1999
|
3d95d65000 |
[AMD][RDNA4] Fix gfx12 (RDNA 4 / Wave32) related CI issues (#2313)
* [AMD][RDNA4] Fix RDNA4 (gfx1201 / Wave32) CI failures
- reduce.h: Replace hardcoded wave64 assumptions with runtime
__builtin_amdgcn_wavefrontsize() in SharedReduceWarp and AllReduce
- scan.h: Add SEG=32 support in InclusiveScan/CumSum/CumMax for RDNA wave32
- finalize_reducer.cc: Use TargetGetWarpSize() instead of hardcoded 64
- carver/arch/__init__.py + rdna.py: Add gfx12 (RDNA4) support alongside gfx11
- wmma_macro_generator.py: Fix k_pack indexing in ldmatrix_a/b for k_pack > 1
These changes correspond to PR #2210 (Fix RDNA4 gfx1201/Wave32 CI Failures),
rebased cleanly onto upstream main (
|
||
|
Lei Wang
|
146a1d3bef |
[TileOP] Add scan operators (#2262)
Add TileLang scan operators |
||
|
|
a84c9b1433 |
[CuTeDSL] Integrate host codegen call sites (#2221)
[CuTeDSL] Integrate host codegen call sites
Drive CuTeDSL launcher generation from lowered host tvm_call_packed call sites so multi-kernel wrappers preserve host IR order and repeated launches can be represented.
Expose separate device, host, and generated module sources through the CuTeDSL adapter and persist them separately in the kernel/autotune caches.
Normalize bare cutedsl targets through TileLang CUDA arch detection so Hopper runners use sm_90a for CuTeDSL compilation instead of the generic sm_90 Target("cuda") result.
Add CuTeDSL host-codegen integration tests covering multi-call launcher ordering, source accessors, cache restore behavior, and bare target architecture normalization.
Co-authored-by: Jayce <jayce.su@enflame-tech.com>
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> |
||
|
Leonard
|
c2a9d9872d |
[Fix][AMD] Fix Roller autotuner for RDNA3 WMMA targets (#2208)
* Tune RDNA WMMA Roller defaults * Fix matmul Roller benchmark callback * lint fix |
||
|
Leonard
|
d863f1d89d |
Add RDNA gfx1151 ROCm target support (#2127)
* Add RDNA gfx1151 ROCm target support * Fix optional annotation handling for gfx1100 * Address CodeRabbit ROCm cleanup feedback * Scope MFMA tests to CDNA targets * Scope gfx950 copy async GEMM test * Reject non-divisible warp layouts * Skip CUDA/CDNA GEMM dtype tests on RDNA * Scope RDNA device model to gfx11 * Use RDNA arch helper in block scoring * Make RDNA tensor intrinsics generation-aware * Exclude gfx1010 from SDOT4 guard * Store RDNA tensor instructions as list * Reject unsupported RDNA carver targets |
||
|
Lei Wang
|
4929ad89cc |
[Refactor] Improve cp.async lowering and add async_copy op (#1887)
* Add CPAsync copy selection/lowering and vectorized ptx_cp_async folding
* [transform] Rename SimplifyCPAsyncSync to OptimizeCPAsyncSync
* ThreadSync: treat ptx_cp_async as async-copy
* Implement async_copy operator and enhance cp.async handling
- Introduced the async_copy operator for explicit asynchronous global-to-shared memory copies, ensuring no automatic waits are inserted.
- Updated the copy operator to normalize copy regions and handle both synchronous and asynchronous copy semantics.
- Enhanced pipeline planning to recognize explicit cp.async calls and bind commit and wait operations to the appropriate stages.
- Added tests to validate the behavior of async_copy and its integration with pipeline planning.
This commit improves the handling of asynchronous memory operations and clarifies the semantics of copy operations in the context of the new async_copy functionality.
* docs: clarify T.copy vs T.async_copy
* style: apply format.sh fixes
* async_copy: improve errors and defer cp.async check
* test: check T.copy vs T.async_copy codegen equivalence on SM80
* copy: allow predicated cp.async for OOB global->shared
* pipeline: schedule cp.async commits with cp.async groups
* lower: preserve manual pipeline statement arity
* pipeline: skip commit-order check when fused with cp.async
* pipeline: enforce cp.async ordering in PipelinePlanning
PipelinePlanning may float pure control intrinsics like ptx_commit_group earlier because they have no buffer regions, which can generate commit-before-cp_async.\n\nAdd an order-repair step that models cp.async group semantics (cp_async* -> commit, commit(group_i) -> cp_async(group_{i+1})) and wait->consumer constraints, then applies a stable topological sort to produce a valid software_pipeline_order.\n\nAlso adds a regression test to assert cp.async is ordered before commit when split into separate blocks.
* Enhance cp.async synchronization handling in pipeline planning
- Introduced logic to relax loop wait conditions and split epilogue waits for cp.async operations, improving overlap between consumer phases.
- Added a new method to identify and manage consumer phases based on async memory reads, ensuring correct synchronization.
- Updated tests to validate the new behavior, confirming that wait groups are correctly relaxed and inserted as needed in the optimized pipeline.
This commit refines the handling of cp.async synchronization, enhancing performance and correctness in the generated code.
* lint fix
* Refactor cp.async copy handling and improve vectorization logic
- Removed the unused IsCPAsyncElementBytes function to streamline the code.
- Enhanced the CPAsyncStoreRewriter class with a new method to track vectorized lanes during statement rewriting.
- Updated CheckCPAsyncCopy to remove the buffer_oob parameter, simplifying its signature.
- Added logging for vectorized loop iterations to aid in debugging and performance analysis.
- Improved error handling in LowerCPAsyncCopy to provide clearer warnings when cp.async rewrites fail.
These changes aim to optimize the cp.async copy process and improve code maintainability.
* lint fix
* Update cp.async copy handling to improve rewrite success tracking
- Set the successfully_rewritten_ flag to true when a copy operation is successfully rewritten.
- Adjust the flag to false in cases where the load is not global-like or the buffer is not shared.
- Removed logging for vectorized loop iterations to streamline the code.
These changes enhance the tracking of rewrite success in the cp.async copy process, contributing to better error handling and maintainability.
* fix
* test: cover commit-merge + wait relax
* fix: skip wait relaxation under conditional async copy
* test fix
* Implement async copy support in TileLang
- Introduced a new configuration option `tl.enable_async_copy` to control the lowering of global to shared memory copies to PTX `cp.async`.
- Updated the `CopyNode::LowerCPAsyncCopy` method to handle async copy semantics based on the new configuration.
- Enhanced the `InjectPTXAsyncCopy` pass to rewrite eligible global->shared copies into `cp.async`, including necessary commit and wait statements.
- Added tests to verify the correct behavior of async copy lowering and the impact of the configuration option.
These changes improve the flexibility and performance of memory copy operations in TileLang.
* lint fix
* Refactor async copy handling in TileLang
- Renamed `InjectPTXAsyncCopy` to `LowerPTXAsyncCopy` to better reflect its functionality.
- Updated the `OptimizeForTarget` function to use `LowerPTXAsyncCopy` instead of the deprecated `InjectPTXAsyncCopy`.
- Added a deprecated alias for `InjectPTXAsyncCopy` to maintain backward compatibility.
- Adjusted documentation to clarify the purpose of the new function name.
These changes enhance clarity and maintainability in the async copy transformation process.
* fix
* lint fix
* test fix
* Enhance PTX async copy integration in TileLang
- Added the `InjectPTXAsyncCopy` function to facilitate the injection of PTX `cp.async` operations into the lowering process.
- Updated the `LowerTileOp` function to conditionally apply `InjectPTXAsyncCopy` during parallel loop lowering when async copy is enabled.
- Removed deprecated calls to `LowerPTXAsyncCopy` in the `OptimizeForTarget` function to streamline the optimization process.
These changes improve the handling of async copy operations, ensuring better performance and compatibility with CUDA targets.
* lint fix
* fix
* fix
* fix
* set parallisim to 1 to test the ci
* fix
* Refactor async copy handling in PTX lowering
- Improved the `PrepareCopyIndexInfo` function to clarify the calculation of transfer bytes and ensure valid async copy operations.
- Updated synchronization state management for injected `cp.async` transfers, enhancing clarity and consistency in handling open copy regions and uncommitted transfers.
- Added helper functions for zero value checks and linear offset flattening to streamline index processing.
These changes enhance the robustness and readability of the async copy implementation.
* Refactor nested pipeline functions and update test cases
- Removed unused parameters `order` and `stage` from several functions, simplifying their signatures.
- Updated nested pipeline implementations to use `num_stages=2` for better clarity and performance.
- Adjusted test cases to reflect the changes in function signatures, ensuring they run correctly without the removed parameters.
These modifications enhance code readability and maintainability while ensuring the tests remain functional.
* Refactor async copy statement generation in PTX lowering
- Updated the `MakeCPAsyncStmt` function to `MakeCPAsyncStmtFromLoads`, improving the handling of base indices for vectorized accesses.
- Introduced `ExtractVectorBaseIndices` to streamline the extraction of base indices from vectorized loads and stores.
- Enhanced error handling for predicated async copy operations when base indices cannot be extracted.
These changes improve the clarity and robustness of async copy statement generation in the PTX lowering process.
* lint fix
* Enhance async copy operation handling in CodeGenTileLangCuTeDSL
- Added support for conditional async copy operations in the `VisitExpr_` method, allowing for an optional predicate argument.
- Improved error checking to ensure the correct number of arguments for `tl::ptx_cp_async`.
- Streamlined the generation of async copy statements for both standard and conditional cases.
These updates improve the flexibility and robustness of async copy operations in the code generation process.ix
* lint fix
* disable auto cp async without pipeline.
* lint fix
* Enhance cp.async synchronization handling in TileLang
- Introduced `ThreadSyncAfterWaitGroupInserter` to ensure correct insertion of storage synchronization after `ptx_wait_group` statements, addressing potential visibility issues in shared memory.
- Updated `WarpSpecialized` to utilize `tl.ptx_cp_async_barrier_noinc` for releasing producer-consumer dependencies, preventing incomplete memory observations.
- Added tests to verify the correct insertion of synchronization and the handling of cp.async barriers under warp specialization.
These changes improve the correctness and reliability of async copy operations in the TileLang transformation pipeline.
* fix
* fix
* test fix
* Refactor sparse_mla_fwd_pipelined.py and enhance phase.py synchronization
- Removed unused post-processing function and related code from `sparse_mla_fwd_pipelined.py` to streamline the implementation.
- Added debug print statements in `phase.py` to log the state of the module before and after thread synchronization, aiding in troubleshooting and verification of synchronization behavior.
* Enhance async copy handling and loop synchronization in TileLang
- Introduced new loop annotations for async copy operations, including `parallel_prefer_async` and `parallel_async_without_async_commit_wait`, to improve control over async behavior in parallel loops.
- Updated the `InjectPTXAsyncCopy` function to accommodate the new annotations, ensuring correct async copy injection based on user-defined preferences.
- Enhanced the pipeline planning logic to prioritize cp.async groups by their earliest consumers, improving scheduling efficiency and reducing wait times.
- Added tests to validate the new async copy behaviors and ensure correct synchronization in pipelined loops.
These changes enhance the flexibility and performance of async copy operations in the TileLang transformation pipeline.
* lint fix
* Add Regression Perf for mhc
* clean
* clean
* lint fix
* lint fix
|