Files
tilelang/testing/python/transform
Lei Wang 2ff1dd50dc [Feature] Introduce DecoupleTypeCast pass for mixed-precision vectorization (#1644)
* [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
2026-01-18 01:23:54 +08:00
..