* Fix thread storage synchronization logic in `thread_storage_sync.cc` to correctly identify conflicts between read and write operations based on loop carry conditions.
* lint fix
* Refactor `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to use shared memory for `sorted_token_ids` instead of local memory, improving thread synchronization. Adjust default argument values for M, N, and K in the main function for better testing scenarios.
* Add UniformExprChecker to enforce thread synchronization rules
Introduce the UniformExprChecker class to determine if expressions are uniform across threads, crucial for safe synchronization in conditional statements. Update the TileLangThreadSyncPlanner to hoist synchronization points out of non-uniform if-statements to prevent potential deadlocks. Enhance tests to validate sync hoisting behavior for various non-uniform conditions involving thread indices and shared memory access.
* lint fix
* Enhance `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` with cache disabling and kernel source printing for debugging. Update thread synchronization logic in `thread_storage_sync.cc` to check for runtime-dependent conditions, preventing potential deadlocks by hoisting sync points as necessary.
* Update submodule `tvm` to latest commit and remove deprecated `example_gqa_decode_varlen_logits_paged.py` file. Refactor `example_gqa_decode_varlen_logits.py` to enhance performance and maintainability by removing unused imports and optimizing shared memory usage. Adjust test cases to reflect the removal of the paged example.
* fix
* Enhance thread synchronization logic in `thread_storage_sync.cc` by adding a configurable warp size parameter to `RuntimeDependentConditionChecker` and `TileLangThreadSyncPlanner`. This allows for better adaptability to different target architectures. Update the logic to ensure thread extent is a constant and improve handling of runtime-dependent conditions.
* lint fix
* Refactor thread extent validation in `thread_storage_sync.cc` to use pointer checks instead of optional values. This change improves clarity and ensures that the thread extent is correctly validated as a constant.
* Adjust loop variable constraints in `thread_storage_sync.cc` for loop-carry analysis by modifying the extent calculation. This change ensures valid iteration comparisons by reducing the extent by one, allowing for accurate analysis of loop iterations.
* lint fix
* Fix type annotations for T.reshape and T.view
* Fix issue #1666: reduce_sum with clear=False not accumulating correctly
* address comments and add testcases
* add more tests
---------
Co-authored-by: SiriusNEO <chaofan@deepseek.com>
* support tcgen5mma lower for int8/uint8->int32
* add test
* add template for u8 and fix a bug in codegen
* lint
* fix ai suggestions
* add testing to maint
[Feature] Enhance Simplify Pass with Let Inline and Common Subexpression Elimination
- Added a new configuration option `tilelang_enable_simplify_let_inline` to enable inlining of let statements during simplification.
- Updated the Simplify pass to conditionally inline let statements based on the new configuration.
- Integrated common subexpression elimination into the device code generation process to optimize performance.
- Enhanced test coverage for the Simplify pass, including various scenarios to validate simplification behavior.
* [Feature] Add support for predicated loads and stores in LDGSTG lowering
- Introduced a new configuration option to disable loop unswitching.
- Enhanced the lowering logic to ensure that loads within predicated stores are also treated as predicated, preventing out-of-bounds memory access.
- Added a test case to verify the correct behavior of predicated loads and stores in the transformation pipeline.
- Included debug print statements to visualize the module state before and after the lowering process.
* lint fix
- Updated the loop unswitching pass to support hoisting Let-bound variables alongside their conditions, improving optimization for loops with variable dependencies.
- Modified the CallCheckerExcludingIf class to exclude matching If nodes when checking for CallNode presence, ensuring safety during unswitching.
- Enhanced IfBranchReplacer to replace all instances of If nodes with matching conditions, streamlining the transformation process.
- Added comprehensive tests to validate the new behavior, including scenarios with multiple Let-bound variables and identical conditions across If statements.
* [Feature] Implement Hoist Loop-Invariant If Pass
- Added a new pass to hoist loop-invariant if statements out of loops, improving optimization opportunities.
- Introduced classes for collecting and checking conditions, as well as for rewriting statements.
- Integrated the new pass into the optimization pipeline and provided a corresponding API for usage.
* Add tilelang semantics guide to programming guides (#1740)
* Add CallNodeChecker to Loop Unswitching Pass and integrate tilelang cache disabling in benchmark script
- Introduced CallNodeChecker class to identify CallNode expressions in loop conditions, enhancing loop-invariant checks.
- Updated IsLoopInvariant function to reject conditions containing CallNodes, preventing potential side effects.
- Added tilelang.disable_cache() in benchmark_mha_sink_fwd.py to optimize performance during benchmarking.
* lint fix
* lint fix
* Enhance Loop Unswitching Pass with Let Binding Support
- Added support for Let-bound variables in the WrittenBufferReadChecker to improve buffer read checks.
- Introduced UsesLoopVarThroughLetBindings function to check if conditions depend on loop variables through Let bindings.
- Updated IsLoopInvariant function to account for Let bindings when determining loop invariance.
- Enhanced HoistableIfFinder to track Let bindings for variables bound to BufferLoad expressions.
- Added debug print statements in the OptimizeForTarget function to visualize the module state before and after loop unswitching.
* fix
* lint fix
* Enhance Loop Unswitching with CallNode Exclusion Check
- Added CallCheckerExcludingIf class to ensure function calls outside of hoisted if statements are identified, preventing potential synchronization issues during loop unswitching.
- Updated loop unswitching logic to incorporate the new call checker, enhancing safety and correctness.
- Integrated debug print statements in OptimizeForTarget to visualize module state before and after loop unswitching.
- Disabled tilelang cache in the benchmark script for improved performance.
* Update TVM subproject to latest commit and remove debug print statements in OptimizeForTarget function
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [ThreadSync] Use Z3 for constraint equivalence checking
Replace const_int_bound based range comparison with Z3 SMT solver
for checking thread constraint equivalence in FindConflict.
* [ThreadSync] Clarify conflict detection for constraint equivalence
Added comments to improve understanding of the logic determining whether constraints are equivalent or in conflict in the thread storage synchronization process.
* add test
* Update test_issue_1106 to check syncthreads is outside for loop
* test fix
* [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.
* [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
* 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>
* [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
[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.
* [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
* [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>
[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>
* 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>
* 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
* [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
* 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.
* [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
* 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
* 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
* [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
* [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
* 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
- 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>
* 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.
* [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
* 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>
* 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>
* 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
Lei Wang