Misc. fixes for new distributed partitioning

pytato/distributed/execute.py

@@ -88,7 +88,7 @@ def execute_distributed_partition(
 
     from mpi4py import MPI
 
-    if len(partition.parts) != 1:
+    if any(part.input_name_to_recv_node for part in partition.parts.values()):
         recv_names_tup, recv_requests_tup, recv_buffers_tup = zip(*[
             (name,) + _post_receive(mpi_communicator, recv)
             for part in partition.parts.values()
@@ -100,7 +100,6 @@ def execute_distributed_partition(
         del recv_requests_tup
         del recv_buffers_tup
     else:
-        # Only a single partition, no recv requests exist
         recv_names = []
         recv_requests = []
         recv_buffers = []
@@ -171,8 +170,8 @@ def wait_for_some_recvs() -> None:
             buf = recv_buffers.pop(idx)
 
             # FIXME: pytato shouldn't depend on pyopencl
-            import pyopencl as cl
-            context[name] = cl.array.to_device(queue, buf, allocator=allocator)
+            import pyopencl.array as cl_array
+            context[name] = cl_array.to_device(queue, buf, allocator=allocator)
             recv_names_completed.add(name)
 
     # {{{ main loop

pytato/distributed/partition.py

@@ -357,7 +357,7 @@ def map_distributed_recv(
         recv_id = _recv_to_comm_id(self.local_rank, expr)
 
         if recv_id in self.local_recv_id_to_recv_node:
-            raise ValueError(f"Multiple recveives found for '{recv_id}'")
+            raise ValueError(f"Multiple receives found for '{recv_id}'")
 
         self.local_recv_id_to_recv_node[recv_id] = expr
 
@@ -923,6 +923,10 @@ def get_needed_local_sends(
     comm_batches, comm_id_to_part_id = _make_local_comm_batches(
             local_rank, local_comm_to_needed_local_comms)
 
+    # FIXME?: comm_batches isn't being used for anything
+
+    nparts = max(part_id for part_id in comm_id_to_part_id.values()) + 1
+
     # }}}
 
     # {{{ assign each materialized array to a batch/part
@@ -935,43 +939,66 @@ def get_needed_local_sends(
             for send in lsrdg.local_send_id_to_send_node.values()}
     sent_arrays = frozenset(sent_array_to_send_node)
 
-    # While sent arrays are materialized, we shouldn't be including them
-    # here because we're using sent arrays (which need to be materialized
-    # in order to send them) as anchors to place *other* materialized data
-    # into the batches.
-    materialized_arrays = frozenset(
-        materialized_arrays_collector.materialized_arrays) - sent_arrays
-
     received_arrays = frozenset(lsrdg.local_recv_id_to_recv_node.values())
 
-    sent_array_dep_mapper = SubsetDependencyMapper(sent_arrays)
+    # While sent/received arrays are materialized, we shouldn't be including them
+    # here because we're using sent/received arrays as anchors to place *other*
+    # materialized data into the batches.
+    materialized_arrays = (
+        frozenset(materialized_arrays_collector.materialized_arrays)
+        - sent_arrays
+        - received_arrays)
+
+    materialized_array_dep_mapper = SubsetDependencyMapper(materialized_arrays)
     recvd_array_dep_mapper = SubsetDependencyMapper(received_arrays)
 
-    materialized_ary_to_part_id: Dict[Array, int] = {
+    materialized_ary_to_part_id_start: Dict[Array, int] = {
             ary: max(
-                max(
-                    comm_id_to_part_id[
-                        _send_to_comm_id(local_rank,
-                                         sent_array_to_send_node[sent_array])]
-                    for sent_array in sent_array_dep_mapper(ary)),
-                max(
-                    comm_id_to_part_id[
+                    (comm_id_to_part_id[
                         _recv_to_comm_id(local_rank,
-                                         cast(DistributedRecv, recvd_array))]
-                    for recvd_array in recvd_array_dep_mapper(ary))
-                )
+                                         cast(DistributedRecv, recvd_ary))]
+                    for recvd_ary in recvd_array_dep_mapper(ary)),
+                    default=0)
             for ary in materialized_arrays
             }
 
+    materialized_ary_to_part_id_end: Dict[Array, int] = {
+        ary: nparts
+        for ary in materialized_arrays}
+    for sent_ary in sent_arrays:
+        for ary in materialized_array_dep_mapper(sent_ary):
+            materialized_ary_to_part_id_end[ary] = min(
+                materialized_ary_to_part_id_end[ary],
+                comm_id_to_part_id[
+                    _send_to_comm_id(local_rank,
+                                     sent_array_to_send_node[sent_ary])] + 1)
+
+    assert all(
+            (
+                materialized_ary_to_part_id_start[ary]
+                < materialized_ary_to_part_id_end[ary])
+            for ary in materialized_arrays), \
+        "unable to find suitable part for materialized array"
+
+    materialized_ary_to_part_id: Dict[Array, int] = {
+            ary: materialized_ary_to_part_id_end[ary] - 1
+            for ary in materialized_arrays}
+
     # }}}
 
     sent_ary_to_part_id = {
             sent_ary: comm_id_to_part_id[
                 _send_to_comm_id(local_rank, send_node)]
             for sent_ary, send_node in sent_array_to_send_node.items()}
 
+    recvd_ary_to_part_id = {
+            recvd_ary: comm_id_to_part_id[
+                _recv_to_comm_id(local_rank, recvd_ary)]
+            for recvd_ary in received_arrays}
+
     stored_ary_to_part_id = materialized_ary_to_part_id.copy()
     stored_ary_to_part_id.update(sent_ary_to_part_id)
+    stored_ary_to_part_id.update(recvd_ary_to_part_id)
 
     # {{{ find which materialized arrays have users in multiple parts
     # (and promote them to part outputs)