mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Lei Wang
7171aff6dc
* [Enhancement] Add commit ID to versioning and improve logging initialization * Updated `get_tilelang_version` to include an optional commit ID in the version string. * Enhanced the `TileLangBuilPydCommand` to write the version with commit ID to the VERSION file during the build process. * Introduced a new function `get_git_commit_id` in `version.py` to retrieve the current git commit hash. * Refactored logger initialization in `autotuner/__init__.py` to ensure handlers are set up only once, improving performance and clarity. * Minor fixes in `flatten_buffer.cc` and `kernel_cache.py` for better handling of versioning and logging. * [Refactor] Enhance AutoTuner and JITKernel for improved performance and caching * Refactored the AutoTuner class to include new methods for setting compilation and profiling arguments, enhancing configurability. * Introduced caching mechanisms for tuning results, allowing for faster retrieval of previously computed configurations. * Updated JITKernel to store tuning results, including latency and configuration details, improving the kernel's performance tracking. * Added new methods for generating cache keys and saving/loading results to/from disk, streamlining the tuning process. * Enhanced the overall structure and readability of the autotuning logic, ensuring better maintainability and clarity. * Minor adjustments in related modules to support the new caching and profiling features. * [Refactor] Clean up code formatting and improve readability in AutoTuner and related modules * Consolidated import statements and removed unnecessary line breaks for better readability. * Standardized function argument formatting across the AutoTuner and CompileArgs classes. * Enhanced consistency in the use of whitespace and indentation throughout the codebase. * Minor adjustments in the Profiler and JITKernel classes to improve clarity and maintainability. * Ensured that all changes adhere to the project's coding style guidelines. * [Refactor] Remove redundant type hints in AutoTuner modules * Simplified import statements in `__init__.py` and `param.py` by removing unnecessary duplicate type hints for `Any`. * Improved code readability and maintainability by streamlining type imports across the AutoTuner module. * [Refactor] Update AutoTuner configuration for improved profiling and target detection * Enhanced the AutoTuner configuration across multiple examples by adding `set_profile_args` to better manage profiling settings. * Standardized the use of `target="auto"` in compile arguments to ensure automatic target detection. * Removed redundant target specifications in certain instances to streamline the configuration process. * Improved overall clarity and maintainability of the autotuning logic in various example scripts. * [Refactor] Simplify code formatting and improve readability in example scripts * Consolidated function argument formatting in `benchmark_mla_decode_amd_tilelang.py`, `example_elementwise_add.py`, and `performance.py` for better clarity. * Removed unnecessary line breaks and standardized argument placement across multiple files. * Enhanced overall code readability and maintainability in autotuning examples and performance scripts. * [Refactor] Update JIT decorator usage across multiple files * Removed redundant parameters from the JIT decorator in various benchmark and example scripts, simplifying the code. * Standardized the import of the JIT decorator from `tilelang`, enhancing consistency across the codebase. * Improved overall readability and maintainability by consolidating import statements and cleaning up function definitions. * [Refactor] Standardize JIT decorator formatting across benchmark and example scripts * Simplified the formatting of the JIT decorator in multiple files by removing unnecessary line breaks. * Enhanced code readability and consistency in the usage of the JIT decorator across benchmark and example scripts. * Improved overall maintainability by ensuring uniformity in function definitions and decorator usage.
92 lines
2.9 KiB
Python
92 lines
2.9 KiB
Python
import argparse
|
|
import tilelang.language as T
|
|
from tilelang.autotuner import AutoTuner
|
|
|
|
|
|
def ref_program(A, B):
|
|
return A @ B.T
|
|
|
|
|
|
def get_configs():
|
|
configs = [{
|
|
"block_M": 128,
|
|
"block_N": 128,
|
|
"block_K": 64,
|
|
"num_stages": 2,
|
|
"thread_num": 256,
|
|
"enable_rasteration": True, # keep param name for backward-compat
|
|
}]
|
|
return configs
|
|
|
|
|
|
def run(M, N, K):
|
|
|
|
def kernel(
|
|
block_M=None,
|
|
block_N=None,
|
|
block_K=None,
|
|
num_stages=None,
|
|
thread_num=None,
|
|
enable_rasteration=None,
|
|
):
|
|
dtype = "float16"
|
|
accum_dtype = "float"
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, K), dtype),
|
|
B: T.Tensor((N, K), dtype),
|
|
C: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(
|
|
T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by):
|
|
A_shared = T.alloc_shared((block_M, block_K), dtype)
|
|
B_shared = T.alloc_shared((block_N, block_K), dtype)
|
|
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
C_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
T.use_swizzle(panel_size=10, enable=enable_rasteration)
|
|
T.clear(C_local)
|
|
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
|
T.copy(A[by * block_M, k * block_K], A_shared)
|
|
T.copy(B[bx * block_N, k * block_K], B_shared)
|
|
T.gemm(
|
|
A_shared,
|
|
B_shared,
|
|
C_local,
|
|
transpose_B=True,
|
|
)
|
|
T.copy(C_local, C_shared)
|
|
T.copy(C_shared, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
autotuner = AutoTuner.from_kernel(
|
|
kernel=kernel, configs=get_configs()).set_compile_args(
|
|
out_idx=[-1],
|
|
target="auto",
|
|
).set_profile_args(
|
|
ref_prog=ref_program,)
|
|
return autotuner.run(warmup=3, rep=20)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Parse command-line arguments for matrix dimensions
|
|
parser = argparse.ArgumentParser(description="Autotuned MatMul Benchmark")
|
|
parser.add_argument("--m", type=int, default=16384, help="Matrix dimension M")
|
|
parser.add_argument("--n", type=int, default=16384, help="Matrix dimension N")
|
|
parser.add_argument("--k", type=int, default=16384, help="Matrix dimension K")
|
|
args = parser.parse_args()
|
|
|
|
M, N, K = args.m, args.n, args.k
|
|
|
|
# Compute total floating-point operations to measure throughput
|
|
total_flops = 2 * M * N * K
|
|
|
|
result = run(M, N, K)
|
|
|
|
print(f"Latency: {result.latency}")
|
|
print(f"TFlops: {total_flops / result.latency * 1e-9:.3f}")
|
|
print(f"Config: {result.config}")
|
|
|
|
print(f"Reference TFlops: {total_flops / result.ref_latency * 1e-9:.3f}")
|