replacing set with OrderedSet to avoid control-replication violations (…

…#1054)

replacing set with OrderedSet to avoid control-replication violations

cunumeric/_ufunc/ufunc.py

@@ -17,6 +17,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Sequence, Union
 
 import numpy as np
+from legate.core.utils import OrderedSet
 
 from ..array import check_writeable, convert_to_cunumeric_ndarray, ndarray
 from ..config import BinaryOpCode, UnaryOpCode, UnaryRedCode
@@ -552,7 +553,7 @@ def _find_common_type(
         arrs: Sequence[ndarray], orig_args: Sequence[Any]
     ) -> np.dtype[Any]:
         all_ndarray = all(isinstance(arg, ndarray) for arg in orig_args)
-        unique_dtypes = set(arr.dtype for arr in arrs)
+        unique_dtypes = OrderedSet(arr.dtype for arr in arrs)
         # If all operands are ndarrays and they all have the same dtype,
         # we already know the common dtype
         if len(unique_dtypes) == 1 and all_ndarray:

cunumeric/array.py

@@ -25,14 +25,14 @@
     Literal,
     Optional,
     Sequence,
-    Set,
     TypeVar,
     Union,
     cast,
 )
 
 import numpy as np
 from legate.core import Array, Field
+from legate.core.utils import OrderedSet
 from numpy.core.multiarray import (  # type: ignore [attr-defined]
     normalize_axis_index,
 )
