Commit Graph
13 Commits
Author SHA1 Message Date
Tong WUandLeiWang1999 228c7c0486 [Feature] Support iket profiler for CUDA backend (#2515)
* Add experimental IKET CUDA backend hooks

* Keep IKET APIs under namespace

* Use real Perfetto screenshot for IKET example

* Show IKET events in Perfetto timeline screenshot

* Add IKET validation overhead example

* Remove IKET validation overhead example

* Fix IKET lint issues

* Move IKET integration to CUDA tools

* Add tools documentation and fix workflow regressions

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-07-12 20:03:30 +08:00
b939fa0137 [TIR][IR] Update to use tirx (#2216)
* [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

* Add support for disabling loop unswitching optimization

- Introduced a new configuration option `tl.disable_loop_unswitching` to allow users to disable the loop unswitching optimization.
- Updated the Loop Unswitching pass to check this configuration and return the original function if the option is enabled.
- Added relevant documentation in the PassConfigKey enumeration for clarity.

* update tvm

* Refactor: Update TileLang to use tirx instead of tir

- Removed references to `tvm.tir` and replaced them with `tvm.tirx` across various files, including examples and backend operations.
- Updated target configuration in documentation to reflect the new usage of target config dictionaries instead of CLI-style strings.
- Cleaned up `pyproject.toml` and `load_tvm.cmake` by removing obsolete paths.
- Enhanced examples for dequantization and GEMM operations to utilize the new `tirx` constructs.
- Adjusted various backend operations to ensure compatibility with the new `tirx` namespace.

* pass fix

* Update submodule 'tvm' to latest commit f24c7b19

* Apply formatting fixes

* Update CMakeLists.txt to set HIDE_PRIVATE_SYMBOLS for TileLang and modify target_host initialization in lower.py for compatibility with TVM's target handling.

* Refactor CMakeLists.txt and kernel_cache.py to update target names and library dependencies for TVM integration. Adjusted output targets and library names for better compatibility across platforms.

* Update repair-wheel-command in pyproject.toml to exclude libz3.so with wildcard for better compatibility with dependency packages.

* Fix macOS Z3 dylib lookup

* Move macOS Z3 install-name fix to install step

* Update TVM submodule for runtime import fix

* Update ROCm backend for TVM module factory changes

* Fix build-tree Z3 runtime lookup

* Fix macOS build-tree Z3 lookup

* Support legacy target option strings

* Normalize CUDA target in transform test

* Qualify Metal stream FFI lookup

* Fix PyPI Z3 packaging links

* Update TVM for Windows DLL lookup

* Register Windows dependency DLL directories

* Bump tvm-ffi to match tirx

* tvm update

* Apply lint fixes

* Enhance reduction operations with support for nan propagation and vectorization. Introduce new functions for preferred vectorized size and update existing reduction logic to handle packed operations for bfloat16 and float16 types. Add nan-aware min and max operations in CUDA and ROCm backends, and update related tests to validate functionality.

* Refactor ReduceLowerer: Remove unnecessary buffer allocation logic for packed operations in reduce.h

* Normalize PyPI Z3 dylib before linking

* Drop legacy target option string support

* Fix type reference for barrier annotations in CUDA copy operations. Updated `BufferLoadNode` to `tirx::BufferLoadNode` in `GetBarrier` and `LowerCluster` methods to ensure compatibility with recent changes in the TIR API.

---------

Co-authored-by: Yu Cheng <yu.cheng@pku.edu.cn>
Co-authored-by: Yichen Yan <oraluben@outlook.com>
2026-05-20 15:38:13 +08:00
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
2026-05-08 17:11:34 +08:00
Lei Wang d85e0c6e96 [Enhancement] Improve plot_layout visualization for Layouts (#1811)
* Enhance plot_layout function to support both Fragment and Layout types for visualization. Update parameters for colormap and formats, and introduce helper functions for format parsing and saving plots. Improve documentation for clarity on usage and expected input types.

* lint fix

* Refactor swizzle layout functions to use dedicated layout creators. Replace inline 2D swizzle functions with calls to `make_full_bank_swizzled_layout`, `make_half_bank_swizzled_layout`, and `make_quarter_bank_swizzled_layout` for improved clarity and maintainability in layout generation.

* Remove outdated documentation from layout_swizzle.py and ensure plots are closed after saving in plot_layout.py for better resource management.
2026-02-08 22:37:47 +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 29051439db [Lint] Phaseout Yapf format and embrace ruff format (#1417) 2025-12-12 17:22:18 +08:00
Jiaxing Ding 60567ba3b2 [AMD] Supoort T.gemm_v2 for AMD Backend (#1136) 2025-10-28 13:48:01 +08:00
Jiaxing Ding 6c67a77f2b [Layout] fix plot layout (#890) 2025-09-29 11:29:33 +08:00
Lei Wang f2e99180d5 [Refactor] Phaseout LLVM Dependency by Making it Optional (#247)
* remove llvm build

* [Refactor] Update kernel compilation and profiling in examples

- Replaced `tilelang.lower` with `tilelang.compile` in multiple example scripts to streamline kernel compilation.
- Updated profiling calls to utilize the new `get_profiler` method, enhancing performance measurement consistency.
- Adjusted assertions and benchmarking methods to align with the new profiling structure across various examples, ensuring correctness and clarity in performance evaluations.

* lint fix

* License Update

* [Refactor] Improve code formatting and documentation in CUDA header and HIP runtime files

- Adjusted formatting in `cuda.h` for better readability, including alignment of comments and struct fields.
- Cleaned up whitespace and improved comment clarity in `rt_mod_hip.cc` to enhance code maintainability.

* [Refactor] Enhance formatting and clarity in CUDA header and HIP runtime files

- Improved comment alignment and readability in `cuda.h`.
- Cleaned up whitespace and formatting in `rt_mod_hip.cc` to enhance maintainability.

* lint fix

* lint fix

* lint fix

* lint fix

* fix

* License update

* [Enhancement] Update JITKernel to use artifact for kernel source

- Assigned the generated artifact to `self.artifact` for better management.
- Updated kernel source references to use `artifact.kernel_source` for consistency in execution backend handling.

* lint fix

* Add @tilelang.testing.requires_llvm decorator to vectorization tests

* Enhance setup.py and env.py for library management

- Added functionality to remove original files after copying in CMakeBuild.
- Updated TVM_LIBRARY_PATH in env.py to include the PyPI build library path for better integration.

* Refactor TVM_LIBRARY_PATH assignment for improved readability in env.py

* Refactor CMakeBuild file handling in setup.py

- Added a check to ensure the target library directory exists before copying .so files.
- Improved the logic for creating the target directory and copying files to enhance robustness.

* bugfix

* Rename BuildTLDebug to BuildTileLangCUDAWithoutCompile and update registration. Add @tilelang.testing.requires_llvm decorator to multiple tests for LLVM requirement.

* lint fix

* Enhance TileLang code generation by adding support for device code generation without compilation. Updated `host_codegen` and `device_codegen` functions to include new transformations and registration for `tilelang_hip_without_compile`. Refactored JIT kernel adapters to accommodate host and device modules, improving overall integration and flexibility.

* lint fix

* Add support for C target in device code generation

- Updated `device_codegen_without_compile` to include handling for the C target by registering the `tilelang_cpp` function.

* [Enhancement] Implement auto-clear cache feature based on environment variable

* Added TILELANG_CLEAR_CACHE environment variable to control cache clearing.
* Updated CI workflow to set TILELANG_CLEAR_CACHE during testing.
* Modified cache initialization to clear cache if TILELANG_CLEAR_CACHE is set to true.

* [Refactor] Update kernel invocation and import paths in tests and cache

* Changed kernel invocation in `test_tilelang_kernel_dequantize_gemm.py` to return the result.
* Updated import statements in `test_tilelang_kernel_int4_gemm_mma.py` to use `bitblas` instead of `tilelang`.
* Refactored paths for artifact and parameters in `kernel_cache.py` for better maintainability.

* [Refactor] Clean up whitespace and improve code formatting in kernel_cache.py

* Removed unnecessary blank lines and adjusted spacing for better readability in the KernelCache class.
* Enhanced overall code formatting to align with project standards.

* [Enhancement] Add bfloat16 test case and improve kernel caching logic

* Introduced a new test case for bfloat16 matrix multiplication in `test_tilelang_kernel_gemm_mma_intrinsic.py`.
* Updated `KernelCache` to handle multiple kernel source files and improve error handling during saving and loading.
* Refactored `JITKernel` to support instantiation from a database, enhancing flexibility in kernel management.
* Adjusted `CtypesKernelAdapter` and `CythonKernelAdapter` to utilize the new kernel loading mechanism from the database.
* Improved code formatting and readability across several files.

* lint fix

* Update bfloat16 matrix multiplication test case to use larger dimensions for improved coverage
2025-03-20 20:08:20 +08:00
66RING 6ab29ffc82 Update expired example code. (#196)
Expired example code, update readme.
2025-03-12 20:52:26 +08:00
Yu Cheng e89e8b6cec [Doc] Add MLA Decoding Performance Benchmarks and Documentation (#137)
- Update news and MLA performance benchmark in README.md
- Move performance benchmark and layout images to a dedicated 'figures' directory
- Improve code formatting and image references in documentation
2025-03-04 13:56:03 +08:00
Lei Wang e32311b2db [Debug] Improve Memory Layout Plot (#136)
* Change default log level from WARNING to INFO in TileLang initialization

* Refactor Flash Attention Variable-Length MHA Example with Cython Backend Support

- Update `example_mha_fwd_varlen.py` to use Cython backend for kernel compilation
- Remove unused imports and simplify function signature
- Modify `flashattn` function to handle max sequence length as a separate argument
- Update kernel call to include max sequence length parameter
- Improve code readability and remove commented-out code
- Add print statement to confirm successful assertion

* Refactor code formatting in TileLang lowering and example files

- Improve line breaks and code formatting in `lower.py`, `wrapper.py`, and `tensor.py`
- Simplify line breaks and reduce unnecessary whitespace
- Enhance code readability by adjusting indentation and line breaks
- Update example MHA forward pass script with cleaner tensor initialization

* Update TileLang kernel test with import path changes for MMA layout and macro generator

- Modify import statements in test_tilelang_kernel_dequantize_gemm.py
- Replace bitblas imports with tilelang.intrinsics imports for MMA-related utilities
- Update main function to use tilelang.testing.main()

* Add Block Sparse Attention Examples for TileLang and Triton

- Implement block sparse attention kernels for both TileLang and Triton
- Add utility functions for generating sparse attention masks using top-k and threshold methods
- Support causal and variable-length attention scenarios
- Include test cases for different sequence length configurations
- Demonstrate block-level sparse attention with configurable parameters

* Refactor Block Sparse Attention Examples with Code Style Improvements

- Improve code formatting in block_sparse_attn_tilelang.py and block_sparse_attn_triton.py
- Enhance readability by adjusting line breaks and indentation
- Simplify kernel and function calls with better formatting
- Add whitespace and line break improvements for better code clarity

* Enhance Layout Plotting with Multi-Replication and Dynamic Visualization

- Update plot_layout function to support multiple replications in thread and value mapping
- Improve thread and value mapping to handle replicated layouts
- Dynamically adjust figure size and legend positioning
- Add print statements for saved plot file paths
- Modify example fragment_mma_load_a.py to uncomment and enable warp and block layout plotting
2025-03-04 02:10:25 +08:00
Lei Wang f9b6a92ee8 [Tools] Introduce plot_layout to visualize the fragment layout (#68)
* [Enhancement] Add VectorizeLoop function and update imports for compatibility

* [CI][Test] Improve test cases for vectorization and fix typos in parser comments

* lint fix

* Fix incorrect module reference for VectorizeLoop transformation

* Refactor vectorize_loop transformation by removing unused extent mutation logic

* [Enhancement] Add support for FP8 data types and global barriers in CUDA codegen

* Fix formatting in CUDA FP8 header file for consistency

* Refactor CI workflow to use 'tilelang_ci' virtual environment and update CUDA type printing for better clarity

* Update submodule 'tvm' to latest commit for improved functionality

* Refactor execution backend references from 'dl_pack' to 'dlpack' for consistency and clarity; add apply_simplify function to simplify PrimFunc or IRModule.

* Refactor CUDA code for improved readability; clean up formatting and remove unnecessary whitespace in multiple files.

* Refactor import statement in test_tilelang_kernel_dequantize_gemm.py to use 'tilelang.language' for consistency

* Add CUDA requirements to FP8 test cases and update references for clarity

* Add a blank line for improved readability in test_tilelang_kernel_fp8_gemm_mma.py

* Fix data type in reference result calculation for consistency in test_tilelang_kernel_gemm_mma_intrinsic.py

* Add CUDA requirements and FP8 test cases for matmul and gemv simulations

* Remove debug print statements and use tilelang's testing assertion for result validation in test_tilelang_kernel_gemm_mma_intrinsic.py

* Remove outdated comment regarding FP8 tests in test_tilelang_kernel_gemv_simt.py

* Add BF16 support to matrix multiplication and introduce corresponding test cases

* Add a blank line for improved readability in BF16 GEMM test

* Update acknowledgements in README to include supervision by Zhi Yang at Peking University

* enhance acknowledgement

* Replace tutorial on memory layout optimization with new tutorial on writing high-performance kernels with thread primitives

* Update subproject commit for TVM dependency

* Update subproject commit for TVM dependency

* Add int4_t type and functions for packing char values in CUDA common header

* Add plot_layout example and implement GetForwardVars method in layout classes

* Refactor code for improved readability by adjusting line breaks and formatting in layout and test files

* Fix formatting by removing unnecessary line break in layout.h

* Refactor make_int4 function for improved readability by adjusting parameter formatting
2025-02-10 01:03:01 +08:00