minor fix

Author: kurisu6912

tilelang/language/v2/ast.py

@@ -396,11 +396,18 @@ def flush_binds():
     def visit_Assign(self, node: ast.Assign) -> list[ast.AST]:
         node = self.generic_visit(node)
         rval = node.value
-        stmts = []
-        for target in reversed(node.targets):
-            stmts.extend(self._emit_assign_target(target, rval))
-            rval = target
-        return stmts
+        if len(node.targets) == 1:
+            return self._emit_assign_target(node.targets[0], rval)
+        else:
+            tmp_name = self.get_tmp()
+            tmp_store = ast.Name(tmp_name, ctx=ast.Store())
+            tmp_load = ast.Name(tmp_name, ctx=ast.Load())
+            ast_set_span(tmp_store, node.targets[0])
+            ast_set_span(tmp_load, node.targets[0])
+            stmt = self._emit_assign_target(tmp_store, rval)
+            for target in node.targets:
+                stmt.extend(self._emit_assign_target(target, tmp_load))
+            return stmt
 
     def visit_AugAssign(self, node: ast.AugAssign) -> list[ast.AST]:
         node = self.generic_visit(node)

tilelang/language/v2/builder.py

@@ -29,7 +29,7 @@ def unwrap_expr(expr) -> PrimExpr | int | float:
         expr = tir.BufferLoad(expr, indices=[0])
     elif isinstance(expr, (EqualOp, NotEqualOp)):
         expr = expr.asobject()
-    elif isinstance(expr, tir.IntImm) and expr.dtype == 'int32':
+    elif isinstance(expr, IntImm) and expr.dtype == 'int32':
         expr = expr.value
     return expr
 
@@ -257,10 +257,9 @@ def bind(self, name, value, annot=BaseBuilder.empty):
         return res
 
     def unwrap_value(self, value):
+        value = unwrap_expr(value)
         # handle bx, by = tl.Kernel(128, 128), rval is frame
-        if isinstance(value, tir.meta_var):
-            return value.value
-        elif isinstance(value, tir.frame.IRBuilderFrame):
+        if isinstance(value, tir.frame.IRBuilderFrame):
             return self.enter_frame(value)
         else:
             return value
@@ -295,9 +294,9 @@ def assign_slice(self, lval: Any, sl: slice, value: Any, annot=BaseBuilder.empty
             return super().assign_slice(lval, sl, value)
 
     def aug_assign(self, op, target, aug_value):
-        if isinstance(target, Buffer) and target.scope() == 'local.var':
-            tir.buffer_store(target, eval_op(op, target, aug_value), 0)
-        if isinstance(target, Buffer):
+        if is_var(target):
+            tir.buffer_store(target, eval_op(op, target[0], aug_value), 0)
+        elif isinstance(target, Buffer):
             raise RuntimeError("Augmented assignment is not supported for Buffer")
         else:
             return super().aug_assign(op, target, aug_value)
@@ -370,11 +369,7 @@ def rval(self, name: str, value: Any) -> Any:
                     f"Use immutable variable `{name}` outside its defining region, did you forget **alloc_var**?\n"
                     f"variable `{name}` is defined in frame: {frame}, current frames: {self.frames}."
                 )
-        if isinstance(value, tir.IntImm):
-            return value.value
-        if isinstance(value, Buffer) and value.scope() == 'local.var':
-            return tir.BufferLoad(value, indices=[0])
-        return super().rval(name, value)
+        return unwrap_expr(value)
 
     def arg(self, name, value):
         if self.find_frame_idx(MacroFrame) is not None:

tilelang/language/v2/utils.py

@@ -100,7 +100,7 @@ def get_compiled_object(source: str | ast.AST,
             compiled = disk_compile(source, name)
     except Exception as e:
         source_str = source if isinstance(source, str) else ast.unparse(source)
-        raise RuntimeError(f'Failed to compile source for {name}:\n{source_str}') from e
+        raise RuntimeError(f'Failed to compile source for {name}, Error: {e}:\n{source_str}') from e
     locs = {}
     exec(compiled, globals, locs)
     return locs[name]