Commit Graph
340 Commits
Author SHA1 Message Date
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
Lei Wang 350f987a5b [Refactor] Disable Predicated LDG PTX Lowering by default (#1739)
[Refactor] Update predicated global load/store lowering configuration

- Changed the configuration option from `kDisableLowerLDGSTGPredicated` to `kEnableLowerLDGSTGPredicated` to enable predicated lowering of global load/store operations.
- Updated relevant comments and documentation to reflect the new configuration behavior.
- Adjusted tests to utilize the new configuration option for enabling predicated lowering, ensuring consistency across the codebase.
2026-01-27 20:55:02 +08:00
chaospointer f5525ead06 [BugFix] Fix FP4 related vectorized cast (#1741) 2026-01-27 20:54:40 +08:00
Kuris 7ec5542db9 [BugFix] Remove unnecessary binding in loop variable analysis and add test for issue 1728 (#1735)
* [Fix] Remove unnecessary binding in loop variable analysis and add test for issue 1728

* [Test] Add test
2026-01-27 14:43:28 +08:00
Chaofan Lin a17230e442 [Feature] Support E8M0 related type conversion and vectorized cast (#1731)
* [Feature] Support E8M0 related vectorized cast

* fix

* address comments
2026-01-27 12:21:55 +08:00
Lei Wang 8bd0293740 [Enhancement] Add explicit global memory load/store intrinsics (ldg/stg 32/64/128) (#1717)
* [Enhancement] Add global memory load/store intrinsics for various vector widths (#1710)

* Introduced new intrinsics: ldg32, ldg64, ldg128, ldg256 for loading data from global memory.
* Added corresponding store intrinsics: stg32, stg64, stg128, stg256 for storing data to global memory.
* Updated built-in functions and templates to support these new operations.
* Enhanced documentation for usage examples and type checks in the language module.
* Updated CUDA code generation to utilize new intrinsics for improved performance.

* lint fix

* test fix

* Refactor buffer type annotations in global memory load/store intrinsics

Updated the type annotations for the `ldg` and `stg` functions to use `tir.Buffer` instead of the generic `Buffer` type. This change enhances type specificity and aligns with the TVM type system for better clarity and correctness in memory operations.

* lint fix

* [Enhancement] Introduce LDG/STG lowering options and transform pass

* Added configuration options to enable/disable lowering of non-predicated and predicated global memory loads/stores to ldg/stg intrinsics.
* Implemented the LowerLDGSTG transformation pass to optimize global memory access for better performance.
* Updated tests to require a minimum CUDA compute version for relevant functionality.
* Enhanced documentation for new configuration options and transformation behavior.

* lint fix

* Refactor LDG/STG intrinsic selection logic for improved readability

* Replaced switch-case statements with if-else conditions for determining the appropriate LDG operation based on byte width.
* Enhanced error handling for unsupported byte widths with a clear fatal log message.

* lint fix

* fix

* testfix

* Enhance LDGSTG pass to skip processing for CuTeDSL backend

* Added a check to skip the transformation if the target has the "cutedsl" key.
* Updated tests to require a minimum CUDA compute version of 10.0 for LDG and STG code generation tests.

* Refactor global memory load/store functions to include conditional variants

* Updated the CUDA code generation to support conditional versions of LDG and STG functions, enhancing flexibility in memory operations.
* Modified existing functions to include new conditional variants (e.g., `ldg32_conditional`, `stg32_conditional`) for better control over memory access based on predicates.
* Adjusted comments to reflect the changes in function signatures and usage.

* Update comments and input handling for LDG intrinsics to include address and optional predicate

* Modified comments for `ldg32`, `ldg64`, `ldg128`, and `ldg256` to clarify the function signatures, indicating the use of 'address' and an optional 'predicate'.
* Changed the number of inputs for these LDG functions from 2 to -1 to accommodate the new signature format.

* Add correctness verification to LDG code generation tests

* Introduced assertions to verify the correctness of generated CUDA source for `ldg32`, `ldg64`, `ldg128`, and `ldg256` functions.
* Added checks to ensure the output matches the expected reference values within specified tolerances.
* Updated test cases to enhance validation of LDG intrinsic functionality.

* test fix

* Refactor LDG and STG intrinsics to use unified naming convention

* Updated all instances of `ldg` and `stg` functions to their new names: `load_global` and `store_global`, respectively, across the codebase.
* Modified related test cases to reflect the new function names, ensuring consistency in the testing framework.
* Enhanced comments and documentation to clarify the changes in function signatures and usage.

* Update reference values in LDG256 predicated code generation test to improve accuracy

* lint fix

* Rename test function for clarity and add new test for predicated load/store global intrinsics. Update assertions to reflect new naming convention in generated kernel source.

* lint fix

* subtype cast fix
2026-01-26 12:16:42 +08:00
Lei WangandZhiwen Mo 56dbc9f8f0 [Refactor] Enhance variable substitution in device function generation (#1723)
* [Refactor] Enhance variable substitution in device function generation

- Updated `wrapBodyWithHostSideAssumes` to accept a mapping of host-side variable names to device-side variables, ensuring consistent variable references during SSA conversion.
- Introduced a substitution mechanism for old parameter variables to new ones in the device function, preventing variable sharing between host and device.
- Adjusted buffer declarations to utilize new variable mappings, enhancing the robustness of the device function generation process.

* lint fix

* test fix

* lint fix

---------

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2026-01-23 17:24:37 +08:00
Kuris 5358f5a272 [EagerJIT] Lazy Evaluation of Kernel Body in Eager JIT (#1690) (#1694)
* [Refactor] Support lazy evaluate in eager jit

* add test script

* fix lint error

* canonicalize signature and error of `T.empty`

* Fix comment from genai

* [EagerJIT] Add comment for eager jit stage
2026-01-23 14:33:12 +08:00
Lei WangandZhiwen Mo 4ab369d6ad [Refactor] re-implement vector subtype and its access method (#1722)
[Refactor] Update fp4 data structures and CUDA code generation

- Refactored `fp4_e2_2_t` structure to include methods for accessing and setting low and high bits.
- Updated `PrintVecElemLoad` and `PrintVecElemStore` methods in `codegen_cuda.cc` to reflect changes in member access for `fp4_e2_4_t` and `fp4_e2_2_t`.
- Adjusted test for `tilelang_copy_fp4` to ensure correct input dimensions and output validation based on data type consistency.

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2026-01-23 14:17:54 +08:00
470a19237e [Enhancement] Remove RewriteUnsafeSelect Pass (#1705)
* Remove RewriteUnsafeSelect Pass

* Add test for

* Add test for T.Select

* Add test for T.Select

* Add test for T.Select

* Update test_tilelang_language_select.py

---------

Co-authored-by: wangxiangwen <xiangwen.wang@deepseek.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
2026-01-21 18:12:08 +08:00
Lei Wang 9a9255f5f7 Enhance loop vectorization logic for CallNode handling (#1696)
* Enhance loop vectorization logic and buffer handling

- Refactored vectorization logic to utilize configuration settings for enabling/disabling vectorization.
- Introduced new checks for buffer access invariance in call nodes to improve vectorization accuracy.
- Updated GCD calculations for vector sizes to include call node constraints, enhancing overall performance.
- Added verbose logging for better insights during vectorization planning.
- Introduced tests for vectorization of infinity and bitwise logical operations to validate new functionality.

* lint fix

* fix

* test fix
2026-01-20 03:52:10 +08:00
Lei Wang c9cec62557 [Bugfix][Refactor] Always disable light storage reuse (#1691)
* fix

* lint fix

* Remove print statement in test_tilelang_issue_1678.py

Remove print statement for kernel source in test.
2026-01-19 18:06:35 +08:00
Yichen YanandLeiWang1999 10baa27ce7 [tvm-ffi] Enable tvm-ffi for metal backend (#1289)
* enable tvm-ffi for metal

* copy metal codegen from tvm

* Fix

* upd

* Compile host runtime with torch's mps

* upd

* add current device/hsot code

* Pass basic test

* fmt

* add comment

* no need to change device codegen

* cleanup

* add torch to cache key for metal

* remove hardcoded compile options

* Add `MarkHostMetalContextMutator`

* update codegen

* update and cleanup

* tvm_ffi for metal by default

* upd

* cleanup

* fix

* fix

* test

* fix

* fix test

* fix

* test

* fix

* upd

* upd

* fix

* upd

* add PrintLine helper in c codegen

* use tvm's hook

* Use latest tvm

* test install wheels again

* limit ccache to 200 MB

* refactor

* lint fix

* resolve comments

* fix

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-01-19 13:46:03 +08:00
Lei Wang 2ff1dd50dc [Feature] Introduce DecoupleTypeCast pass for mixed-precision vectorization (#1644)
* [Feature] Introduce DecoupleTypeCast pass to optimize vectorization constraints

- Added a new pass to decouple type cast vectorization constraints, allowing for optimal vectorization in mixed-precision scenarios.
- Implemented transformations for local-to-memory and memory-to-local cases, inserting intermediate local buffers as needed.
- Updated pass configuration to enable verbose output for vectorization planning.
- Enhanced loop vectorization logic to handle complex cases with multiple statements.
- Added tests to validate the new functionality and ensure correctness in various scenarios.

* Refactor verbose output in loop vectorization for improved readability

- Updated the output formatting in the VectorizePlanner to use newline characters instead of std::endl for better performance and clarity.
- Ensured consistent output style across various verbose logging statements related to buffer sizes and vectorization strategies.

* Refactor verbose output in loop vectorization for consistency

- Simplified the output formatting in the verbose logging of vector size to enhance readability by removing unnecessary line breaks.

* use eager style.

* lint fix

* Enhance loop vectorization by introducing GetBufferStrides function

- Added GetBufferStrides function to compute buffer strides when not defined, improving clarity and reusability.
- Updated VectorizeFindGlobalAccess to utilize the new function for stride calculations, simplifying the logic.
- Refactored related code to remove redundant stride calculations and improve overall readability.
- Introduced a test case to ensure no transformation occurs when different dtypes are only in if-then-else conditions, validating the new behavior.

* Refactor loop vectorization logic and enhance target support checks

- Updated the logical condition in TargetIsSm100 to use the correct logical AND operator.
- Introduced TargetSupportVectorize256 function to encapsulate vectorization support checks for specific architectures.
- Refactored VectorizeFindGlobalAccess to VectorizeFindMemoryAccess for improved clarity and added support for shared memory access checks.
- Updated vectorization logic to utilize the new support checks, enhancing overall performance and maintainability.
- Cleaned up redundant logging statements and improved code formatting for better readability.

* Update TargetSupportVectorize256 to simplify architecture check

- Modified the architecture check in TargetSupportVectorize256 to return true for architectures 100 and above, removing the upper limit of 110. This change enhances compatibility with newer architectures while maintaining existing functionality.

* lint fix

* Refactor tilelang JIT configurations and optimize GEMM calls in fused MoE example

- Updated JIT pass configurations to use enum keys for better clarity.
- Removed unused parameters `k_pack` and `coalesced_width` from the `moe_forward_tilelang_routed` function.
- Simplified GEMM calls by removing the `k_pack` argument, enhancing readability.
- Adjusted swizzle usage to remove the `enable` flag for consistency.
- Added a call to `tilelang.disable_cache()` in the main function to improve performance.

* lint fix

* Update TVM submodule and enhance vectorization logging

- Updated the TVM submodule to the latest commit for improved functionality.
- Added logging in the VectorizePlanner to provide insights during the planning phase of for loops.
- Ensured GCD calculation for vector size aligns with loop extent for better vectorization accuracy.
- Enabled verbose logging for the vectorize planner in the test configuration to aid debugging.

* lint fix

* performance fix

* fix

* fix

* fix

* fix

* bench fix
2026-01-18 01:23: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
silentCoder-devandLeiWang1999 1d71bed2db [Feature] Reimplement Threadsync with ConstrVisitor (#1631)
* reimplement ThreadSync pass

* fix bugs

* remove old storage_access

* remove deadcode for debugging

* format

* add WhileOp & expose MakeGuard in constr_visitor

* bugfix

* bugfix

* bugfix

* bugfix

* trigger ci

* add try for checking

* bugfix

* update tvm with better analyzer

* update tvm

* Log for debugging in thread_storage_sync

* format

* consider let node

* typo

* remove debugging log

* consider while op in thread_sync & fix bug sbout letnode

* add testing for thread_sync issues

* update testing for thread_sync

* format

* consider T.Ramp expr

* remove useless header

* format

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-01-15 18:59:50 +08:00
Lei Wang 2d8d3676ed [CUDA] Enhance Broadcast Codegen for Symbolic Value (#1669)
* Enhance CUDA code generation for BroadcastNode by implementing compile-time constant folding and runtime broadcasting for various lane configurations. Improved handling for 4-bit and 8-bit integer types, ensuring correct replication and type casting in output expressions. This update increases performance and correctness in CUDA kernel generation.

* add test

* lint fix

* fix
2026-01-14 16:02:31 +08:00
Tong WU 732971a8de [Feature] Support cp.reduce.async.bulk.tensor (#1667)
* support cp.reduce.async.bulk.tensor and add test

* Refactor flash attention example by removing unnecessary layout annotations

* support swizzle layout for tma reduce

* auto swizzle for non-1d tma atomic add

* upd example and test

* lint

* typo

* add constraint for test

* Refactor CUDA data type mapping by moving the to_CUtensorMapDataType function to utils.cc and utils.h, while removing redundant definitions from atomic_add.cc and copy.cc.

* lint
2026-01-14 14:39:46 +08:00
Kuris c3c8831db6 [EagerJIT] Add Support for Parameter Only Kernel Compilation (#1664)
* [Fix] Refactor type hint extraction logic in DSLMutator for better clarity and handling of annotations

* [Refactor] Remove redundant tensor creation in loop layout tests and update kernel compilation parameters
2026-01-13 16:22:49 +08:00