mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Lei Wang
6e439530b5
* 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.