Commit Graph
81 Commits
Author SHA1 Message Date
Lei Wang 4956b5835f ThreadSync: avoid barriers between atomic ops (#1852) 2026-02-15 21:01:37 +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 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 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
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
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
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
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
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
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
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 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 a56212de10 [Bugfix] Introduce a flag to avoid unnecessary broadcast hoist and enable for let stmt (#1638)
* [Enhancement] Implement hoisting of broadcast values in Let statements and update tests for validation

* lint fix

* [Enhancement] Update hoisting logic for broadcast values in Let statements and add a new test case for nested BufferStore broadcasts

* lint fix
2026-01-08 17:56:27 +08:00
silentCoder-devandkurisu6912 b5be1a16ab [Feature] Add hoist_broadcast_values pass (#1606)
* add hoist_broadcast_values pass

* only hoist intImm or floatImm

* refactor

* add test for hoist_broadcast_values

* add a test for issue 1601

* add transform test for hoist_broadcast_value

---------

Co-authored-by: kurisu6912 <227995639+kurisu6912@users.noreply.github.com>
2026-01-07 15:16:25 +08:00
Kuris 5d806298de [Feat] Add strong checker to detect data racing in T.Parallel (#1615)
* add verifier for parallel loop data racing

* fix lint error

* fix typo

* fix lint error

* fix bug in metal test
2026-01-06 17:31:24 +08:00
Lei Wang 6fd0bd31fe [Refactor] Use access_ptr instead of buffer and offsets for cp async params (#1590)
* refactor params of ptx_cp_aync

* Enhancement: Introduce ptx_cp_async intrinsic and improve async copy handling

- Added the `ptx_cp_async` intrinsic for asynchronous memory copy operations in TileLang, supporting both predicated and non-predicated versions.
- Updated the code generation for `ptx_cp_async` to handle offsets correctly and added inline functions for simplified usage.
- Enhanced the `inject_ptx_async_copy` transformation to support predicated vectorized copies.
- Updated example scripts to disable caching and print kernel sources for better debugging and performance insights.

* Remove tilelang.disable_cache() calls from example scripts for cleaner execution

* Refactor comments and formatting in CUDA code generation and PTX async copy transformation

- Updated comments for clarity in `codegen_cuda.cc` regarding `ptx_cp_async` arguments.
- Improved formatting in `inject_ptx_async_copy.cc` for better readability of access pointer creation and logging.
- Added a blank line in `op.py` for consistency in the `ptx_cp_async` function definition.

* Remove unused ptx_cp_async handling from CUDA code generation

- Eliminated the implementation for `ptx_cp_async` and `tl::ptx_cp_async` from `codegen_cuda.cc`, streamlining the code by removing unnecessary checks and assembly generation.
- This change focuses on cleaning up the codebase and improving maintainability by removing deprecated functionality.
2026-01-03 19:00:55 +08:00
Lei Wang 27db71f5bb [Host] Provide post process to customize host code and enhance nullable check (#1562)
* [Enhancement] Add C host post-processing callback and improve nullable buffer handling

- Introduced a new callback for C host code generation that allows for post-processing of generated code before it is wrapped into a CSourceModule.
- Enhanced the ArgBinder to log shape variable sources and ensure safe binding of symbolic shape variables, preventing potential segmentation faults when dealing with nullable buffers.
- Added a regression test to verify that a single buffer with a symbolic shape variable must be non-null, ensuring robustness against null inputs.

* [Refactor] Clean up formatting and improve readability in C host code generation and ArgBinder

- Adjusted indentation in the C host code generation to enhance clarity.
- Reformatted comments and code structure in ArgBinder for better readability, ensuring consistent style throughout the file.
- Minor whitespace adjustments to maintain code consistency.
2025-12-29 16:36:40 +08:00
Lei Wang f6db2014a3 [ArgBinder] Enhance shape variable handling and assertions (#1467)
* feat(arg_binder): enhance shape variable handling and assertions

- Implemented special handling for comparing if_then_else expressions to simplify conditions involving NULL checks.
- Added methods to set shared shape variables and finalize deferred bindings, generating cascading if_then_else expressions and runtime assertions for non-NULL buffers.
- Updated the binding logic to defer shape variable bindings for shared variables, ensuring proper handling across multiple nullable buffers.

* refactor(arg_binder): clean up shape variable handling and remove unused code

- Removed deprecated methods for setting shared shape variables and finalizing deferred bindings, streamlining the argument binding process.
- Simplified the logic for handling shape values in the `BindDLTensor` function, ensuring immediate binding for normal shape variables.
- Enhanced clarity by eliminating unnecessary comments and code related to cascading if_then_else expressions for shared variables.

* refactor(arg_binder): enhance DLTensor binding with improved shape handling

- Replaced the single `BindDLTensor` method with `BindDLTensors` to support multiple buffers, improving flexibility in handling DLTensor bindings.
- Introduced a two-pass approach for shape variable handling, allowing for better management of symbolic dimensions and null checks.
- Updated the logic to assert non-null conditions at runtime and utilize cascaded if_then_else expressions for shape retrieval, enhancing robustness.
- Removed deprecated code and streamlined the binding process for clarity and maintainability.

* fix(test_nullable_buffer_params): improve formatting and consistency in test output

- Updated string formatting for better readability in the `test_nullable_shared_shape` function.
- Ensured consistent use of double quotes for string literals.
- Added a missing newline at the end of the file for proper formatting.

* refactor(arg_binder): simplify allocation size calculation in BindDLTensors

- Streamlined the calculation of allocation size by replacing a lambda function with a direct loop, enhancing readability and maintainability.
- Improved clarity in the null check message for data pointers, ensuring better understanding of the binding process.

* Remove debug prints from phase.py

Removed debug print statements after MakePackedAPI transformation.
2025-12-19 11:19:04 +08:00
Lei Wang c750fb8a98 [Enhancement] Update examples and tests for improved type handling functionality (#1448)
* [Enhancement] Update examples and tests for improved type handling and functionality

- Enhanced various example scripts to support new data types and improve compatibility with PyTorch.
- Updated tests across multiple modules to ensure correct functionality with the latest changes in type handling.
- Refactored code in examples to streamline operations and improve clarity, particularly in tensor operations and memory management.
- Added comprehensive tests for new features and fixed existing issues related to type conversions and buffer handling.

* [Refactor] Update accumulation data type to float32 across examples

- Changed accumulation data type from "float" to T.float32 in multiple example scripts to ensure consistency and improve numerical stability.
- This update affects various modules including flash attention, GEMM analysis, convolution, and deepseek MLA examples, enhancing type handling across the board.

* [Refactor] Standardize data type usage across benchmark scripts

- Updated data type definitions in benchmark scripts to use T.float16 and T.float32 consistently, enhancing clarity and type handling.
- Adjusted dtype assignments in matmul functions and configuration setups to align with the new standard.
- Improved overall code consistency and maintainability by ensuring uniform data type usage across various modules.

* [Refactor] Standardize data type usage in templates and scripts

- Updated data type definitions in various templates and scripts to use string representations (e.g., "float16", "int32") instead of T.float16 and T.int32 for improved consistency and clarity.
- Enhanced overall code maintainability by ensuring uniform data type usage across multiple modules, including convolution, elementwise operations, and matrix multiplication templates.
- This change aims to streamline type handling and improve compatibility with existing workflows.

* [Refactor] Standardize data type usage in examples and benchmarks

- Updated data type definitions in various example and benchmark scripts to use T.float16 and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in kernel functions and configuration setups to align with the new standard.
- Improved overall code consistency by ensuring uniform data type usage across multiple modules, including attention mechanisms, matrix multiplication, and GEMM examples.

* [Refactor] Import dtypes from language.v2 module

- Added import statement for dtypes from the language.v2 module to enhance type handling and maintain consistency across the codebase.
- This change aims to streamline data type management and improve overall code clarity.

* fix

* [Refactor] Standardize data type usage across scripts

- Updated data type definitions in various scripts to use string representations (e.g., "float16", "int8") instead of T.float16 and T.int8 for improved consistency and clarity.
- Adjusted dtype assignments in functions and configuration setups to align with the new standard, enhancing overall code maintainability.
- This change affects multiple modules, including benchmark and attention mechanisms, ensuring uniform data type usage throughout the codebase.

* [Refactor] Update data type handling for consistency and clarity

- Changed string representations of data types in the Hint class to use T.float32 and T.int32 for improved consistency.
- Added new data types "int4" and "int16" to the dtypes module, enhancing type support across the codebase.
- Updated function signatures and assertions in the lop3 and mxfp modules to utilize the new data types, ensuring uniformity in type handling.
- This refactor aims to streamline data type management and improve overall code clarity and maintainability.

* [Enhancement] Improve data type handling and error messaging

- Introduced a mapping for canonical data types to their display strings, enhancing clarity in type representation.
- Updated the dtype creation logic to utilize the new mapping, ensuring more intuitive handling of string inputs.
- Refined error messages in the lop3 module to provide clearer feedback on invalid source formats, improving debugging and user experience.

* [Fix] Correct boolean flag in GEMM SP test case

- Updated the boolean flag in the test_gemm_sp_sm90 function to ensure proper functionality in the test case.
- This change enhances the accuracy of the test and aligns it with expected behavior for the GEMM SP implementation.

* [Refactor] Standardize data type usage across scripts

- Updated data type definitions in various scripts to use T.float16 and T.bfloat16 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in function signatures and argument parsing to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change affects multiple modules, including benchmarks and examples, improving overall code consistency and readability.

* [Refactor] Standardize data type usage in various modules

- Updated data type assignments in multiple scripts to utilize T.float32, T.int8, and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted function signatures and parameter types across benchmarks, examples, and tests to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change improves overall code consistency and readability, impacting modules related to matrix multiplication, GEMM, and tensor operations.

* [Refactor] Update argument parsing for data types in benchmarks

- Changed argument parsing for data types in benchmark_matmul_intrinsic.py and benchmark_matmul_sp.py to use string representations ("float16", "int8", "float") instead of T.float16 and T.float.
- This update enhances consistency in data type handling across benchmark scripts, improving clarity and maintainability.

* [Refactor] Update data type handling in benchmark and example scripts

- Changed data type arguments in benchmark and example scripts to use string representations ("float16") instead of T.float16 for improved consistency.
- Updated function signatures and argument parsing to align with the new standard, enhancing clarity and maintainability across the codebase.
- This change affects multiple modules related to attention mechanisms and tensor operations, ensuring uniform data type usage throughout the examples.

* [Refactor] Fix data type conversion in multiple scripts

- Corrected the usage of the data type conversion method from dtype..as_torch() to dtype.as_torch() across various benchmark and example scripts.
- This change enhances consistency in data type handling and improves code readability, impacting modules related to attention mechanisms and tensor operations.

* [Refactor] Update float8 data type usage across multiple scripts

- Changed instances of T.float8_e4m3 to T.float8_e4m3fn in various benchmark, example, and test scripts to ensure consistency in data type handling.
- This update enhances clarity and maintainability across the codebase, particularly in modules related to matrix multiplication and tensor operations.

* [Refactor] Enhance float8 data type handling in CUDA code generation

- Updated the handling of float8 data types in the CUDA code generation to include additional float8 variants, improving type conversion logic.
- Adjusted conditions to ensure proper type checks for float8 conversions, enhancing clarity and maintainability in the codebase.
- Modified layout inference to streamline float8 type checks, ensuring consistency across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.

* [Refactor] Streamline float8 data type handling in CUDA and related modules

- Enhanced float8 data type handling in CUDA code generation by refining type conversion logic and ensuring consistent type checks.
- Updated layout inference for float8 types to improve clarity and maintainability across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.

* [Refactor] Remove unnecessary cache disabling in float8 example script

- Eliminated the call to tilelang.disable_cache() in example_group_per_split_token_cast_to_fp8.py to streamline the code.
- This change enhances clarity and maintainability of the example script without affecting its functionality.

* [Refactor] Update data type usage in debug print tests

- Changed the argument for dtype in the test_debug_print_buffer function from a string representation to the corresponding T.bool type.
- This update enhances consistency in data type handling within the test suite, improving clarity and maintainability.

* lint fix

* Update function parameter types from `str` to `T.dtype` for improved type safety in attention sink and related examples

* Refactor `gemv_alloc_reducer` function signature for improved readability by formatting parameters across multiple lines.
2025-12-17 14:41:25 +08:00
Kuris 9c21586b54 [Feat] Integrate Z3 in TVM Arith Analyzer (#1367) 2025-12-17 11:22:55 +08:00
Lei Wang 00dd738835 [CUDA] Add read-only parameter annotation for CUDA codegen (#1416)
* [Enhancement] Add read-only parameter annotation for CUDA codegen

* Introduced the `AnnotateReadOnlyParams` transformation to annotate read-only handle parameters in PrimFuncs, enabling the generation of `const` qualifiers in CUDA codegen.
* Updated `PrintFunctionSignature` and `AddFunction` methods to utilize the new attribute `tl.readonly_param_indices`, enhancing performance by allowing read-only cache loads.
* Modified the optimization pipeline to include the new annotation step, improving the overall efficiency of the code generation process.

* lint fix

* [Dependency] Update apache-tvm-ffi version to >=0.1.3

* Updated the version of apache-tvm-ffi in pyproject.toml, requirements.txt, and requirements-dev.txt to ensure compatibility with the latest features and fixes.
* Made adjustments in CUDA and HIP template files to use `const` qualifiers for global pointer parameters, enhancing code safety and clarity.

* lint fix

* [Enhancement] Refactor ReadWriteMarker for improved parameter handling

* Updated the ReadWriteMarker class to accept a set of parameter or data variables, enhancing its ability to track written variables.
* Introduced a new method, ResolveDataVarFromPtrArg, to resolve underlying buffer data from pointer-like arguments, improving accuracy in identifying written variables.
* Modified the MarkReadOnlyParams function to gather handle parameters and their corresponding buffer data variables, streamlining the process of determining read-only parameters.
* Enhanced the logic for identifying written variables to account for aliased data variables, ensuring comprehensive tracking of modifications.

* lint fix

* Update tma_load function to use const qualifier for global memory pointer

* Changed the parameter type of gmem_ptr in the tma_load function from void* to void const* to enhance type safety and clarity in memory operations.
* This modification ensures that the function correctly handles read-only global memory pointers, aligning with best practices in CUDA programming.

* Remove commented-out code and reorder transformations in OptimizeForTarget function for clarity

* Refactor buffer marking logic in annotate_read_only_params.cc to improve accuracy in identifying written variables. Update OptimizeForTarget function to reorder transformations for better clarity.
2025-12-14 00:33:00 +08:00
Lei Wang 3546e2ee08 [Atomic] Use ptr for atomicAdd dst instead of reference (#1425)
* [Enhancement] Update AtomicAdd function signature to accept pointer to destination

* Modified AtomicAdd in CUDA to take a pointer instead of a reference for the destination argument.
* Updated related code in atomicadd_vectorize.cc to ensure compatibility with the new signature.
* Adjusted Python interface in atomic.py to pass the destination by pointer, aligning with device function requirements.

* [Enhancement] Refactor AtomicAddRet function signature to accept pointer

* Updated AtomicAddRet in both CUDA and HIP to take a pointer instead of a reference for the address argument, improving consistency with the AtomicAdd function.
* Adjusted the implementation to ensure proper reinterpretation of the address type for atomic operations.

* lint fix

* [Enhancement] Refactor AtomicAddNode::MakeSIMTLoop to use destination pointer

* Updated the MakeSIMTLoop function to build a pointer to the destination element using tvm_access_ptr instead of loading the destination value directly.
* Simplified the handling of source and destination predicates, improving clarity and maintainability of the code.
* Ensured compatibility with the new pointer-based approach for atomic operations.

* lint fix

* test fix

* lint fix
2025-12-14 00:14:31 +08:00
Lei Wang 29051439db [Lint] Phaseout Yapf format and embrace ruff format (#1417) 2025-12-12 17:22:18 +08:00
Lei Wang f2858fa11d [Enhancement] Refactor inflight computing to support dynamic pipeline extents (#1399)
* [Build] Update CMake configuration for tilelang_cython_wrapper installation

- Adjusted output directories for the tilelang_cython_wrapper to ensure that development builds place the extension in build/lib.
- Updated installation paths to place the extension in tilelang/lib within the wheel, improving organization and avoiding potential conflicts with other modules.
- Modified the internal library path exposure in env.py to prevent shadowing of common module names, enhancing compatibility and usability in user projects.

* [Build] Standardize output directories for tilelang libraries

- Set output directories for both tilelang and tilelang_module libraries to "${CMAKE_BINARY_DIR}/lib" for consistency in development builds.
- This change enhances organization and ensures that all build artifacts are located in a unified directory structure.

* [Refactor] Update TVM subproject and enhance pipeline loop handling

- Updated the TVM subproject to commit 90581fe9e5287bbcf1844ad14255a1e1e8cdf7f0.
- Added new fields to `PipelineAnnotation` and `RewrittenBlockInfo` structures to track original statement indices and improve async state management.
- Refactored `EmitImpl` and `PopulateWaitCounts` methods to enhance clarity and functionality, including better handling of commit groups and wait counts.
- Simplified access index calculations and strengthened analyzer constraints for loop bounds.

* [Cleanup] Remove license block and unused includes from inject_pipeline.cc

- Eliminated the Apache license block from the top of the file to streamline the code.
- Removed unused include directives for memory and stringstream to enhance code clarity and reduce unnecessary dependencies.

* [Refactor] Enhance transformation pipeline and test execution

- Added an additional Simplify transformation in the InjectSoftwarePipeline to improve optimization.
- Updated the test file to call `test_trival_pipeline()` directly, commenting out the previous main execution for better test isolation.
2025-12-10 22:04:18 +08:00
ConvolutedDog fac0400680 [Feat] Extend LegalizeNegativeIndex to support buffer store stmts (#1339)
This commit enhances the LegalizeNegativeIndex transformation pass to handle
both buffer load and store operations with negative indices and adds some
test cases.
2025-11-26 14:02:09 +08:00
Tong WUandcoderabbitai[bot] 3ab93cd76b [Enhancement] Keep max score attention across blocks in FlashAttention for better numerical stablity (#1269)
* Implement max score retention across blocks in FlashAttention for improved stability

* fix manual pipeline parameters

* Update examples/flash_attention/example_gqa_fwd_varlen.py

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

* fix typo

* more

* fix a previous typo

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
2025-11-17 21:53:19 +08:00
5f202fe5c1 [Language] Initial version of tilelang frontend v2 (#1120)
* tilelang frontend v2

* syntax sugar: defining a local var by annotation

* [Refactor] fix type linting warning like `T.float32`

* Add tl.local_var_init for new tl.float32

* allow passing default argument as function annotation

* allow default arguments as annotation

* fix lint error

* minor fix

* [Refactor] refactor tilelang.jit and tilelang.autotune

* minor fix

* minor fix

* minor fix

* fix metal get function name

* add par_compile impl and tests

* Type consistency on tvm datatype
1. isinstance(tl.float32, tvm.DataType) == True
2. Allow `tl.float32` as function annotations
3. Allow `tl.float32` as argument to be passed to `tl.alloc` or other functions

* fix lint error

* add more warning in frontend

* update tvm version

* Minor fix on tvm_ffi annotations

* add document and examples

* fix lint error

* Simplify index calculations in example_chunk_o_bwd.py

Refactor index calculations for dg_last_fragment assignment.

* minor fix

* lint fix

---------

Co-authored-by: Lei Wang <leiwang1999@outlook.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
2025-11-03 18:15:56 +08:00
Lei Wang d99853b665 [Language] Add Correctness and performance check scripts for V2 (#1174)
* fix

* lint fix

* fix

* lint fix

* fix

* upd
2025-11-03 02:51:59 +08:00
Lei WangandZhiwen Mo aef0a6bb07 [Language] Expose T.warpgroup_fence_operand for nvcc code motion (#986)
* remove debug print

* pipeline fix

* use the correct buffer access scope

* rs support

* warp warpgroup_fence_operand

* fix

* fp8 dtype ptx enhance

* mma fix

* TCGEN05 Interface

* tcgen05 support

* rebase

* update

* Enhance TCGEN05 support by adding new intrinsic operations and descriptors. Introduced `ptx_tcgen05_mma_ts` for tensor-memory to shared-memory instructions and `tcgen05_mma_arrive` for signaling barrier completion. Updated existing descriptors and code generation logic to accommodate these changes, ensuring compatibility with new instruction sets. Refactored related allocation functions and improved handling of shared memory descriptors.

* lint fix

* Refactor buffer reference handling in CUDA code generation and update test execution in tilelang. Ensure default annotations for unrolling are set correctly in TIR IR module.

* wgmma fix

---------

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2025-11-03 00:50:31 +08:00
Lei Wang 4efd2d2d74 [CI] allow dirty workspace for format.sh and introduce loop carry thread sync unit test (#1153)
* atomic_fix

* atomic_fix

* mem fix

* lint fix

* add some comments

* fix

* fix

* lint fix

* handle async copy

* lint fix

* lint fix
2025-10-29 13:17:23 +08:00
bb8b3cd790 [Enhancement] Update async intrinsic handling in inject_fence_proxy (#1068)
* [Enhancement] Update async intrinsic handling in inject_fence_proxy

* Added support for wgmma async intrinsics in IsAsyncIntrinsic function.
* Changed handling of unknown externs to treat them as Generic instead of Async, improving accuracy in proxy kind determination.

* test fix

* Update testing/python/transform/test_tilelang_transform_inject_fence_proxy.py

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

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
2025-10-21 01:05:18 +08:00
Lei Wang a7730272e4 [Language] Recommend using T.dynamic instead of T.symbolic (#1076)
* recommend using T.dynamic instead of T.symbolic

* lint fix

* lint fix
2025-10-20 19:48:31 +08:00
Chaofan Linandpre-commit-ci[bot] 72111642a6 [Refactor] Refactor Pass LegalizeSafeMemoryAccess to support recursive load/store rewrite (#1050)
* [Refactor] Refactor Pass  to support recursive load/store rewrite

* lint

* recursive collect conds for call_extern

* fix name

* [Lint]: [pre-commit.ci] auto fixes [...]

* lint

* [Lint]: [pre-commit.ci] auto fixes [...]

* lint

* [Lint]: [pre-commit.ci] auto fixes [...]

* address comment

* rename pad_value to safe_value

* lint

* add oob store test

* [Lint]: [pre-commit.ci] auto fixes [...]

* fix

* fix

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2025-10-17 20:56:01 +08:00
Lei Wang ddfaac36c1 [Refactor] Refactor Pass InjectFenceProxy and expose some warp group primitives in frontend (#977)
* • InjectFenceProxy docs and tests

  - annotate proxy fence injector with context comments for async/generic detection
  - add compiler internals doc covering the pass mechanics and link it in docs index
  - repair fence proxy test by fixing descriptor init usage and fence counter logic

* do not consider call_extern as async.

* doc update.

* reduce test size for sparse mla
2025-10-11 16:46:48 +08:00
Lei Wang f8ae600ce4 [Bugfix] Do not force inline let stmt (#947)
* remove debug print

* Remove inline let expressions from the LowerAndLegalize function in phase.py

* add test

* Update sparse MLA examples to support SKV adjustment and correctness checks

- Changed SKV parameter from 32768 to 8192 in sparse MLA backward and forward tests.
- Added check_correctness parameter to test functions for validation of outputs.
- Updated test cases to reflect new SKV values and correctness checks.

* reduce test shape

* Update documentation structure and refactor main function parameters in example_fusedmoe_tilelang.py

- Added a new section for compiler internals in the documentation.
- Refactored the main function in example_fusedmoe_tilelang.py to accept parameters for hidden dimensions, expert configurations, and batch/sequence sizes, improving flexibility and readability.

* Update buffer access checks in merge_shared_memory_allocations.cc

- Changed the condition for buffer access from less than (<) to less than or equal to (<=) to allow access at the same scope level.
- Adjusted the logic for determining the access level when touching buffers to ensure correct handling of scope levels.

* lint fix

* Support pipeline with LetStmt

* lint fix

* • Fix LowerTileOp let handling to avoid LetInline dependency

  - inline let-bound BufferLoad nodes via resolver helpers and structured return
  - remap layouts/buffers using original data vars and only rewrite when needed
  - update pipeline planner to understand let-bound address_of buffers
  - document the new inline behaviour in docs/let_inline_fix.md

* fix for wgmma pipeline with let binding

* lint fix

* test fix

* reduce smem usage.

* let binding enhancement

* fix for dpgm

* fix simplify

* lint fix

* use tilelang.Simplify instead of tir.Simplify

* • Add TL_FORCE_LET_INLINE pass config and gate eager LetInline usage

  - register the new config in builtin headers/registration
  - add helper to pipeline enabling LetInline based on pass context
  - document LetStmt inlining controls and usage
2025-10-10 16:20:22 +08:00
Zhiwen Mo f58bcd43f1 [SM100] Add sm100 GEMM layouts and tcgen05 support (#887)
* update sm100 related utcmma, tmem, ld/st256 in src
* update sm100 related utcmma, tmem, ld/st256 in tilelang
* Remove deprecated GEMM examples and related README documentation for SM100 architecture support
* Update GEMM implementation to replace UTCMMA with TCGEN5MMA across relevant files
* Remove gemm_umma.py example and update README to reflect TCGEN5MMA terminology changes
* Update README.md for gemm_sm100 example by removing outdated API sections and streamlining documentation
* Update README and source files to reflect TCGEN5.MMA terminology changes
* Refactor CUDA GEMM header for improved readability
2025-09-28 16:49:39 +08:00
Yichen Yan 5e52952201 [Lint] Add ruff config to check for useless spaces (#807)
* update lint config

* Remove spaces for blank line

* update
2025-09-13 21:06:25 +08:00
Lei Wang 91a7bb2be6 [TileOp] Introduce a experimental python defined T.gemm_v2 (#793)
* Refactor GEMM and GEMM-SP operations to enhance clarity and maintainability

- Removed deprecated prime factorization functions from `gemm.cc` and `gemm_sp.cc`.
- Introduced a new `GemmWarpPolicy` class to manage warp policy attributes and methods, improving encapsulation.
- Updated reflection methods to include the new policy structure, ensuring proper registration and introspection capabilities.
- Enhanced `GetArchInt` function in `utils.cc` for better readability and type safety.
- Added new `gemm_v2` function in `gemm.py` for improved GEMM operation with additional parameters and checks.

* Refactor GEMM and frontend legalize operations for improved clarity and functionality

- Updated `gemm_py.h` to include the correct header for GEMM operations.
- Renamed `FrontendLegalizer` class to `LetInliner` and updated related methods to reflect this change, enhancing code clarity.
- Modified the pass function from `FrontendLegalize` to `LetInline` for better alignment with its purpose.
- Updated test cases to utilize the new `gemm_v2` function and adjusted the testing framework for improved output and clarity.
- Removed obsolete test file `test_tilelang_transform_frontend_legalize.py` to streamline the test suite.
- Enhanced the `LowerAndLegalize` function to utilize the new `LetInline` pass, improving the overall transformation process.

* Enhance CUDA code generation and testing for GEMM operations

- Added indentation printing in `codegen_cuda.cc` for improved assembly code formatting.
- Updated `test_tilelang_tilelibrary_gemm.py` to include additional GEMM test cases and shared memory allocation with specified scope.
- Introduced new `matmul_sr` and `run_gemm_sr` functions for GEMM operations with shared and fragment memory layouts.
- Refactored layout inference in `mma_macro_generator.py` to improve clarity and correctness in shared memory handling.
- Enhanced `gemm/__init__.py` to support new GEMM operation combinations and layout inference logic.

These changes improve the clarity, functionality, and testing coverage of GEMM operations in the TileLang framework.

* Refactor GEMM layout and testing for improved clarity and functionality

- Updated `gemm_layouts.cc` to enhance the layout generation logic for transposed and non-transposed GEMM operations.
- Renamed and modified functions in `test_tilelang_tilelibrary_gemm.py` to reflect changes in GEMM function signatures and improve test coverage.
- Introduced new GEMM operation combinations in `gemm/__init__.py` to support additional layouts and configurations.
- Enhanced layout inference in `mma_layout.py` and `mma_macro_generator.py` for better handling of shared memory layouts.

These changes improve the clarity, functionality, and testing coverage of GEMM operations in the TileLang framework.

* Refactor GEMM layout and Python integration for improved functionality

- Updated `gemm_layouts.cc` to correct the order of layout replication and repetition for transposed and non-transposed GEMM operations.
- Enhanced `gemm_py.cc` to handle block realization more robustly, ensuring correct assignment of global symbols and block attributes.
- Refactored `inject_pipeline.cc` to streamline buffer read/write region handling, improving clarity and maintainability.
- Cleaned up test cases in `test_tilelang_tilelibrary_gemm.py` by removing unnecessary print statements and adjusting function calls for better test execution flow.

These changes enhance the clarity, functionality, and robustness of GEMM operations and their testing in the TileLang framework.

* Refactor GEMM layout and testing for improved clarity and functionality

- Updated `gemm_layouts.cc` to enhance layout generation logic for transposed and non-transposed GEMM operations.
- Improved block realization handling in `gemm_py.cc` for better assignment of global symbols.
- Streamlined buffer read/write region handling in `inject_pipeline.cc` for clarity.
- Enhanced test cases in `test_tilelang_tilelibrary_gemm.py` by adjusting function calls and adding new GEMM operation combinations.

These changes improve the clarity, functionality, and robustness of GEMM operations and their testing in the TileLang framework.

* tfloat32 support.

* lint fix

* lint fix

* Refactor shared memory allocation in GEMM tests

- Removed unnecessary scope specification in shared memory allocation for matrices A and B in `test_tilelang_tilelibrary_gemm.py`.
- This change simplifies the allocation process and aligns with the updated GEMM function signatures.
2025-09-11 03:26:20 +08:00
Wenhao Xie ff35fc0881 [Bugfix] Address PassContext contamination from CI and fix incorrect rewrites in warp specialized pass (#767)
* fix ci and pass bug

* fix

* try

* lint
2025-08-28 23:46:26 +08:00
Lei Wang 556d411e99 [Typo] Remove disable_cache in some tests (#755)
* Update test parameters and remove debug print statement

- Adjusted test cases in `test_tilelang_dynamic_symbolic_bench.py` to use smaller matrix sizes (1024x1024) for improved performance and quicker execution.
- Removed a debug print statement from `phase.py` to clean up the code and enhance clarity.

* Refactor loop stack management in warp_specialized_rewriter

- Introduced a new `LoopInfo` struct to encapsulate loop variable details, including `loop_var`, `extent`, and `min`, enhancing clarity and maintainability.
- Updated the `loop_stack_` to utilize `LoopInfo` instead of a pair, improving type safety and readability.
- Adjusted linear index calculations to account for the new structure, ensuring correct behavior in loop transformations.

* Remove unused `torch.backends` import and `tilelang.disable_cache()` calls from multiple test files to enhance code clarity and maintainability.
2025-08-25 01:04:50 +08:00
Lei Wang 6b125028f6 [Refactor] Merge ThreadPartialSync and ThreadStorageSync (#741)
* Remove `thread_partial_sync.cc` and refactor `thread_storage_sync.cc` to streamline synchronization handling. Introduce `thread_sync_types.h` for thread-bound key definitions and reserved named barriers. Update related logic in `ThreadSyncInserter` and `TileLangThreadSync` for improved clarity and efficiency.

* Remove `sync_thread_partial` references and related documentation from the codebase. Update CUDA and HIP code generation files to eliminate calls to the removed function. Refactor `__sync_thread_partial` to `sync_thread_partial` in CUDA common header for consistency.

* Remove unused import of `bulk_copy.h` in `codegen_hip.cc` to enhance code clarity and maintainability.

* Add import of `bulk_copy.h` in `codegen_hip.cc` to support new functionality.

* typo fix

* Update data type in reduce_sum tests from float16 to float32 for consistency and clarity. Remove redundant dtype tests and streamline run functions. Enhance reshape kernel compilation with pass configurations to address shared memory layout issues.

* lint fix

* test fix

* Enhance CI configuration by adding verbose output to pip install command for better visibility during installation.

* use ninja instead of make

* Add CMake configuration step for Ninja build system in setup.py

* Update pyproject.toml to include additional build dependencies: build, torch, tox, auditwheel, patchelf, and ninja.

* Enhance CI configuration by adding verbose output to pytest commands for improved test visibility.

* Update pyproject.toml to add Cython as a build dependency. Enhance thread storage synchronization in thread_storage_sync.cc by introducing new thread variable handling and improving index disjointness checks.

* Update data type in cumulative sum tests from float16 to float32 for consistency. Modify run_cumsum function to utilize the updated dtype and enhance result validation with assertions. Adjust test cases accordingly.

* Refactor storage access handling by introducing buffer data mapping in TileLangStorageAccessVisitor. Enhance access entry structure to include pointer access flag. Update thread storage synchronization to accommodate new buffer data mappings. Adjust quickstart example to print kernel source for debugging purposes.

* Refactor linear index conversion in TileLangStorageAccessVisitor to utilize the analyzer for simplification. Update buffer index calculations to ensure consistent simplification of range expressions.

* bugfix

* Refactor buffer index calculation in TileLangStorageAccessVisitor to simplify access handling. Removed unused buffer mapping logic, ensuring consistent buffer index generation with a default ramp.

* Refactor TileLangStorageAccessVisitor to replace buffer indices with buffer ranges for improved pointer access handling. Update AccessEntry structure to include buffer_ranges and adjust thread storage synchronization logic to account for pointer access conflicts.

* Refactor thread storage synchronization to replace 'shared.dyn' with 'shared' for consistency in memory allocation. Update related test cases to reflect this change and ensure proper functionality.
2025-08-23 17:07:10 +08:00
Lei Wang 5c11d245ae [Refactor] Merge bulk copy into copy and improve layout inference for bulk copy (#746)
* [Refactor] Merge bulk copy into copy and refactor layout inference for bulk copy

* Deleted the `bulk_copy` operator implementation and its header file as it is no longer needed.
* Introduced a new function `cuTensorMapType()` to return the data type for CUDA tensor mapping.
* Updated related files to reflect these changes, ensuring that the codebase remains clean and maintainable.

* lint fix

* Fix typos in intrinsic names and remove unused print statement in block_sparse_attn_tilelang.py. Updated references from `ptx_ldmatirx` to `ptx_ldmatrix` across multiple files for consistency.

* remove bulk copy

* Refactor copy and atomic add operations to support TMA lower configuration

- Updated `GetCopyInst` to accept a `disable_tma_lower` parameter, allowing for conditional usage of TMA in bulk load/store operations.
- Modified `Lower` method in `Copy` to incorporate the new TMA configuration.
- Refactored `AtomicAdd::Lower` to streamline layout inference and vectorization logic.
- Removed unused `disable_tma_lower` field from `LowerArgs` structure for clarity.
- Enhanced atomic add vectorization by replacing the buggy implementation with a more robust loop vectorization approach.

* Enhance TMA bulk copy logic in `LowerBulkCopy` method

- Added a condition to set `desc.swizzle` to `CU_TENSOR_MAP_SWIZZLE_NONE` when `shared_layout` matches `linear_layout`, improving clarity in layout handling.
- Updated warning log to provide more detailed information about fallback scenarios, including source and destination buffer names and shapes, enhancing debugging capabilities.

* lint fix

* Remove fallback logging for non-swizzled global layout in `LowerBulkCopy` method to streamline the bulk copy logic. This change enhances code clarity by eliminating unnecessary warning messages related to inner box dimensions.

* Enhance reshape kernel compilation in `run_reshape` and `run_reshape_smem_1d_2_2d` functions

- Updated the `tl.compile` method to include `pass_configs` that disable TMA lower and warp specialization, addressing shared memory layout transformation limitations.
- Added TODO comments to indicate the need for further improvements in shared memory handling.

* Update `native_sparse_attention` function to include TMA configuration options

- Added `pass_configs` to the JIT decorator to disable TMA lower and warp specialization, addressing potential issues with shared memory layout transformations.
- Updated comments to clarify modifications in tensor shapes for inference, specifically setting `q` sequence length to 1.

* Refactor JIT decorator formatting in `native_sparse_attention` function

- Improved readability by reformatting the JIT decorator parameters for `native_sparse_attention`, ensuring consistent style across the codebase.
- No functional changes were made; this update focuses on code clarity and maintainability.

* Enhance thread management and logging in TileLang compilation

- Added a method to check if printing is enabled during compilation, improving control over logging behavior.
- Updated the JIT kernel class to utilize the new method for logging compilation status, ensuring consistent and clear output.
- Added comments to clarify the purpose of changes and improve code readability.

* Add warp specialization scope and refactor register management in TileLang

- Introduced a new constant `kWarpSpecializationScope` in `builtin.h` for better attribute management.
- Removed the `SetMaxNRegCollector` class and its related logic from `warp_specialized_rewriter.cc`, streamlining the warp specialization process.
- Added functions `annotate_producer_reg_dealloc` and `annotate_consumer_reg_alloc` in `builtin.py` to facilitate register management.
- Implemented `AnnotateWarpGroupRegAlloc` in `__init__.py` to inject register allocation calls into warp-specialized functions, enhancing the overall register handling in the compilation process.

* Refactor test for InjectSetMaxNReg pass in TileLang

- Improved readability by restructuring conditional checks and assertions in the test cases.
- Enhanced clarity in the collection of `set_max_nreg` calls by simplifying the logic.
- Ensured consistent formatting and spacing throughout the test functions for better maintainability.

* Enhance bulk copy and store checks in `Copy` class

- Updated scope validation for source and destination tensors in `CheckBulkLoad` and `CheckBulkStore` methods to include both `shared.dyn` and `shared` as valid options.
- Modified `CheckLDSMCopy` and `CheckSTSMCopy` methods to accommodate the new scope validation, ensuring compatibility with shared memory configurations.
- Improved logging in `LowerBulkCopy` to provide clearer warnings regarding unsupported swizzle layouts, including source and destination names for better debugging.

* lint fix
2025-08-22 14:41:34 +08:00
Lei WangandChenggang Zhao 1b308baf9a [Language] Introduce StridedTensor to support non contigious torch inputs (#722)
* Update submodule 'tvm' to commit e11521e6936a827efa334588d29571fbb4620107

* Support strided tensors

* Refactor target attribute helper functions for improved clarity

* No code changes made in proxy.py and setup.py

* lint fix

* lint fix via gemini

* lint fix

* test fix

* test fix

* lint fix

* Update wrapper.py

* test fix

* Enhance test for InjectSoftwarePipeline by adding LowerOpaqueBlock transformation and updating expected function signature to use match_buffer for better clarity.

* lint fix

---------

Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com>
2025-08-18 00:53:29 +08:00