Delay the eliminate_load_store pass until the hir2ir stage

This is in preparation for implementing closures/lambdas. But this also
fixes an edge case with flattening of if-else statements with a constant
condition (see test_constfold.py).

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

src/cuda/tile/_compile.py

@@ -16,7 +16,7 @@
 import tempfile
 import threading
 import traceback
-from typing import Callable, Optional
+from typing import Callable, Optional, Any, Set
 import zipfile
 
 from cuda.tile._cext import get_compute_capability, TileContext, default_tile_context
@@ -25,10 +25,11 @@
 from cuda.tile._exception import (
     TileCompilerError,
     TileCompilerExecutionError,
-    TileCompilerTimeoutError,
+    TileCompilerTimeoutError, TileValueError, TileTypeError
 )
 from cuda.tile._ir import ir, hir
-from cuda.tile._ir.ir import bind_kernel_arguments
+from cuda.tile._ir.ir import Argument
+from cuda.tile._ir.typing_support import typeof_pyval, get_constant_value
 from cuda.tile._passes.ast2hir import get_function_hir
 from cuda.tile._passes.code_motion import hoist_loop_invariants
 from cuda.tile._passes.eliminate_assign_ops import eliminate_assign_ops
@@ -74,9 +75,9 @@ def wrapper(*args, **kwargs):
 
 def _get_final_ir(pyfunc, args, tile_context) -> ir.Block:
     ir_ctx = ir.IRContext(tile_context)
-    func_hir: hir.Block = get_function_hir(pyfunc, ir_ctx, call_site=None)
+    func_hir: hir.Function = get_function_hir(pyfunc, ir_ctx, call_site=None)
 
-    ir_args = bind_kernel_arguments(func_hir.params, args, get_constant_annotations(pyfunc))
+    ir_args = _bind_kernel_arguments(func_hir.param_names, args, get_constant_annotations(pyfunc))
     func_body = hir2ir(func_hir, ir_args, ir_ctx)
     eliminate_assign_ops(func_body)
     dead_code_elimination_pass(func_body)
@@ -96,6 +97,34 @@ def _get_final_ir(pyfunc, args, tile_context) -> ir.Block:
     return func_body
 
 
+def _bind_kernel_arguments(param_names: tuple[str, ...],
+                           args: tuple[Any, ...],
+                           constant_args: Set[str]) -> tuple[Argument, ...]:
+    # TODO: unify this logic with dispatcher from c extension
+    # Refactor "extract_cuda_args" to return type descriptor
+    # that can be wrapped as IR Type for type inference.
+    if len(args) != len(param_names):
+        msg = f"Expected {len(param_names)} arguments, got {len(args)}"
+        raise TileValueError(msg)
+
+    ir_args = []
+    for param_name, arg_value in zip(param_names, args, strict=True):
+        const_val = None
+        is_const = param_name in constant_args
+        ty = typeof_pyval(arg_value, kernel_arg=not is_const)
+        if is_const:
+            try:
+                const_val = get_constant_value(arg_value)
+            except TileTypeError:
+                raise TileTypeError(
+                    f"Argument `{param_name}` is a constexpr, "
+                    f"but the value is not a supported constant.")
+        ir_args.append(Argument(type=ty,
+                                is_const=is_const,
+                                const_value=const_val))
+    return tuple(ir_args)
+
+
 def _log_mlir(bytecode_buf):
     try:
         from cuda.tile_internal import _internal_cext

src/cuda/tile/_ir/hir.py

@@ -17,7 +17,7 @@
 import enum
 from dataclasses import dataclass
 from textwrap import indent
-from typing import Any, Set
+from typing import Any, Set, Mapping
 
 from cuda.tile._exception import Loc
 from cuda.tile._ir.ir import Var
@@ -95,6 +95,14 @@ def jump_str(self):
         return f"{self.jump._value_}{results_str}  # Line {self.jump_loc.line}"
 
 
+@dataclass
+class Function:
+    body: Block
+    param_names: tuple[str, ...]
+    param_locs: tuple[Loc, ...]
+    frozen_globals: Mapping[str, Any]
+
+
 @dataclass
 class _OperandFormatter:
     blocks: list["Block"]
