Commit Graph
359 Commits
Author SHA1 Message Date
Lei WangandZhiwen Mo ccfc127d91 [Feature] Add CUDA intrinsic for isfinite operation (#1925)
* [Feature] Add CUDA intrinsic for isfinite operation

- Introduced DispatchCUDAIsFinite function to handle the isfinite intrinsic for CUDA, enabling support for checking finite values in floating-point computations.
- Registered the new intrinsic with TVM under the operation name "tir.isfinite", enhancing the CUDA intrinsic capabilities.

Updated submodule reference for TVM to the latest commit.

* merge cluster mod

---------

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2026-03-12 17:21:01 +08:00
Lei WangandZhiwen Mo bdbb32c99b Support ptr-table grouped GEMM kernels (#1923)
Add ptr-table grouped GEMM support

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2026-03-12 12:51:09 +08:00
xyyy1420 ff1b8cc56c [Feature] Add T.gemm support for CPU target (#1904)
* [Feature] Add T.gemm support for CPU target

- Add GemmInst::kCPU enum and TargetIsCPU() helper for CPU target detection
- Add CPU branch in getGemmInst() for both GemmNode and GemmPyNode
- Add CPU early return in GemmNode::Lower() and InferLayout() (gemm_v1 no-op)
- Add GemmScalar: triple-loop scalar gemm implementation (gemm_v2 path)
- Add CPU dispatch in GemmPy._get_implementation_class()
- Add pytest tests for codegen and correctness (NN/TN/NT/TT, multiple sizes)

* rename kCPU to kScalar for consistency with existing GemmInst naming
2026-03-10 18:02:40 +08:00
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
2026-03-06 17:35:16 +08:00
xyyy1420 26ac2fe45e [BugFix] Add vector type definitions to common.h for CPU codegen (#1901)
This is the follow-up fix mentioned in #1899 for the remaining float4
compilation issue.

The CPU C codegen emits vector types (e.g. float4, int32_t4) when
VectorizeLoop vectorizes T.copy operations, but common.h had no
definitions for these types, causing g++ compilation failures.

Add a generic vec_type<T,N> template and type aliases for all
combinations that PrintType() can emit (float/double/half/intN_t,
lanes 2/4/8/16).

Update test_matmul_with_copy_cython to verify end-to-end correctness
now that both the wrapper and vector type issues are resolved.
2026-03-06 01:19:13 +08:00
xyyy1420 00267e43aa [BugFix] add target context and avoid redundant re-lowering in TLCPUSourceWrapper (#1899)
fix(cpu): skip redundant re-lowering in TLCPUSourceWrapper when mods already exist

TLCPUSourceWrapper.parse_source_information() unconditionally called
get_annotated_mod() to re-lower the IR, even when device_mod and
host_mod were already passed in from tilelang.lower(). This caused a
crash when the kernel uses T.copy (which generates T.Parallel loops),
because VectorizePlanner requires Target::Current() during lowering.

Fix: check if device_mod/host_mod already exist before re-lowering,
matching the pattern used by TLCUDASourceWrapper. Also add self.target
to the context manager for the fallback path.

Added test_matmul_with_copy_cython to reproduce the issue.
2026-03-05 13:38:05 +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 WangandZhiwen Mo df5ade9992 Refactor CUDA version checks for compute 9.0 (#1893)
* Refactor CUDA version checks in test files to use equality comparison for compute version 9.0

- Updated multiple test files to change `requires_cuda_compute_version_ge(9, 0)` to `requires_cuda_compute_version_eq(9, 0)`.
- Removed deprecated test functions in `test_example_blocksparse_attention.py`.
- Simplified configuration generation in `test_tilelang_autotune.py` by removing the `with_roller` parameter.
- Ensured consistency across various examples and tests regarding CUDA version requirements.

* revert unexpected changes

---------

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2026-03-03 18:28:11 +08:00
HaleandTong WU eb2755750b [CUDA] Support tcgen5mma gemm ts (#1866)
* feat: add mma_ts init

* feat: mma_ts valid

* fix: tcgen05.st & .ld mismatch on 16-bit pack/unpack, cause function mistake in gemm

* fix: diable unpack for bf16's tcgen05.st, enforce each tmem cell pack with 2 val, to be consistent with mma_ts's behaviour

* test: add test example in testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py

* docs: add docstrings to gemm_base,gemm_tcgen05 and test files

* style: apply clang-format and ruff format

* refactor: merge ld/st meta, remove example, rename getTcgen05Meta to getTcgen05MetaLd_*

* Restore TODO for tcgen05.cp support in copy.cc

Added a TODO comment for future support of tcgen05.cp in conjunction with LowerTmemCopy.

* fix: swap Ramp/PrimExpr isinstance order in access_ptr_from

---------

Co-authored-by: Tong WU <109033598+Rachmanino@users.noreply.github.com>
2026-02-28 16:33:13 +08:00
Tong WU 2e854b068e [Feature] Support cluster launch, query, synchronization and barrier operations (#1874)
* support cluster launch for tvm_ffi and cython backend

* lint

* add test

* lint

* update tvm

* Add cluster synchronization operations and related functions

- Introduced new built-in operations for cluster synchronization: `cluster_arrive_relaxed`, `cluster_arrive`, `cluster_wait`, `cluster_sync`, and `block_rank_in_cluster`.
- Updated the CUDA code generator to handle these new operations.
- Added corresponding Python bindings and documentation for the new cluster functions.
- Included necessary header files for cluster operations in the CUDA code generation process.

* add test for intrinsics

* lint

* remove legacy code

* Add cluster barrier support and related functionality

- Introduced `ptx_arrive_cluster_barrier` built-in for arriving at cluster barriers.
- Added `alloc_cluster_barrier` function for allocating cluster barrier buffers.
- Updated CUDA code generation to handle the new cluster barrier operations.
- Enhanced existing functions to support shared and cluster barrier scopes.
- Added tests for cluster barrier functionality in the TileLang testing suite.

* lint

* fix typo

* add storage sync for cluster
2026-02-28 14:17:47 +08:00
Chaofan Lin 7fcd62b0ba [Enhancement] Clarify the semantic rule of copy operator and add shape mismatched tests (#1883)
* [Enhancement] Clarify the semantic rule of copy operator and add shape mismatched tests

* remove debug msg

* fix after rebase

* fix

* fix
2026-02-27 19:49:38 +08:00
Taoyu Zhu 0fa78e235f [AMD] Fix gfx950 ci and add 16x16x32_bf16/fp16 instructions support (#1878)
* add

* add

* add

* zz

* zz

* tmp not right

* right

* right

* right

* right

* right

* clear

* pre-commit lint

* pre-commit lint

* pre-commit lint

* pre-commit lint

* pre-commit lint

* fix commit

* rerun ci?

* add tvm ffi target support for cdna4

* add tvm ffi target support for cdna4

* add tvm ffi target support for cdna4

* add tvm ffi target support for cdna4

* tmp 1 process to run rocm cis

* restore

* fix
2026-02-25 19:46:46 +08:00
Lei Wang 808b0cefcc [Feature] Introduce higher-dimensional gemm layout support (#1798)
* [Refactor] Update swizzle layout functions to accept Buffer objects directly

* Refactor swizzle layout functions to take Buffer parameters instead of stride and continuous values, improving usability and consistency.
* Introduce helper functions to extract stride and continuous dimensions from Buffer objects.
* Update related functions and documentation to reflect the new parameter types.
* Enhance layout inference logic to merge swizzle layouts based on Buffer objects.
* Add new reshape functionality to the Layout class for better shape management.

* lint fix

* use layout.expand instead of reshape to lift a swizzled layout

* lint fix
2026-02-25 15:50:58 +08:00
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
2026-02-25 15:19:36 +08:00
Lei Wang 9bd62f5b5f [Layout] Add Layout.expand to lift a layout into higher dimensions and improve repeat errors (#1876)
* layout: add Layout.expand and improve repeat errors

* revert unexpected changes
2026-02-25 12:54:17 +08:00
Lei Wang a22fa61dfe layout: add Layout.repeat for tiling atom layouts (#1875)
layout: add Layout.repeat via FFI
2026-02-24 22:45:10 +08:00
Lei Wang 7ab2b2be92 [BugFix] ArgBinder: relax shared-shape binding for unused nullable buffers (#1870)
* [BugFix] tvm_ffi: handle unused nullable buffers with shared symbolic shapes

* [BugFix] ArgBinder: relax shared-shape binding for unused nullable buffers
2026-02-24 11:43:55 +08:00
Lei Wang 9f259544d4 [Refactor] Refactor Pass InjectFenceProxy (#1863)
* tir: add T.cdiv alias for T.ceildiv

* docs: Update InjectFenceProxy documentation and enhance code comments

- Clarified the description of `tl.InjectFenceProxy` to specify the transition between generic and async proxy operations.
- Improved explanations of the pass's functionality, including state tracking and the handling of TMA store synchronization.
- Added details on the new `ProxyStateSet` class and its role in managing proxy states.
- Updated usage instructions for proxy hints to include new options for custom operations.
- Enhanced test coverage for handling unknown external calls and proxy hint overrides.

* fix

* enhance

* fix

* remove tl.proxy_hint

* refactor tma_store_arrive and tma_store_wait

* refactor

* fix

* InjectFenceProxy: hoist fence out of pure-async loops

* InjectFenceProxy: hoist fence for if/while pure-async regions

* Remove LowerTileOp transform test
2026-02-23 16:43:17 +08:00
Gabriel Wuandyuxic 391cf5b6b9 [Enhancement] GEMM V2 on SM90/SM100 CuTeDSL backend (#1855)
* fix(cutedsl): add rsqrt/rsqrtf support in codegen

- Add rsqrt/rsqrtf to fastmath function mapping
- Implement ieee_frsqrt using tl.rsqrt
- Export rsqrt from contrib/cutedsl/math.py

Fixes: norm/test_rms_norm.py

* fix(cutedsl): add atomic_add support and fix shared.dyn access

- Add atomic_add_elem_op and atomic_add_ret_elem_op handling
- Exclude shared.dyn from direct indexing to avoid out-of-bounds access
  (shared.dyn shape is set to (1,) but accessed with arbitrary indices)

Fixes: gemv/test_example_gemv.py

* fix(cutedsl): add builtin::ptx_ldmatrix support

Implement builtin::ptx_ldmatrix by reusing tl.ptx_ldmatrix_x* functions.

* feat(cutedsl): add ptx_mma support with all MMA types

Add complete PTX MMA intrinsic support for CuTeDSL backend:

- Implement ptx_mma codegen in codegen_cutedsl.cc
- Add GetVarPtr_ helper for local buffer pointer extraction
- Create ptx_mma.py with inline PTX assembly wrappers for:
  - FP16: m16n8k16 -> f16/f32
  - BF16: m16n8k16 -> f32
  - INT8/UINT8: m16n8k32 -> i32
  - INT4/UINT4: m16n8k32 -> i32
  - TF32: m16n8k4, m16n8k8 -> f32
  - FP64: m8n8k4 -> f64
  - FP8 e4m3/e5m2: m16n8k32 -> f16/f32
- Fix BroadcastNode to emit typed values for make_filled_tensor
- Update ptx_ldmatrix to use GetVarPtr_ for local buffers

* feat(cutedsl): comprehensive atomic operations support

- Add atomic.py module with full atomic operations:
  - AtomicAdd: float32 (nvvm.atomicrmw), float16 (inline PTX with i16 bitcast),
    int32/int64 (nvvm.atomicrmw)
  - AtomicAddx2/x4: vectorized float32 via PTX atom.global.add.v2/v4.f32
  - AtomicMax/Min: int32/int64 (nvvm.atomicrmw), float32 (CAS loop via inline PTX)
  - AtomicLoad/Store: all memory orderings via inline PTX (relaxed/acquire/release/seq_cst)
- Add codegen support for atomic_load/store/max/min/addx2/addx4 in codegen_cutedsl.cc
- Move AtomicAdd from __init__.py to atomic.py, extend with float16/int64 support

* feat(cutedsl): add fabsf/fabs support in codegen and math module

* feat(cutedsl): add wait_wgmma support via nvvm.wgmma_wait_group_sync_aligned

* feat(cutedsl): support i32 min/max

* fix(cutedsl): handle FP8 narrow precision alignment in codegen

CuTeDSL requires all FP8 vector operations to be 32-bit aligned, but
TIR generates float8x2 (16-bit) vectors. Additionally, MLIR vector
extractelement on FP8 types fails with unrealized_conversion_cast.

Fix CastNode, BufferLoadNode, and BufferStoreNode to:
- Pad unaligned FP8 casts to aligned width (e.g. f32x4 -> fp8x4)
- Return rmem tensor names (not .load()) so downstream uses rmem
  element access instead of MLIR vector extractelement
- Use padded aligned stores for local rmem, scalar element stores
  for global/shared memory (to handle non-aligned offsets)

* fix(cutedsl): use pointer arithmetic for shared.barrier buffer access

_Pointer type from tl.alloc_smem() does not support subscript [i],
but supports pointer arithmetic (ptr + i). This fixes warp_specialize
tests that use mbarrier.

- GetBufferRef_: use (vid + index) for shared.barrier instead of vid[index]
- BufferLoadNode: skip .load() suffix for shared.barrier (returns pointer, not value)

Fixes 3/4 warp_specialize tests.

* fix(cutedsl): handle single coordinate in tma_load

When tma_load receives a single Int32 coordinate (for 1D tensors) instead
of a tuple, wrap it in a tuple before iterating. This fixes the error:
'Int32' object is not iterable

Fixes flash_decoding/test_example_flash_decoding.py

* feat(cutedsl): add decode_i4u/s_to_f16 and warp primitives

- Add quantize.py with decode_i4u_to_f16 and decode_i4s_to_f16
  using inline PTX (lop3.b32 + sub.f16x2)
- Add warp.py with __activemask, __shfl_down_sync, __shfl_sync
  using CuTeDSL's shuffle_sync functions
- Update codegen to map decode_* and __* functions to tl.* namespace

Fixes dequant_gemv_fp16xint4 test.

* feat(cutedsl): add decode_fp4_to_bf16_twiddling

Implement FP4 to BF16 conversion using bit twiddling technique.
- Use prmt.b32 for endianness handling
- Use mul.bf16x2 for scaling
- Pack output with correct high/low BF16 layout

Fixes test_example_dequant_gemm_bf16_mxfp4_hopper.

* fix(cutedsl): support reinterpret for arbitrary expressions and fix shift type promotion

- reinterpret: handle non-BufferLoad args with tl.bitcast()
- shift_right/shift_left: wrap result in dtype() to preserve type
- CuTeDSL Int8>>4 returns Int32 (Python promotion), CUDA keeps int8_t
- Add bitcast() function using llvm.bitcast MLIR op

* feat: disable v2->v1 fallback

* fix(cutedsl): TensorSSA for store (Select/scalar) and TF32 MMA dispatch

- ArithValue: emit tl.where() for Select and wrap scalar store RHS in
  make_filled_tensor().load() so CuTeDSL Tensor.store() gets TensorSSA.
- TF32: accept fp32/f32/float32 for m16n8k4 and m16n8k8 in ptx_mma dispatch
  (e.g. deepseek_mhc).

* feat(cutedsl): implement CumSum1D and CumSum2D for cumulative sum operations

- Add CumSum1D and CumSum2D classes in reduce.py based on CUDA templates
- Add __shfl_up_sync wrapper in warp.py using CuTeDSL's shuffle_sync_up
- Implement warp-level prefix sum (forward/reverse) using shuffle operations
- Use @cute.jit decorator and nested if for DSL-compatible control flow
- Fixes gdn/test_example_cumsum_compilation test (7/8 gdn tests now pass)

* feat(cutedsl): implement tma_reduce for TMA atomic add reduction

* fix(cutedsl): handle if_then_else as tl.where and unify Select/where arm dtypes

- Override builtin::if_then_else in CallNode to emit tl.where() instead of
  Python ternary (which yields ArithValue, not TensorSSA)
- Use make_rmem_tensor + element assignment instead of make_filled_tensor
  (which creates unsupported vector<1xT> types for int)
- Unify arm dtypes in both SelectNode and if_then_else: when true/false arms
  differ (e.g. BFloat16 vs Float32), cast both to common type via .to()
- Fix BufferStore: skip make_filled_tensor wrapping for Select/if_then_else
  values (already TensorSSA); extract [0] for scalar element assignment
- Fix GetBufferRef_: extract Ramp base for make_tensor_at_offset offset
  parameter to avoid emitting tuple offsets

Fixes: cast/group_per_split_token, topk, deepseek_mhc, flash varlen

* fix(cutedsl): convert non-integer offsets to Int64 in make_tensor_at_offset

Complex arithmetic on cutlass.Int64 values (involving tl.min, //, %) can
produce Float64 ArithValue that CuTeDSL Pointer.__add__ does not accept.
Detect non-integer offsets via is_integer() and wrap with cutlass.Int64().

Fixes blocksparse_attention varlen_indice and varlen_mask (6/6 passed).

* fix(cutedsl): fix narrow unaligned load/store for sub-32-bit types

Three issues in the narrow unaligned (e.g. uint8x2, fp8x2) code paths:

1. Load from shared/global: using aligned_lanes (4) elements caused OOB
   near buffer boundaries and MisalignedAddress for non-aligned offsets.
   Fix: load exactly value_lanes (2) elements with no cute.assume.

2. Load from rmem: cute.assume(offset, divby=aligned_lanes) silently
   truncates offsets not divisible by aligned_lanes (e.g. offset=6 with
   divby=4 → 4). Fix: drop cute.assume for rmem (no hardware constraint).

3. Store to rmem: padded (aligned_lanes,) .store() writes beyond
   value_lanes elements, causing overlapping writes and OOB at tensor
   boundaries. Also triggers CuTeDSL's 32-bit alignment check for FP8.
   Fix: use element-by-element assignment (vid[i]=val).

Fixes dequantize_gemm w4a8 (cudaErrorMisalignedAddress → correct results).

* fix(cutedsl): handle signless int types in tl.min/tl.max fallback

* fix(cutedsl): handle _Pointer src_values in AtomicAddx2/x4

* fix(cutedsl): remove double PrintIndent in cp_async_gs_conditional codegen

* fix(example): support SM100 in attention_sink bwd configs

Change get_bwd_configs() to use >= 90 instead of == 90 so that SM100
(Blackwell) uses the same config as SM90 (Hopper): block_M=128,
block_N=32, num_stages=2, threads=256.

* feat: implement WGMMA V2 (ss + rs) for CuTeDSL backend

Implement explicit WGMMA descriptor management ops in the CuTeDSL backend,
replacing LOG(FATAL) stubs with working codegen for all 8 WGMMA-related ops.

New Python files:
- typing.py: WGMMA type map (string dtype -> nvvm.WGMMATypes)
- gemm_V2.py: GmmaDescriptor class, initialize/increase descriptor,
  warpgroup sync ops, wgmma_ss (via nvvm.wgmma_mma_async),
  wgmma_rs (via PTX inline asm)

Key findings and fixes:
- nvvm.mma_smem_desc does not correctly pack the layout_type/swizzle
  field in the WGMMA descriptor; replaced with manual bit packing
- wgmma_rs implemented using llvm.inline_asm with dynamically generated
  PTX and constraint strings based on M/N/K Constexpr parameters

C++ codegen changes (codegen_cutedsl.cc):
- warpgroup_arrive/commit_batch/wait/fence_operand
- ptx_wgmma_ss (15 args) and ptx_wgmma_rs (14 args)
- initialize_wgmma_descriptor (arg reorder: TIR->Python)
- increase_descriptor_offset

All GEMM tests (4/4) and flash attention tests (14/14) pass.

* fix(cutedsl): replace break with guard variable pattern for CuTeDSL compatibility

* feat(cutedsl): implement tcgen05 (SM100/Blackwell) MMA support

Add CuTeDSL codegen and Python runtime support for tcgen05 MMA operations:

- New gemm_tcgen05.py: Tcgen05SmemDescriptor, descriptor initialization,
  MMA variants (ss/ws_ss/ts), mbarrier arrive, TMEM alloc/dealloc/load
- All MMA and commit instructions guarded by elect_one_sync via inline PTX
  (@q predicate on individual instructions, not on block scopes)
- TMEM load uses recursive power-of-2 splitting with direct MLIR emission
- C++ codegen: replace 6 LOG(FATAL) stubs with proper code generation for
  tcgen05 MMA/descriptor/TMEM ops, handle descriptor scope allocation
- Fix shared.barrier pointer access (skip .iterator for alloc_smem results)
- Fix mbarrier_wait: use blocking PTX loop instead of single try_wait call

* fix(cutedsl): correct tcgen05_ld 64b/128b/256b PTX variants and cleanup

- Fix tcgen05_ld_32dp64bNx/128bNx/256bNx: use correct PTX instructions
  (16x64b/16x128b/16x256b) with two-half decomposition instead of
  hardcoded 32x32b
- Generalize _emit_tmem_ld_segment to accept ptx_type and regs_per_x
- Remove dead expression in ptx_mma.py (c_ptr + a_offset)
- Remove unused GemmCuTeDSL and is_cutedsl_target imports

* fix(cutedsl): review cleanup — warp shuffle, mbar timeout, atomic addr space, dead code

- Fix warp.py: correct mask_and_clamp formulas to match CUDA semantics
  for __shfl_down_sync, __shfl_up_sync, __shfl_sync (width param was ignored)
- Fix mbar.py: use timeout_ns parameter instead of hardcoded 0x989680
- Fix atomic.py: use generic address space (ld.f32, atom.cas.b32) instead of
  .global for AtomicMax/AtomicMin CAS loops, supporting shared memory pointers
- Add ICHECK for non-unit Ramp stride in codegen_cutedsl.cc GetBufferRef_
- Remove unreachable duplicate elif gemm_inst.is_tcgen5mma() branch
- Remove unused _to_ir_value helper from ptx_mma.py

* fix(cutedsl): skip BulkLoad1D for CuTeDSL + fix uint broadcast type mismatch

Two fixes for dequant/groupedgemm_bf16_mxfp4_hopper on CuTeDSL:

1. Skip BulkLoad1D/BulkStore1D for CuTeDSL target in copy.cc:
   cp_async_bulk_shared_cluster_global (raw 1D TMA) combined with WGMMA
   in the same kernel triggers a ptxas ICE in the NVPTX backend.
   Falls through to descriptor-based BulkLoad/BulkStore instead.

2. Use signed int type for BroadcastNode fill values in codegen_cutedsl.cc:
   CuTeDSL/MLIR normalizes unsigned integer tensor loads to signed types
   (Uint8 pointer -> i8 elements). BroadcastNode fill values now use
   matching signed types to avoid Int8/Uint8 mismatch in tl.where().

* fix(cutedsl): support fp16/bf16 in AtomicAddx2 via atom.add.noftz.v2.f16/bf16

* ci(cutedsl): add CuTeDSL examples to CI pipeline

Add a new GitHub Actions step that runs all examples with
TILELANG_TARGET=cutedsl. Known failures (sparse ops) are marked as
xfail and unsupported tests (convolution, stream-K, flash_decoding FA3)
are auto-skipped via pytest_collection_modifyitems in conftest.py.
The step uses continue-on-error: true initially until stabilized.

* chore: fix pre-commit

* fix(adapter): use name-based fallback for tir.Var dynamic symbolic lookup

tir.Var uses object identity for hash/eq, so when param_shapes (from
the original PrimFunc) and dynamic_symbolic_map (from the transformed
PrimFunc in ir_module) contain different tir.Var objects with the same
name, direct dict lookup fails with KeyError.

Add a secondary name-based index (_dynamic_symbolic_name_map) and a
_lookup_dynamic_symbolic() helper that falls back to name matching.
This fixes KeyError: batch under pytest-xdist parallel execution.

Affected adapters: cutedsl, nvrtc. (cython and tvm_ffi already had
a str()-based workaround.)

* chore: remove dead GemmCuTeDSL class and update stale V1 comments

- Delete `gemm_cutedsl.py` (no references anywhere in codebase)
- Remove outdated comment in `gemm_op.py`
- Update `run_local_ci_test.sh` comment: CuTeDSL no longer requires V1
- Keep `TILELANG_USE_GEMM_V1` env var as a debugging escape hatch

* refactor(cutedsl): reorganize contrib/cutedsl module structure

- Rename gemm_V1.py/gemm_V2.py to snake_case (gemm_v1.py/gemm_v2.py)
- Merge typing.py (19 lines) and mbar.py (53 lines) into utils.py and cpasync.py
- Extract utility functions from __init__.py into new utils.py module
- Slim __init__.py to pure re-exports
- Add __all__ to all submodules for explicit public API documentation

* fix(ci): use .empty() instead of .size() >= 1 in pipeline_planning.cc

* fix: coderabbit comments

* fix(cutedsl): correct WGMMA result register count for f16/bf16 output

wgmma_ss/wgmma_rs used M*N//128 for result registers, which is only
correct for 32-bit output types. For f16/bf16, each i32 register holds
2 packed elements, halving the count. Derive from cutlass type .width
instead of hardcoding.

Also update barrier test assertion to match codegen's pointer-arithmetic
syntax: (barriers + i) instead of barriers[i].

* refactor(cutedsl): templatize ptx_mma.py with factory function

Replace 13 near-identical @dsl_user_op PTX MMA functions (~840 lines)
with a single _make_ptx_mma factory and a config table (~120 lines).
Also correctly interpolates a_layout/b_layout into the PTX mnemonic
instead of hardcoding row.col.

* fix(cutedsl): fix NaN CAS loop, wgmma_rs int accum, and stale pipeline state

- atomic.py: Use integer-domain comparison (setp.ne.b32) in AtomicMax/Min
  CAS loops to prevent infinite spin when values are NaN
- gemm_v2.py: Use "=r"/i32 constraints for wgmma_rs when C_dtype is s32
  (int8 GEMM) instead of hardcoded "=f"/f32
- reduce.py: Add assert H <= 32 guard for CumSum2D dim=0 column-wise path
- pipeline_planning.cc: Clear pending tcgen05 state unconditionally after
  arrive to prevent stale entries from leaking to future arrives

* fix(cutedsl): fix ptx_mma brace escaping and wgmma_rs DSL compatibility

- ptx_mma.py: Double-brace register lists in PTX template so
  .format(a_layout=, b_layout=) doesn't interpret {$0, ...} as
  format fields (KeyError regression from templatization)
- gemm_v2.py: Replace if/else branches inside DSL-traced loops with
  ternary expressions to avoid CuTeDSL "operands type change inside
  dynamic if" error
- ci.yml: Remove continue-on-error from cutedsl CI, reduce maxfail
  to 3, increase parallelism to 4

* feat(cutedsl): implement ptx_mma_sp (sparse MMA) codegen and runtime

* fix(cutedsl): guard saturate param, move CuTeDSL CI step after core tests

- ptx_mma/ptx_mma_sp: raise NotImplementedError if saturate=True
- warpgroup_fence_operand: add comment explaining no-op rationale
- ci.yml: move CuTeDSL examples after CUDA/ROCm/Metal test steps
  so a CuTeDSL failure no longer blocks core tests

* feat(cutedsl): implement 16 previously unsupported ops

Add support for thread index queries (get_lane_idx, get_warp_idx,
get_warp_idx_sync, get_warp_group_idx), exp10 math function,
warp-level reductions (sum, max, min, bitand, bitor), and
IEEE-754 compliant arithmetic with explicit rounding modes
(add, sub, mul, fma, rcp, sqrt, div).

Reduces unsupported op count from 26 to 10. The remaining ops
are either high-complexity (tma_load_im2col, tl_gemm_sp),
require wrapper changes (sync_grid), or are legacy/unused.

* fix(cutedsl): wrap non-int32 integer literals with CuTeDSL type constructors

CuTeDSL's as_numeric() converts bare Python ints to Int32, causing
"Type mismatch" errors when storing to narrower tensors (e.g. Int16).
Override IntImmNode to emit cutlass.Int16(0) instead of bare 0 for
non-int32 literals. Also add a BufferStore narrowing guard as safety net.

Fixes compress_kernel Int32→Int16 error in gemm_sp/custom_compress.

* feat(cutedsl): implement sync_grid with cooperative kernel launch

Add grid-level synchronization via a software barrier using
llvm.mlir.global counter + inline PTX atomic spin-wait. The CuTeDSL
wrapper detects sync_grid usage and automatically switches to
cuLaunchCooperativeKernel to guarantee all blocks are resident.

- grid_sync.py: module-global counter + PTX barrier protocol
- wrapper.py: CPP_COOPERATIVE_KERNEL_LAUNCH_TEMPLATE
- codegen_cutedsl.cc: emit tl.sync_grid() instead of LOG(FATAL)

---------

Co-authored-by: yuxic <yuxic@nvidia.com>
2026-02-23 01:22:47 +08:00
Yichen Yan fbbda0f22b [Codegen] Metal codegen on Linux (#1857) 2026-02-19 11:19:17 +08:00
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>
2026-02-16 00:34:10 +08:00
Lei Wang 4956b5835f ThreadSync: avoid barriers between atomic ops (#1852) 2026-02-15 21:01:37 +08:00
48a8f4a6ce [Feature] Hierarchical reduction and warp reduction intrinsics support (#1762)
* [Feature] Add new hierarchical reduce with less workspace size and new ptx

* Lint

* Add intrinsic functions and more types

* Fix sm100a redux instruction

* Fix data type

* [Test] Add tests for reduce operations

* [Refactor] Unify AllReduce barrier mechanism via policy template

Replace the duplicated run()/run_hopper() methods in AllReduce with a
single run() that accepts a Barrier policy template parameter
(SyncThreadsBarrier or NamedBarrier<N>). Extract the shared inter-warp
reduction logic into a warp_inter_reduce helper, and split the dispatch
into private hierarchical_reduce/butterfly_reduce methods.

Update codegen in reduce.cc and finalize_reducer.cc to emit
NamedBarrier<all_threads> for SM >= 90 targets instead of the old
all_threads + run_hopper pattern.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Add handling for non-float types in warp_reduce function

This update introduces a new case in the warp_reduce function to cast non-float types to float using static_cast. This enhancement improves type flexibility and ensures compatibility with a broader range of input types during reduction operations.

* [Refactor] Remove hierarchical reduce because of performance regression

* [BugFix] Buffer for reduction

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-14 01:55:17 +08:00
Lei WangandZhiwen Mo 6d10d1602c [CUDA][Feature] Add packed FP32x2 math intrinsics and auto vectorized support (#1839)
* [Feature] Add packed FP32x2 math intrinsics and support for CUDA/HIP

This commit introduces new packed FP32x2 math operations: fadd2, fmul2, and fma2, which leverage PTX instructions on supported architectures. The changes include:

- Definitions and implementations of fadd2, fmul2, and fma2 in both CUDA and HIP codegen files.
- New Python API functions for these operations, ensuring they validate input types and handle fallbacks for unsupported architectures.
- Documentation updates to reflect the new intrinsics in the math language module.

These enhancements improve performance for vectorized floating-point operations in TileLang.

* revert hip changes

* annotate cuda only

---------

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2026-02-12 12:50:30 +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
Lei Wang e666d2d3cc [Refactor] Treat local.var as local buffers when deciding vectorization for stable actions (#1835)
lint fix
2026-02-11 16:47:53 +08:00
Lei WangandCopilot f8b4d9e677 [Bugfix] Fix threadIdx variable lookup by thread_tag instead of position in ThreadSync (#1829)
* Refactor OptimizeForTarget function by removing debug print statements and updating module state visualization

* Add test for sync hoisting in non-uniform if within loop using shared memory

This commit introduces a new test case to verify the correct behavior of sync hoisting when a non-uniform if statement is present inside a loop that utilizes shared memory. The test ensures that the synchronization occurs before the if statement, confirming the expected transformation in the module's intermediate representation.

* fix

* Refactor thread variable handling in thread_storage_sync.cc

This commit removes the optional thread variable lookup by tag and replaces it with direct indexing based on the expected position of threadIdx variables. This change addresses a bug where the last three elements of env_threads_ were incorrectly assumed to always correspond to threadIdx.x/y/z, improving the accuracy of thread variable access in the synchronization logic.

* Update testing/python/transform/test_tilelang_transform_thread_sync.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-02-10 01:03:01 +08:00
Lei Wang c99879ad12 [Bugfix] Fix ast builder error for value -= 1 (#1825)
* Enhance augmented assignment methods in BaseBuilder and Builder classes

- Updated `aug_assign` method to include an optional `name` parameter for better handling of variable names during augmented assignments.
- Improved handling of immutable variables in the `Builder` class to treat augmented assignments as re-bindings, providing warnings for potential issues.
- Added comments to clarify the changes and ensure user expectations are met regarding variable assignments.

* lint fix
2026-02-09 17:30:22 +08:00
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.
2026-02-09 13:26:17 +08:00
Lei Wang e9d056954c [BugFix] LoopUnswitching: gate non-trivial else behind PassConfig (#1816)
* [BugFix] LoopUnswitching: gate non-trivial else behind PassConfig

* lint fix
2026-02-09 11:20:20 +08:00
Kuris c47d6df379 [BugFix] Missing Recursive Loop Var Checking in Loop Unswitching (#1801)
* [Fix] Update loop unswitching logic to handle multiple let bindings and add corresponding test case

* remove debug print
2026-02-06 17:45:28 +08:00
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
2026-02-06 12:50:22 +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 b01cb93586 [Refactor] Reorganize ParallelOp code structure and move ProveFragmentContains to layout utils (#1779)
* [Feature] Implement ProveFragmentContains Function for Fragment Thread Validation

- Added the ProveFragmentContains function to check if the threads accessing elements of a smaller fragment are a subset of those accessing a larger fragment.
- This function ensures valid access when transitioning from a smaller to a larger fragment layout.
- Updated layout.cc and utils.cc to incorporate this new functionality, enhancing the layout validation process.
- Removed the previous implementation of ProveFragmentContains from parallel.cc to streamline the codebase.

* fix

* Refactor ParallelOpNode Layout Handling

- Removed the initial DeReplicate attempt from InferLayout to streamline layout inference.
- Added DeReplicate logic to ComputeLoopLayoutFromBuffer to reduce replication when validating layout candidates.
- Updated test cases to disable caching and ensure proper functionality of loop layout kernels.

* fix

* Refactor Test Cases for Loop Layout

- Removed caching disablement and print statements from the loop layout identity test for cleaner output.
- Updated the main execution block to directly call the testing framework, enhancing test execution flow.
2026-02-06 12:28:58 +08:00
Lei Wang 191d879ddd [Bugfix] Fix thread storage sync conflict detection for loop carry write-after-read (#1781)
* Fix thread storage synchronization logic in `thread_storage_sync.cc` to correctly identify conflicts between read and write operations based on loop carry conditions.

* lint fix

* Refactor `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to use shared memory for `sorted_token_ids` instead of local memory, improving thread synchronization. Adjust default argument values for M, N, and K in the main function for better testing scenarios.

* Add UniformExprChecker to enforce thread synchronization rules

Introduce the UniformExprChecker class to determine if expressions are uniform across threads, crucial for safe synchronization in conditional statements. Update the TileLangThreadSyncPlanner to hoist synchronization points out of non-uniform if-statements to prevent potential deadlocks. Enhance tests to validate sync hoisting behavior for various non-uniform conditions involving thread indices and shared memory access.

* lint fix

* Enhance `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` with cache disabling and kernel source printing for debugging. Update thread synchronization logic in `thread_storage_sync.cc` to check for runtime-dependent conditions, preventing potential deadlocks by hoisting sync points as necessary.

* Update submodule `tvm` to latest commit and remove deprecated `example_gqa_decode_varlen_logits_paged.py` file. Refactor `example_gqa_decode_varlen_logits.py` to enhance performance and maintainability by removing unused imports and optimizing shared memory usage. Adjust test cases to reflect the removal of the paged example.

* fix

* Enhance thread synchronization logic in `thread_storage_sync.cc` by adding a configurable warp size parameter to `RuntimeDependentConditionChecker` and `TileLangThreadSyncPlanner`. This allows for better adaptability to different target architectures. Update the logic to ensure thread extent is a constant and improve handling of runtime-dependent conditions.

* lint fix

* Refactor thread extent validation in `thread_storage_sync.cc` to use pointer checks instead of optional values. This change improves clarity and ensures that the thread extent is correctly validated as a constant.

* Adjust loop variable constraints in `thread_storage_sync.cc` for loop-carry analysis by modifying the extent calculation. This change ensures valid iteration comparisons by reducing the extent by one, allowing for accurate analysis of loop iterations.

* lint fix
2026-02-05 03:47:26 +08:00
ShaobinChenandSiriusNEO 215f436b05 [BugFix] Fix reduce_sum with clear=False not accumulating correctly (#1778)
* Fix type annotations for T.reshape and T.view

* Fix issue #1666: reduce_sum with clear=False not accumulating correctly

* address comments and add testcases

* add more tests

---------

Co-authored-by: SiriusNEO <chaofan@deepseek.com>
2026-02-04 11:32:40 +08:00
Chaofan Lin 378a8f2573 [Refactor] Unify the usage of cast-related operators (#1757)
* [Refactor] Unify the usage of cast-related operators

* reinterpret auto detect
2026-02-02 18:47:18 +08:00
Tong WU 727d108315 [Feature] Support tcgen5mma lowering for .kind::i8 (#1764)
* support tcgen5mma lower for int8/uint8->int32

* add test

* add template for u8 and fix a bug in codegen

* lint

* fix ai suggestions

* add testing to maint
2026-02-02 17:12:27 +08:00
Lei Wang 7bf43447d5 [Feature] Add PassConfig for Controlling Let Statement Inlining in Simplify Pass (#1769)
[Feature] Enhance Simplify Pass with Let Inline and Common Subexpression Elimination

- Added a new configuration option `tilelang_enable_simplify_let_inline` to enable inlining of let statements during simplification.
- Updated the Simplify pass to conditionally inline let statements based on the new configuration.
- Integrated common subexpression elimination into the device code generation process to optimize performance.
- Enhanced test coverage for the Simplify pass, including various scenarios to validate simplification behavior.
2026-02-02 15:48:14 +08:00
Lei Wang c4748da9d8 [Bugfix] Add predicate to loads inside predicated stores in LowerLDGSTG pass (#1767)
* [Feature] Add support for predicated loads and stores in LDGSTG lowering

- Introduced a new configuration option to disable loop unswitching.
- Enhanced the lowering logic to ensure that loads within predicated stores are also treated as predicated, preventing out-of-bounds memory access.
- Added a test case to verify the correct behavior of predicated loads and stores in the transformation pipeline.
- Included debug print statements to visualize the module state before and after the lowering process.

* lint fix
2026-02-02 13:59:58 +08:00
Lei Wang 5f60c09060 [Feature] Enhance Loop Unswitching with Let Binding and Condition Handling (#1766)
- Updated the loop unswitching pass to support hoisting Let-bound variables alongside their conditions, improving optimization for loops with variable dependencies.
- Modified the CallCheckerExcludingIf class to exclude matching If nodes when checking for CallNode presence, ensuring safety during unswitching.
- Enhanced IfBranchReplacer to replace all instances of If nodes with matching conditions, streamlining the transformation process.
- Added comprehensive tests to validate the new behavior, including scenarios with multiple Let-bound variables and identical conditions across If statements.
2026-02-02 13:28:33 +08:00
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>
2026-02-01 00:32:51 +08:00
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
2026-01-30 17:30:25 +08:00
Lei Wang 07d75bfc73 [Enhancement] Add dynamic symbolic constraints support for Profiler benchmarking (#1753)
* [Enhancement] Add dynamic symbolic constraints support in Profiler inputs and benchmarking

- Updated `_get_inputs` method to accept dynamic symbolic constraints for kernel parameters.
- Introduced `_substitute_dynamic_symbols` method to handle substitution of dynamic shapes with concrete values.
- Modified `benchmark` method to utilize dynamic symbolic constraints when provided, enhancing flexibility for benchmarking kernels with dynamic shapes.

* [BugFix] Correct error handling in dynamic symbolic constraints test

- Updated the test for missing constraints in the profiler to raise a ValueError instead of using an assertion. This change ensures that the error is properly raised and caught, improving the clarity of the test's intent and behavior.

* Remove copyright and license comments

Removed copyright and license comments from the test file.
2026-01-30 16:31:24 +08:00
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
2026-01-30 11:53:31 +08:00
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
2026-01-29 23:59:24 +08:00
Tong WU 3fbf562cd2 [Feature] Support message-only debug print (#1755)
* support msg-only debug print

* lint
2026-01-29 15:37:36 +08:00
Hubert Lu a55a82302b [AMD] Fix ROCm FP8 dtype selection and MFMA support on gfx942/gfx950 (#1743)
* Fix gfx950 FP8 E4M3 selection

* Support FP8 E5M2 through BF8 MFMA intrinsics

* Use determine_fp8_type and determine_torch_fp8_type instead
2026-01-29 12:07:30 +08:00
fe5bf3569d [Refactor] Enhance T.alloc_barrier with new features and deprecate legacy mbarrier related intrinsics (#1733)
* enhance `T.alloc_barrier` and let `T.wait/arrive` only accept buffers

* lint

* Refactor mbarrier handling and update related documentation

- Updated the `T.alloc_barrier` function to accept an array of arrive counts instead of individual thread counts.
- Modified the `get_mbarrier` function documentation to clarify its purpose and usage.
- Enhanced the `CreateMBarrierBuffer` function to create mbarrier buffers with shared storage scope.
- Updated tests to reflect changes in mbarrier allocation and usage.
- Adjusted comments in various files to improve clarity regarding mbarrier functionality.

* lint

* fix test

* tcgen5 check input

* lint

* fix name conflict

* add name check

* apply ai suggestions

* fix shared mem disalignment and missing fence for mbarrier init

* fix cutedsl codegen

* bugfix

* add test for cutedsl codegen

* Update testing/python/jit/test_tilelang_jit_cutedsl.py

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

* Remove tilelang.disable_cache() call

Remove cache disabling from the test file.

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
2026-01-28 15:20:13 +08:00
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
2026-01-28 10:56:25 +08:00