Commit Graph
15 Commits
Author SHA1 Message Date
jpwang aa7df8670e [BugFix] Fix fp16/bf16 T.atomic_max/atomic_min silently corrupting fp32 values (#2780)
* [BugFix] Convert value to dst dtype in fp16/bf16 atomic max/min

* Make the linter happy
2026-07-27 20:28:20 +08:00
morlutoandLeiWang1999 ab1d2df419 [BugFix] Support BufferRegion destinations in atomic_addx2 return_prev (#2753)
* fix(atomic): support sliced x2 return destinations

* fix(atomic): derive vector return dtype structurally

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-07-24 14:45:21 +08:00
24a023c636 [BugFix][Transform] Never classify side-effecting binds as replayable (atomics re-executed at every use site since v0.1.11) (#2651)
* [Transform] Never classify side-effecting binds as replayable

IsReplayableScalarBind (introduced in #2245) decides whether a scalar
bind may be dissolved and its value re-evaluated at every use site
(IfStmtBinding, default-on) or replayed across pipeline stages
(PipelinePlanning, InjectSoftwarePipeline). The predicate checked the
dtype and read/write disjointness, but never the value's effect kind,
so a bind like

    pos = T.atomic_add(counter[0], 1, return_prev=True)

used at two sites (bound check + store index) was inlined into both use
sites, executing the atomic once per use site instead of exactly once
and silently corrupting compaction/histogram/top-k kernels since
v0.1.11. The atomic intrinsic is correctly registered kOpaque; its
pointer argument enters the expression through tvm_access_ptr, which
the access collector records as a read, so the write-disjointness check
passed and the bind was declared replayable.

Reject any bind whose value has SideEffect > kReadState, mirroring the
effect guard StmtSimplifier already applies to let inlining. Read-only
values remain replayable; read consistency is enforced by the existing
write-set disjointness check.

The shipped examples/deepseek_v32/topk_selector.py uses exactly this
pattern and is miscompiled on v0.1.11/v0.1.12; its test only printed
recall and could not fail, so it also gains a recall assertion.

Regression tests:
- transform: IfStmtBinding must keep a bind whose value is an atomic
  RMW returning the previous value.
- language (CUDA): compaction kernel using atomic_add(return_prev=True)
  must contain exactly one AtomicAddRet and produce a permutation with
  counter == N.

* [Transform] Validate scalar bind replay safety

Track access-pointer write masks during bind and pipeline analysis, and reject unsafe precomputed replay annotations during pipeline injection.

Add regressions for atomic target reads and side-effecting replay masks.

* Remove recall calculation from top-k selector

Removed recall calculation and assertion for top-k selection.

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
2026-07-22 17:43:01 +08:00
LyscoriaandSiriusNEO 88e007d4d7 [BugFix] Guard T.atomic_addx4 return type for sliced destinations (#2590)
[BugFix] Guard T.atomic_addx4 return type computation for sliced destinations

Co-authored-by: SiriusNEO <chaofan@deepseek.com>
2026-07-21 17:48:36 +08:00
Chennes 4f442a395e [BugFix] Reject invalid atomic load and store memory orders (#2666)
T.atomic_load and T.atomic_store previously accepted every order in the shared memory-order map, including operation-illegal combinations that trigger device-side assertions in libcu++.

Validate memory orders against operation-specific allow lists before lowering, report invalid values with a frontend ValueError, and update the documentation and regression coverage accordingly.

Fixes #2574
2026-07-16 14:22:20 +08:00
Xuyang LiandLeiWang1999 c5e5307621 [BugFix] Return previous value for scalar atomic_min/atomic_max (#2672)
* [BugFix] Return previous value for atomic min/max/add (scalar + vectorized)

`atomic_min`/`atomic_max` (scalar) and `atomic_addx2`/`atomic_addx4`
(vectorized) typed their call with a non-void return under
`return_prev=True` but always selected the void `*_elem_op`, so codegen
emitted an empty binding (`prev = ;`) and nvcc failed.

Select the `*_ret_elem_op` variant under `return_prev`, as `atomic_add`
already does:
- min/max: the ret ops, codegen handlers, and device templates
  (AtomicMaxRet/AtomicMinRet) already existed.
- addx2/addx4 (float32): register the ret ops (builtin.cc/.h) and add
  codegen handlers emitting AtomicAddx2Ret/AtomicAddx4Ret (already in
  the device headers) as expressions; the return is typed as the packed
  2-/4-lane vector so the store matches.
- addx2 (fp16/bf16): the ret store LHS is `uint1`; codegen wraps the
  half2/bf162 result with `tl::to_uint1`.
- addx4 (fp16/bf16): add AtomicAddx4Ret(half_t*)/(bfloat16_t*) device
  overloads that run two AtomicAddx2Ret halves and return them packed as
  uint2 (a fp16x4 is stored as uint2).

The 16-bit x4 path is per-pair atomic (two half2 atomics), not
whole-quad single-atomic -- hardware has no vector atomic wider than
`.f16x2` for 16-bit types. This matches the existing fp32-x4 scalar
fallback, and the docstring only promises "quad-width atomic addition".
Adds tests across scalar/x2/x4 and float32/fp16/bf16.

Fixes #2667

* minor fix for missing conditiions

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-07-16 13:48:57 +08:00
Tong WU aca581263a [Feature] Expose multiple CUDA intrinsics (#2473)
* [CUDA] Add T.__fns intrinsic for find-nth-set bit

Expose CUDA __fns in TileLang for direct lookup of the k-th active lane in a 32-bit bitmask, complementing the existing T.__ffs helper.

* Add builtin intrinsics for shared ldst and atomics
2026-07-02 17:29:54 +08:00
Jayce Suanddingsg 39ade83ff2 [BugFix][CUDA] Use PTX v4 atomics for fp16/bf16 atomic_addx4 (#2492)
Use PTX v4 atomics for 16-bit CUDA atomic_addx4

The initial fp16/bf16 atomic_addx4 fix avoided the generic float4 path by splitting each operation into two AtomicAddx2 calls. That fixed correctness, but it still used the simulated implementation on Hopper where PTX provides native vector atomics.

Add sm90+ half_t and bfloat16_t AtomicAddx4 overloads backed by the PTX vector forms atom.global.v4.f16.add.noftz and atom.global.v4.bf16.add.noftz, including the existing memory-order variants. Keep the x2 fallback for pre-sm90 targets because PTX vector f16/bf16 atomics require sm90 or newer.

Keep the CUDA language regression focused on fp16/bf16 atomic_addx4 codegen and offset coverage.

Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
2026-07-02 16:21:02 +08:00
Chennes 9f76d522ed [BugFix] Fix vectorized atomic_add dtype mismatch reinterpret (#2414)
Vectorized atomic_add into an fp16/bf16 destination with an fp32 source
bit-reinterpreted the source float2 as half2/__nv_bfloat162 instead of
numerically converting it, silently producing wrong results. This hit
both the auto-vectorized T.atomic_add path (AtomicAddx2 value form) and
the explicit T.atomic_addx2 API (pointer form).

Add float2 / float* converting overloads (round-to-nearest) to ToHalf2
and ToBfloat162 so the vectorized helper converts instead of
reinterpreting the source bytes. Add mixed-dtype regression tests for
both entry points.

Fixes #2382
2026-06-19 23:21:35 +08:00
ZihaoMu 7bf8de13ca [ROCm] Try to fix ROCm CI error (#2179)
* Fix ROCm tvm-ffi DLPack handling

* Fix CuTeDSL FP4 torch dtype map

* Fix CuTeDSL FP4 packed stores

* Mark DeepSeek V4 FP4 CuTeDSL xfail

* Keep ROCm DLPack PR focused

* Revert "Keep ROCm DLPack PR focused"

This reverts commit bd7ebea864.

* Split FP4 changes from ROCm DLPack PR

* Clarify ROCm FP8 tvm-ffi runtime args

* Gate ROCm FP8 tvm-ffi runtime args

* Keep ROCm DLPack PR focused

* Drop FP8 fallback diff from ROCm CI PR
2026-05-12 15:13:34 +08:00
Chaofan Lin c214f0ae0e [BugFix] Fix type mismatch when lowering to AtomicAddx2 template (#1898)
* [BugFix] Fix type mismatch when lowering to AtomicAddx2 template

* remove debug output

* add tests
2026-03-04 20:44:12 +08:00
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.
2026-02-11 22:51:59 +08:00
Chaofan Lin 4349b2c50b [Feature] Support passing PrimExpr value in tile-level atomic operation (#1796)
* [Feature] Support passing PrimExpr value in tile-level atomic operation

* fix after rebase

* address comments

* fix tvm ver

* fix
2026-02-06 12:41:54 +08:00
Lei Wang 5feb22535c [Refactor] Move AtomicAdd Vectorization to VectorizeLoop Pass (#1677)
* move to vectorize_loop

* fix

* lint fix

* Refactor atomic add functions to improve type handling and reduce redundancy

- Introduced utility functions ToFloat2 and ToFloat4 for type conversion, enhancing the flexibility of atomic add operations.
- Updated AtomicAddx2 and AtomicAddx4 functions to accept generic value types, streamlining the code and improving readability.
- Removed outdated function overloads and consolidated logic for atomic addition, ensuring consistency across implementations.
- Enhanced memory order handling in atomic operations for better performance and clarity.

* add related test

* Remove atomic add vectorization implementation

- Deleted the atomicadd_vectorize.cc and atomicadd_vectorize.h files, which contained the automatic vectorization pass for atomic add operations.
- Removed references to the deleted files in atomic_add.cc and loop_partition.cc to clean up the codebase and eliminate unused dependencies.
2026-01-16 15:16:16 +08:00
Lei Wang b27fb9286a [Feature] Atomic Reduction Operations and Vectorization Enhancement (#1676)
* [Feature] Implement atomic reduction operations and enhance atomic add functionality

- Added support for atomic max and min operations with corresponding classes and methods.
- Introduced vectorization capabilities for atomic add operations based on data type and target architecture.
- Refactored atomic add intrinsic calls to improve clarity and consistency.
- Enhanced layout inference for atomic operations and integrated new utility functions for target architecture checks.
- Updated tests to validate the new atomic operations and their vectorization behavior.

* Refactor atomic operation code for improved readability and consistency

- Cleaned up whitespace and formatting in atomic add and reduce implementations for better code clarity.
- Enhanced comments in CUDA code generation for atomic operations to improve understanding.
- Updated function signatures and layout definitions for consistency across atomic operation files.
- Ensured proper alignment of code for better maintainability and readability.

* lint fix

* Enhance loop vectorization and layout handling across multiple operations

- Updated function signatures in atomic add, atomic reduce, copy, and fill operations to include layout map parameters for improved layout handling.
- Refactored vectorization logic to utilize layout maps, ensuring better performance and compatibility with various buffer layouts.
- Enhanced the LowerParallelLoop function to accept layout maps, facilitating more efficient loop transformations.
- Added checks for buffer contiguity in vectorization processes to ensure correctness when using layout maps.
- Updated tests to validate the new layout handling and vectorization behavior.

* Add atomic operation checks in legalize_safe_memory_access

- Introduced a new method to identify atomic operations within the legalize_safe_memory_access transformation.
- Updated the VisitStmt_ function to handle both CallExtern and atomic operations, ensuring recursive condition collection for these cases.
- Enhanced comments for clarity on the handling of atomic operations in the context of memory access legality.

* Add atomic operations for HIP code generation

- Implemented support for various atomic operations including atomic add, atomic load, atomic store, atomic max, and atomic min in the HIP code generation.
- Enhanced the handling of atomic operations to include optional memory order parameters.
- Improved code readability with added comments explaining the purpose of each atomic operation.
- Ensured consistency with existing atomic operation implementations in the codebase.

* Refactor AtomicOpBaseNode and related classes for improved clarity

- Changed the GetElemOpStatic method to a virtual GetElemOp method in AtomicOpBaseNode, enhancing polymorphism.
- Updated AtomicAddNode, AtomicMaxNode, and AtomicMinNode to override the new GetElemOp method.
- Removed unnecessary template parameters from AtomicOpBaseNode, simplifying the class structure.
- Cleaned up includes in atomic_reduce.cc to remove unused dependencies, improving code organization.

* Enhance argument validation for atomic operations

- Added checks in AtomicAdd, AtomicMax, and AtomicMin constructors to ensure at least two arguments are provided, improving error handling and user feedback.
- Removed the unused LowerTMA method declaration from the AtomicAdd class, streamlining the codebase.

* lint fix

* fix

* test fix

* performance fix

* remove useless comments
2026-01-16 00:03:01 +08:00