Simplify undefined variable in IR

Remove the mutable _undefined flag from Var and instead use InvalidType
directly at creation sites (scope lookup, store_invalid, aggregate
placeholders) to represent variables with partially defined type.

DCE no longer mutates Var state — instead of _mark_unused_vars_as_undefined, a
new _replace_pruned_with_dummy function inserts MakeDummy() op
for carried variables whose defining ops were pruned, using body vars as
type fallbacks when the pruned var has InvalidType.

This eliminates the special bytecode helpers (get_value_allow_undefined,
get_value_or_zero, undefined_value), allowing Loop/Continue/Break to
use plain get_value since all vars are now guaranteed to have
defining ops.

Downstream passes (alias analysis, code motion) are
simplified to use structural checks (map membership, dict.get)
instead of querying the removed flag.

Signed-off-by: Jay Gu <jagu@nvidia.com>

Author: haijieg

src/cuda/tile/_ir/ir.py

@@ -41,18 +41,18 @@ def __init__(self, config: TileContextConfig, tileiras_version: BytecodeVersion)
         self.tileiras_version: BytecodeVersion = tileiras_version
 
     #  Make a Var with a unique name based on `name`.
-    def make_var(self, name: str, loc: Loc, undefined: bool = False) -> Var:
+    def make_var(self, name: str, loc: Loc) -> Var:
         var_name = name
         while var_name in self._all_vars:
             var_name = f"{name}.{next(self._counter_by_name[name])}"
         self._all_vars[var_name] = name
-        return Var(var_name, loc, self, undefined)
+        return Var(var_name, loc, self)
 
     def make_var_like(self, var: Var) -> Var:
-        return self.make_var(self.get_original_name(var.name), var.loc, var.is_undefined())
+        return self.make_var(self.get_original_name(var.name), var.loc)
 
-    def make_temp(self, loc: Loc, undefined: bool = False) -> Var:
-        return self.make_var(f"${next(self._temp_counter)}", loc, undefined=undefined)
+    def make_temp(self, loc: Loc) -> Var:
+        return self.make_var(f"${next(self._temp_counter)}", loc)
 
     def get_original_name(self, var_name: str) -> str:
         return self._all_vars[var_name]
@@ -154,11 +154,10 @@ def as_tuple(self) -> tuple["Var", ...]:
 
 
 class Var:
-    def __init__(self, name: str, loc: Loc, ctx: IRContext, undefined: bool = False):
+    def __init__(self, name: str, loc: Loc, ctx: IRContext):
         self.name = name
         self.loc = loc
         self.ctx = ctx
-        self._undefined = undefined
 
     def try_get_type(self) -> Optional[Type]:
         return self.ctx.typemap.get(self.name)
@@ -173,9 +172,6 @@ def get_type_allow_invalid(self) -> Type:
         try:
             return self.ctx.typemap[self.name]
         except KeyError:
-            if self._undefined:
-                return InvalidType(f"Use of potentially undefined variable"
-                                   f" `{self.get_original_name()}`", loc=Loc.unknown())
             raise TileInternalError(f"Type of variable {self.name} not found")
 
     def set_type(self, ty: Type, force: bool = False):
@@ -213,12 +209,6 @@ def set_loose_type(self, ty: Type, force: bool = False):
             assert self.name not in self.ctx._loose_typemap
         self.ctx._loose_typemap[self.name] = ty
 
-    def is_undefined(self) -> bool:
-        return self._undefined
-
-    def set_undefined(self):
-        self._undefined = True
-
     def get_original_name(self) -> str:
         return self.ctx.get_original_name(self.name)
 

src/cuda/tile/_ir/ops.py

@@ -25,6 +25,7 @@
     TupleValue, make_aggregate, RangeValue, BoundMethodValue, ArrayValue, ConstantState,
     ListValue, ClosureValue, MemoryEffect, attribute, operand, BlockRestriction
 )
