Virtualize unflatten_aggregate() via ImplRegistry

This gets rid of MakeTensorView ops in cuda.lang.

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

Author: gbonik

experimental/cuda-lang/src/cuda/lang/_compile.py

@@ -129,7 +129,7 @@ def get_function_ir(
     if constant_mask is None:
         constant_mask = [False] * len(signature.parameters)
     parameter_names = function.signature.parameters.keys()
-    with ir.TileBuilder(ctx, function.body.loc) as builder:
+    with ir.TileBuilder(ctx, function.body.loc) as builder, cuda_lang_impl_registry.as_current():
         params = _create_kernel_parameters(
             signature.parameters,
             constant_mask,
@@ -138,8 +138,7 @@ def get_function_ir(
             ctx
         )
         canonicalize_parameters(params, builder)
-        with cuda_lang_impl_registry.as_current():
-            hir2ir(function, params.aggregate_vars, ctx)
+        hir2ir(function, params.aggregate_vars, ctx)
     func_body = ctx.make_block("entry", function.body.loc)
     func_body.params = sum((vars for vars, _ in params.nonconstant_flat_vars), ())
     func_body.extend(builder.ops)

experimental/cuda-lang/src/cuda/lang/_passes/ir2mlir/pass_definition.py

@@ -282,7 +282,6 @@ def _get_mlir_comparison_op(
 # These operations have aggregate results. The RHS's elements are stored to
 # the LHS's when lowering Assign operations and are no-ops at the MLIR level.
 _NOOP_LOWERINGS = frozenset([
-    ops.MakeTensorView,
     ops.ReinterpretPointerAsArray,
 ])
 

experimental/cuda-lang/test/passes/test_flatten_cfg.py

@@ -22,7 +22,6 @@ def test_kernel(A):
         else:
             A[0] = 0
 
-    # BEFORE: A{{.+}}: Array[int32,(?):(?)] = make_tensor_view
     # BEFORE: $[[ITEM:[0-9]+]]: Tile[int32,()] = load_pointer
     # BEFORE: $[[ITEM_CASTED:[0-9]+]]: Tile[bool_,()] = tile_astype(x=$[[ITEM]])
     # BEFORE: if(cond=$[[ITEM_CASTED]])
@@ -37,7 +36,6 @@ def test_kernel(A):
     filecheck(str(body), get_source(), ("BEFORE",))
 
     # AFTER: ^entry({{.+}}):
-    # AFTER:   A{{.+}}: Array[int32,(?):(?)] = make_tensor_view
     # AFTER:   $[[ITEM:[0-9]+]]: Tile[int32,()] = load_pointer
     # AFTER:   $[[ITEM_CASTED:[0-9]+]]: Tile[bool_,()] = tile_astype(x=$[[ITEM]])
     # AFTER:   cond_br $[[ITEM_CASTED]]: Tile[bool_,()] ^then() ^else()

src/cuda/tile/_compile.py

@@ -270,12 +270,12 @@ def get_final_ir(self, signature_index: int) -> ir.Block:
                                   tileiras_version=self.bytecode_version,
                                   typing_hooks=_TileTypingHooks())
             with ir.Builder(ir_ctx, self._func_hir.body.loc) as ir_builder:
-                params = _create_kernel_parameters(sig.parameters,
-                                                   self.ann_func.constant_parameter_mask,
-                                                   param_names,
-                                                   self._func_hir.param_locs,
-                                                   ir_ctx)
                 with tile_impl_registry.as_current():
+                    params = _create_kernel_parameters(sig.parameters,
+                                                       self.ann_func.constant_parameter_mask,
+                                                       param_names,
+                                                       self._func_hir.param_locs,
+                                                       ir_ctx)
                     hir2ir(self._func_hir, params.aggregate_vars, ir_ctx)
 
             func_body = ir.Block(ir_ctx, self._func_hir.body.loc)

src/cuda/tile/_ir/aggregate_support.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Sequence, Iterator
+
+from cuda.tile._ir.ir import Var, Builder
+from cuda.tile._ir.op_impl import ImplRegistry
+from cuda.tile._ir.type import Type, InvalidType
+
+
+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())
+        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)
+            ret.extend(items)
+    return tuple(ret)
+
+
+def flatten_aggregate_types(types: Sequence[Type]) -> tuple[Type, ...]:
+    ret = []
+    for ty in types:
+        ret.extend(ty.flatten_aggregate())
+    return tuple(ret)
+
+
+def unflatten_aggregates(flattened: tuple[Var, ...],
+                         nominal: Sequence[Type], actual: Sequence[Type]) -> tuple[Var, ...]:
+    it = iter(flattened)
+    ret = tuple(_maybe_unflatten_aggregate(it, n, a) for n, a in zip(nominal, actual, strict=True))
+    assert next(it, None) is None
+    return ret
+
+
+def _maybe_unflatten_aggregate(flattened_iter: Iterator[Var], nominal: Type, actual: Type) -> Var:
+    if not nominal.is_aggregate():
+        return next(flattened_iter)
+    return _unflatten_proper_aggregate(flattened_iter, nominal, actual, result_var=None)
+
+
+def expand_aggregate_var(var: Var) -> tuple[Var, ...]:
+    item_types = tuple(var.get_type().flatten_aggregate())
+    ret = tuple(var.ctx.make_var(f"{var.get_original_name()}_{i}", var.loc)
+                for i in range(len(item_types)))
+    for item, item_ty in zip(ret, item_types, strict=True):
+        item.set_type(item_ty)
+    return ret
+
+
+def flatten_block_parameters(vars: Sequence[Var]) -> list[tuple[Var, ...]]:
+    ret = []
+    for v in vars:
+        ty = v.get_type_allow_invalid()
+        if ty.is_aggregate():
+            flattened_vars = expand_aggregate_var(v)
+            ret.append(flattened_vars)
+            it = iter(flattened_vars)
+            _unflatten_proper_aggregate(it, ty, ty, v)
+            assert next(it, None) is None
+        else:
+            ret.append((v,))
+    return ret
+
+
+def _unflatten_proper_aggregate(flattened_iter: Iterator[Var], nominal: Type, actual: Type,
+                                result_var: Var | None) -> Var:
+    nominal_item_types = nominal.aggregate_item_types()
+    if isinstance(actual, InvalidType):
+        # Pop values from the iterator and throw them out
+        for _ in nominal_item_types:
+            next(flattened_iter)
+        builder = Builder.get_current()
+        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
+                  in zip(nominal_item_types, actual.aggregate_item_types(), strict=True))
+    val = nominal.make_aggregate_value(items)
+
+    impl = ImplRegistry.get_current().unflatten_aggregate_implementations.get(type(nominal))
+    if impl is None:
+        return Builder.get_current().make_aggregate(val, nominal, result_var=result_var)
+    else:
+        return impl(val, nominal, result_var)

