• Lei Wang's avatar
    [Language] Enhance alias to support blockwise memory load (#261) · 927e50d9
    Lei Wang authored
    * [Enhancement] Introduce caching control and frame management in TileLang
    
    - Added cache control functions (`enable_cache`, `disable_cache`, `is_cache_enabled`) in `env.py` to manage kernel caching behavior.
    - Updated `kernel_cache.py` to utilize the cache state, preventing unnecessary kernel compilation when caching is disabled.
    - Introduced a new `frame.py` module to manage LetFrame instances, including a stack for variable-value mapping and enhanced frame management.
    - Updated imports in various modules to accommodate new caching and frame functionalities, improving overall organization and clarity.
    
    * [Refactor] Clean up and enhance caching and frame management in TileLang
    
    - Added spacing for improved readability in `env.py` and `frame.py`.
    - Refactored `LetFrame` class to enhance clarity in buffer region assignment.
    - Ensured consistent formatting and organization across caching control and frame management functions.
    
    * [Feature] Add matrix multiplication functionality in TileLang
    
    - Introduced a new test file `test_tilelang_language_alias.py` that implements a matrix multiplication function using TileLang's primitives.
    - The `matmul` function defines a kernel for performing tile-level GEMM operations, with support for customizable block sizes and data types.
    - Added a `run_matmul` function to compile and execute the kernel, along with a test function to validate the implementation.
    - Updated `gemm.py` to allow `tir.Buffer` or `tir.Var` as valid argument types for the `gemm` function, enhancing flexibility in argument handling.
    
    * [Refactor] Improve formatting and readability in test_tilelang_language_alias.py
    
    - Adjusted spacing and alignment in the `matmul` and `run_matmul` functions for better readability.
    - Cleaned up unnecessary blank lines and ensured consistent formatting throughout the file.
    - Enhanced overall code clarity without altering functionality.
    927e50d9
test_tilelang_language_alias.py 1.84 KB