@@ -118,6 +126,8 @@ def __call__(self, x: Operand) -> str:
 # ==================================
 
 def if_else(cond, then_block, else_block, /): ...
-def loop(body, iterable, /, *initial_values): ...  # infinite if `iterable` is None
+def loop(body, iterable, /): ...  # infinite if `iterable` is None
 def build_tuple(*items): ...  # Makes a tuple (i.e. returns `items`)
 def identity(x): ...   # Identity function (i.e. returns `x`)
+def store_var(name, value, /): ...  # Store into a named variable
+def load_var(name, /): ...  # Load from a named variable

src/cuda/tile/_ir/ir.py

@@ -18,11 +18,8 @@
     List, Optional, Dict, Tuple, Set, Any, TYPE_CHECKING, Sequence, Iterator
 )
 from .type import Type, InvalidType
-from .typing_support import typeof_pyval, get_constant_value, loose_type_of_pyval
 from cuda.tile._exception import (
-    TileTypeError,
-    TileValueError,
-    Loc, TileInternalError
+    TileTypeError, Loc, TileInternalError, TileSyntaxError
 )
 from .._cext import TileContext
 
@@ -308,28 +305,122 @@ def finalize_loopvar_type(self, body_var: Var):
 class LoopInfo:
     var_states: tuple[LoopVarState, ...]
     is_for_loop: bool
+    stored_names: tuple[str, ...]
+    flatten: bool = False
 
 
-def get_innermost_loop() -> LoopInfo | None:
-    return Builder.get_current().loop_info
+@dataclass
+class IfElseInfo:
+    result_phis: tuple[PhiState, ...]
+    stored_names: tuple[str, ...]
+    flatten: bool = False
+    flattened_results: tuple[Var, ...] = ()
+    have_end_branch: bool = False
 
 
 @contextmanager
 def nested_block(name: str, loc: Loc, params: Sequence[Var] = (),
-                 loop_info: LoopInfo | None = None):
+                 loop_info: LoopInfo | None = None,
+                 if_else_info: IfElseInfo | None = None):
     prev_builder = Builder.get_current()
     block = Block(prev_builder.ir_ctx, params=params, name=name, loc=loc)
     new_loop_info = loop_info or prev_builder.loop_info
-    with Builder(prev_builder.ir_ctx, loc, new_loop_info) as builder:
+    new_if_else_info = if_else_info or prev_builder.if_else_info
+    scope = prev_builder.scope
+    with Builder(prev_builder.ir_ctx, loc, scope, new_loop_info, new_if_else_info) as builder, \
+            scope.local.enter_branch():
         yield block
     block.extend(builder.ops)
 
 
+class LocalScope:
+    def __init__(self,
+                 all_locals: Set[str],
+                 ir_ctx: IRContext,
+                 parent: Optional["LocalScope"] = None):
+        self._all_locals = all_locals
+        self._ir_ctx = ir_ctx
+        self._map = dict()
+        self._parent = parent
+
+    def is_local_name(self, name: str):
+        current = self
+        while current is not None:
+            if name in current._all_locals:
+                return True
+            current = current._parent
+        return False
+
+    def redefine(self, name: str, loc: Loc) -> Var:
+        var = self._ir_ctx.make_var(name, loc)
+        self._map[name] = var
+        return var
+
+    def __getitem__(self, name: str):
+        var = self._lookup(name)
+        if var is None:
+            raise TileSyntaxError(f"Undefined variable {name} used")
+        return var
+
+    def get(self, name: str, loc: Loc):
+        var = self._lookup(name)
+        if var is None:
+            return self._ir_ctx.make_var(name, loc, undefined=True)
+        else:
+            return var
+
+    def _lookup(self, name: str) -> Optional[Var]:
+        seen = set()
+        current = self
+        while current is not None:
+            var = current._map.get(name)
+            if var is not None:
+                return var
+            # Sanity check, should not reach here.
+            if id(current) in seen:
+                raise RuntimeError("Cycle detected in Scope chain")
+            seen.add(id(current))
+            current = current._parent
+        return None
+
+    @contextmanager
+    def enter_branch(self):
+        old = self._map
+        self._map = _OverlayDict(old)
+        try:
+            yield
+        finally:
+            self._map = old
+
+
+class _OverlayDict:
+    def __init__(self, orig_dict: dict):
+        self._orig = orig_dict
+        self._overlay = dict()
+
+    def get(self, key):
+        value = self._overlay.get(key)
+        return self._orig.get(key) if value is None else value
+
+    def __setitem__(self, key, value):
+        self._overlay[key] = value
+
+
+@dataclass
+class Scope:
+    local: LocalScope
+    frozen_globals: Mapping[str, Any]
+
+
 class Builder:
