Decouple HIR from the IRContext

Instead of ir.Var, HIR now uses a separate hir.Value. It is numbered,
not named, and is unique per function rather than per IRContext. This
allows us to convert AST to HIR once and reuse it multiple times.
We use this property to add an LRU cache to get_function_hir().

As a result, we lose the mapping between HIR and IR using variable
names, which is occasionally useful for debugging. Hopefully, line
numbers should suffice for this purpose. Also, there is no name
mapping to the AST anyway.

Giving up on this name-based mapping also simplifies the logic in
_dispatch_call(): we no longer need to remap the variable names and can
directly use the variable returned by the op implementation. We also get
rid of two uses of assign(): there is now only one use remaining in
store_var_impl().

Signed-off-by: Greg Bonik <[email protected]>

Author: gbonik

src/cuda/tile/_compile.py

@@ -73,9 +73,9 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def _get_final_ir(pyfunc, args, tile_context) -> ir.Block:
+def _get_final_ir(pyfunc, args, tile_context) -> ir.Function:
     ir_ctx = ir.IRContext(tile_context)
-    func_hir: hir.Function = get_function_hir(pyfunc, ir_ctx, entry_point=True)
+    func_hir: hir.Function = get_function_hir(pyfunc, entry_point=True)
 
     ir_args = _bind_kernel_arguments(func_hir.param_names, args, get_constant_annotations(pyfunc))
     func_body = hir2ir(func_hir, ir_args, ir_ctx)
@@ -94,7 +94,7 @@ def _get_final_ir(pyfunc, args, tile_context) -> ir.Block:
 
     split_loops(func_body)
     dead_code_elimination_pass(func_body)
