diff --git a/examples/gemm_sp/example_custom_compress.py b/examples/gemm_sp/example_custom_compress.py index 7f18523b7..7b93f2a77 100644 --- a/examples/gemm_sp/example_custom_compress.py +++ b/examples/gemm_sp/example_custom_compress.py @@ -1,5 +1,3 @@ -# Copyright (c) Tile-AI Corporation. -# Licensed under the MIT License. import argparse import tilelang diff --git a/examples/gemm_sp/example_gemm_sp.py b/examples/gemm_sp/example_gemm_sp.py index 708bc7231..10f524adb 100644 --- a/examples/gemm_sp/example_gemm_sp.py +++ b/examples/gemm_sp/example_gemm_sp.py @@ -1,5 +1,3 @@ -# Copyright (c) Tile-AI Corporation. -# Licensed under the MIT License. import argparse import tilelang diff --git a/tilelang/language/__init__.py b/tilelang/language/__init__.py index bf2a144a3..97f838581 100644 --- a/tilelang/language/__init__.py +++ b/tilelang/language/__init__.py @@ -58,12 +58,12 @@ alloc_tcgen05_instr_desc, # noqa: F401 empty, # noqa: F401 ) -from .copy import copy, c2d_im2col # noqa: F401 +from .copy_op import copy, c2d_im2col # noqa: F401 from tilelang.tileop.base import GemmWarpPolicy # noqa: F401 -from .gemm import gemm, gemm_v1, gemm_v2 # noqa: F401 +from .gemm_op import gemm, gemm_v1, gemm_v2 # noqa: F401 from .experimental.gemm_sp import gemm_sp, gemm_sp_v2 # noqa: F401 -from .fill import fill, clear # noqa: F401 -from .reduce import ( +from .fill_op import fill, clear # noqa: F401 +from .reduce_op import ( reduce, # noqa: F401 reduce_max, # noqa: F401 reduce_min, # noqa: F401 @@ -81,7 +81,7 @@ warp_reduce_bitand, # noqa: F401 warp_reduce_bitor, # noqa: F401 ) -from .print import print, device_assert # noqa: F401 +from .print_op import print, device_assert # noqa: F401 from .customize import ( atomic_max, # noqa: F401 atomic_min, # noqa: F401 diff --git a/tilelang/language/atomic.py b/tilelang/language/atomic.py index b63b48c56..a801f75f4 100644 --- a/tilelang/language/atomic.py +++ b/tilelang/language/atomic.py @@ -1,6 +1,4 @@ -# Copyright (c) Tile-AI Corporation. -# Licensed under the MIT License. -"""Atomic operations for tilelang.""" +"""Atomic operations exposed on the TileLang language surface.""" from __future__ import annotations diff --git a/tilelang/language/builtin.py b/tilelang/language/builtin.py index 60739e611..2932656ca 100644 --- a/tilelang/language/builtin.py +++ b/tilelang/language/builtin.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""Builtin operations exposed on the TileLang language surface.""" from __future__ import annotations diff --git a/tilelang/language/copy.py b/tilelang/language/copy_op.py similarity index 98% rename from tilelang/language/copy.py rename to tilelang/language/copy_op.py index 1bc84a530..0b55c410c 100644 --- a/tilelang/language/copy.py +++ b/tilelang/language/copy_op.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""Copy operations exposed on the TileLang language surface.""" from __future__ import annotations from typing import Literal diff --git a/tilelang/language/customize.py b/tilelang/language/customize.py index e2f4b1c8a..ae4e754f7 100644 --- a/tilelang/language/customize.py +++ b/tilelang/language/customize.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""Some customized operations frequently used in tensor programming, exposed on the TileLang language surface.""" from __future__ import annotations import tilelang.language as T diff --git a/tilelang/language/fastmath.py b/tilelang/language/fastmath.py index 0146f53ac..c77fad34c 100644 --- a/tilelang/language/fastmath.py +++ b/tilelang/language/fastmath.py @@ -1,3 +1,5 @@ +"""Fast math operations exposed on the TileLang language surface.""" + from tvm import tir diff --git a/tilelang/language/fill.py b/tilelang/language/fill_op.py similarity index 97% rename from tilelang/language/fill.py rename to tilelang/language/fill_op.py index af301c264..a093a8459 100644 --- a/tilelang/language/fill.py +++ b/tilelang/language/fill_op.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""Fill operations exposed on the TileLang language surface.""" from __future__ import annotations from tvm import tir diff --git a/tilelang/language/gemm.py b/tilelang/language/gemm_op.py similarity index 98% rename from tilelang/language/gemm.py rename to tilelang/language/gemm_op.py index 6f650470d..e2bda2b94 100644 --- a/tilelang/language/gemm.py +++ b/tilelang/language/gemm_op.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""GEMM (General Matrix Multiplication) operators exposed on the TileLang language surface.""" from __future__ import annotations from tilelang.tileop.base import GemmWarpPolicy diff --git a/tilelang/language/kernel.py b/tilelang/language/kernel.py index 625531b38..73f7ed949 100644 --- a/tilelang/language/kernel.py +++ b/tilelang/language/kernel.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""Kernel launching language interface in TileLang.""" from __future__ import annotations from collections import deque diff --git a/tilelang/language/logical.py b/tilelang/language/logical.py index fb4b88a6e..66f0a2e2b 100644 --- a/tilelang/language/logical.py +++ b/tilelang/language/logical.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""Logical operations exposed on the TileLang language surface.""" from __future__ import annotations diff --git a/tilelang/language/loop.py b/tilelang/language/loop.py index f28f097cb..4fbd4e9f8 100644 --- a/tilelang/language/loop.py +++ b/tilelang/language/loop.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""Loop related language interfaces in TileLang.""" from __future__ import annotations from typing import Any @@ -175,5 +175,12 @@ def unroll( return UnrollForWithStep(start, stop, step, annotations=annotations) -Serial = serial -Unroll = unroll +# "Serial" and "Unroll" are aliases of "T.serial" and "T.unroll". We use uppercase to emphasize that they are tile-level loops. + + +def Serial(*args, **kwargs): + return serial(*args, **kwargs) + + +def Unroll(*args, **kwargs): + return unroll(*args, **kwargs) diff --git a/tilelang/language/math_intrinsics.py b/tilelang/language/math_intrinsics.py index 7a6104c74..6dfb617e5 100644 --- a/tilelang/language/math_intrinsics.py +++ b/tilelang/language/math_intrinsics.py @@ -1,3 +1,5 @@ +"""Common math intrinsics exposed on the TileLang language surface.""" + from tvm import tir diff --git a/tilelang/language/print.py b/tilelang/language/print_op.py similarity index 100% rename from tilelang/language/print.py rename to tilelang/language/print_op.py diff --git a/tilelang/language/proxy.py b/tilelang/language/proxy.py index 7807a4669..b739de6b5 100644 --- a/tilelang/language/proxy.py +++ b/tilelang/language/proxy.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""Buffer/Tensor proxy in TileLang.""" from __future__ import annotations diff --git a/tilelang/language/reduce.py b/tilelang/language/reduce_op.py similarity index 99% rename from tilelang/language/reduce.py rename to tilelang/language/reduce_op.py index eb464f6af..9db56df0d 100644 --- a/tilelang/language/reduce.py +++ b/tilelang/language/reduce_op.py @@ -1,4 +1,4 @@ -"""The language interface for tl programs.""" +"""Reduce operations exposed on the TileLang language surface.""" from __future__ import annotations from tvm import tir diff --git a/tilelang/quantize/quantization.py b/tilelang/quantize/quantization.py index 13552f674..74a545f25 100644 --- a/tilelang/quantize/quantization.py +++ b/tilelang/quantize/quantization.py @@ -16,8 +16,7 @@ # specific language governing permissions and limitations # under the License. # -# Copyright (c) Tile-AI Corporation. -# Licensed under the MIT License. + # The code below is mostly copied from mlc.ai quantization.py in mlc-llm. # pylint: disable=invalid-name,missing-function-docstring,unused-variable """TIR computation utilities for quantization."""