mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
056ccfc267f03498d8dcc79f6efd2aba2bae8abd
11
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Lei Wang
|
a797e5120f |
[Refactor] Move backend-specific GEMM implementations and transforms into backend directories (#2165)
* [Refactor] Move backend-specific GEMM implementations and transforms into backend directories Restructure the codebase so that each backend (cpu, cuda, rocm) owns its GEMM implementations, sparse GEMM implementations, and transform passes under a consistent op/ subdirectory layout. The shared GEMM registry and base classes remain in tileop/ as the platform-agnostic dispatch layer. - Move gemm/gemm_sp registries from backend/ into tileop/ as registry.py - Move CUDA GEMM impls (mma, mma_sm70, wgmma, tcgen05) into backend/cuda/op/gemm/ - Move CUDA sparse GEMM impl into backend/cuda/op/gemm_sp/ - Move CPU GEMM impl (scalar) into backend/cpu/op/gemm/ - Move ROCm GEMM impls (mfma, wmma) into backend/rocm/op/gemm/ - Move CUDA-specific transform passes from src/transform/ into src/backend/cuda/transform/ - Move CUDA runtime sources from src/runtime/ into src/backend/cuda/ - Remove dead backend-importing wrappers from transform/__init__.py - Update phase.py and tests to import CUDA transforms from their canonical location Each backend now has a symmetric op/ directory structure. Adding a new backend no longer requires modifying shared transform or tileop modules. * Refactor Python backend package layout * Remove deprecated intrinsic implementations and related tests for matrix multiplication and element-wise addition. This includes the deletion of files for `benchmark_matmul_intrinsic.py`, `example_tilelang_gemm_fp8_intrinsic.py`, and associated test files, streamlining the codebase by eliminating unused components. * Move CUDA transform passes back to common transform |
||
|
Lei WangandZhiwen Mo
|
f8dc61c02c |
test: reduce CI runtime for slow Python suites (#1932)
* test: trim CI runtime for slow Python suites * remove useless test * Update TVM submodule to latest commit and remove manual device assert test from CI. Replace with a call to tilelang.testing.main(). * lint fix * fix * revert test for eager jit --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei Wang
|
4929ad89cc |
[Refactor] Improve cp.async lowering and add async_copy op (#1887)
* Add CPAsync copy selection/lowering and vectorized ptx_cp_async folding
* [transform] Rename SimplifyCPAsyncSync to OptimizeCPAsyncSync
* ThreadSync: treat ptx_cp_async as async-copy
* Implement async_copy operator and enhance cp.async handling
- Introduced the async_copy operator for explicit asynchronous global-to-shared memory copies, ensuring no automatic waits are inserted.
- Updated the copy operator to normalize copy regions and handle both synchronous and asynchronous copy semantics.
- Enhanced pipeline planning to recognize explicit cp.async calls and bind commit and wait operations to the appropriate stages.
- Added tests to validate the behavior of async_copy and its integration with pipeline planning.
This commit improves the handling of asynchronous memory operations and clarifies the semantics of copy operations in the context of the new async_copy functionality.
* docs: clarify T.copy vs T.async_copy
* style: apply format.sh fixes
* async_copy: improve errors and defer cp.async check
* test: check T.copy vs T.async_copy codegen equivalence on SM80
* copy: allow predicated cp.async for OOB global->shared
* pipeline: schedule cp.async commits with cp.async groups
* lower: preserve manual pipeline statement arity
* pipeline: skip commit-order check when fused with cp.async
* pipeline: enforce cp.async ordering in PipelinePlanning
PipelinePlanning may float pure control intrinsics like ptx_commit_group earlier because they have no buffer regions, which can generate commit-before-cp_async.\n\nAdd an order-repair step that models cp.async group semantics (cp_async* -> commit, commit(group_i) -> cp_async(group_{i+1})) and wait->consumer constraints, then applies a stable topological sort to produce a valid software_pipeline_order.\n\nAlso adds a regression test to assert cp.async is ordered before commit when split into separate blocks.
* Enhance cp.async synchronization handling in pipeline planning
- Introduced logic to relax loop wait conditions and split epilogue waits for cp.async operations, improving overlap between consumer phases.
- Added a new method to identify and manage consumer phases based on async memory reads, ensuring correct synchronization.
- Updated tests to validate the new behavior, confirming that wait groups are correctly relaxed and inserted as needed in the optimized pipeline.
This commit refines the handling of cp.async synchronization, enhancing performance and correctness in the generated code.
* lint fix
* Refactor cp.async copy handling and improve vectorization logic
- Removed the unused IsCPAsyncElementBytes function to streamline the code.
- Enhanced the CPAsyncStoreRewriter class with a new method to track vectorized lanes during statement rewriting.
- Updated CheckCPAsyncCopy to remove the buffer_oob parameter, simplifying its signature.
- Added logging for vectorized loop iterations to aid in debugging and performance analysis.
- Improved error handling in LowerCPAsyncCopy to provide clearer warnings when cp.async rewrites fail.
These changes aim to optimize the cp.async copy process and improve code maintainability.
* lint fix
* Update cp.async copy handling to improve rewrite success tracking
- Set the successfully_rewritten_ flag to true when a copy operation is successfully rewritten.
- Adjust the flag to false in cases where the load is not global-like or the buffer is not shared.
- Removed logging for vectorized loop iterations to streamline the code.
These changes enhance the tracking of rewrite success in the cp.async copy process, contributing to better error handling and maintainability.
* fix
* test: cover commit-merge + wait relax
* fix: skip wait relaxation under conditional async copy
* test fix
* Implement async copy support in TileLang
- Introduced a new configuration option `tl.enable_async_copy` to control the lowering of global to shared memory copies to PTX `cp.async`.
- Updated the `CopyNode::LowerCPAsyncCopy` method to handle async copy semantics based on the new configuration.
- Enhanced the `InjectPTXAsyncCopy` pass to rewrite eligible global->shared copies into `cp.async`, including necessary commit and wait statements.
- Added tests to verify the correct behavior of async copy lowering and the impact of the configuration option.
These changes improve the flexibility and performance of memory copy operations in TileLang.
* lint fix
* Refactor async copy handling in TileLang
- Renamed `InjectPTXAsyncCopy` to `LowerPTXAsyncCopy` to better reflect its functionality.
- Updated the `OptimizeForTarget` function to use `LowerPTXAsyncCopy` instead of the deprecated `InjectPTXAsyncCopy`.
- Added a deprecated alias for `InjectPTXAsyncCopy` to maintain backward compatibility.
- Adjusted documentation to clarify the purpose of the new function name.
These changes enhance clarity and maintainability in the async copy transformation process.
* fix
* lint fix
* test fix
* Enhance PTX async copy integration in TileLang
- Added the `InjectPTXAsyncCopy` function to facilitate the injection of PTX `cp.async` operations into the lowering process.
- Updated the `LowerTileOp` function to conditionally apply `InjectPTXAsyncCopy` during parallel loop lowering when async copy is enabled.
- Removed deprecated calls to `LowerPTXAsyncCopy` in the `OptimizeForTarget` function to streamline the optimization process.
These changes improve the handling of async copy operations, ensuring better performance and compatibility with CUDA targets.
* lint fix
* fix
* fix
* fix
* set parallisim to 1 to test the ci
* fix
* Refactor async copy handling in PTX lowering
- Improved the `PrepareCopyIndexInfo` function to clarify the calculation of transfer bytes and ensure valid async copy operations.
- Updated synchronization state management for injected `cp.async` transfers, enhancing clarity and consistency in handling open copy regions and uncommitted transfers.
- Added helper functions for zero value checks and linear offset flattening to streamline index processing.
These changes enhance the robustness and readability of the async copy implementation.
* Refactor nested pipeline functions and update test cases
- Removed unused parameters `order` and `stage` from several functions, simplifying their signatures.
- Updated nested pipeline implementations to use `num_stages=2` for better clarity and performance.
- Adjusted test cases to reflect the changes in function signatures, ensuring they run correctly without the removed parameters.
These modifications enhance code readability and maintainability while ensuring the tests remain functional.
* Refactor async copy statement generation in PTX lowering
- Updated the `MakeCPAsyncStmt` function to `MakeCPAsyncStmtFromLoads`, improving the handling of base indices for vectorized accesses.
- Introduced `ExtractVectorBaseIndices` to streamline the extraction of base indices from vectorized loads and stores.
- Enhanced error handling for predicated async copy operations when base indices cannot be extracted.
These changes improve the clarity and robustness of async copy statement generation in the PTX lowering process.
* lint fix
* Enhance async copy operation handling in CodeGenTileLangCuTeDSL
- Added support for conditional async copy operations in the `VisitExpr_` method, allowing for an optional predicate argument.
- Improved error checking to ensure the correct number of arguments for `tl::ptx_cp_async`.
- Streamlined the generation of async copy statements for both standard and conditional cases.
These updates improve the flexibility and robustness of async copy operations in the code generation process.ix
* lint fix
* disable auto cp async without pipeline.
* lint fix
* Enhance cp.async synchronization handling in TileLang
- Introduced `ThreadSyncAfterWaitGroupInserter` to ensure correct insertion of storage synchronization after `ptx_wait_group` statements, addressing potential visibility issues in shared memory.
- Updated `WarpSpecialized` to utilize `tl.ptx_cp_async_barrier_noinc` for releasing producer-consumer dependencies, preventing incomplete memory observations.
- Added tests to verify the correct insertion of synchronization and the handling of cp.async barriers under warp specialization.
These changes improve the correctness and reliability of async copy operations in the TileLang transformation pipeline.
* fix
* fix
* test fix
* Refactor sparse_mla_fwd_pipelined.py and enhance phase.py synchronization
- Removed unused post-processing function and related code from `sparse_mla_fwd_pipelined.py` to streamline the implementation.
- Added debug print statements in `phase.py` to log the state of the module before and after thread synchronization, aiding in troubleshooting and verification of synchronization behavior.
* Enhance async copy handling and loop synchronization in TileLang
- Introduced new loop annotations for async copy operations, including `parallel_prefer_async` and `parallel_async_without_async_commit_wait`, to improve control over async behavior in parallel loops.
- Updated the `InjectPTXAsyncCopy` function to accommodate the new annotations, ensuring correct async copy injection based on user-defined preferences.
- Enhanced the pipeline planning logic to prioritize cp.async groups by their earliest consumers, improving scheduling efficiency and reducing wait times.
- Added tests to validate the new async copy behaviors and ensure correct synchronization in pipelined loops.
These changes enhance the flexibility and performance of async copy operations in the TileLang transformation pipeline.
* lint fix
* Add Regression Perf for mhc
* clean
* clean
* lint fix
* lint fix
|
||
|
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 |
||
|
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 |
||
|
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 |
||
|
Chaofan Lin
|
dff10e5283 | [BugFix] Complete vectorized loading for common dtypes (#1536) | ||
|
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.
|
||
|
Lei Wang
|
29051439db | [Lint] Phaseout Yapf format and embrace ruff format (#1417) | ||
|
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 |
||
|
LJC00118andpre-commit-ci[bot]
|
35cf8885ef |
[Enhancement] Remove constraint requiring last dimension stride to be 1 (#1040)
* remove last dimension stride must be 1 constraint * add vectorize test * minor fix * [Lint]: [pre-commit.ci] auto fixes [...] --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> |