src/cuda/tile/_ir/op_impl.py

@@ -59,6 +59,7 @@ class ImplRegistry:
     def __init__(self):
         self.op_implementations = dict()
         self._overloaded_implementations = defaultdict(dict)
+        self.unflatten_aggregate_implementations = dict()
 
     @staticmethod
     def get_current() -> "ImplRegistry":
@@ -79,6 +80,7 @@ def update(self, source: "ImplRegistry"):
         self.op_implementations.update(source.op_implementations)
         for stub, overloads in source._overloaded_implementations.items():
             self._overloaded_implementations[stub].update(overloads)
+        self.unflatten_aggregate_implementations.update(source.unflatten_aggregate_implementations)
 
     def overload_dispatcher(self, stub, *, fixed_args: Sequence[Any] = ()):
         """
@@ -215,6 +217,12 @@ def _have_overload_matching_first_param(self, stub: Callable, first_param: Any)
         return any(predicates[0](first_param)
                    for _priority, predicates, _impl in candidates.values())
 
+    def unflatten_aggregate_impl(self, aggregate_type_class: type[Type]):
+        def decorate(func):
+            self.unflatten_aggregate_implementations[aggregate_type_class] = func
+            return func
+        return decorate
+
 
 def _predicate_from_overload_pattern(pattern):
     if pattern == WILDCARD:

src/cuda/tile/_ir/ops.py

@@ -6,7 +6,7 @@
 import operator
 from dataclasses import dataclass
 from typing import (
-    Literal, Sequence, Tuple, Optional, Any, List, Callable, Iterator, Iterable,
+    Literal, Sequence, Tuple, Optional, Any, List, Callable, Iterable,
 )
 
 from typing_extensions import override
@@ -21,6 +21,8 @@
     PhiState, LoopVarState, make_aggregate, ConstantState, MemoryEffect, attribute, operand,
     BlockRestriction, add_operation_variadic,
 )
+from .aggregate_support import flatten_block_parameters, expand_aggregate_var, \
+    flatten_aggregate_types, flatten_aggregates, unflatten_aggregates
 from .arithmetic_ops import reshape, broadcast_to, astype, compare_tensorlike, \
     binary_bitwise_tensorlike, bitwise_shift_tensorlike, binary_arithmetic_tensorlike, \
     compare_tensorlike_raw, where, binary_bitwise_tensorlike_raw, where_raw, TileReshape, \
@@ -1218,113 +1220,6 @@ def generate_bytecode(self, ctx: "BytecodeContext"):
         return bc.encode_MakePartitionViewOp(ctx.builder, pv_ty, tv)
 
 
-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())
-        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)
-            ret.extend(items)
-    return tuple(ret)
-
-
-def flatten_aggregate_types(types: Sequence[Type]) -> tuple[Type, ...]:
-    ret = []
-    for ty in types:
-        ret.extend(ty.flatten_aggregate())
-    return tuple(ret)
-
-
-def unflatten_aggregates(flattened: Tuple[Var, ...],
-                         nominal: Sequence[Type], actual: Sequence[Type]) -> tuple[Var, ...]:
-    it = iter(flattened)
-    ret = tuple(_maybe_unflatten_aggregate(it, n, a) for n, a in zip(nominal, actual, strict=True))
-    assert next(it, None) is None
-    return ret
-
-
-def _maybe_unflatten_aggregate(flattened_iter: Iterator[Var], nominal: Type, actual: Type) -> Var:
-    if not nominal.is_aggregate():
-        return next(flattened_iter)
-    return _unflatten_proper_aggregate(flattened_iter, nominal, actual, result_var=None)
-
-
-def expand_aggregate_var(var: Var) -> Tuple[Var, ...]:
-    item_types = tuple(var.get_type().flatten_aggregate())
-    ret = tuple(var.ctx.make_var(f"{var.get_original_name()}_{i}", var.loc)
-                for i in range(len(item_types)))
-    for item, item_ty in zip(ret, item_types, strict=True):
-        item.set_type(item_ty)
-    return ret
-
-
-def flatten_block_parameters(vars: Sequence[Var]) -> list[tuple[Var, ...]]:
-    ret = []
-    for v in vars:
-        ty = v.get_type_allow_invalid()
-        if ty.is_aggregate():
-            flattened_vars = expand_aggregate_var(v)
-            ret.append(flattened_vars)
-            it = iter(flattened_vars)
-            _unflatten_proper_aggregate(it, ty, ty, v)
-            assert next(it, None) is None
-        else:
-            ret.append((v,))
-    return ret
-
-
-def _unflatten_proper_aggregate(flattened_iter: Iterator[Var], nominal: Type, actual: Type,
-                                result_var: Var | None) -> Var:
-    nominal_item_types = nominal.aggregate_item_types()
-    if isinstance(actual, InvalidType):
-        # Pop values from the iterator and throw them out
-        for _ in nominal_item_types:
-            next(flattened_iter)
-        builder = Builder.get_current()
-        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
-                  in zip(nominal_item_types, actual.aggregate_item_types(), strict=True))
-    val = nominal.make_aggregate_value(items)
-
-    builder = Builder.get_current()
-    if isinstance(nominal, ArrayTy):
-        assert isinstance(val, ArrayValue)
-        base_ptr = val.base_ptr
-        shape = tuple(assume_bounded(x, 0, None) for x in val.shape)
-
-        all_strides = []
-        dynamic_strides = []
-        for x, s in zip(val.strides, nominal.strides, strict=True):
-            if s is None:
-                x = assume_bounded(x, 0, None)
-                dynamic_strides.append(x)
-            all_strides.append(x)
-
-        operands = dict(base_ptr=base_ptr, shape=shape, dynamic_strides=tuple(dynamic_strides))
-        ret = builder.add_operation(MakeTensorView, nominal, operands, result_var)
-        ret.set_aggregate(ArrayValue(base_ptr, shape, tuple(all_strides)))
-        return ret
-    elif isinstance(nominal, ListTy):
-        assert isinstance(val, ListValue)
-        operands = dict(base_ptr=val.base_ptr, length=val.length)
-        ret = builder.add_operation(MakeListView, nominal, operands, result_var)
-        ret.set_aggregate(val)
-        return ret
-    else:
-        return builder.make_aggregate(val, nominal, result_var=result_var)
-
-
 @dataclass(eq=False)
 class TileNumBlocks(Operation, opcode="tile_num_blocks"):
     axis: int = attribute()
@@ -3921,6 +3816,35 @@ def store_advanced_impl(array: Var, indices: Var, tile: Var,
                   latency=latency_val, allow_tma=allow_tma_val)
 
 
+@tile_impl_registry.unflatten_aggregate_impl(ArrayTy)
+def _unflatten_aggregate_array_impl(val: ArrayValue, ty: ArrayTy, result_var: Var):
+    assert isinstance(val, ArrayValue)
+    base_ptr = val.base_ptr
+    shape = tuple(assume_bounded(x, 0, None) for x in val.shape)
+
+    all_strides = []
+    dynamic_strides = []
+    for x, s in zip(val.strides, ty.strides, strict=True):
+        if s is None:
+            x = assume_bounded(x, 0, None)
+            dynamic_strides.append(x)
+        all_strides.append(x)
+
+    operands = dict(base_ptr=base_ptr, shape=shape, dynamic_strides=tuple(dynamic_strides))
+    ret = Builder.get_current().add_operation(MakeTensorView, ty, operands, result_var)
+    ret.set_aggregate(ArrayValue(base_ptr, shape, tuple(all_strides)))
+    return ret
+
+
+@tile_impl_registry.unflatten_aggregate_impl(ListTy)
+def _unflatten_aggregate_list_impl(val: ListValue, ty: ListTy, result_var: Var):
+    assert isinstance(val, ListValue)
+    operands = dict(base_ptr=val.base_ptr, length=val.length)
+    ret = Builder.get_current().add_operation(MakeListView, ty, operands, result_var)
+    ret.set_aggregate(val)
+    return ret
+
+
 def _add_dummy_op_to_invalid_vars(vars: Sequence[Var],
                                   actual_types: Sequence[Type]) -> tuple[Var, ...]:
     return tuple(add_operation(MakeDummy, actual)