Commit Graph
316 Commits
Author SHA1 Message Date
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
Lei Wang 4fb96a0f56 [Bugfix] Fix variable scoping issue in InjectSoftwarePipeline for transitive LetStmt dependencies (#1657)
* [Enhancement] Update global load/store functions for CUDA compatibility (#1652)

Refactor the `ld_global_256` and `st_global_256` functions to support both CUDA versions above 12.9 and earlier versions. This change ensures that 256-bit loads and stores are handled correctly across different CUDA versions, improving performance and compatibility. The implementation now uses two 128-bit loads/stores for older versions, enhancing the robustness of the codebase.

* Update comments in global load/store functions for CUDA compatibility

Clarified comments in `ld_global_256` and `st_global_256` functions to indicate that the fallback for CUDA versions below 12.9 may have performance regressions. This change enhances code readability and provides better context for developers working with different CUDA versions.

* Update submodule and enhance LetStmt handling in inject_pipeline.cc

- Updated the TVM submodule to the latest commit.
- Improved the handling of LetStmt in the inject_pipeline.cc file to account for transitive dependencies on loop variables, ensuring correct variable substitution in rewritten blocks.
- Adjusted test_tilelang_issue_1263.py to remove unnecessary jit decorator and updated the kernel compilation process with specific pass configurations.

* lint fix

* revert tvm

* remove unused test

* test fix
2026-01-13 14:21:12 +08:00
Lei Wang 9936636438 [Enhancement] Improve unroll loop functionality for dynamic extent and corresponding test case (#1654)
* Add unroll loop functionality and corresponding test case

- Introduced a new `UnrollLoop` function in the transform module to unroll loops based on various configuration options.
- Added a test case in `test_tilelang_language_unroll.py` to validate the behavior of `T.unroll` with only the extent parameter, ensuring correct kernel generation with unroll pragmas.

* Refactor unroll kernel implementation and update test case

- Changed the kernel function in `test_tilelang_language_unroll.py` to use a new `unroll_kernel` function that compiles and returns the output tensor, improving clarity and structure.
- Updated the `OptimizeForTarget` function in `phase.py` to ensure the `UnrollLoop` transformation is applied correctly, maintaining consistency in optimization phases.

* lint fix

* lint fix
2026-01-12 22:33:17 +08:00
Kuris 326be4dd3f [Feat] Allow dangling producer in wasp pipeline planning (#1263) (#1647)
* [Fix] Handle cases where producers lack corresponding consumers in synchronization logic

* add test for pipeline without consumer

* fix lint error
2026-01-09 22:14:30 +08:00
Lei WangandClaude Opus 4.5 01d651a311 [Refactor] Move dtypes.py from eager to language and add bits/bytes properties (#1646)
- Move `tilelang/language/eager/dtypes.py` to `tilelang/language/dtypes.py`
  since dtype definitions are general-purpose, not eager-specific
- Add `bytes` property to dtype (alias for `itemsize`) for convenience
- Add type hints for `bits` and `bytes` properties in TYPE_CHECKING block
- Create `tilelang/dtypes.py` for direct import via `from tilelang.dtypes import ...`
- Update all import paths accordingly

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
2026-01-09 21:53:52 +08:00
Lei Wang 409677dfb3 [Bugfix] Fix pdl related intrin handling to avoid strict annotation codegen (#1650)
bugfix
2026-01-09 17:31:33 +08:00
Lei Wang 6e439530b5 [Refactor] Unify @jit and @lazy_jit into a single @jit decorator (#1632)
* Refactor: Replace @tilelang.lazy_jit with @tilelang.jit in examples and tests

Updated all instances of the @tilelang.lazy_jit decorator to @tilelang.jit in the lazyjit example notebooks and related test files. This change aligns the code with the new JIT compilation approach, enhancing consistency across the codebase. Additionally, removed the lazy_jit import from the module initialization to streamline the API.

* Enhancement: Improve layout annotation handling in T.copy and parallel loops

Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify layout hints for the outermost parallel loop. Enhanced layout annotation handling in CopyNode and AtomicAddNode classes to ensure compatibility with SIMT operations. Added tests to validate the functionality of loop layout annotations, improving robustness and clarity in layout management for parallel loops.

* Add out_idx parameter to @jit and validate it's only used in lazy mode

* lintfix

* Refactor: centralize mode inference in JITImpl.__call__ before parse_args

* test fix

* Enhancement: Introduce JITNoBuilderError for eager mode checks in T.Kernel and T.const functions. Ensure proper error handling when Builder is unavailable, improving robustness in JIT context management.

* lint fix

* lint fix

* Enhancement: Update example_triton_sparse_gqa_decode_varlen_indice.py to improve performance metrics output by adding average FLOPS calculations. Modify layout_inference.cc to refine buffer layout inference checks by ensuring fragment buffers are considered.

* lint fix

* Enhancement: Refactor JIT implementation to centralize mode initialization and improve error handling for eager mode. Update matmul tests to utilize both Cython and TVM FFI JIT kernels, ensuring consistency in output validation.

* Refactor: Remove unused simplify_prim_func decorator from example_gemm_intrinsics.py and related test files. Enhance layout.cc to simplify integer set analysis for multi-threaded compilation, ensuring constant integer outputs. Update lower_tile_op.cc to enforce constant integer checks in layout output shapes. Improve error handling in builder.py for JIT context management.

* remove claude.md

* Refactor: Clean up code formatting in layout.cc, lower_tile_op.cc, and proxy.py for improved readability. Add missing import in test_tilelang_kernel_int4_gemm_mma.py and ensure consistent spacing in builder.py.

* Refactor: Replace logical conditions with tir_all in complex expressions within test_arith_hard.py for improved clarity and consistency. Update main execution flow to call test_hard_prove instead of tilelang.testing.main.

* Enhancement: Add @tilelang.testing.requires_cuda decorator to test functions in test_tilelang_issue_1549.py and test_tilelang_issue_1601.py to ensure CUDA requirements are explicitly stated for relevant tests.

* lint fix

* Refactor: Simplify output shape calculation in layout.cc by directly setting max value from ist.max(). Update test_valid_loop in test_tilelang_fragment_loop_checker.py to disable certain tests and ensure proper execution flow. Enhance error handling in builder.py to correctly manage PrimFunc return types.

* Refactor: Update boolean condition in test_boolop to use Or and Not for improved clarity and consistency in logical expressions.

* Refactor: Update test execution flow in test_tilelang_fragment_loop_checker.py to call tilelang.testing.main. Modify JIT implementation to handle callable functions more generically. Enhance TirTemplate class in builder.py to return an empty tuple if matcher is None, improving robustness.

* Refactor JIT function handling: Replace LazyJITFunc with JITFunc across the codebase, updating related function signatures and logic to support eager JIT compilation mode. This change enhances the clarity and functionality of JIT operations.

* Update JITImpl documentation: Clarify output tensor index attribute description and remove outdated comment regarding implicit TIR construction.

* Refactor TileLang to transition from v2 to eager language module. Update imports across multiple files to replace references to the deprecated v2 module with the new eager module, ensuring compatibility with the latest JIT compilation features. Remove obsolete v2 module files to streamline the codebase.

* Enhance JIT compilation support: Update JIT function handling to improve compatibility with eager execution. Refactor related function signatures and logic, ensuring seamless integration with the latest compilation features. This change aims to streamline JIT operations and enhance overall code clarity.
2026-01-09 15:41:39 +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
Taoyu Zhuandzhutaoyu b914318b46 [BugFix] Correct index_map selection for transposed A matrix in MFMA Layout with k_dim==4 and open rocm-ci for gemmsr (#1627)
* fix(rocm): correct index_map selection for transposed A matrix in MFMA when k_dim=4

Previously, when k_dim=4, the index_map always used non-transposed layout
for matrix A regardless of the transposed flag. This caused precision
issues for transposed GEMM operations on ROCm.

Re-enable the previously skipped test cases for trans_A=True with float dtype.

* also open int8 gemm

---------

Co-authored-by: zhutaoyu <tazhu@amd.com>
2026-01-08 13:08:29 +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
Lei Wang 358f899951 [Bugfix] Do not consider local.var as local buffer during LowerTileOP (#1628)
* [Refactor] Simplify IsLocalBuffer checks in lower_tile_op.cc by removing redundant boolean argument

* lint fix
2026-01-07 14:51:14 +08:00
silentCoder-dev a756074c5c [Feature] add T.sync_warp & T.shfl_sync; change extern pdl into intrin (#1614)
* add T.sync_warp & T.shfl_sync

* add requires_cuda for warp_sync test

* define intrin tl.sync_warp

* define intrin tl.pdl_trigger & tl.pdl_sync
2026-01-06 23:54:11 +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
Chaofan Lin 1b0022015f [BugFix] Fix some bugs in lowering ParallelOp and VectorizeLoop (#1607)
* fix

* [BugFix] Fix some bugs in lowering ParallelOp and VectorizeLoop

* remove debug output
2026-01-06 16:52:12 +08:00
silentCoder-dev 7cfd087aea [Feature] Add more curand operations & support vectorization (#1582)
* add more curand operations & support vectorization

* fix typo about bitwidth

* merge random float into a single operation
2026-01-06 02:20:01 +08:00
Gongen-Ali 4e53b50750 [Bugfix] Open Rocm ci test and fix some bugs. (#1443)
* Fix some bugs on ci and open rocm ci test

* fix pre-commit

* fix build issue on hip

* refactor on hip

* pre-commit fix

* fix test with pdl

* fix pdl test
2026-01-06 02:18:29 +08:00
Lei Wang e3c9a581bd [Enhancement][Subtype] Enhance symbolic shape/stride handling for subtype (#1599)
* [Enhancement] Improve shape and stride binding for subbyte types in ArgBinder

* Updated the shape binding logic in ArgBinder to handle subbyte types (bits < 8) by introducing a packing factor for logical shapes.
* Enhanced the handling of strides for subbyte types, ensuring correct mapping between logical and runtime strides.
* Refactored the binding process to accommodate both symbolic and constant dimensions, improving the robustness of the shape binding mechanism.
* Removed obsolete test file `test_ir_kernel_frame.py` as part of cleanup efforts.

* lint fix
2026-01-04 17:06:37 +08:00
9fd6ea7d3b [Feat] PDL Support (#1494)
* [Feat] PDL Support

* remove unused comments/codes

* fix mismatch funcname

* Resolve duplicate config issue for AI review

* fix typo error && clang-tidy error

* chore: retrigger CI

* change file comment

* rename test_tilelang_jit_ctypes.py to test_tilelang_jit_cython.py

* remove ctype backend

* Move pdl attributes to header and remove redundant VisitStmt_

* chore: trigger CI

* modify the cuda kernel with pdl due to nvcc's bug

* throw an error when invoking __ldg with pdl due to nvcc's bug

* Throw an error when pdl is not supported

* fix checking about pdl

* add comments about MarkCudaSyncCalls

* remove pdl support with compute_capability<90

* ruff check

* fix nvrtc about pdl

* update tvm & extend pdl support for tvm_ffi & add test for tvm_ffi

* ensure that kUseDynamicSharedMemoryTag is the last tag in launch_param_tags

* fix bug introduced by branch merging

* [Feature] Add get_host_source method to CythonKernelAdapter

* Introduced a new method, get_host_source, in the CythonKernelAdapter class to return the source code of the host function. This addition enhances the functionality of the adapter by providing direct access to the host kernel source.

---------

Co-authored-by: silentCoder-dev <silentcoder@foxmail.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-01-04 17:03:37 +08:00
silentCoder-dev 23c446da6a [Fix] Add register to read A ptr in test_tilelang_language_cooperative.py (#1593)
* add register to read A ptr

* add randomize initial value for tensor

* remove manual boundary checking
2026-01-04 15:45:38 +08:00
Lei Wang 946f611c13 [Typo] Rename loop layout annotation test(#1596) 2026-01-04 13:40:21 +08:00
Lei Wang 5685c56a45 [Layout] Support annotating loop layout in frontend (#1579)
* [Enhancement] Improve layout annotation handling in parallel loops

* Updated the handling of layout annotations in parallel loops to ensure they are only attached to the outermost loop. This change clarifies the rationale behind layout management, allowing for better reasoning and transformation of nested parallel regions.
* Introduced new parameters in the `Parallel` function to accept layout annotations directly, enhancing usability and validation during layout inference.
* Refactored the layout inference logic to prioritize annotated layouts over inferred ones, improving the efficiency of layout management in parallel operations.

This update aims to strengthen the integrity of layout annotations and streamline the handling of nested parallel loops in the system.

* lint fix

* lint fix

* [Enhancement] Update test and builder for improved debugging and functionality

* Added print statements in `test_tilelang_language_let.py` to output the `main` function and the source of the compiled module, aiding in debugging.
* Modified the `bind_immutable` method in `builder.py` to include logging of the bound value type, enhancing traceability during the binding process.
* Commented out the call to `tilelang.testing.main()` in favor of directly invoking the test function, streamlining the testing workflow.

* replicate example

* [Enhancement] Add loop layout support to T.copy and improve parallel loop handling

* Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify a layout hint for the outermost parallel loop in SIMT copies.
* Enhanced the handling of layout annotations in the CopyNode class to enforce compatibility with SIMT copy operations, ensuring that layout hints are only applied when appropriate.
* Added tests to validate the functionality of loop layout annotations in copy operations, improving the robustness of the layout management system.

* [Enhancement] Refine loop annotation handling in AtomicAdd and Copy operations

* Updated the handling of loop annotations in both AtomicAddNode and CopyNode classes to ensure that coalesced width and parallel loop layout annotations are only applied to the outermost loop.
* This change improves the clarity and correctness of loop management in SIMT operations, aligning with the recent enhancements made to layout annotation handling in parallel loops.
2026-01-04 12:57:16 +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
silentCoder-dev e1138ad0ca [Feature] add kUseCooperativeLaunch tag for tvm_ffi (#1572)
* add kUseCooperativeLaunch tag

* add test for sync_grid in cooperative launch

* add test for sync_grid in cooperative launch
2025-12-31 11:11:43 +08:00
Lei Wang 0fa16b4c41 [Enhancement] Improve equality checks in layout nodes and fragment validation (#1573)
* [Enhancement] Improve equality checks in layout nodes and fragment validation

* Enhanced the IsEqual method in LayoutNode and FragmentNode to include detailed comparisons of forward mappings, ensuring accurate equality checks.
* Introduced a new parameter in ProveFragmentContains to validate physical indices when checking fragment containment, improving correctness in layout validation.
* Removed obsolete test file related to layout inference.

This update aims to strengthen the integrity of layout comparisons and fragment validations in the system.

* lint fix
2025-12-31 00:36:29 +08:00
Lei Wang 0f9bbd759c [Parallel][Infer] Free-mode chooses minimal replication between buffer-based and PlanLoopPartition (#1559)
* [Enhancement] Improve layout inference in ParallelOpNode

* Enhanced the layout inference mechanism in ParallelOpNode to utilize two strategies: compute_loop_layout_from_buffer and PlanLoopPartition, selecting the one that minimizes replication while ensuring compatibility.
* Updated the logic to choose the best candidate layout based on replication size and containment checks.
* Refactored the HasKnownLayoutAnchor function to clarify its purpose in prioritizing buffer layouts.
* Added a new test case to validate the layout inference behavior, ensuring the correct fragments are generated in the output.

This update aims to optimize layout inference for parallel operations, improving performance and resource utilization.

* lint fix

* bug fix and refactor

* lint fix
2025-12-30 13:28:54 +08:00