diff --git a/pytato/distributed/execute.py b/pytato/distributed/execute.py index bf464dcf5..3f31482bd 100644 --- a/pytato/distributed/execute.py +++ b/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 diff --git a/pytato/distributed/partition.py b/pytato/distributed/partition.py index 856fa9535..ee8ae6463 100644 --- a/pytato/distributed/partition.py +++ b/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,34 +939,51 @@ 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 = { @@ -970,8 +991,14 @@ def get_needed_local_sends( _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)