mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
Lei Wang
c99879ad12
* Enhance augmented assignment methods in BaseBuilder and Builder classes - Updated `aug_assign` method to include an optional `name` parameter for better handling of variable names during augmented assignments. - Improved handling of immutable variables in the `Builder` class to treat augmented assignments as re-bindings, providing warnings for potential issues. - Added comments to clarify the changes and ensure user expectations are met regarding variable assignments. * lint fix
35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
import tilelang.language as T
|
|
from tilelang.language.eager.ast import mutate
|
|
|
|
|
|
def test_augassign_ast_no_placeholder_collision_for_value():
|
|
def demo(
|
|
A: T.Tensor[(1,), T.int32],
|
|
B: T.Tensor[(1,), T.int32],
|
|
):
|
|
value = A[0]
|
|
value -= 1
|
|
B[0] = value
|
|
|
|
ir_gen = mutate(demo)
|
|
# Regression test: the AugAssign lowering previously used placeholder names
|
|
# that collided with common user variable names like `value`, producing:
|
|
# value = __tb.aug_assign('Sub', 1, 1)
|
|
assert "__tb.aug_assign('Sub', value, 1, name='value')" in ir_gen.source
|
|
assert "__tb.aug_assign('Sub', 1, 1)" not in ir_gen.source
|
|
|
|
|
|
def test_augassign_immutable_var_is_lowered_as_rebind():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor[(1,), T.int32],
|
|
B: T.Tensor[(1,), T.int32],
|
|
):
|
|
x = A[0]
|
|
x -= 1
|
|
B[0] = x
|
|
|
|
# Just building the PrimFunc is sufficient; pre-fix this raised because
|
|
# augmented assignment on immutable `Var` was rejected.
|
|
assert main is not None
|