+from .type import PointerTy
 from . import hir
 from .hir import ResolvedName
 from .op_impl import (
@@ -102,8 +103,8 @@ def body_vars(self) -> tuple[Var, ...]:
     @override
     def generate_bytecode(self, ctx: BytecodeContext) -> tuple[bc.Value, ...]:
         types = tuple(x.get_type() for x in self.body_vars)
-        initial_values = [ctx.get_value_allow_undefined(input_var, ty)
-                          for input_var, ty in zip(self.initial_values, types, strict=True)]
+        initial_values = [ctx.get_value(input_var)
+                          for input_var in self.initial_values]
         result_type_ids = [typeid(ctx.type_table, ty) for ty in types]
 
         if self.is_for_loop:
@@ -224,9 +225,7 @@ async def loop_impl(body: hir.Block, iterable: Var):
         ty = body_var.get_type_allow_invalid()
         is_valid = not isinstance(ty, InvalidType)
         mask.append(is_valid)
-        if not is_valid:
-            body_var.set_undefined()
-        elif not was_valid and ty.is_aggregate():
+        if not was_valid and is_valid and ty.is_aggregate():
             # The initial variable is invalid but the loop variable is preserved,
             # and the loop variable is aggregate. In this case, `flat_body_vars[i]`
             # will contain a single variable (previously of InvalidType,
@@ -246,23 +245,28 @@ async def loop_impl(body: hir.Block, iterable: Var):
     new_body.params = tuple(body_params)
     valid_var_types = tuple(v.get_type() for v, is_valid in zip(body_vars, mask, strict=True)
                             if is_valid)
+    flat_var_types = flatten_aggregate_types(valid_var_types)
 
     # Update Continue/Break statements
     for jump_info in loop_info.jumps:
         values = tuple(out
                        for out, is_valid in zip(jump_info.outputs, mask, strict=True)
                        if is_valid)
         flat_values = flatten_aggregates(values, valid_var_types)
+        # For undefined break/continue value, add a MakeDummy op as its producer
+        flat_values = _add_dummy_op_to_invalid_vars(flat_values, flat_var_types)
         assert len(flat_values) == len(all_flattened_body_vars)
         jump_info.jump_op.values = flat_values
 
     # Create the loop Operation
-    valid_initial_values = tuple(v for v, is_valid in zip(initial_values, mask, strict=True)
+    valid_initial_values = tuple(v for v, is_valid
+                                 in zip(initial_values, mask, strict=True)
                                  if is_valid)
     flat_initial_values = flatten_aggregates(valid_initial_values, valid_var_types)
+    # For undefined initial value, add a MakeDummy op as its producer
+    flat_initial_values = _add_dummy_op_to_invalid_vars(flat_initial_values, flat_var_types)
     assert len(flat_initial_values) == len(all_flattened_body_vars)
 
-    flat_var_types = flatten_aggregate_types(valid_var_types)
     if range_ty is None:
         start = stop = step = None
     else:
@@ -297,7 +301,8 @@ async def loop_impl(body: hir.Block, iterable: Var):
     for body_var, state, local_idx, is_valid in zip(body_vars, var_states, stored_locals, mask,
                                                     strict=True):
         if not is_valid:
-            store_undefined(local_idx, body_var.get_type_allow_invalid(), state.result_phi.last_loc)
+            store_invalid(local_idx, body_var.get_type_allow_invalid(),
+                          state.result_phi.last_loc)
 
     # Do this check at the end because this may be an automatically inserted loop
     # around the helper function's body.
@@ -465,14 +470,16 @@ async def if_else_impl(cond: Var, then_block: hir.Block, else_block: hir.Block)
         assert num_explicit_results == 1
         ret = all_results[0]
         if ret is None:
-            ret = builder.ir_ctx.make_temp(builder.loc, undefined=True)
+            assert isinstance(result_phis[0].ty, InvalidType)
+            ret = builder.ir_ctx.make_temp(builder.loc)
+            ret.set_type(result_phis[0].ty)
 
     # Update the scope for stored named
     for res_var, phi, local_idx in zip(all_results[num_explicit_results:],
                                        result_phis[num_explicit_results:],
                                        stored_locals, strict=True):
         if res_var is None:
-            store_undefined(local_idx, phi.ty, phi.last_loc)
+            store_invalid(local_idx, phi.ty, phi.last_loc)
         else:
             store_var(local_idx, res_var, phi.last_loc)
 
@@ -486,9 +493,7 @@ class Continue(Operation, opcode="continue", terminator=True):
 
     @override
     def generate_bytecode(self, ctx: BytecodeContext) -> tuple[()]:
-        next_values = [ctx.get_value_allow_undefined(var, ctx.typeof(body_var))
-                       for var, body_var
-                       in zip(self.values, ctx.innermost_loop.body_vars, strict=True)]
+        next_values = [ctx.get_value(var) for var in self.values]
         bc.encode_ContinueOp(ctx.builder, next_values)
         return ()
 
@@ -517,11 +522,7 @@ class Break(Operation, opcode="break", terminator=True):
 
     @override
     def generate_bytecode(self, ctx: BytecodeContext) -> tuple[()]:
-        # body_vars is not a typo. We use body variables because they always contain the actual
-        # types of the loop variables, whereas result variables may have an InvalidType.
-        output_values = [ctx.get_value_allow_undefined(var, ctx.typeof(body_var))
-                         for var, body_var
-                         in zip(self.values, ctx.innermost_loop.body_vars, strict=True)]
+        output_values = [ctx.get_value(var) for var in self.values]
         bc.encode_BreakOp(ctx.builder, output_values)
         return ()
 
@@ -616,6 +617,27 @@ def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
         return ctx.constant(self.value, ctx.typeof(self.result_var))
 
 
+@dataclass(eq=False)
+class MakeDummy(Operation, opcode="make_dummy"):
+    """Placeholder value inserted for undefined variable in loop.
+
+    The use case for undefined variables is to represent loop's
+    initial_values or continue/break's next_values during type inference or
+    post dead code elimination.
+    """
+
+    @override
+    def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
+        ty = ctx.typeof(self.result_var)
+        if isinstance(ty, TokenTy):
+            return bc.encode_MakeTokenOp(ctx.builder, ctx.type_table.Token)
+        if isinstance(ty, TileTy) and isinstance(ty.dtype, PointerTy):
+            int_ty = TileTy(dtype=datatype.int64, shape=ty.shape)
+            const = ctx.constant(0, int_ty)
+            return bc.encode_IntToPtrOp(ctx.builder, typeid(ctx.type_table, ty), const)
+        return ctx.constant(0, ty)
+
+
 def loosely_typed_const(value: Any,
                         ty: Optional[Type] = None,
                         loose_ty: Optional[Type] = None) -> Var:
@@ -1583,8 +1605,6 @@ def _to_string_rhs(self) -> str:
 def assign(value: Var, res: Var) -> None:
     Builder.get_current().append_verbatim(Assign(value=value, result_vars=(res,), loc=res.loc))
     res.ctx.copy_type_information(value, res)
-    if value.is_undefined():
-        res.set_undefined()
 
 
 @impl(hir.identity)
@@ -1776,8 +1796,12 @@ def flatten_aggregates(vars: Sequence[Var], types: Sequence[Type]) -> tuple[Var,
     ret = []
     for x, ty in zip(vars, types, strict=True):
         item_types = tuple(ty.flatten_aggregate())
-        if isinstance(x.get_type_allow_invalid(), InvalidType):
-            ret.extend(x.ctx.make_temp(x.loc, undefined=True) for _ in item_types)
+        x_ty = x.get_type_allow_invalid()
+        if isinstance(x_ty, InvalidType):
+            for _ in item_types:
+                t = x.ctx.make_temp(x.loc)
+                t.set_type(x_ty)
+                ret.append(t)
         else:
             items = tuple(x.flatten_aggregate())
             assert len(items) == len(item_types)
@@ -1837,10 +1861,10 @@ def _unflatten_proper_aggregate(flattened_iter: Iterator[Var], nominal: Type, ac
         # Pop values from the iterator and throw them out
         for _ in nominal_item_types:
             next(flattened_iter)
-        # Return an undefined variable. It is OK that we don't create an aggregate value for it --
-        # any use of it should be invalid anyway.
         builder = Builder.get_current()
-        return builder.ir_crx.make_temp(builder.loc, undefined=True)
+        t = builder.ir_ctx.make_temp(builder.loc)
+        t.set_type(actual)
+        return t
 
     items = tuple(_maybe_unflatten_aggregate(flattened_iter, item_nominal, item_actual)
                   for item_nominal, item_actual
@@ -3991,10 +4015,10 @@ def store_var(local_idx: int, value: Var, loc: Loc | None = None):
     assign(value, new_var)
 
 
-def store_undefined(local_idx: int, ty: Type, loc: Loc | None = None):
+def store_invalid(local_idx: int, ty: Type, loc: Loc | None = None):
+    assert isinstance(ty, InvalidType)
     scope = Scope.get_current()
     new_var = scope.local.redefine(local_idx, loc or Builder.get_current().loc)
-    new_var.set_undefined()
     new_var.set_type(ty)
 
 
@@ -4014,8 +4038,6 @@ def load_var_impl(name):
     if rn.depth >= 0:
         ret = scope.local_scopes[rn.depth][rn.index]
         ret.get_type()  # Trigger an InvalidType check
-        if ret.is_undefined():
-            raise TileSyntaxError(f"Undefined variable {name} used")
         return ret
     elif rn.index >= 0:
         val = scope.func_hir.frozen_global_values[rn.index]
@@ -4206,3 +4228,11 @@ def sym2var(x: Any) -> Var:
 
     x = get_constant_value(x)
     return loosely_typed_const(x)
+
+
+def _add_dummy_op_to_invalid_vars(vars: Sequence[Var],
+                                  actual_types: Sequence[Type]) -> tuple[Var, ...]:
+    return tuple(add_operation(MakeDummy, actual)
+                 if isinstance(v.get_type_allow_invalid(), InvalidType)
+                 else v
+                 for v, actual in zip(vars, actual_types, strict=True))

src/cuda/tile/_ir/scope.py

@@ -13,6 +13,7 @@
 from cuda.tile._ir import hir
 from cuda.tile._ir.hir import ResolvedName
 from cuda.tile._ir.ir import Operation, Var, IRContext
+from cuda.tile._ir.type import InvalidType
 
 
 @dataclass
@@ -71,7 +72,9 @@ def get(self, index: int, loc: Loc):
         assert index >= 0
         var = self._map[index]
         if var is None:
-            return self._ir_ctx.make_var(self._local_names[index], loc, undefined=True)
+            name = self._local_names[index]
+            var = self._ir_ctx.make_var(name, loc)
+            var.set_type(InvalidType(f"Use of potentially undefined variable `{name}`", loc=loc))
         return var
 
     @contextmanager

src/cuda/tile/_ir2bytecode.py

@@ -313,9 +313,6 @@ def get_constant_or_default(self, var: Var, default=None):
     def get_value(self, var: Var) -> bc.Value:
         return self._value_map[var.name]
 
-    def get_value_allow_undefined(self, var: Var, ty: Type) -> bc.Value:
-        return self.undefined_value(ty) if var.is_undefined() else self.get_value(var)
-
     def get_optional_value(self, var: Var) -> Optional[bc.Value]:
         if var.name in self._constants and self._constants[var.name] is None:
             return None
@@ -365,16 +362,6 @@ def constant_tuple(self, value, ty: Type) -> Tuple[bc.Value, ...]:
                         for item_ty, item_val in zip(ty.value_types, value, strict=True)), ())
         return self.constant(value, ty),
 
-    def undefined_value(self, ty: Type) -> bc.Value:
-        if isinstance(ty, TokenTy):
-            return bc.encode_MakeTokenOp(self.builder, typeid(self.type_table, ty))
-
-        if isinstance(ty, TileTy) and isinstance(ty.dtype, PointerTy):
-            const = self.constant(0, TileTy(dtype=datatype.int64, shape=ty.shape))
-            return bc.encode_IntToPtrOp(self.builder, typeid(self.type_table, ty), const)
-
-        return self.constant(0, ty)
-
     def index_tuple(self, index: tuple[Var, ...]) -> Tuple[bc.Value, ...]:
         i32_tile_ty = self.type_table.tile(self.type_table.I32, ())
         item_types = tuple(x.get_type() for x in index)

src/cuda/tile/_passes/alias_analysis.py

@@ -83,9 +83,6 @@ def finalize(self):
 def _propagate(alias_tracker: _AliasTracker,
                src: Var,
                dst: Var):
-    if src.is_undefined():
-        alias_tracker[src.name] = frozenset()
-
     src_aliases = alias_tracker[src.name]
     dst_aliases = alias_tracker.get(dst.name, frozenset())
     alias_tracker[dst.name] = dst_aliases | src_aliases

src/cuda/tile/_passes/code_motion.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-from cuda.tile._ir.ir import Block, Var, MemoryEffect
+from cuda.tile._ir.ir import Block, MemoryEffect
 from cuda.tile._ir.ops import Loop, IfElse, Continue, Break, EndBranch, Return, TileReduce, TileScan
 
 from dataclasses import dataclass
@@ -93,10 +93,10 @@ def _hoist(block: Block, stack: list[_StackItem], def_depth: dict[str, int], is_
 
             inputs = op.initial_values if isinstance(op, Loop) else op.xs
             for v in inputs:
-                depinfo.update(_get_def_depth(def_depth, v), depth)
+                depinfo.update(def_depth[v.name], depth)
             depinfo.update(body_res.min_depth, depth)
         elif isinstance(op, IfElse):
-            depinfo.update(_get_def_depth(def_depth, op.cond), depth)
+            depinfo.update(def_depth[op.cond.name], depth)
             for branch in (op.then_block, op.else_block):
                 branch_res = _hoist(branch, stack, def_depth, False)
                 depinfo.update(branch_res.min_depth, depth)
@@ -115,12 +115,12 @@ def _hoist(block: Block, stack: list[_StackItem], def_depth: dict[str, int], is_
         elif isinstance(op, EndBranch):
             depinfo.must_stay = True
             for v in op.outputs:
-                depinfo.update(_get_def_depth(def_depth, v), depth)
+                depinfo.update(def_depth[v.name], depth)
         else:
             # "Pure" operation without any nested blocks, side effects and jumps.
             assert len(op.nested_blocks) == 0
             for v in op.all_inputs():
-                depinfo.update(_get_def_depth(def_depth, v), depth)
+                depinfo.update(def_depth[v.name], depth)
 
         target_depth = depth
         if depinfo.must_stay:
@@ -140,12 +140,3 @@ def _hoist(block: Block, stack: list[_StackItem], def_depth: dict[str, int], is_
     stack.pop()
     block[:] = new_block.detach_all()
     return ret
-
-
-def _get_def_depth(def_depth: dict[str, int], var: Var) -> int:
-    try:
-        return def_depth[var.name]
-    except KeyError:
-        pass
-    assert var.is_undefined(), var.name
-    return 0