Skip to content

[Fix] Fix buffer re-import typo in tilelang.languge #1214

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LeiWang1999 merged 2 commits into from
Nov 10, 2025
Merged

[Fix] Fix buffer re-import typo in tilelang.languge #1214

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ca8e74d
Fix Buffer re-import typo in tilelang.langugage
kurisu6912 Nov 10, 2025
32bdb07
fix lint error
kurisu6912 Nov 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions testing/python/issue/test_tilelang_issue_1198.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
import tilelang.testing
import tilelang.language as T


def test_issue_1198():

@T.prim_func
def foo(x: T.Buffer([
32,
], "int32")):
pass


if __name__ == '__main__':
tilelang.testing.main()
25 changes: 14 additions & 11 deletions tilelang/language/builtin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from tvm import DataType, tir
from tvm.runtime import convert
from typing import Any
from tvm.tir import PrimExpr, Var, Call, Buffer, BufferLoad
from tvm.tir import PrimExpr, Var, Call, BufferLoad

_IS_HIP_AVAILABLE = check_hip_availability()

Expand Down Expand Up @@ -430,7 +430,7 @@ def shuffle_elect(thread_extent: int) -> PrimExpr:
return tir.call_intrin("bool", tir.op.Op.get("tl.tl_shuffle_elect"), thread_extent)


def warpgroup_fence_operand(buffer_or_ptr: Buffer | PrimExpr,
def warpgroup_fence_operand(buffer_or_ptr: tir.Buffer | PrimExpr,
offset: int | PrimExpr = 0,
num_regs: int | PrimExpr | None = None,
dtype: str | None = None):
Expand All @@ -456,7 +456,7 @@ def warpgroup_fence_operand(buffer_or_ptr: Buffer | PrimExpr,
if isinstance(buffer_or_ptr, BufferLoad):
raise TypeError("Expected a buffer handle or pointer expression, got BufferLoad.")

if isinstance(buffer_or_ptr, Buffer):
if isinstance(buffer_or_ptr, tir.Buffer):
data_ptr = buffer_or_ptr.data
inferred_dtype = buffer_or_ptr.dtype
if dtype is not None and dtype != inferred_dtype:
Expand Down Expand Up @@ -599,18 +599,19 @@ def sync_grid():


def initialize_wgmma_descriptor(
descriptor: Buffer,
descriptor: tir.Buffer,
start_address: PrimExpr,
layout_type_: int = 0,
leading_byte_offset: int = 0,
stride_byte_offset: int = 0,
) -> PrimExpr:
"""Initialize a WGMMA/UTCMMA shared-memory descriptor."""

if not isinstance(descriptor, (BufferLoad, Buffer)):
if not isinstance(descriptor, (BufferLoad, tir.Buffer)):
raise TypeError("Descriptor must be a tvm.tir.Buffer or tvm.tir.BufferLoad.")

if isinstance(descriptor, Buffer) and (len(descriptor.shape) != 1 or descriptor.shape[0] != 1):
if isinstance(descriptor, tir.Buffer) and (len(descriptor.shape) != 1 or
descriptor.shape[0] != 1):
raise ValueError("Descriptor must be a 1D buffer of size 1.")

descriptor = descriptor if isinstance(descriptor, BufferLoad) else tir.BufferLoad(
Expand All @@ -629,7 +630,7 @@ def initialize_wgmma_descriptor(


def initialize_tcgen05_descriptor(
descriptor: Buffer,
descriptor: tir.Buffer,
start_address: PrimExpr,
leading_byte_offset: int,
stride_byte_offset: int,
Expand All @@ -639,10 +640,11 @@ def initialize_tcgen05_descriptor(
) -> PrimExpr:
"""Initialize a TCGEN05 shared-memory descriptor."""

if not isinstance(descriptor, (BufferLoad, Buffer)):
if not isinstance(descriptor, (BufferLoad, tir.Buffer)):
raise TypeError("Descriptor must be a tvm.tir.Buffer or tvm.tir.BufferLoad.")

if isinstance(descriptor, Buffer) and (len(descriptor.shape) != 1 or descriptor.shape[0] != 1):
if isinstance(descriptor, tir.Buffer) and (len(descriptor.shape) != 1 or
descriptor.shape[0] != 1):
raise ValueError("Descriptor must be a 1D buffer of size 1.")

descriptor = descriptor if isinstance(descriptor, BufferLoad) else tir.BufferLoad(
Expand Down Expand Up @@ -673,10 +675,11 @@ def increase_descriptor_offset(descriptor: PrimExpr, offset: PrimExpr) -> PrimEx
Returns:
PrimExpr: A handle representing the modified descriptor.
"""
if not isinstance(descriptor, (BufferLoad, Buffer)):
if not isinstance(descriptor, (BufferLoad, tir.Buffer)):
raise TypeError("Descriptor must be a tvm.tir.Buffer or tvm.tir.BufferLoad.")

if isinstance(descriptor, Buffer) and len(descriptor.shape) != 1 or descriptor.shape[0] != 1:
if isinstance(descriptor, tir.Buffer) and len(
descriptor.shape) != 1 or descriptor.shape[0] != 1:
raise ValueError("Descriptor must be a 1D buffer of size 1.")
Comment on lines +678 to 683
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Fix operator precedence issue in condition

Lines 681-682 have a logic error due to operator precedence. The condition:

if isinstance(descriptor, tir.Buffer) and len(descriptor.shape) != 1 or descriptor.shape[0] != 1:

Evaluates as: (isinstance(descriptor, tir.Buffer) and len(descriptor.shape) != 1) or (descriptor.shape[0] != 1)

This could cause an AttributeError when descriptor is a BufferLoad, as the second part of the or evaluates without the isinstance guard.

Apply this diff to fix the precedence:

-    if isinstance(descriptor, tir.Buffer) and len(
-            descriptor.shape) != 1 or descriptor.shape[0] != 1:
+    if isinstance(descriptor, tir.Buffer) and (len(descriptor.shape) != 1 or descriptor.shape[0] != 1):

Additionally, consider extracting the error message:

-    if not isinstance(descriptor, (BufferLoad, tir.Buffer)):
-        raise TypeError("Descriptor must be a tvm.tir.Buffer or tvm.tir.BufferLoad.")
+    DESCRIPTOR_TYPE_ERROR = "Descriptor must be a tvm.tir.Buffer or tvm.tir.BufferLoad."
+    if not isinstance(descriptor, (BufferLoad, tir.Buffer)):
+        raise TypeError(DESCRIPTOR_TYPE_ERROR)
🧰 Tools
🪛 Ruff (0.14.3)

679-679: Avoid specifying long messages outside the exception class

(TRY003)

681-682: Parenthesize a and b expressions when chaining and and or together, to make the precedence clear

Parenthesize the and subexpression

(RUF021)

683-683: Avoid specifying long messages outside the exception class

(TRY003)

🤖 Prompt for AI Agents 
In tilelang/language/builtin.py around lines 678 to 683, the conditional has an
operator precedence bug causing the descriptor.shape[0] check to run for
non-Buffer objects; change the condition to guard both shape checks behind the
isinstance test by grouping the shape checks with parentheses (e.g. if
isinstance(descriptor, tir.Buffer) and (len(descriptor.shape) != 1 or
descriptor.shape[0] != 1): raise ValueError(...)), and extract the ValueError
message into a small named variable or constant above the check so the message
is declared once and reused.


descriptor = descriptor if isinstance(descriptor, BufferLoad) else tir.BufferLoad(
Expand Down
Loading