Handle some param-domain complexity

Author: inducer

namedisl/expression_like.py

@@ -387,7 +387,10 @@ def from_piece_and_aff(piece: Set, aff: Aff) -> PwAff:
         if not piece.space.order_equal(aff.space.as_set_space()):
             raise ValueError("spaces don't match")
 
-        return PwAff(isl.PwAff.alloc(piece._obj, aff._obj), aff.space)
+        if aff.space.dim(DimType.in_):
+            return PwAff(isl.PwAff.alloc(piece._obj, aff._obj), aff.space)
+        else:
+            return PwAff(isl.PwAff.alloc(piece._obj.params(), aff._obj), aff.space)
 
     def var_pw_aff(self: PwAff, name: str) -> PwAff:
         """Return a :class:`PwAff` that evaluates to *name* in the space of *self*."""
@@ -419,7 +422,10 @@ def where(self,
                 self.space)
         self_a, rhs_a = align_two(self, rhs)
         from .set_like import Set
-        return Set(func(self_a._obj, rhs_a._obj), self_a.space.as_set_space())
+        res_set = func(self_a._obj, rhs_a._obj)
+        return Set(
+            res_set.from_params() if res_set.is_params() else res_set,
+            self_a.space.as_set_space())
 
     def eq_set(self, rhs: int | PwAff) -> Set: return self.where("=", rhs)
     def ne_set(self, rhs: int | PwAff) -> Set: return self.where("!=", rhs)
@@ -447,7 +453,9 @@ def get_pieces(self) -> list[tuple[Set, Aff]]:
         set_space = self.space.as_set_space()
         from .set_like import Set
         return [
-            (Set(set, set_space), Aff(aff, self.space))
+            (Set(set.from_params() if set.is_params() else set,
+                set_space),
+                Aff(aff, self.space))
             for set, aff in self._obj.get_pieces()
         ]
 
@@ -456,7 +464,10 @@ def coalesce(self) -> PwAff:
 
     def get_aggregate_domain(self) -> Set:
         from .set_like import Set
-        return Set(self._obj.get_aggregate_domain(), self.space.as_set_space())
+        agg_domain = self._obj.get_aggregate_domain()
+        return Set(
+            agg_domain.from_params() if agg_domain.is_params() else agg_domain,
+            self.space.as_set_space())
 
     def union_max(self, other: PwAff) -> PwAff:
         self_a, other_a = align_two(self, other)

namedisl/set_like.py

@@ -303,6 +303,11 @@ class _NamedIslSetLike(_NamedIslSetOrMapLike[IslSetLikeT]):
     active_dim_types: ClassVar[frozenset[DimType]] = frozenset(
         {DimType.param, DimType.out})
 
+    if __debug__:
+        def __post_init__(self) -> None:
+            assert not self._obj.is_params()
+            return super().__post_init__()
+
     def is_bounded(self) -> bool:
         return self._obj.is_bounded()
 
@@ -399,6 +404,11 @@ def make_basic_set(src: isl.BasicSet) -> BasicSet: ...
 
 def make_basic_set(src: str | isl.BasicSet, ctx: isl.Context | None = None) -> BasicSet:
     obj = isl.BasicSet(src, ctx) if isinstance(src, str) else src
+
+    if obj.is_params():
+        # shield the user from 'param domain' complexity
+        obj = obj.from_params()
+
     return BasicSet(obj, Space.from_isl(obj, BasicSet.active_dim_types))
 
 
@@ -486,6 +496,11 @@ def make_set(src: isl.Set) -> Set: ...
 
 def make_set(src: isl.Set | str, ctx: isl.Context | None = None) -> Set:
     obj = isl.Set(src, ctx) if isinstance(src, str) else src
+
+    if obj.is_params():
+        # shield the user from 'param domain' complexity
+        obj = obj.from_params()
+
     return Set(obj, Space.from_isl(obj, Set.active_dim_types))