legate.sparse tracking branch [DO NOT MERGE]

legate/core/__init__.py

@@ -52,6 +52,8 @@
     PartitionByRestriction,
     PartitionByImage,
     PartitionByImageRange,
+    PartitionByPreimage,
+    PartitionByPreimageRange,
     EqualPartition,
     PartitionByWeights,
     IndexPartition,

legate/core/_legion/__init__.py

@@ -38,6 +38,8 @@
     PartitionByRestriction,
     PartitionByImage,
     PartitionByImageRange,
+    PartitionByPreimage,
+    PartitionByPreimageRange,
     EqualPartition,
     PartitionByWeights,
     PartitionByDomain,
@@ -88,6 +90,8 @@
     "PartitionByDomain",
     "PartitionByImage",
     "PartitionByImageRange",
+    "PartitionByPreimage",
+    "PartitionByPreimageRange",
     "PartitionByRestriction",
     "PartitionByWeights",
     "PartitionFunctor",

legate/core/_legion/future.py

@@ -343,7 +343,7 @@ def from_list(
             context,
             domain,
             points,
-            futures,
+            futures_,
             num_futures,
             False,
             0,

legate/core/_legion/partition.py

@@ -109,6 +109,14 @@ def get_root(self) -> Region:
         """
         return self.parent.get_root()
 
+    @property
+    def disjoint(self) -> bool:
+        return self.index_partition.disjoint
+
+    @property
+    def complete(self) -> bool:
+        return self.index_partition.complete
+
 
 class IndexPartition:
     _logical_handle: Any
@@ -260,3 +268,15 @@ def get_root(self) -> IndexSpace:
         Return the root IndexSpace in this tree.
         """
         return self.parent.get_root()
+
+    @property
+    def disjoint(self) -> bool:
+        return legion.legion_index_partition_is_disjoint(  # type: ignore
+            self.runtime, self.handle
+        )
+
+    @property
+    def complete(self) -> bool:
+        return legion.legion_index_partition_is_complete(  # type: ignore
+            self.runtime, self.handle
+        )

legate/core/_legion/partition_functor.py

@@ -17,11 +17,20 @@
 from typing import TYPE_CHECKING, Any, Union
 
 from .. import ffi, legion
+from . import FieldID
 from .future import FutureMap
 from .geometry import Point
 
 if TYPE_CHECKING:
-    from . import FieldID, IndexPartition, IndexSpace, Rect, Region, Transform
+    from . import (
+        FieldID,
+        IndexPartition,
+        IndexSpace,
+        Partition,
+        Rect,
+        Region,
+        Transform,
+    )
 
 
 class PartitionFunctor:
@@ -103,7 +112,7 @@ class PartitionByImage(PartitionFunctor):
     def __init__(
         self,
         region: Region,
-        part: IndexPartition,
+        part: Partition,
         field: Union[int, FieldID],
         mapper: int = 0,
         tag: int = 0,
@@ -148,7 +157,7 @@ class PartitionByImageRange(PartitionFunctor):
     def __init__(
         self,
         region: Region,
-        part: IndexPartition,
+        part: Partition,
         field: Union[int, FieldID],
         mapper: int = 0,
         tag: int = 0,

legate/core/constraints.py

@@ -17,7 +17,7 @@
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any, Iterator, Optional, Protocol, Union
 
-from .partition import Restriction
+from .partition import ImagePartition, Replicate, Restriction
 
 if TYPE_CHECKING:
     from .partition import PartitionBase
@@ -219,6 +219,79 @@ def unknowns(self) -> Iterator[PartSym]:
             yield unknown
 
 
+class Image(Expr):
+    def __init__(
+        self,
+        source_store: Store,
+        dst_store: Store,
+        src_part_sym: Expr,
+        mapper: int,
+        range: bool = False,
+        functor: Any = ImagePartition,
+        disjoint: bool = True,
+        complete: bool = True,
+    ):
+        self._source_store = source_store
+        self._dst_store = dst_store
+        self._src_part_sym = src_part_sym
+        self._mapper = mapper
+        self._range = range
+        self._functor = functor
+        self._disjoint = disjoint
+        self._complete = complete
+
+    def subst(self, mapping: dict[PartSym, PartitionBase]) -> Expr:
+        return Image(
+            self._source_store,
+            self._dst_store,
+            self._src_part_sym.subst(mapping),
+            self._mapper,
+            range=self._range,
+            functor=self._functor,
+            disjoint=self._disjoint,
+            complete=self._complete,
+        )
+
+    @property
+    def ndim(self) -> int:
+        return self._src_part_sym.ndim
+
+    def reduce(self) -> Lit:
+        expr = self._src_part_sym.reduce()
+        assert isinstance(expr, Lit)
+        part = expr._part
+        if isinstance(part, Replicate):
+            return Lit(part)
+        return Lit(
+            self._functor(
+                self._source_store,
+                part,
+                self._mapper,
+                range=self._range,
+                disjoint=self._disjoint,
+                complete=self._complete,
+            )
+        )
+
+    def unknowns(self) -> Iterator[PartSym]:
+        for unknown in self._src_part_sym.unknowns():
+            yield unknown
+
+    def equals(self, other: object) -> bool:
+        return (
+            isinstance(other, Image)
+            and self._source_store == other._source_store
+            and self._dst_store == other._dst_store
+            # Careful! Overloaded equals operator.
+            and self._src_part_sym is other._src_part_sym
+            and self._range == other._range
+            and self._mapper == other._mapper
+            and self._functor == other._functor
+            and self._disjoint == other._disjoint
+            and self._complete == other._complete
+        )
+
+
 class Constraint:
     pass
 

legate/core/operation.py

@@ -27,9 +27,9 @@
 import legate.core.types as ty
 
 from . import Future, FutureMap, Rect
-from .constraints import PartSym
+from .constraints import Image, PartSym
 from .launcher import CopyLauncher, FillLauncher, TaskLauncher
-from .partition import REPLICATE, Weighted
+from .partition import REPLICATE, ImagePartition, Weighted
 from .shape import Shape
 from .store import Store, StorePartition
 from .utils import OrderedSet, capture_traceback_repr
@@ -157,6 +157,12 @@ def get_all_stores(self) -> OrderedSet[Store]:
         result.update(store for (store, _) in self._reductions)
         return result
 
+    def get_all_modified_stores(self) -> OrderedSet[Store]:
+        result: OrderedSet[Store] = OrderedSet()
+        result.update(self._outputs)
+        result.update(store for (store, _) in self._reductions)
+        return result
+
     def add_alignment(self, store1: Store, store2: Store) -> None:
         """
         Sets an alignment between stores. Equivalent to the following code:
@@ -216,6 +222,36 @@ def add_broadcast(
         part = self._get_unique_partition(store)
         self.add_constraint(part.broadcast(axes=axes))
 
+    # add_image_constraint adds a constraint that the image of store1 is
+    # contained within the partition of store2.
+    def add_image_constraint(
+        self,
+        store1: Store,
+        store2: Store,
+        range: bool = False,
+        functor: Any = ImagePartition,
+        disjoint: bool = True,
+        complete: bool = True,
+    ) -> None:
+        self._check_store(store1)
+        self._check_store(store2)
+        # TODO (rohany): We only support point (and rect types if range) here.
+        #  It seems like rects should be added to legate.core's type system
+        #  rather than an external type system to understand this then.
+        part1 = self._get_unique_partition(store1)
+        part2 = self._get_unique_partition(store2)
+        image = Image(
+            store1,
+            store2,
+            part1,
+            self._context.mapper_id,
+            range=range,
+            functor=functor,
+            disjoint=disjoint,
+            complete=complete,
+        )
+        self.add_constraint(image <= part2)
+
     def add_constraint(self, constraint: Constraint) -> None:
         """
         Adds a partitioning constraint to the operation
@@ -851,14 +887,17 @@ def __init__(
             op_id=op_id,
         )
         self._launch_domain: Rect = launch_domain
-        self._input_projs: list[Union[ProjFn, None]] = []
+        # TODO (rohany): The int here is an explicit ID.
+        self._input_projs: list[Union[ProjFn, None, int]] = []
         self._output_projs: list[Union[ProjFn, None]] = []
         self._reduction_projs: list[Union[ProjFn, None]] = []
 
         self._input_parts: list[StorePartition] = []
         self._output_parts: list[StorePartition] = []
         self._reduction_parts: list[tuple[StorePartition, int]] = []
 
+        self._scalar_future_maps: list[FutureMap] = []
+
     @property
     def launch_ndim(self) -> int:
         return self._launch_domain.dim
@@ -876,7 +915,7 @@ def _check_arg(arg: Union[Store, StorePartition]) -> None:
     def add_input(
         self,
         arg: Union[Store, StorePartition],
-        proj: Optional[ProjFn] = None,
+        proj: Optional[Union[ProjFn, int]] = None,
     ) -> None:
         """
         Adds a store as input to the task
@@ -918,10 +957,9 @@ def add_output(
         self._check_arg(arg)
         if isinstance(arg, Store):
             if arg.unbound:
-                raise NotImplementedError(
-                    "Unbound store cannot be used with "
-                    "manually parallelized task"
-                )
+                self._unbound_outputs.append(len(self._outputs))
+                self._outputs.append(arg)
+                return
             if arg.kind is Future:
                 self._scalar_outputs.append(len(self._outputs))
                 self._outputs.append(arg)
@@ -1006,6 +1044,16 @@ def launch(self, strategy: Strategy) -> None:
                 part.store, req, tag=0, read_write=can_read_write
             )
 
+        for fm in self._scalar_future_maps:
+            launcher.add_future_map(fm)
+
+        # Add all unbound stores.
+        for store_idx in self._unbound_outputs:
+            store = self._outputs[store_idx]
+            fspace = self.context.runtime.create_field_space()
+            field_id = fspace.allocate_field(store.type)
+            launcher.add_unbound_output(store, fspace, field_id)
+
         self._add_scalar_args_to_launcher(launcher)
 
         launcher.set_can_raise_exception(self.can_raise_exception)