Rename TupleConstraint.elements to items

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

Author: gbonik

cext/tile_kernel.cpp

@@ -1691,12 +1691,12 @@ static PyPtr parse_param_constraint(ConstantCursor& cursor,
                                     Cursor<RefPtr<LeafAnnotationNode>>* annotation_cursor) {
     ParameterKind pk = param_cursor->next();
     if (pk == ParameterKind::TupleBegin) {
-        PyPtr elements_list = steal(PyList_New(0));
-        if (!elements_list) return {};
+        PyPtr items_list = steal(PyList_New(0));
+        if (!items_list) return {};
         while (param_cursor->peek() != ParameterKind::TupleEnd) {
-            PyPtr elem = parse_param_constraint(cursor, param_cursor, annotation_cursor);
-            if (!elem) return {};
-            if (PyList_Append(elements_list.get(), elem.get()) < 0) return {};
+            PyPtr item = parse_param_constraint(cursor, param_cursor, annotation_cursor);
+            if (!item) return {};
+            if (PyList_Append(items_list.get(), item.get()) < 0) return {};
         }
         ParameterKind tuple_end = param_cursor->next();
         CHECK(tuple_end == ParameterKind::TupleEnd);
@@ -1705,7 +1705,7 @@ static PyPtr parse_param_constraint(ConstantCursor& cursor,
         if (!signature_module) return {};
         PyPtr constraint_class = getattr(signature_module, "TupleConstraint");
         if (!constraint_class) return {};
-        return steal(PyObject_CallOneArg(constraint_class.get(), elements_list.get()));
+        return steal(PyObject_CallOneArg(constraint_class.get(), items_list.get()));
     }
     LeafAnnotationNode* annotation = annotation_cursor->next().get();
     return parse_element_constraint(cursor, pk, *annotation);

src/cuda/tile/_compile.py

@@ -177,21 +177,21 @@ def _create_parameter(
 
     if isinstance(constraint, TupleConstraint):
         if isinstance(annotation, LeafAnnotationNode):
-            item_nodes = [annotation] * len(constraint.elements)
+            item_nodes = [annotation] * len(constraint.items)
         elif isinstance(annotation, HomogeneousTupleNode):
-            item_nodes = [annotation.each] * len(constraint.elements)
+            item_nodes = [annotation.each] * len(constraint.items)
         elif isinstance(annotation, HeterogeneousTupleNode):
-            if len(annotation.items) != len(constraint.elements):
+            if len(annotation.items) != len(constraint.items):
                 raise _make_constraint_error(
-                        f"Received a tuple of length {len(constraint.elements)}"
+                        f"Received a tuple of length {len(constraint.items)}"
                         f" but the annotation implies length {len(annotation.items)}.",
                         path)
             item_nodes = annotation.items
         else:
             assert False
 
         item_vars = []
-        for i, (item, node) in enumerate(zip(constraint.elements, item_nodes, strict=True)):
+        for i, (item, node) in enumerate(zip(constraint.items, item_nodes, strict=True)):
             item_var = var.ctx.make_var(var.name + f"_{i}", var.loc)
             _create_parameter(item, node, path.tuple_item(i), item_var, nonconstant_flat_vars)
             item_vars.append(item_var)

src/cuda/tile/_passes/dataflow_analysis.py

@@ -68,21 +68,21 @@ def _register_leaf_param(state, constraint: ArrayConstraint | ScalarConstraint,
 
 def _register_tuple_params(state, constraint: TupleConstraint, flat_params, offset: int,
                            alias_set_mapper) -> int:
-    for elem in constraint.elements:
-        if isinstance(elem, (ArrayConstraint, ScalarConstraint)):
-            n = 1 + 2 * elem.ndim if isinstance(elem, ArrayConstraint) else 1
-            _register_leaf_param(state, elem, flat_params[offset:offset + n], alias_set_mapper)
+    for item in constraint.items:
+        if isinstance(item, (ArrayConstraint, ScalarConstraint)):
+            n = 1 + 2 * item.ndim if isinstance(item, ArrayConstraint) else 1
+            _register_leaf_param(state, item, flat_params[offset:offset + n], alias_set_mapper)
             offset += n
-        elif isinstance(elem, TupleConstraint):
-            offset = _register_tuple_params(state, elem, flat_params, offset, alias_set_mapper)
-        elif isinstance(elem, ListConstraint):
-            assert isinstance(elem.element, ArrayConstraint)
+        elif isinstance(item, TupleConstraint):
+            offset = _register_tuple_params(state, item, flat_params, offset, alias_set_mapper)
+        elif isinstance(item, ListConstraint):
+            assert isinstance(item.element, ArrayConstraint)
             base_ptr, size_var = flat_params[offset], flat_params[offset + 1]
             state.tracker.update(base_ptr,
-                                 DataPredicate(alias_set=alias_set_mapper(elem.alias_groups),
+                                 DataPredicate(alias_set=alias_set_mapper(item.alias_groups),
                                                div_by=1,
-                                               may_alias_internally=elem.elements_may_alias))
-            elt_predicates = _get_array_predicates(elem.element, alias_set_mapper)
+                                               may_alias_internally=item.elements_may_alias))
+            elt_predicates = _get_array_predicates(item.element, alias_set_mapper)
             state.list_array_tracker.update(base_ptr,
                                             _AggregatePredicate(dict(enumerate(elt_predicates))))
             state.set_always_true(size_var)

src/cuda/tile/compilation/_name_mangling.py

@@ -395,17 +395,17 @@ def _demangle_list_constraint(cursor: _Cursor,
 
 def _mangle_tuple_constraint(constraint: TupleConstraint, alias_group_map: dict[str, int],
                              cconv: CallingConvention) -> str:
-    # Format: {count}{elem0_mangling}{elem1_mangling}...
-    return f"{len(constraint.elements)}" + "".join(
-        _mangle_constraint(e, alias_group_map, cconv) for e in constraint.elements)
+    # Format: {count}{item0_mangling}{item1_mangling}...
+    return f"{len(constraint.items)}" + "".join(
+        _mangle_constraint(e, alias_group_map, cconv) for e in constraint.items)
 
 
 def _demangle_tuple_constraint(cursor: _Cursor,
                                alias_group_demangler: _AliasGroupDemangler,
                                cconv: CallingConvention) -> TupleConstraint:
     count = int(cursor.expect("[0-9]+", "Expected element count"))
-    elements = [_demangle_constraint(cursor, alias_group_demangler, cconv) for _ in range(count)]
-    return TupleConstraint(elements)
+    items = [_demangle_constraint(cursor, alias_group_demangler, cconv) for _ in range(count)]
+    return TupleConstraint(items)
 
 
 def _mangle_dtype(dtype: DType):

src/cuda/tile/compilation/_signature.py

@@ -227,25 +227,19 @@ class TupleConstraint:
     Describes a tuple kernel parameter.
 
     Args:
-        elements: Per-element parameter constraints. Elements may be
-            :class:`ScalarConstraint`, :class:`ArrayConstraint`,
-            :class:`ConstantConstraint`, :class:`ListConstraint`, or nested
-            :class:`TupleConstraint`.
+        items: Per-item constraints.
     """
-    elements: "tuple[ParameterConstraint, ...]"
-
-    def __init__(
-            self,
-            elements: "Sequence[ParameterConstraint]",
-    ):
-        for i, e in enumerate(elements):
-            if not isinstance(e, (ScalarConstraint, ArrayConstraint, ConstantConstraint,
-                                  TupleConstraint, ListConstraint)):
+    items: "tuple[ParameterConstraint, ...]"
+
+    def __init__(self,
+                 items: "Sequence[ParameterConstraint]"):
+        items = tuple(items)
+        for i, item_constraint in enumerate(items):
+            if not isinstance(item_constraint, ParameterConstraint):
                 raise TypeError(
-                    f"TupleConstraint element {i} must be a ScalarConstraint,"
-                    f" ArrayConstraint, ConstantConstraint, TupleConstraint, or ListConstraint,"
-                    f" got {type(e).__name__}")
-        object.__setattr__(self, "elements", tuple(elements))
+                    f"TupleConstraint item #{i} must be a ParameterConstraint,"
+                    f" got {type(item_constraint).__name__}")
+        object.__setattr__(self, "items", items)
 
 
 @dataclass(frozen=False, eq=False)
@@ -445,7 +439,7 @@ def _collect_alias_groups(parameters: Sequence[ParameterConstraint]
             yield p, p.alias_groups
             yield from _collect_alias_groups([p.element])
         elif isinstance(p, TupleConstraint):
-            yield from _collect_alias_groups(list(p.elements))
+            yield from _collect_alias_groups(p.items)
 
 
 def _check_optional_int(i: int, val, param_name: str):