Unverified Commit 1a3a64fb authored by Chaofan Lin's avatar Chaofan Lin Committed by GitHub
Browse files

[Language] Make TL scripts friendly to Python syntax highlights (#1466)

* Language] Make TL scripts friendly to Python syntax highlights

* add comments

* fix submodule
parent f6db2014
# Copyright (c) Tile-AI Corporation.
# Licensed under the MIT License.
import argparse import argparse
import tilelang import tilelang
......
# Copyright (c) Tile-AI Corporation.
# Licensed under the MIT License.
import argparse import argparse
import tilelang import tilelang
......
...@@ -58,12 +58,12 @@ from .allocate import ( ...@@ -58,12 +58,12 @@ from .allocate import (
alloc_tcgen05_instr_desc, # noqa: F401 alloc_tcgen05_instr_desc, # noqa: F401
empty, # 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 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 .experimental.gemm_sp import gemm_sp, gemm_sp_v2 # noqa: F401
from .fill import fill, clear # noqa: F401 from .fill_op import fill, clear # noqa: F401
from .reduce import ( from .reduce_op import (
reduce, # noqa: F401 reduce, # noqa: F401
reduce_max, # noqa: F401 reduce_max, # noqa: F401
reduce_min, # noqa: F401 reduce_min, # noqa: F401
...@@ -81,7 +81,7 @@ from .reduce import ( ...@@ -81,7 +81,7 @@ from .reduce import (
warp_reduce_bitand, # noqa: F401 warp_reduce_bitand, # noqa: F401
warp_reduce_bitor, # 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 ( from .customize import (
atomic_max, # noqa: F401 atomic_max, # noqa: F401
atomic_min, # noqa: F401 atomic_min, # noqa: F401
......
# Copyright (c) Tile-AI Corporation. """Atomic operations exposed on the TileLang language surface."""
# Licensed under the MIT License.
"""Atomic operations for tilelang."""
from __future__ import annotations from __future__ import annotations
......
"""The language interface for tl programs.""" """Builtin operations exposed on the TileLang language surface."""
from __future__ import annotations from __future__ import annotations
......
"""The language interface for tl programs.""" """Copy operations exposed on the TileLang language surface."""
from __future__ import annotations from __future__ import annotations
from typing import Literal from typing import Literal
......
"""The language interface for tl programs.""" """Some customized operations frequently used in tensor programming, exposed on the TileLang language surface."""
from __future__ import annotations from __future__ import annotations
import tilelang.language as T import tilelang.language as T
......
"""Fast math operations exposed on the TileLang language surface."""
from tvm import tir from tvm import tir
......
"""The language interface for tl programs.""" """Fill operations exposed on the TileLang language surface."""
from __future__ import annotations from __future__ import annotations
from tvm import tir from tvm import tir
......
"""The language interface for tl programs.""" """GEMM (General Matrix Multiplication) operators exposed on the TileLang language surface."""
from __future__ import annotations from __future__ import annotations
from tilelang.tileop.base import GemmWarpPolicy from tilelang.tileop.base import GemmWarpPolicy
......
"""The language interface for tl programs.""" """Kernel launching language interface in TileLang."""
from __future__ import annotations from __future__ import annotations
from collections import deque from collections import deque
......
"""The language interface for tl programs.""" """Logical operations exposed on the TileLang language surface."""
from __future__ import annotations from __future__ import annotations
......
"""The language interface for tl programs.""" """Loop related language interfaces in TileLang."""
from __future__ import annotations from __future__ import annotations
from typing import Any from typing import Any
...@@ -175,5 +175,12 @@ def unroll( ...@@ -175,5 +175,12 @@ def unroll(
return UnrollForWithStep(start, stop, step, annotations=annotations) return UnrollForWithStep(start, stop, step, annotations=annotations)
Serial = serial # "Serial" and "Unroll" are aliases of "T.serial" and "T.unroll". We use uppercase to emphasize that they are tile-level loops.
Unroll = unroll
def Serial(*args, **kwargs):
return serial(*args, **kwargs)
def Unroll(*args, **kwargs):
return unroll(*args, **kwargs)
"""Common math intrinsics exposed on the TileLang language surface."""
from tvm import tir from tvm import tir
......
"""The language interface for tl programs.""" """Buffer/Tensor proxy in TileLang."""
from __future__ import annotations from __future__ import annotations
......
"""The language interface for tl programs.""" """Reduce operations exposed on the TileLang language surface."""
from __future__ import annotations from __future__ import annotations
from tvm import tir from tvm import tir
......
...@@ -16,8 +16,7 @@ ...@@ -16,8 +16,7 @@
# specific language governing permissions and limitations # specific language governing permissions and limitations
# under the License. # 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. # The code below is mostly copied from mlc.ai quantization.py in mlc-llm.
# pylint: disable=invalid-name,missing-function-docstring,unused-variable # pylint: disable=invalid-name,missing-function-docstring,unused-variable
"""TIR computation utilities for quantization.""" """TIR computation utilities for quantization."""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment