Merge branch 'main' of https://github.com/tile-ai/tilelang into v2_1106

Author: LeiWang1999

.github/workflows/ci.yml

@@ -129,13 +129,6 @@ jobs:
           echo "UV_CACHE_DIR=${XDG_CACHE_HOME}/uv" | tee -a "${GITHUB_ENV}"
           echo "PRE_COMMIT_HOME=${XDG_CACHE_HOME}/pip/.pre-commit" | tee -a "${GITHUB_ENV}"
 
-      - name: Set environment (GitHub-hosted runners)
-        if: ${{ !startsWith(matrix.runner.name, 'self-hosted') }}
-        run: |
-          # Enable ccache on GitHub-hosted runners to speed up builds
-          echo "CMAKE_C_COMPILER_LAUNCHER=ccache" | tee -a "${GITHUB_ENV}"
-          echo "CMAKE_CXX_COMPILER_LAUNCHER=ccache" | tee -a "${GITHUB_ENV}"
-
       # Do not use ccache on self-hosted runners, as it will download/upload caches which is slow.
       # Self-hosted runners usually have more CPU power to compile without ccache.
       - name: Setup ccache (GitHub-hosted runners)
@@ -144,8 +137,13 @@ jobs:
         uses: hendrikmuhs/ccache-action@v1
         with:
           create-symlink: true
-          key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ matrix.python-version }}-${{ matrix.runner.name }}-${{ matrix.runner.toolkit }}
           evict-old-files: "7d"
+          append-timestamp: false
+          key: ${{ runner.os }}-${{ runner.arch }}-${{ matrix.runner.toolkit }}-${{ hashFiles('**/*.cc') }}
+          restore-keys: |
+            ${{ runner.os }}-${{ runner.arch }}-${{ matrix.runner.toolkit }}-${{ hashFiles('**/*.cc') }}
+            ${{ runner.os }}-${{ runner.arch }}-${{ matrix.runner.toolkit }}
+            ${{ runner.os }}-${{ runner.arch }}
 
       - name: Set environment (CUDA)
         if: contains(matrix.runner.toolkit, 'CUDA')

.github/workflows/dist.yml

@@ -44,11 +44,11 @@ jobs:
     runs-on: macos-latest
     timeout-minutes: 30
     env:
-      # NO_VERSION_LABEL disables embedding the toolchain / git commit hash in version metadata.
+      # `NO_VERSION_LABEL=ON` disables embedding the toolchain / git commit hash in version metadata.
       # Otherwise, the version of the SDist has a git hash suffix (e.g., 0.1.0+gitabcdef12),
       # but the package built from the SDist has no way to get the git hash (it is not a git repo),
       # leading to inconsistent versions between SDist and built packages (+gitabcdef12 vs. +gitunknown).
-      NO_VERSION_LABEL: 'OFF'
+      NO_VERSION_LABEL: 'ON'
 
     steps:
       - name: Checkout repository
@@ -72,18 +72,20 @@ jobs:
         uses: hendrikmuhs/ccache-action@v1
         with:
           create-symlink: true
-          key: ccache-${{ runner.os }}-${{ runner.arch }}
           evict-old-files: "7d"
+          append-timestamp: false
+          key: sdist-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('**/*.cc') }}
+          restore-keys: |
+            sdist-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('**/*.cc') }}
+            sdist-${{ runner.os }}-${{ runner.arch }}
+            ${{ runner.os }}-${{ runner.arch }}
 
       - name: Test SDist buildable
         run: |
           TEMP_DIR="$(mktemp -d -t tilelang-sdist-test)"
           cp -r dist "${TEMP_DIR}/dist"
-          uv venv --seed "${TEMP_DIR}/venv"
-          source "${TEMP_DIR}/venv/bin/activate"
           cd "${TEMP_DIR}"
-          python3 -m pip install --upgrade pip setuptools wheel
-          python3 -m pip install -v dist/*.tar.gz
+          uv pip install -v dist/*.tar.gz
           python3 -c "import tilelang; print(tilelang.__version__)"
 
       - name: Upload SDist
@@ -125,14 +127,19 @@ jobs:
           fetch-depth: 1
           submodules: recursive
 
-      # NB: CIBW builds wheels in containers on Linux
-      - name: Setup ccache (macOS only)
-        if: runner.os == 'macOS'
+      - name: Setup ccache
         uses: hendrikmuhs/ccache-action@v1
         with:
           create-symlink: true
-          key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ matrix.python-version }}-${{ matrix.target.toolkit }}
           evict-old-files: "7d"
+          append-timestamp: false
+          key: wheel-${{ runner.os }}-${{ runner.arch }}-${{ matrix.target.toolkit }}-${{ hashFiles('**/*.cc') }}
+          restore-keys: |
+            wheel-${{ runner.os }}-${{ runner.arch }}-${{ matrix.target.toolkit }}-${{ hashFiles('**/*.cc') }}
+            wheel-${{ runner.os }}-${{ runner.arch }}-${{ matrix.target.toolkit }}
+            wheel-${{ runner.os }}-${{ runner.arch }}
+            ${{ runner.os }}-${{ runner.arch }}-${{ matrix.target.toolkit }}
+            ${{ runner.os }}-${{ runner.arch }}
 
       - name: Set CIBW_BUILD
         run: |
