mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
2d8d3676eda18bd3d8e6fa783399ff96d3cd4ded
8
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Kuris
|
c3c8831db6 |
[EagerJIT] Add Support for Parameter Only Kernel Compilation (#1664)
* [Fix] Refactor type hint extraction logic in DSLMutator for better clarity and handling of annotations * [Refactor] Remove redundant tensor creation in loop layout tests and update kernel compilation parameters |
||
|
Lei Wang
|
6e439530b5 |
[Refactor] Unify @jit and @lazy_jit into a single @jit decorator (#1632)
* Refactor: Replace @tilelang.lazy_jit with @tilelang.jit in examples and tests Updated all instances of the @tilelang.lazy_jit decorator to @tilelang.jit in the lazyjit example notebooks and related test files. This change aligns the code with the new JIT compilation approach, enhancing consistency across the codebase. Additionally, removed the lazy_jit import from the module initialization to streamline the API. * Enhancement: Improve layout annotation handling in T.copy and parallel loops Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify layout hints for the outermost parallel loop. Enhanced layout annotation handling in CopyNode and AtomicAddNode classes to ensure compatibility with SIMT operations. Added tests to validate the functionality of loop layout annotations, improving robustness and clarity in layout management for parallel loops. * Add out_idx parameter to @jit and validate it's only used in lazy mode * lintfix * Refactor: centralize mode inference in JITImpl.__call__ before parse_args * test fix * Enhancement: Introduce JITNoBuilderError for eager mode checks in T.Kernel and T.const functions. Ensure proper error handling when Builder is unavailable, improving robustness in JIT context management. * lint fix * lint fix * Enhancement: Update example_triton_sparse_gqa_decode_varlen_indice.py to improve performance metrics output by adding average FLOPS calculations. Modify layout_inference.cc to refine buffer layout inference checks by ensuring fragment buffers are considered. * lint fix * Enhancement: Refactor JIT implementation to centralize mode initialization and improve error handling for eager mode. Update matmul tests to utilize both Cython and TVM FFI JIT kernels, ensuring consistency in output validation. * Refactor: Remove unused simplify_prim_func decorator from example_gemm_intrinsics.py and related test files. Enhance layout.cc to simplify integer set analysis for multi-threaded compilation, ensuring constant integer outputs. Update lower_tile_op.cc to enforce constant integer checks in layout output shapes. Improve error handling in builder.py for JIT context management. * remove claude.md * Refactor: Clean up code formatting in layout.cc, lower_tile_op.cc, and proxy.py for improved readability. Add missing import in test_tilelang_kernel_int4_gemm_mma.py and ensure consistent spacing in builder.py. * Refactor: Replace logical conditions with tir_all in complex expressions within test_arith_hard.py for improved clarity and consistency. Update main execution flow to call test_hard_prove instead of tilelang.testing.main. * Enhancement: Add @tilelang.testing.requires_cuda decorator to test functions in test_tilelang_issue_1549.py and test_tilelang_issue_1601.py to ensure CUDA requirements are explicitly stated for relevant tests. * lint fix * Refactor: Simplify output shape calculation in layout.cc by directly setting max value from ist.max(). Update test_valid_loop in test_tilelang_fragment_loop_checker.py to disable certain tests and ensure proper execution flow. Enhance error handling in builder.py to correctly manage PrimFunc return types. * Refactor: Update boolean condition in test_boolop to use Or and Not for improved clarity and consistency in logical expressions. * Refactor: Update test execution flow in test_tilelang_fragment_loop_checker.py to call tilelang.testing.main. Modify JIT implementation to handle callable functions more generically. Enhance TirTemplate class in builder.py to return an empty tuple if matcher is None, improving robustness. * Refactor JIT function handling: Replace LazyJITFunc with JITFunc across the codebase, updating related function signatures and logic to support eager JIT compilation mode. This change enhances the clarity and functionality of JIT operations. * Update JITImpl documentation: Clarify output tensor index attribute description and remove outdated comment regarding implicit TIR construction. * Refactor TileLang to transition from v2 to eager language module. Update imports across multiple files to replace references to the deprecated v2 module with the new eager module, ensuring compatibility with the latest JIT compilation features. Remove obsolete v2 module files to streamline the codebase. * Enhance JIT compilation support: Update JIT function handling to improve compatibility with eager execution. Refactor related function signatures and logic, ensuring seamless integration with the latest compilation features. This change aims to streamline JIT operations and enhance overall code clarity. |
||
|
Lei Wang
|
946f611c13 | [Typo] Rename loop layout annotation test(#1596) | ||
|
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. |
||
|
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 |
||
|
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) | ||
|
Lei Wang
|
2957afcabf |
[Enhancement] Improve iterator handling in layout utilities and parallel operations (#1221)
* [Enhancement] Improve iterator handling in layout utilities and parallel operations * Added a new function, DivideUnusedIterators, to detect per-iterator gaps in fused index expressions, enhancing the accuracy of unused iterator detection. * Updated CompleteBufferFragment to prefer direct inversion for bijective index mappings and introduced a fallback mechanism for non-bijective cases, improving layout inversion robustness. * Added a new test for layout inference in fused kernels to ensure correct compilation and execution without layout inversion failures. * lint fix |