Commit Graph
283 Commits
Author SHA1 Message Date
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
Kuris b702299b38 [LazyJIT] Move Type Annotations to Function Body (#1480)
* [LazyJIT] Refactor to move outside annotation into function

* [SyntaxSugar] Allow write expression bind before match buffer

* [SyntaxSugar] Allow decoupled argument annotation

* [SyntaxSugar] Add `T.empty` and return

* Fix typos

* fix lint error

* [Testing] Update lazy_jit tests

* fix typo

* add assertion and error message

* fix typo

* fix double pop of __tune_params

* fix double pop error
2025-12-29 17:46:45 +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
Gabriel Wu 72ce848be0 [Feat] Adapt gemm v2 for cutedsl backend (#1544)
* feat: adapt gemm v2 for cutedsl backend

* fix: ruff
2025-12-27 13:07:16 +08:00
Jinjie LiuandZihua Wu 9ff7c523a9 [Enhancement] Refactor KernelCache to use inheritance-based design (#1483)
* refactor kernel_cache

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* set key to auto when execution_backend is None

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* fix type error on _load_kernel_source override

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* put kernel cache subclasses in adapter

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* remove auto in execution_backend

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* remove _save_kernel_source_code_to_disk for NVRTCKernelCache

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* refactor all cutedsl if statements

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* remove unused kernel_cubin_path

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* add func arguments

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* remove ctypes

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* update ci.yaml to test kernel cache

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* fix bugs after #1512 merged

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* test: fix kernel cache test

* fix: ruff

* fix: ruff

* fix: coderabbit

* fix: coderabbit

* fix: ruff

* fix: MMA

* fix: coderabbit

* fix: codex

* fix: coderabbit

* fix: coderabbit

---------

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
Co-authored-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
2025-12-26 15:41:22 +08:00
Chaofan Lin dff10e5283 [BugFix] Complete vectorized loading for common dtypes (#1536) 2025-12-25 22:28:53 +08:00
KurisandLeiWang1999 d5d959ef7e [Fix] Add support for non-var complement arithmetic computation (#1374) (#1533)
* [Fix] Add complement for non-var IterExpr

* fix lint error

* lint fix

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2025-12-25 21:25:01 +08:00
Lei WangandZhiwen Mo bea40bdd0e [Dtype] Improve host codegen handling for subtype (#1517)
* fp4 related update, require_cu13

* Enhance CUDA type conversion handling and optimize dtype management

- Updated CUDA vectorized cast functions to ensure proper handling of float16, float32, bfloat16, and float8 conversions, adding checks for bit sizes.
- Refactored dtype conversion logic in `cuda_fp4.h` to utilize `cudaRoundZero` for improved accuracy in floating-point conversions.
- Introduced a new method in `KernelParam` to convert TVM DataType to TileLang dtype.
- Adjusted argument binding logic in `arg_binder.cc` to allow for better subtype matching based on total bit counts.
- Enhanced dtype handling in `dtypes.py` to accommodate new float4_e2m1fn types and ensure compatibility with PyTorch.

This update aims to improve type safety and conversion accuracy across the codebase.

* lint fix

* lint fix

* typo fix

---------

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2025-12-24 17:06:39 +08:00
Lei Wang 09385e7deb [Refactor] Support auto swizzling for tma store and phaseout related layout annotations (#1509)
* Remove unnecessary swizzled layout annotations from various attention sink examples and kernel analysis script for improved clarity and performance.

* Remove unnecessary swizzled layout annotations from various examples to enhance code clarity and performance.

* lint fix

* Uncomment main test execution for vectorized cast

* Add swizzled layout annotation for B_shared in dequantize GEMM example

This change introduces a layout annotation for the B_shared tensor in the example_dequant_gemm_bf16_mxfp4_hopper.py file, enhancing the memory layout optimization for better performance during matrix multiplication.

* Remove unnecessary cache disabling in GQA sink example for improved clarity and performance.

* lint fix

* Refactor layout functions to support row-major linear layout for any dimension

- Renamed `makeGemmLayoutLinear` to `makeLinearLayout` and updated its implementation to handle arbitrary dimensions.
- Updated related function calls in `gemm_layouts.cc` and `layout.cc` to use the new layout function.
- Enhanced layout inference in `CumSumOpNode` to enforce linear layout for shared buffers in strict mode.

* Refactor `make_linear_layout` to accept a single argument and support arbitrary dimensions

- Updated the function signature to take a `Buffer`, `BufferLoad`, or `BufferRegion` directly.
- Simplified the implementation by removing argument checks and directly obtaining the shape from the buffer info.
- Enhanced the documentation to clarify the function's purpose and usage.

* skip callback test

* Remove example_dequant_gemm_bf16_mxfp4_hopper_tma.py file, eliminating unused code related to dequantization GEMM example.

* typo fix
2025-12-24 14:22:35 +08:00
Kuris c7e8cab883 [Testing] Add Memory Leak Test (#1516)
* fix TL_ENABLE_PTXAS_VERBOSE_OUTPUT is not print in tvm-ffi

* add memory leak test in tvm-ffi

* fix lint error

* fix typos
2025-12-24 14:19:39 +08:00
Lei Wang 11f122e6a8 [Refactor] Phaseout execution_backend ctypes (#1510)
* Refactor execution backend options by removing 'ctypes' from the list of supported backends across multiple files. Update related documentation and tests to reflect this change, ensuring consistency in the autotuning and JIT compilation processes.

* minor fix
2025-12-23 19:52:16 +08:00
174fbe1908 [Enhancement] Refactor CUDA vectorized cast generation and remove unsupported FP8 type (#1474)
* Refactor CUDA vectorized cast generation and remove unsupported FP8 type

* test fix

* lint fix

* Refactor CUDA vectorized cast function naming for clarity

* Add support for float4_e2m1fn type conversions in CUDA vectorized casts

- Implemented conversions between float4_e2m1fn and float32, half2, and float2 in utils.cc and cuda_fp4.h.
- Updated test_tilelang_language_vectorized_cast.py to validate new conversions and ensure correctness.
- Enhanced dtype conversion in dtypes.py to handle float4_e2m1fn appropriately, logging a warning for unsupported types in PyTorch.

* Enhance vectorized cast tests for new data types

- Added tests for vectorized casting of float8 and float4 data types, ensuring compatibility with CUDA compute versions.
- Refactored existing test functions to improve clarity and organization, separating tests for different data types.
- Updated parameterization to include additional test cases for new conversions.

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2025-12-22 22:44:03 +08:00
Kuris 5acaab7605 [Language] Remove ConstIf Frame for better meta programming (#1496) 2025-12-22 17:20:26 +08:00
Lei Wang a874e4e823 [Refactor] Phaseout PassConfig kDisableDynamicTailSplit and kDynamicAlignment as they are legacy (#1486)
* [Cleanup] Remove dynamic shape example and related tests

* Deleted the dynamic shape example script `example_dynamic.py` and its corresponding test file `test_example_dynamic.py` to streamline the codebase.
* Removed unused dynamic tail split and dynamic alignment configurations from `builtin.h` and `pass_config.py`.
* Cleaned up the dynamic shape testing files to eliminate redundancy and improve maintainability.

* build fix
2025-12-21 16:12:24 +08:00
Lei Wang 7e8d1f8273 [Enhancement] Enhance let binding handling in layout inference and warp specialized pass (#1484)
* [Feature] Add FullyReplicated Fragment Layout and Enhance Layout Inference

* Introduced a new static method `FullyReplicated` in the `Fragment` class to create fully replicated fragment layouts, ensuring all threads hold identical copies of the buffer.
* Updated `CopyNode` to collect fragment layouts and mark them as fully replicated during layout inference.
* Enhanced `ParallelOpNode` to expand let bindings for fragment buffer accesses, improving layout inference accuracy.
* Added documentation for new methods and updated existing methods to support the new layout features.

* lint fix

* Remove debug logging statements from layout inference process to streamline output and improve performance.
2025-12-21 03:26:12 +08:00
silentCoder-dev 95e3b5a716 [Refactor] Remove triton dependence in testing & move triton baseline into examples (#1470)
* remove triton dependence in testing & move triton baseline into example

* use ceildiv and handles arbitrary M correctly for triton
2025-12-19 12:34:39 +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
silentCoder-dev f06726031e [Refactor] Rename test for curand & add triton baseline in test_tilelang_language_rand.py (#1464)
* rename test for curand & add triton baseline

* add a comment for calling T.rng_rand() four times

* refactor tilelang&triton kernel

* Add boundary checks for M not divisible by 128
2025-12-19 11:08:53 +08:00
7248a810d9 feat(cutedsl): add CuTeDSL backend (#1421)
* feat: CuTeDSL backend

* fix: clang-tidy

* fix: clang-format

* fix: ci

* fix: revert example gemm fp8

* fix: remove duplicate code

* fix: switch-case

* fix: fp16 silence

* fix: TVM IR print

* fix: useless tir

* fix: clang-format

* fix: remove tilelang/contrib/cutedsl/.gitignore

* fix: use hexfloat

* fix: gsym guard

* fix: unknown storage sync type

* fix: string literal

* fix: add args guard

* fix: name hint dedup

* fix: better find_kernel_by_pattern

* fix: set libpath for from_database path

* fix: guard buffer.strides

* fix: from guard

* fix: eviction guard

* fix: use thread local tma descs

* fix: ruff

* fix: drop tma_init_cpp

* fix: exc_info

* fix: negative unmatch early return

* fix: rename postproc func and add test

* fix: handle fast math according to pass config

* fix: dyn_sym parse

* fix: wrap_forward

* fix: use tvm_ffi.libinfo instead of cli

* fix: keep signature

* fix: C++ string safety

* fix: mark tma_store_add as unsupported

* fix: tvm version

* resolve ldsm and cpasync issues.

* fix: minor fixes

* fix: parse signature using ast

* fix: guard global_addr

* fix: create tempfile only when necessary

* fix: use logger.execption for exceptions

* fix: guard lib_path and host_func

* fix: remove tma_cpp_init and add timeout for cpp compile

* add timeout for mbarrier_wait.

* fix: _load_kernel_from_disk signature

* resolve codegen issues.

* fix: logger.exception

* add comment for div_by=1

* merge

* fix: reserve cutlass,cute,tl

* fix: guard tma_store

* fix: allow int64 offset in make_tensor_at_offset

* fix: guard barrier

* fix: add comments for div_by=16

* fix: div_by=1 issue

* delete div_by when offset is 0

* use tl.make_tensor when offset is 0

* fix: explicitly check cutedsl target

* fix: use param.torch_dtype()

---------

Co-authored-by: yuxic <yuxic@nvidia.com>
Co-authored-by: Yong <yong@local>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2025-12-18 11:55:18 +08:00
silentCoder-devandcoderabbitai[bot] cae06edd06 [Language]Adds a random number generation capability through curand_kernel (#1461)
* add curand.{curand_init, curand}

* run format.sh

* add default value for curand_init & add test for curand

* Update testing/python/language/test_rand.py

Remove unused thread binding

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

* remove unused library

* enable tilelang cache for testing

* run format.sh

* Revert "run format.sh"

This reverts commit 5afaff782f.

* Revert "enable tilelang cache for testing"

This reverts commit c277a43e77.

* Revert "remove unused library"

This reverts commit 568ad20611.

* run format.sh

* ensure FreshName for __philox_state

* ensure FreshName for __philox_state

* change the return type of T.rng_init

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
2025-12-18 11:03:20 +08:00
Kuris 91cf7966e7 [Analyzer] Fix floordiv & floormod bug in z3 prover (#1458)
* fix floordiv & floormod in z3 prover

* fix lint error
2025-12-17 18:29:52 +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
Lei Wang 0814b17129 [Language] Introduce T.annotate_restrict_buffers (#1428)
* [Enhancement] Introduce non-restrict parameter support in code generation

- Added a new PrimFunc-level attribute `tl.non_restrict_params` to specify handle Vars that should not be marked with the restrict qualifier during code generation.
- Updated `CodeGenTileLangCPP`, `CodeGenTileLangCUDA`, and `CodeGenTileLangHIP` to handle non-restrict parameters, ensuring proper treatment of overlapping buffer aliases.
- Implemented a new annotation function `annotate_restrict_buffers` to facilitate the marking of buffer parameters as non-restrict.
- Enhanced the `SplitHostDevice` transformation to propagate non-restrict parameters from host to device functions.
- Added a new transform function `HoistNonRestrictParams` to manage non-restrict parameters effectively.

* [Enhancement] Improve HoistNonRestrictParams transformation

- Updated the HoistNonRestrictParams function to recursively collect all `tl.non_restrict_params` annotations from nested blocks, enhancing flexibility in annotation placement.
- Introduced a new NonRestrictCollector class to manage the collection and deduplication of non-restrict parameters.
- Modified the SplitHostDevice transformation to remove the non-restrict attribute from the host-side PrimFunc after propagation to device kernels.
- Adjusted the LowerAndLegalize function to directly apply the HoistNonRestrictParams transformation without exception handling, streamlining the process.

* [Refactor] Simplify non-restrict parameter handling in code generation

- Removed unnecessary normalization logic and associated data structures from `CodeGenTileLangCPP`, `CodeGenTileLangCUDA`, and `CodeGenTileLangHIP`.
- Streamlined the handling of non-restrict parameters by directly inserting them into the `non_restrict` set, improving code clarity and maintainability.
- Updated conditional checks to eliminate redundant checks against normalized names, enhancing performance and readability.

* [Dependency] Update TVM subproject to latest commit 68aa8461

- Updated the TVM subproject to the latest commit, ensuring compatibility with recent changes and improvements.
- Refactored non-restrict parameter handling in `CodeGenTileLangCPP`, `CodeGenTileLangCUDA`, and `CodeGenTileLangHIP` to enhance code clarity and maintainability.
- Adjusted the `SplitHostDevice` transformation to streamline the propagation of non-restrict parameters.

* fix
2025-12-17 12:44:19 +08:00
Kuris 9c21586b54 [Feat] Integrate Z3 in TVM Arith Analyzer (#1367) 2025-12-17 11:22:55 +08:00
Kurisandcoderabbitai[bot] 0b6336b583 [Refactor] Use pytest.mark.parameterize to speedup parallel testing (#1447)
* Refactor GEMM tests to use parameterized pytest fixtures

- Converted multiple test cases for GEMM operations in `test_tilelang_tilelibrary_gemm_sp.py` to use `pytest.mark.parametrize` for better maintainability and readability.
- Similar refactoring applied to `test_tilelang_tilelibrary_gemm_sp_v2.py`, consolidating test cases for `run_gemm_ss`, `run_gemm_rs`, `run_gemm_sr`, and `run_gemm_rr` into parameterized tests.
- This change reduces code duplication and enhances the clarity of test configurations.

* Update testing/python/amd/test_tilelang_gemm_mfma_preshuffle.py

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

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
2025-12-16 15:17:32 +08:00
Lei WangandZhiwen Mo dda451269c [Refactor] Reduce direct dependency on PyTorch due to its limited type support (#1444)
* [Enhancement] Update KernelParam to use tvm.DataType directly and add torch_dtype conversion method

- Changed dtype in KernelParam from torch.dtype to tvm.DataType to support a wider range of data types and prevent information loss during conversions.
- Added a new method, torch_dtype, to convert tvm.DataType back to torch.dtype for tensor creation.
- Updated various adapters to utilize the new torch_dtype method for parameter type conversion during initialization.

* [Enhancement] Refactor CUDA type handling and add support for FP4 and FP8 types

- Renamed functions for clarity: GetFP8Type, GetFP6Type, and GetFP4Type are now GetTileLangFP8Type, GetTileLangFP6Type, and GetTileLangFP4Type respectively.
- Enhanced FP4 type handling to support additional lane sizes (2, 4, 8, 16, 32, 64).
- Updated CUDA code generation to include new FP8 and FP4 types, ensuring proper type handling in PrintType and related functions.
- Introduced new structures for FP8 types in cuda_fp8.h to facilitate better memory management and type packing.
- Added methods in KernelParam and tensor utilities to recognize and handle float4 types, improving compatibility with PyTorch.
- Enhanced logging for debugging purposes in various CUDA functions to track type handling and memory operations more effectively.

* lint fix

* Remove unnecessary logging statements from CUDA code generation and delete obsolete matrix multiplication test file.

* [Enhancement] Add support for FP4 and FP8 types in CUDA code generation

- Enhanced PrintVecElemLoad and PrintVecElemStore functions to handle new FP4 types.
- Updated arg_binder to allow float4 to match int8 at runtime, improving compatibility with PyTorch.
- Modified loop_vectorize to account for buffer dtype lanes in vectorization calculations.
- Refactored tensor type mapping to support new float4 and float8 types, ensuring correct type handling in tensor operations.
- Added tests for FP4 and FP8 copy operations to validate functionality and integration with existing workflows.

---------

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2025-12-16 14:26:07 +08:00
DayuxiaoshuiandLeiWang1999 869f021b2b [Feature] Support region as input of T.cumsum (#1426)
* [Feature] Support region as input of T.cumsum

- Extend T.cumsum to accept BufferRegion and BufferLoad inputs in addition to Buffer
- This enables operations on buffer slices/regions like:
  T.cumsum(InputG_fragment[i * chunk_size:(i + 1) * chunk_size], dim=0)
- Update cumsum_fragment to handle region inputs properly
- Add comprehensive tests for 1D and 2D region inputs including normal and reverse modes

Fixes #879

* Fix formatting and add docstring for cumsum_fragment

- Add comprehensive docstring for cumsum_fragment function
- Format code according to ruff style guidelines

* Fix CodeRabbit review issues

- Fix negative dimension bounds check (dim < -len(shape) instead of dim <= -len(shape))
- Add src/dst shape compatibility validation for out-of-place cumsum
- Update copy() type annotation to accept BufferRegion as dst parameter
- Fix test in-place mutation issues by using out-of-place cumsum operations
- Add non-divisible size test cases for tail region coverage

* Fix out-of-bounds access in region tests

- Add bounds clamping using T.min() for chunk_end calculations
- Prevents accessing beyond tensor bounds for non-divisible sizes
- Matches reference implementation behavior
- Fixes both 1D and 2D region test cases

* Fix region test: use simple slice expressions instead of T.min()

- Remove T.min() which cannot be used directly in slice indices
- Use chunk_start + chunk_size form instead
- Rely on system's automatic bounds checking for non-divisible sizes
- Update comments to reflect this approach

* Fix cumsum region: use region extents in lowering and update tests for shared memory

* Simplify fragment scope check using is_fragment()

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2025-12-15 20:08:45 +08:00
Xiangwen Wang bcae814ed1 Enhance vectorized conversion support (#1438) 2025-12-15 20:07:58 +08:00
Lei Wang 89521e632e [Refactor] Phase out the primitives folder since its design has been merged into tileop (#1429)
* Phase out primitives

* revert changes

* Refactor GemmWarpPolicy method signature for clarity

Updated the `from_warp_partition` method in the `GemmWarpPolicy` class to return the type `GemmWarpPolicy` instead of a string, enhancing type safety and clarity in the codebase. Removed an unnecessary blank line for improved readability.

* fix
2025-12-15 11:39:04 +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
Xiangwen Wang e84b24bc6a [Enhancement] Improve vectorization invariant check (#1398)
* Improve loop vectorize

* Improve loop vectorize

* Improve loop vectorize

* Improve loop vectorize

* Improve loop vectorize

* Add some vectorize tests and comments
2025-12-12 17:19:33 +08:00
Lei Wang 6f67da84bc [Enhancement] Introduce T.__ldg (#1414)
* [Enhancement] Add __ldg intrinsic for CUDA read-only cache loads

* Introduced the __ldg intrinsic to enable explicit read-only cached loads from global memory in CUDA.
* Updated the corresponding documentation and added support in both CUDA and HIP code generation.
* Enhanced the Python interface for __ldg to accept BufferLoad and Buffer types, improving usability.

* [Enhancement] Update formatting and linting rules in pyproject.toml; minor test adjustment

* Added new formatting rules in pyproject.toml to enforce consistent code style, including hanging indents and argument splitting.
* Updated test_tilelang_language_intrinsics_codegen.py to improve readability by adding a blank line before the main execution block.
* Refactored error messages in builtin.py for better clarity and consistency, ensuring proper formatting in function definitions and raising ValueErrors.

* lint fix
2025-12-12 16:30:56 +08:00