-    def __init__(self, ctx: IRContext, loc: Loc, loop_info: LoopInfo | None = None):
+    def __init__(self, ctx: IRContext, loc: Loc, scope: Scope,
+                 loop_info: LoopInfo | None = None,
+                 if_else_info: IfElseInfo | None = None):
         self.ir_ctx = ctx
+        self.scope = scope
         self.is_terminated = False
         self.loop_info = loop_info
+        self.if_else_info = if_else_info
         self._loc = loc
         self._ops = []
         self._entered = False
@@ -363,6 +454,10 @@ def add_operation(self, op_class,
     def ops(self) -> list[Operation]:
         return self._ops
 
+    @property
+    def loc(self) -> Loc:
+        return self._loc
+
     def append_verbatim(self, op: Operation):
         self._ops.append(op)
 
@@ -384,6 +479,24 @@ def change_loc(self, loc: Loc):
         finally:
             self._loc = old_loc
 
+    @contextmanager
+    def change_if_else_info(self, new_info: IfElseInfo):
+        old = self.if_else_info
+        self.if_else_info = new_info
+        try:
+            yield
+        finally:
+            self.if_else_info = old
+
+    @contextmanager
+    def change_loop_info(self, new_info: LoopInfo):
+        old = self.loop_info
+        self.loop_info = new_info
+        try:
+            yield
+        finally:
+            self.loop_info = old
+
     def __enter__(self):
         assert not self._entered
         self._prev_builder = _current_builder.builder
@@ -685,69 +798,8 @@ def __str__(self) -> str:
         return self.to_string()
 
 
-def bind_kernel_arguments(params: Tuple[Var, ...],
-                          args: Tuple[Any, ...],
-                          constant_args: Set[str]) -> Tuple[Argument, ...]:
-    # TODO: unify this logic with dispatcher from c extension
-    # Refactor "extract_cuda_args" to return type descriptor
-    # that can be wrapped as IR Type for type inference.
-    if len(args) != len(params):
-        msg = f"Expected {len(params)} arguments, got {len(args)}"
-        raise TileValueError(msg)
-
-    ir_args = []
-    for param, arg_value in zip(params, args):
-        const_val = None
-        is_const = param.name in constant_args
-        ty = typeof_pyval(arg_value, kernel_arg=not is_const)
-        loose_type = ty
-        if is_const:
-            try:
-                const_val = get_constant_value(arg_value)
-            except TileTypeError:
-                raise TileTypeError(
-                    f"Argument {param.name} is a constexpr, "
-                    f"but the value is not a supported constant.")
-            loose_type = loose_type_of_pyval(arg_value)
-        ir_args.append(Argument(type=ty,
-                                loose_type=loose_type,
-                                is_const=is_const,
-                                const_value=const_val))
-    return tuple(ir_args)
-
-
+@dataclass
 class Argument:
-    def __init__(self,
-                 type: Type,
-                 loose_type: Type,
-                 is_const: bool = False,
-                 const_value: Any = None):
-        self._type = type
-        self._loose_type = loose_type
-        self._is_const = is_const
-        self._const_value = const_value
-
-    @property
-    def is_const(self) -> bool:
-        return self._is_const
-
-    @property
-    def const_value(self) -> Any:
-        return self._const_value
-
-    @property
-    def type(self) -> Type:
-        return self._type
-
-    @property
-    def loose_type(self) -> Type:
-        return self._loose_type
-
-    def __eq__(self, value: object) -> bool:
-        if not isinstance(value, Argument):
-            return False
-        return (
-            self.type == value.type and
-            self.is_const == value.is_const and
-            self.const_value == value.const_value
-        )
+    type: Type
+    is_const: bool = False
+    const_value: Any = None