Make _get_final_ir easier to use without real arguments

- Don't require a TileContext, take a TileContextConfig instead;
- Take a list of ir.KernelArgument descriptors instead of actual args.

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

Author: gbonik

src/cuda/tile/_compile.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
+import inspect
 import math
 import re
 from dataclasses import dataclass
@@ -16,19 +17,19 @@
 import tempfile
 import threading
 import traceback
-from typing import Callable, Optional, Any, Set
+from typing import Callable, Optional, Any, Set, Sequence
 import zipfile
 
 from cuda.tile._cext import get_compute_capability, TileContext, default_tile_context
 from cuda.tile._compiler_options import CompilerOptions
 from cuda.tile._const_utils import get_constant_annotations
+from cuda.tile._context import TileContextConfig
 from cuda.tile._exception import (
     TileCompilerError,
     TileCompilerExecutionError,
     TileCompilerTimeoutError, TileValueError, TileTypeError
 )
 from cuda.tile._ir import ir, hir
-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
@@ -73,13 +74,13 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def _get_final_ir(pyfunc, args, tile_context) -> ir.Function:
-    ir_ctx = ir.IRContext(tile_context)
+def _get_final_ir(pyfunc,
+                  args: Sequence[ir.KernelArgument],
+                  config: TileContextConfig) -> ir.Function:
     func_hir: hir.Function = get_function_hir(pyfunc, entry_point=True)
 
-    ir_args = _bind_kernel_arguments(tuple(func_hir.signature.parameters),
-                                     args, get_constant_annotations(pyfunc))
-    func_body = hir2ir(func_hir, ir_args, ir_ctx)
+    ir_ctx = ir.IRContext(config)
+    func_body = hir2ir(func_hir, args, ir_ctx)
     eliminate_assign_ops(func_body)
     dead_code_elimination_pass(func_body)
 
@@ -100,7 +101,7 @@ def _get_final_ir(pyfunc, args, tile_context) -> ir.Function:
 
 def _bind_kernel_arguments(param_names: tuple[str, ...],
                            args: tuple[Any, ...],
-                           constant_args: Set[str]) -> tuple[Argument, ...]:
+                           constant_args: Set[str]) -> tuple[ir.KernelArgument, ...]:
     # 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.
@@ -120,9 +121,7 @@ def _bind_kernel_arguments(param_names: tuple[str, ...],
                 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))
+        ir_args.append(ir.KernelArgument(type=ty, is_const=is_const, const_value=const_val))
     return tuple(ir_args)
 
 
@@ -181,7 +180,9 @@ def compile_tile(pyfunc,
                  args,
                  compiler_options: CompilerOptions,
                  context: TileContext = default_tile_context) -> TileLibrary:
-    func_ir = _get_final_ir(pyfunc, args, context)
+    param_names = tuple(inspect.signature(pyfunc).parameters.keys())
+    ir_args = _bind_kernel_arguments(param_names, args, get_constant_annotations(pyfunc))
+    func_ir = _get_final_ir(pyfunc, ir_args, context.config)
 
     if 'CUTILEIR' in context.config.log_keys:
         code = (f"==== CuTile IR for {func_ir.name}==== \n\n"

src/cuda/tile/_ir/ir.py

@@ -22,20 +22,21 @@
     TileTypeError, Loc, TileInternalError
 )
 from .._cext import TileContext
+from .._context import TileContextConfig
 
 if TYPE_CHECKING:
     from cuda.tile._ir2bytecode import BytecodeContext
 
 
 class IRContext:
-    def __init__(self, tile_ctx: TileContext):
+    def __init__(self, config: TileContextConfig):
         self._all_vars: Dict[str, str] = {}
         self._counter_by_name: Dict[str, Iterator[int]] = defaultdict(itertools.count)
         self._temp_counter = itertools.count()
         self.typemap: Dict[str, Type] = dict()
         self.constants: Dict[str, Any] = dict()
         self._loose_typemap: Dict[str, Type] = dict()
-        self.tile_ctx: TileContext = tile_ctx
+        self.config: TileContext = config
         self._aggregate_values: Dict[str, Any] = dict()
 
     #  Make a Var with a unique name based on `name`.
@@ -822,7 +823,7 @@ class Function:
 
 
 @dataclass