@@ -147,6 +154,11 @@ jobs:
             echo "CUDA_VERSION=${CUDA_VERSION}" | tee -a "${GITHUB_ENV}"
           fi
 
+          if [[ "${{ runner.os }}" == "Linux" ]]; then
+            HOST_CCACHE_DIR="$(ccache --get-config cache_dir)"
+            echo "CIBW_BEFORE_BUILD_LINUX=yum install -y ccache && ccache -o cache_dir=/host${HOST_CCACHE_DIR}" | tee -a "${GITHUB_ENV}"
+          fi
+
       - name: Build wheels
         uses: pypa/cibuildwheel@v3.2
         with:

.gitignore

@@ -104,4 +104,4 @@ cmake-build/
 cmake-build-*/
 
 # Git version for sdist
-_git_commit.txt
+.git_commit.txt

CMakeLists.txt

@@ -41,7 +41,16 @@ endif()
 
 find_program(CCACHE_PROGRAM ccache)
 if(CCACHE_PROGRAM)
-  message(STATUS "Using ccache: ${CCACHE_PROGRAM}")
+  message(STATUS "Using ccache: ${CCACHE_PROGRAM} with base_dir=${CMAKE_SOURCE_DIR}")
+  if(APPLE)
+    # Passing configs like `ccache base_dir=/xxx cc ...` is supported
+    # (likely) since ccache 4.x, which has been provided by homebrew.
+    # Our Linux builder image (manylinux2014 & manylinux_2_28) still
+    # provides ccache 3.x and do not support this form.
+    # `cibuildwheel` uses fixed folder on Linux (`/project`) as working directory,
+    # so cache would work without setting `base_dir`.
+    set(CCACHE_PROGRAM "${CCACHE_PROGRAM};base_dir=${CMAKE_SOURCE_DIR}")
+  endif()
   set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}" CACHE STRING "C compiler launcher")
   set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}" CACHE STRING "CXX compiler launcher")
   set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}" CACHE STRING "CUDA compiler launcher")

testing/python/language/test_tilelang_language_frontend_v2.py

@@ -3,6 +3,8 @@
 import torch
 import tilelang.testing
 import tvm
+from tvm.script.ir_builder.base import IRBuilderFrame
+from tvm.tir.expr import IntImm, Var
 
 
 def test_argument():
@@ -273,6 +275,43 @@ def foo() -> T.Tensor((128,), T.float32):
     assert isinstance(foo, T.PrimFunc)
 
 
+def test_serial_for_with_step():
+
+    @tilelang.jit(out_idx=-1)
+    @T.prim_func
+    def test_stepped_serial(A: T.Tensor((10,), T.int32)):
+        with T.Kernel(1) as _:
+            for i in range(0, 10, 2):
+                T.device_assert(0 <= i < 10 and i % 2 == 0, "i out of range")
+                A[i] = 1.0
+            for i in range(1, 10, 2):
+                T.device_assert(1 <= i < 10 and i % 2 == 1, "i out of range")
+                A[i] = 2.0
+
+    ker = test_stepped_serial()
+    res = ker()
+    ref = torch.tensor([1, 2, 1, 2, 1, 2, 1, 2, 1, 2], dtype=torch.int32, device='cuda')
+    assert torch.all(res == ref), f"Expected {ref}, but got {res}"
+
+    @tilelang.jit(out_idx=-1)
+    @T.prim_func
+    def test_serial_step_neg(A: T.Tensor((10,), T.int32)):
+        with T.Kernel(1) as _:
+            for i in range(10, 0, -1):
+                T.device_assert(0 < i <= 10, "i out of range")
+                A[10 - i] = i
+
+    ker = test_serial_step_neg()
+    res = ker()
+    ref = torch.tensor([10, 9, 8, 7, 6, 5, 4, 3, 2, 1], dtype=torch.int32, device='cuda')
+    assert torch.all(res == ref), f"Expected {ref}, but got {res}"
+
+    assert isinstance(T.serial(1, 10, 1), IRBuilderFrame)
+    assert isinstance(T.serial(1, 10, IntImm('int32', 1)), IRBuilderFrame)
+    assert not isinstance(T.serial(1, 10, Var('tmp', 'int32')), IRBuilderFrame)
+    assert not isinstance(T.serial(10, -1, -1), IRBuilderFrame)
+
+
 def test_swap_logic():
 
     @tilelang.jit