@@ -90,7 +90,7 @@ def add_boilerplate(
       parameter (if present), to cuNumeric ndarrays.
     * Convert the special "where" parameter (if present) to a valid predicate.
     """
-    keys = set(array_params)
+    keys = OrderedSet(array_params)
     assert len(keys) == len(array_params)
 
     def decorator(func: Callable[P, R]) -> Callable[P, R]:
@@ -100,11 +100,11 @@ def decorator(func: Callable[P, R]) -> Callable[P, R]:
 
         # For each parameter specified by name, also consider the case where
         # it's passed as a positional parameter.
-        indices: Set[int] = set()
+        indices: OrderedSet[int] = OrderedSet()
         where_idx: Optional[int] = None
         out_idx: Optional[int] = None
         params = signature(func).parameters
-        extra = keys - set(params)
+        extra = keys - OrderedSet(params)
         assert len(extra) == 0, f"unknown parameter(s): {extra}"
         for idx, param in enumerate(params):
             if param == "where":
@@ -2435,7 +2435,7 @@ def _diag_helper(
         else:
             assert axes is not None
             N = len(axes)
-            if len(axes) != len(set(axes)):
+            if len(axes) != len(OrderedSet(axes)):
                 raise ValueError(
                     "axes passed to _diag_helper should be all different"
                 )

cunumeric/coverage.py

@@ -36,6 +36,7 @@
 )
 
 from legate.core import track_provenance
+from legate.core.utils import OrderedSet
 from typing_extensions import Protocol
 
 from .runtime import runtime
@@ -62,7 +63,7 @@ def filter_namespace(
     omit_names: Optional[Container[str]] = None,
     omit_types: tuple[type, ...] = (),
 ) -> dict[str, Any]:
-    omit_names = omit_names or set()
+    omit_names = omit_names or OrderedSet()
     return {
         attr: value
         for attr, value in ns.items()
@@ -330,7 +331,7 @@ def clone_class(
     """
 
     class_name = f"{origin_class.__module__}.{origin_class.__name__}"
-    clean_omit_names = set() if omit_names is None else omit_names
+    clean_omit_names = OrderedSet() if omit_names is None else omit_names
 
     def _clone_class(cls: type) -> type:
         missing = filter_namespace(

cunumeric/deferred.py

@@ -36,6 +36,7 @@
 import legate.core.types as ty
 import numpy as np
 from legate.core import Annotation, Future, ReductionOp, Store
+from legate.core.utils import OrderedSet
 from numpy.core.numeric import (  # type: ignore [attr-defined]
     normalize_axis_tuple,
 )
@@ -96,7 +97,7 @@ def auto_convert(
     """
     Converts all named parameters to DeferredArrays.
     """
-    keys = set(thunk_params)
+    keys = OrderedSet(thunk_params)
     assert len(keys) == len(thunk_params)
 
     def decorator(func: Callable[P, R]) -> Callable[P, R]:
@@ -107,7 +108,7 @@ def decorator(func: Callable[P, R]) -> Callable[P, R]:
         # For each parameter specified by name, also consider the case where
         # it's passed as a positional parameter.
         params = signature(func).parameters
-        extra = keys - set(params)
+        extra = keys - OrderedSet(params)
         assert len(extra) == 0, f"unknown parameter(s): {extra}"
         indices = {idx for (idx, param) in enumerate(params) if param in keys}
 
@@ -1429,16 +1430,16 @@ def fft(
             task.add_scalar_arg(kind.type_id, ty.int32)
             task.add_scalar_arg(direction.value, ty.int32)
             task.add_scalar_arg(
-                len(set(axes)) != len(axes)
+                len(OrderedSet(axes)) != len(axes)
                 or len(axes) != input.ndim
                 or tuple(axes) != tuple(sorted(axes)),
                 ty.bool_,
             )
             for ax in axes:
                 task.add_scalar_arg(ax, ty.int64)
 
-            if input.ndim > len(set(axes)):
-                task.add_broadcast(input, axes=set(axes))
+            if input.ndim > len(OrderedSet(axes)):
+                task.add_broadcast(input, axes=OrderedSet(axes))
             else:
                 task.add_broadcast(input)
             task.add_constraint(p_output == p_input)
@@ -1502,9 +1503,9 @@ def contract(
 
         # Sanity checks
         # no duplicate modes within an array
-        assert len(lhs_modes) == len(set(lhs_modes))
-        assert len(rhs1_modes) == len(set(rhs1_modes))
-        assert len(rhs2_modes) == len(set(rhs2_modes))
+        assert len(lhs_modes) == len(OrderedSet(lhs_modes))
+        assert len(rhs1_modes) == len(OrderedSet(rhs1_modes))
+        assert len(rhs2_modes) == len(OrderedSet(rhs2_modes))
         # no singleton modes
         mode_counts: Counter[str] = Counter()
         mode_counts.update(lhs_modes)

cunumeric/utils.py

@@ -22,6 +22,7 @@
 
 import legate.core.types as ty
 import numpy as np
+from legate.core.utils import OrderedSet
 
 from .types import NdShape
 
@@ -194,8 +195,8 @@ def check_axes(a_axes: Axes, b_axes: Axes) -> None:
             len(a_axes) != len(b_axes)
             or len(a_axes) > a_ndim
             or len(b_axes) > b_ndim
-            or len(a_axes) != len(set(a_axes))
-            or len(b_axes) != len(set(b_axes))
+            or len(a_axes) != len(OrderedSet(a_axes))
+            or len(b_axes) != len(OrderedSet(b_axes))
             or any(ax < 0 for ax in a_axes)
             or any(ax < 0 for ax in b_axes)
             or any(ax >= a_ndim for ax in a_axes)
@@ -211,8 +212,14 @@ def check_axes(a_axes: Axes, b_axes: Axes) -> None:
     b_modes = list(ascii_uppercase[:b_ndim])
     for a_i, b_i in zip(a_axes, b_axes):
         b_modes[b_i] = a_modes[a_i]
-    a_out = [a_modes[a_i] for a_i in sorted(set(range(a_ndim)) - set(a_axes))]
-    b_out = [b_modes[b_i] for b_i in sorted(set(range(b_ndim)) - set(b_axes))]
+    a_out = [
+        a_modes[a_i]
+        for a_i in sorted(OrderedSet(range(a_ndim)) - OrderedSet(a_axes))
+    ]
+    b_out = [
+        b_modes[b_i]
+        for b_i in sorted(OrderedSet(range(b_ndim)) - OrderedSet(b_axes))
+    ]
 
     return (a_modes, b_modes, a_out + b_out)
 

tests/integration/test_einsum.py

@@ -19,6 +19,7 @@
 
 import numpy as np
 import pytest
+from legate.core.utils import OrderedSet
 from utils.comparisons import allclose
 from utils.generators import mk_0to1_array, permutes_to
 
@@ -54,8 +55,8 @@ def gen_operand(
         return
     # If we've hit the limit on distinct modes, only use modes
     # appearing on the same operand
-    if len(op) == dim_lim - 1 and len(set(op)) >= mode_lim:
-        for m in sorted(set(op)):
+    if len(op) == dim_lim - 1 and len(OrderedSet(op)) >= mode_lim:
+        for m in sorted(OrderedSet(op)):
             op.append(m)
             yield from gen_operand(used_modes, dim_lim, mode_lim, op)
             op.pop()
@@ -82,7 +83,7 @@ def gen_expr(
     if opers is None:
         opers = []
     if cache is None:
-        cache = set()
+        cache = OrderedSet()
     # The goal here is to avoid producing duplicate expressions, up to
     # reordering of operands and alpha-renaming, e.g. the following
     # are considered equivalent (for the purposes of testing):
@@ -108,7 +109,9 @@ def gen_expr(
     dim_lim = len(opers[-1]) if len(opers) > 0 else MAX_OPERAND_DIM
     # Between operands of the same length, put those with the most distinct
     # modes first.
-    mode_lim = len(set(opers[-1])) if len(opers) > 0 else MAX_OPERAND_DIM
+    mode_lim = (
+        len(OrderedSet(opers[-1])) if len(opers) > 0 else MAX_OPERAND_DIM
+    )
     for op in gen_operand(used_modes, dim_lim, mode_lim):
         opers.append(op)
         yield from gen_expr(opers, cache)
@@ -187,7 +190,7 @@ def mk_input_that_broadcasts_to(lib, tgt_shape):
     # just one of them to 1. Consider the operation 'aab->ab': (10,10,11),
     # (10,10,1), (1,1,11), (1,1,1) are all acceptable input shapes, but
     # (1,10,11) is not.
-    tgt_sizes = list(sorted(set(tgt_shape)))
+    tgt_sizes = list(sorted(OrderedSet(tgt_shape)))
     res = []
     for mask in product([True, False], repeat=len(tgt_sizes)):
         tgt2src_size = {

tests/integration/utils/contractions.py

@@ -15,6 +15,7 @@
 
 import numpy as np
 from legate.core import LEGATE_MAX_DIM
+from legate.core.utils import OrderedSet
 
 import cunumeric as num
 
@@ -38,7 +39,7 @@ def gen_output_default(lib, modes, a, b):
 
 def gen_shapes(a_modes, b_modes):
     yield ((5,) * len(a_modes), (5,) * len(b_modes))
-    for mode_to_squeeze in set(a_modes + b_modes):
+    for mode_to_squeeze in OrderedSet(a_modes + b_modes):
         a_shape = tuple((1 if m == mode_to_squeeze else 5) for m in a_modes)
         b_shape = tuple((1 if m == mode_to_squeeze else 5) for m in b_modes)
         yield (a_shape, b_shape)
@@ -104,7 +105,10 @@ def gen_output_of_various_types(lib, modes, a, b):
 
 def _test(name, modes, operation, gen_inputs, gen_output=None, **kwargs):
     (a_modes, b_modes, out_modes) = modes
-    if len(set(a_modes) | set(b_modes) | set(out_modes)) > LEGATE_MAX_DIM:
+    if (
+        len(OrderedSet(a_modes) | OrderedSet(b_modes) | OrderedSet(out_modes))
+        > LEGATE_MAX_DIM
+    ):
         # Total number of distinct modes can't exceed maximum Legion dimension,
         # because we may need to promote arrays so that one includes all modes.
         return