-class Argument:
+class KernelArgument:
     type: Type
     is_const: bool = False
     const_value: Any = None

src/cuda/tile/_passes/hir2ir.py

@@ -4,14 +4,14 @@
 import inspect
 import sys
 from contextlib import contextmanager
-from typing import Any
+from typing import Any, Sequence
 
 from .ast2hir import get_function_hir
 from .. import TileTypeError
 from .._coroutine_util import resume_after, run_coroutine
 from .._exception import Loc, TileSyntaxError, TileInternalError, TileError, TileRecursionError
 from .._ir import hir, ir
-from .._ir.ir import Var, IRContext, Argument, BoundMethodValue, ClosureValue
+from .._ir.ir import Var, IRContext, BoundMethodValue, ClosureValue, KernelArgument
 from .._ir.op_impl import op_implementations
 from .._ir.ops import loosely_typed_const, end_branch, return_, continue_, \
     break_, flatten_block_parameters, store_var
@@ -23,14 +23,13 @@
 MAX_RECURSION_DEPTH = 1000
 
 
-def hir2ir(func_hir: hir.Function,
-           args: tuple[Argument, ...],
-           ir_ctx: IRContext) -> ir.Block:
+def hir2ir(func_hir: hir.Function, args: Sequence[KernelArgument], ir_ctx: IRContext) -> ir.Block:
     # Run as a coroutine using a software stack, so that we don't exceed Python's recursion limit.
     return run_coroutine(_hir2ir_coroutine(func_hir, args, ir_ctx))
 
 
-async def _hir2ir_coroutine(func_hir: hir.Function, args: tuple[Argument, ...], ir_ctx: IRContext):
+async def _hir2ir_coroutine(func_hir: hir.Function, args: Sequence[KernelArgument],
+                            ir_ctx: IRContext):
     scope = _create_scope(func_hir, ir_ctx, call_site=None, parent_scopes=())
     aggregate_params = [
         scope.local.redefine(local_idx, param_loc)
@@ -50,7 +49,7 @@ async def _hir2ir_coroutine(func_hir: hir.Function, args: tuple[Argument, ...],
 
             await _dispatch_hir_block_inner(func_hir.body, ir_builder)
         except Exception as e:
-            if 'CUTILEIR' in ir_ctx.tile_ctx.config.log_keys:
+            if 'CUTILEIR' in ir_ctx.config.log_keys:
                 highlight_loc = e.loc if hasattr(e, 'loc') else None
                 ir_str = "\n".join(op.to_string(highlight_loc=highlight_loc)
                                    for op in ir_builder.ops)
@@ -93,7 +92,7 @@ async def _dispatch_hir_block_inner(block: hir.Block, builder: ir.Builder):
         with _wrap_exceptions(loc), builder.change_loc(loc):
             _dispatch_hir_jump(block, scope)
     except Exception:
-        if 'CUTILEIR' in builder.ir_ctx.tile_ctx.config.log_keys:
+        if 'CUTILEIR' in builder.ir_ctx.config.log_keys:
             hir_params = ", ".join(p.name for p in block.params)
             hir_lines = [str(c) for c in block.calls]
             hir_lines.append(block.jump_str())

test/test_dce.py

@@ -2,19 +2,26 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import torch
-
 import cuda.tile as ct
+from cuda.tile._ir.ir import KernelArgument
 from cuda.tile._ir.ops import Loop, Continue, Break
 from cuda.tile._ir import ir
 from cuda.tile._compile import _get_final_ir
 from cuda.tile._cext import default_tile_context
+from cuda.tile._ir.type import ArrayTy, TupleTy, SizeTy
 
 
 def get_ir(func) -> ir.Block:
-    x = torch.zeros(10, device="cuda")
-    ir = _get_final_ir(func, (x,), default_tile_context)
-    print(ir)
+    x = KernelArgument(type=ArrayTy(ct.int32,
+                                    shape=TupleTy((SizeTy(),)),
+                                    strides=TupleTy((SizeTy(1,),)),
+                                    elements_disjoint=True,
+                                    base_ptr_div_by=None,
+                                    stride_div_by=(None,),
+                                    shape_div_by=(None,)),
+                       is_const=False,
+                       const_value=None)
+    ir = _get_final_ir(func, (x,), default_tile_context.config)
     return ir.body