• Lei Wang's avatar
    [Refactor] refactor `tilelang.jit` to support a faster and more flexible kernel cache (#501) · 25a50f1a
    Lei Wang authored
    * [Refactor] Update JIT kernel functions and streamline GEMM tests
    
    * Renamed and refactored matmul and run_gemm functions to matmul_kernel_jit and run_gemm_kernel_jit for clarity.
    * Removed redundant JIT decorator from the matmul function, ensuring it is applied only to the kernel function.
    * Updated test function names to reflect changes in the kernel functions, enhancing consistency and readability.
    * Cleaned up commented-out code and unnecessary imports to improve overall code quality.
    
    * Update main function call in GEMM test to use tilelang testing framework
    
    * Update README and example scripts to include JIT decorator comments
    
    * Added comments in README.md and various example scripts to indicate the use of the @tilelang.jit decorator for returning torch functions.
    * Removed redundant comments that previously instructed to add the decorator, streamlining the documentation and improving clarity.
    
    * Update GEMM test parameters for improved performance
    
    * Set num_stages to 0 and adjusted matrix dimensions in test functions to enhance performance and consistency across GEMM tests in test_tilelang_kernel_gemm.py.
    25a50f1a
README.md 11.6 KB