3 Commits

Include renames

Author	SHA1	Message	Date
Lei Wang	99ee74b343	[CUDA] Support int4 T.gemm (#2063) ... * Support int4 T.gemm and direct packed cp.async lowering * Keep cp.async vectorization as a late fallback * Change tl.ptx_cp_async to use num_elems semantics * fix test	2026-04-20 16:35:43 +08:00
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	2026-03-06 17:35:16 +08:00
Lei Wang	4ff81c7d40	[Refactor] Introduce T.access_of to combine T.address_of and access_ptr (#1827) ... * [Refactor] Update atomic operations to use access_ptr for buffer access * Replaced instances of address_of with access_ptr in atomic operations to improve clarity and maintainability. * Enhanced atomic operation implementations in various files, including atomic_add, atomic_max, and atomic_min. * Updated related logic in loop vectorization and tile language to support new access patterns. * Added a utility function to extract destination scalar dtype for atomic operations, supporting both BufferLoad and tvm_access_ptr. * Improved kernel cache logic to include a build-stamp for the TileLang runtime library, ensuring cache validity during development. * lint fix * Fix access_ptr layout rewrite and atomic vectorization * lint fix * Refactor atomic operations to remove const qualifiers from pointer parameters in atomicAddx2 and ToFloat functions. This change enhances type flexibility and consistency across atomic operations. * Refactor buffer access logic in lower_tile_op and vectorize_loop to enhance clarity and maintainability. Removed unused code and improved offset calculations for tile access. Updated atomic operation handling to streamline dtype extraction and ensure correct vectorization behavior. * Add logging for loop layout and lowered loop in atomic_add; implement tvm_access_ptr handling in loop_vectorize and lower_tile_op This commit introduces logging statements to track the loop layout and the result of the lowered loop in the AtomicAddNode. Additionally, it implements a new method to handle tvm_access_ptr calls in the loop_vectorize and lower_tile_op files, enhancing the handling of memory access patterns and ensuring correct vectorization behavior. * fix * lint fix * fix * Update ToFloat2 and ToFloat4 functions to use const qualifiers for pointer parameters in atomic.h; modify test script to call test_atomic_addx2_float instead of main. * Refactor tvm_access_ptr handling in loop_vectorize to improve argument validation and buffer resolution logic. Ensure at least three arguments are provided and streamline the search for matching buffers in layout_map_. Update vector size calculations to enhance clarity and maintainability. * Add handling for tvm_access_ptr in vectorize_loop to support offset substitution and buffer creation This commit introduces logic to handle tvm_access_ptr calls, allowing for the substitution of loop variables in offsets and the creation of dummy buffers. The new method ensures correct base address calculation and enhances the overall handling of memory access patterns in vectorization. * Add handling for tvm_access_ptr offset validation in loop_vectorize This commit enhances the loop_vectorize functionality by adding checks for the element offset in tvm_access_ptr calls. It ensures that offsets are invariant within the vector boundary, improving the robustness of vectorization logic and memory access patterns. * fix * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update improves the robustness of vectorization logic and ensures offsets remain invariant within vector boundaries. * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update ensures offsets remain invariant within vector boundaries, improving memory access robustness and vectorization logic. * fix * Add tl.access_ptr support in TileLang This commit introduces the `tl.access_ptr` operation, which carries pointer access metadata in the frontend. It allows for better handling of buffer loads by retaining necessary information for downstream analysis. The new operation is integrated into various components, including the loop vectorization and memory access legalizations, ensuring compatibility with existing functionality while enhancing the robustness of memory access patterns. * lint fix * fix * Add support for tl.access_ptr in atomic operations This commit updates the AtomicAddNode and AtomicOpBaseNode implementations to utilize the new tl.access_ptr operation for building access pointers to destination elements. It enhances the handling of buffer loads by incorporating type checks for index data types, improving the robustness of memory access patterns. Additionally, it removes unnecessary logging statements in the lowering process, streamlining the codebase. * lint fix * Refactor buffer load extraction in vectorize_loop to check for non-empty arguments in access_ptr calls. This change improves the robustness of buffer load handling by ensuring valid argument presence before processing. * Remove unused fallback logic in KernelCache and clean up import statements in test_tilelang_language_access_ptr.py for improved code clarity. * Refactor memory access in BitNet and dequantization examples to use tl.access_ptr This update replaces T.address_of calls with T.access_ptr in multiple files, enhancing memory access handling by specifying read and write permissions. This change improves code clarity and aligns with recent enhancements in pointer access management.	2026-02-11 22:51:59 +08:00