tilelang/language/__init__.py

@@ -23,9 +23,7 @@
     SharedBuffer,  # noqa: F401
     LocalBuffer,  # noqa: F401
 )
-from .parallel import Parallel  # noqa: F401
-from .pipeline import Pipelined  # noqa: F401
-from .persistent import Persistent  # noqa: F401
+from .loop import serial, Parallel, Persistent, Pipelined  # noqa: F401
 from .frame import has_let_value, get_let_value  # noqa: F401
 from .math_intrinsics import *  # noqa: F401
 from .kernel import (

tilelang/language/loop.py

@@ -0,0 +1,111 @@
+"""The language interface for tl programs."""
+from __future__ import annotations
+
+from typing import Any
+from tvm import tir
+from tvm.tir import IntImm
+import tvm.script.ir_builder.tir as tb_tir
+from .v2.builder import SerialForWithStep
+from tilelang import _ffi_api
+
+
+def Parallel(*extents: tir.PrimExpr, coalesced_width: int | None = None):
+    """Tools to construct nested parallel for loop.
+       This can be used to create element-wise tensor expression.
+
+    Parameters
+    ----------
+    extents : PrimExpr
+        The extents of the iteration.
+
+    coalesced_width : Optional[int]
+        The coalesced width of the parallel loop.
+
+    Returns
+    -------
+    res : frame.ForFrame
+        The ForFrame.
+    """
+    annotations: dict[str, Any] = {}
+    if coalesced_width is not None:
+        annotations.update({"coalesced_width": coalesced_width})
+    return _ffi_api.Parallel(extents, annotations)  # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def Persistent(
+    domain: list[tir.PrimExpr],
+    wave_size: tir.PrimExpr,
+    index: tir.PrimExpr,
+    group_size: tir.PrimExpr | None = 8,
+):
+    """Tools to construct persistent for loop.
+
+    Parameters
+    ----------
+    domain : List[tir.PrimExpr]
+        The list of dominators.
+    wave_size : int
+        The wave size.
+    index : int
+        The tile index in one wave.
+    group_size : tir.PrimExpr
+        The group size.
+    """
+    return _ffi_api.Persistent(domain, wave_size, index, group_size)
+
+
+def Pipelined(
+    start: tir.PrimExpr,
+    stop: tir.PrimExpr = None,
+    num_stages: int = 0,
+    order: list[int] | None = None,
+    stage: list[int] | None = None,
+    sync: list[list[int]] | None = None,
+    group: list[list[int]] | None = None,
+):
+    """Tools to construct pipelined for loop.
+
+    Parameters
+    ----------
+    start : PrimExpr
+        The minimum value of iteration.
+    stop : PrimExpr
+        The maximum value of iteration.
+    num_stages : int
+        The max number of buffer used between pipeline producers and consumers.
+        if num_stages is 0, pipeline will not be enabled.
+    Returns
+    -------
+    res : frame.ForFrame
+        The ForFrame.
+    """
+    if stop is None:
+        stop = start
+        start = IntImm(start.dtype, 0) if hasattr(start, "dtype") else 0
+    if order is None:
+        order = []
+    if stage is None:
+        stage = []
+    if sync is None:
+        sync = []
+    if group is None:
+        group = []
+    # type: ignore[attr-defined] # pylint: disable=no-member
+    return _ffi_api.Pipelined(start, stop, num_stages, order, stage, sync, group)
+
+
+def serial(start: tir.PrimExpr,
+           stop: tir.PrimExpr | None = None,
+           step: tir.PrimExpr | None = None,
+           *,
+           annotations: dict[str, Any] | None = None):
+    step_is_one = False
+    step_is_one |= isinstance(step, int) and step == 1
+    step_is_one |= isinstance(step, IntImm) and step.value == 1
+    if step is None or step_is_one:
+        return tb_tir.serial(start, stop, annotations=annotations)
+    else:
+        if stop is None:
+            stop = start
+            start = IntImm(start.dtype, 0) if hasattr(start, "dtype") else 0
+        return SerialForWithStep(start, stop, step, annotations=annotations)