-    return func_body
+    return ir.Function(func_body, func_hir.desc.name, func_hir.body.loc)
 
 
 def _bind_kernel_arguments(param_names: tuple[str, ...],
@@ -143,7 +143,7 @@ def _log_mlir(bytecode_buf):
     print(f"Lowering\n==== TILEIR MLIR module ====\n\n{text}", file=sys.stderr)
 
 
-def _compiler_crash_dump(func_body,
+def _compiler_crash_dump(func_ir: ir.Function,
                          bytecode_generator,
                          error_msg,
                          compiler_flags,
@@ -161,13 +161,13 @@ def _compiler_crash_dump(func_body,
         bytecode_generator(writer, anonymize_debug_attr=True)
 
     artifacts = {
-        f"{func_body.name}.bytecode": bytes(bytecode_buf),
-        f"{func_body.name}.cutileir": f"{func_body.to_string(include_loc=False)}\n",
+        f"{func_ir.name}.bytecode": bytes(bytecode_buf),
+        f"{func_ir.name}.cutileir": f"{func_ir.to_string(include_loc=False)}\n",
         "debug_info.txt": debug_info,
     }
 
     timestamp = datetime.datetime.now().timestamp()
-    zip_filename = os.path.abspath(f"crash_dump_{func_body.name}_{timestamp}.zip")
+    zip_filename = os.path.abspath(f"crash_dump_{func_ir.name}_{timestamp}.zip")
     print(f"Dumping crash artifacts to {zip_filename}\n", file=sys.stderr)
 
     with zipfile.ZipFile(zip_filename, "w") as z:
@@ -180,17 +180,17 @@ def compile_tile(pyfunc,
                  args,
                  compiler_options: CompilerOptions,
                  context: TileContext = default_tile_context) -> TileLibrary:
-    func_body = _get_final_ir(pyfunc, args, context)
+    func_ir = _get_final_ir(pyfunc, args, context)
 
     if 'CUTILEIR' in context.config.log_keys:
-        code = (f"==== CuTile IR for {func_body.name}==== \n\n"
-                f"{func_body.to_string(include_loc=False)}\n\n")
+        code = (f"==== CuTile IR for {func_ir.name}==== \n\n"
+                f"{func_ir.to_string(include_loc=False)}\n\n")
         print(f'\n{code}', file=sys.stderr)
 
     sm_arch = get_sm_arch()
 
     bytecode_generator = functools.partial(generate_bytecode_for_kernel,
-                                           func_body, compiler_options, sm_arch)
+                                           func_ir, compiler_options, sm_arch)
 
     bytecode_buf = bytearray()
     with bc.write_bytecode(num_functions=1, buf=bytecode_buf) as writer:
@@ -202,9 +202,9 @@ def compile_tile(pyfunc,
     if CUDA_TILE_DUMP_BYTECODE is not None:
         if not os.path.exists(CUDA_TILE_DUMP_BYTECODE):
             os.makedirs(CUDA_TILE_DUMP_BYTECODE)
-        base_filename = os.path.basename(func_body.loc.filename.split(".")[0])
+        base_filename = os.path.basename(func_ir.loc.filename.split(".")[0])
         path = os.path.join(CUDA_TILE_DUMP_BYTECODE,
-                            f"{base_filename}.ln{func_body.loc.line}.cutile")
+                            f"{base_filename}.ln{func_ir.loc.line}.cutile")
         print(f"Dumping TILEIR bytecode to file: {path}", file=sys.stderr)
         with open(path, "wb") as f:
             f.write(bytecode_buf)
@@ -216,9 +216,9 @@ def compile_tile(pyfunc,
             mlir_text = bytecode_to_mlir_text(bytecode_buf)
             if not os.path.exists(CUDA_TILE_DUMP_TILEIR):
                 os.makedirs(CUDA_TILE_DUMP_TILEIR)
-            base_filename = os.path.basename(func_body.loc.filename.split(".")[0])
+            base_filename = os.path.basename(func_ir.loc.filename.split(".")[0])
             path = os.path.join(
-                CUDA_TILE_DUMP_TILEIR, f"{base_filename}.ln{func_body.loc.line}.cuda_tile.mlir"
+                CUDA_TILE_DUMP_TILEIR, f"{base_filename}.ln{func_ir.loc.line}.cuda_tile.mlir"
             )
             print(f"Dumping TILEIR MLIR module to file:{path}", file=sys.stderr)
             with open(path, "w") as f:
@@ -228,7 +228,7 @@ def compile_tile(pyfunc,
                   "This is currently not a public feature.", file=sys.stderr)
 
     # Compile MLIR module and generate cubin
-    with tempfile.NamedTemporaryFile(suffix='.bytecode', prefix=func_body.name,
+    with tempfile.NamedTemporaryFile(suffix='.bytecode', prefix=func_ir.name,
                                      dir=context.config.temp_dir, delete=False) as f:
         f.write(bytecode_buf)
         f.flush()
@@ -238,12 +238,12 @@ def compile_tile(pyfunc,
                                        timeout_sec=context.config.compiler_timeout_sec)
         except TileCompilerError as e:
             if context.config.enable_crash_dump:
-                _compiler_crash_dump(func_body, bytecode_generator, e.message,
+                _compiler_crash_dump(func_ir, bytecode_generator, e.message,
                                      e.compiler_flags, e.compiler_version)
 
             raise e
 
-    return TileLibrary(func_body.name, cubin_file, bytecode_buf, func_body)
+    return TileLibrary(func_ir.name, cubin_file, bytecode_buf, func_ir.body)
 
 
 # Adapter between compile_tile() and kernel/TileDispatcher

src/cuda/tile/_ir/hir.py

@@ -15,27 +15,58 @@
 
 
 import enum
+import threading
 from dataclasses import dataclass
 from textwrap import indent
 from typing import Any, Set, Mapping
 
 from cuda.tile._exception import Loc, FunctionDesc
-from cuda.tile._ir.ir import Var
+
+
+@dataclass(frozen=True)
+class Value:
+    id: int
+
+    def __str__(self):
+        return f"%{self.id}"
+
+
+# Cache Value objects for reuse
+_value_cache = []
+_value_cache_lock = threading.Lock()
+
+
+def make_value(id: int) -> Value:
+    try:
+        # Fast path
+        return _value_cache[id]
+    except IndexError:
+        pass
+
+    if id >= 2000:
+        # Don't cache too many objects
+        return Value(id)
+
+    # Add 100 objects at a time to avoid triggering the slow path every time
+    with _value_cache_lock:
+        _value_cache.extend([Value(j) for j in range(len(_value_cache), id + 100)])
+    return _value_cache[id]
 
 
 # An "Operand" is a value that can be used as a function's argument, or as the function itself.
-# There are two kinds of Operands: variables and constants. Using a `Var` instance as an Operand
-# signals that this Operand is a variable, e.g. a result of a previous call or a kernel parameter.
-# An object of any other type means that it is an immediate constant.
-Operand = Var | Any
+# There are two kinds of Operands:
+#    - Using a `Value` instance as an Operand signals that this Operand is a result
+#      of a previous call, or a kernel parameter.
+#    - An object of any other type means that it is an immediate constant.
+Operand = Value | Any
 
 
 ModuleType = type(enum)
 
 
 @dataclass
 class Call:
-    results: tuple[Var, ...]
+    result: Value | None
     callee: Operand
     args: tuple[Operand, ...]
     kwargs: tuple[tuple[str, Operand], ...]
@@ -45,24 +76,16 @@ def __str__(self):
         opfmt = _OperandFormatter([])
         loc_str = f"  # Line {self.loc.line}"
         if self.callee is identity:
-            return f"{_lhs_var_str(self.results[0])} = {opfmt(self.args[0])}{loc_str}"
+            return f"{self.result} = {opfmt(self.args[0])}{loc_str}"
+        lhs_str = "" if self.result is None else f"{self.result} = "
         callee_str = opfmt(self.callee)
-        results_str = ", ".join(_lhs_var_str(r) for r in self.results)
-        lhs_str = f"{results_str} = " if results_str else ""
         args_and_kwargs = (*(opfmt(a) for a in self.args),
                            *(f"{k}={opfmt(v)}" for k, v in self.kwargs))
         args_str = ", ".join(args_and_kwargs)
         blocks_str = "".join(indent(f"\n{b}", "    ") for b in opfmt.blocks)
         return f"{lhs_str}{callee_str}({args_str}){loc_str}{blocks_str}"
 
 
-def _lhs_var_str(var: Var):
-    ty = var.try_get_type()
-    if ty is None:
-        return var.name
-    return f"{var.name}: {ty}"
-
-
 class Jump(enum.Enum):
     END_BRANCH = "end_branch"
     CONTINUE = "continue"
@@ -72,26 +95,27 @@ class Jump(enum.Enum):
 
 @dataclass
 class Block:
-    name: str
-    params: tuple[Var, ...]
+    block_id: int
+    params: tuple[Value, ...]
     calls: list[Call]
-    results: tuple[Operand, ...]
+    have_result: bool
+    result: Operand
     jump: Jump | None
     jump_loc: Loc
     stored_names: Set[str]
     loc: Loc
 
     def __str__(self):
-        params_str = ", ".join(p.name for p in self.params)
+        params_str = ", ".join(str(p) for p in self.params)
         calls_str = "".join(f"\n{c}" for c in self.calls)
         if self.jump is not None:
             calls_str += "\n" + self.jump_str()
         calls_str = indent(calls_str, "    ")
-        return f"^{self.name}({params_str}):{calls_str}"
+        return f"^{self.block_id}({params_str}):{calls_str}"
 
     def jump_str(self):
         opfmt = _OperandFormatter([])
-        results_str = ",".join(f" {opfmt(r)}" for r in self.results)
+        results_str = "" if self.result is None else opfmt(self.result)
         return f"{self.jump._value_}{results_str}  # Line {self.jump_loc.line}"
 
 
@@ -102,20 +126,21 @@ class Function:
     param_names: tuple[str, ...]
     param_locs: tuple[Loc, ...]
     frozen_globals: Mapping[str, Any]
+    value_id_upper_bound: int
 
 
 @dataclass
 class _OperandFormatter:
     blocks: list["Block"]
 
     def __call__(self, x: Operand) -> str:
-        if isinstance(x, Var):
-            return x.name
+        if isinstance(x, Value):
+            return str(x)
         elif isinstance(x, ModuleType):
             return str(f"<mod:{x.__name__}>")
         elif isinstance(x, Block):
             self.blocks.append(x)
-            return f"^{x.name}"
+            return f"^{x.block_id}"
         elif callable(x):
             return f"<fn:{x.__name__}>"
         else: