From ce2eddeda2e2a1cd08645df5c2fbf4235aada7e2 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 28 Nov 2022 12:27:10 -0600 Subject: [PATCH 1/9] fix import error --- pytato/distributed/execute.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pytato/distributed/execute.py b/pytato/distributed/execute.py index bf464dcf5..b74a30a6e 100644 --- a/pytato/distributed/execute.py +++ b/pytato/distributed/execute.py @@ -171,8 +171,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 From edf82b3922632fec5e80c582c27efbe4d7bcb42a Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 28 Nov 2022 12:29:31 -0600 Subject: [PATCH 2/9] post receives even if len(partition.parts) == 1 a part might not do any work before receiving --- pytato/distributed/execute.py | 26 ++++++++++---------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/pytato/distributed/execute.py b/pytato/distributed/execute.py index b74a30a6e..84fd2e16a 100644 --- a/pytato/distributed/execute.py +++ b/pytato/distributed/execute.py @@ -88,22 +88,16 @@ def execute_distributed_partition( from mpi4py import MPI - if len(partition.parts) != 1: - recv_names_tup, recv_requests_tup, recv_buffers_tup = zip(*[ - (name,) + _post_receive(mpi_communicator, recv) - for part in partition.parts.values() - for name, recv in part.input_name_to_recv_node.items()]) - recv_names = list(recv_names_tup) - recv_requests = list(recv_requests_tup) - recv_buffers = list(recv_buffers_tup) - del recv_names_tup - del recv_requests_tup - del recv_buffers_tup - else: - # Only a single partition, no recv requests exist - recv_names = [] - recv_requests = [] - recv_buffers = [] + recv_names_tup, recv_requests_tup, recv_buffers_tup = zip(*[ + (name,) + _post_receive(mpi_communicator, recv) + for part in partition.parts.values() + for name, recv in part.input_name_to_recv_node.items()]) + recv_names = list(recv_names_tup) + recv_requests = list(recv_requests_tup) + recv_buffers = list(recv_buffers_tup) + del recv_names_tup + del recv_requests_tup + del recv_buffers_tup context: Dict[str, Any] = input_args.copy() From f268289b8a9ebd8dfc1ef2a75aba6a55610b9197 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 28 Nov 2022 12:36:47 -0600 Subject: [PATCH 3/9] add fixme --- pytato/distributed/partition.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pytato/distributed/partition.py b/pytato/distributed/partition.py index 856fa9535..0a82ac772 100644 --- a/pytato/distributed/partition.py +++ b/pytato/distributed/partition.py @@ -923,6 +923,8 @@ 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 + # }}} # {{{ assign each materialized array to a batch/part From 94c16c7a2b6506451cbb96b1ad97f91d5b0d13b9 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 28 Nov 2022 12:47:03 -0600 Subject: [PATCH 4/9] also assign part ID to materialized arrays that have no send/recv dependencies --- pytato/distributed/partition.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/pytato/distributed/partition.py b/pytato/distributed/partition.py index 0a82ac772..33a946b88 100644 --- a/pytato/distributed/partition.py +++ b/pytato/distributed/partition.py @@ -925,6 +925,8 @@ def get_needed_local_sends( # 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 @@ -952,15 +954,17 @@ def get_needed_local_sends( materialized_ary_to_part_id: Dict[Array, int] = { ary: max( max( - comm_id_to_part_id[ + (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)), + default=nparts-1), 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)) + for recvd_array in recvd_array_dep_mapper(ary)), + default=nparts-1) ) for ary in materialized_arrays } From 88b8ff8a26e9c1f8422f25a4daff4937860a0b03 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 28 Nov 2022 13:56:40 -0600 Subject: [PATCH 5/9] fix typo --- pytato/distributed/partition.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/distributed/partition.py b/pytato/distributed/partition.py index 33a946b88..e544787d8 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 From ea63005957f3942e84f3ce670a649953a999f5eb Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 28 Nov 2022 14:50:01 -0600 Subject: [PATCH 6/9] fix conditional conditional needs to be here because zip won't return the right number of variables if there aren't any receives; checking if #parts == 1 is not sufficient though, need to actually check for recvs --- pytato/distributed/execute.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/pytato/distributed/execute.py b/pytato/distributed/execute.py index 84fd2e16a..3f31482bd 100644 --- a/pytato/distributed/execute.py +++ b/pytato/distributed/execute.py @@ -88,16 +88,21 @@ def execute_distributed_partition( from mpi4py import MPI - recv_names_tup, recv_requests_tup, recv_buffers_tup = zip(*[ - (name,) + _post_receive(mpi_communicator, recv) - for part in partition.parts.values() - for name, recv in part.input_name_to_recv_node.items()]) - recv_names = list(recv_names_tup) - recv_requests = list(recv_requests_tup) - recv_buffers = list(recv_buffers_tup) - del recv_names_tup - del recv_requests_tup - del recv_buffers_tup + 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() + for name, recv in part.input_name_to_recv_node.items()]) + recv_names = list(recv_names_tup) + recv_requests = list(recv_requests_tup) + recv_buffers = list(recv_buffers_tup) + del recv_names_tup + del recv_requests_tup + del recv_buffers_tup + else: + recv_names = [] + recv_requests = [] + recv_buffers = [] context: Dict[str, Any] = input_args.copy() From ee18e54c5cfa9e7bcc5cac4a5fff7dbfece8f3b0 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 28 Nov 2022 15:21:54 -0600 Subject: [PATCH 7/9] attempt to fix materialized_ary_to_part_id computation --- pytato/distributed/partition.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/pytato/distributed/partition.py b/pytato/distributed/partition.py index e544787d8..9769ef170 100644 --- a/pytato/distributed/partition.py +++ b/pytato/distributed/partition.py @@ -951,24 +951,35 @@ def get_needed_local_sends( sent_array_dep_mapper = SubsetDependencyMapper(sent_arrays) recvd_array_dep_mapper = SubsetDependencyMapper(received_arrays) - materialized_ary_to_part_id: 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)), - default=nparts-1), + materialized_ary_to_part_id_range: Dict[Array, Tuple(int, int)] = { + ary: ( max( (comm_id_to_part_id[ _recv_to_comm_id(local_rank, cast(DistributedRecv, recvd_array))] for recvd_array in recvd_array_dep_mapper(ary)), - default=nparts-1) + default=0), + min( + (comm_id_to_part_id[ + _send_to_comm_id(local_rank, + sent_array_to_send_node[sent_array])] + 1 + for sent_array in sent_array_dep_mapper(ary)), + default=nparts), ) for ary in materialized_arrays } + assert all( + start < end + for start, end + in materialized_ary_to_part_id_range.values()), \ + "unable to find suitable part for materialized array" + + materialized_ary_to_part_id: Dict[Array, int] = { + ary: end - 1 + for ary, (_, end) + in materialized_ary_to_part_id_range.items()} + # }}} sent_ary_to_part_id = { From b80875f6f491bd6f973f97423fe24e6162793e1e Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 28 Nov 2022 16:13:33 -0600 Subject: [PATCH 8/9] attempt to fix materialized_ary_to_part_id computation (again) --- pytato/distributed/partition.py | 43 ++++++++++++++++++--------------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/pytato/distributed/partition.py b/pytato/distributed/partition.py index 9769ef170..f5f499829 100644 --- a/pytato/distributed/partition.py +++ b/pytato/distributed/partition.py @@ -948,37 +948,40 @@ def get_needed_local_sends( received_arrays = frozenset(lsrdg.local_recv_id_to_recv_node.values()) - sent_array_dep_mapper = SubsetDependencyMapper(sent_arrays) + materialized_array_dep_mapper = SubsetDependencyMapper(materialized_arrays) recvd_array_dep_mapper = SubsetDependencyMapper(received_arrays) - materialized_ary_to_part_id_range: Dict[Array, Tuple(int, int)] = { - ary: ( - max( + materialized_ary_to_part_id_start: Dict[Array, int] = { + ary: max( (comm_id_to_part_id[ _recv_to_comm_id(local_rank, - cast(DistributedRecv, recvd_array))] - for recvd_array in recvd_array_dep_mapper(ary)), - default=0), - min( - (comm_id_to_part_id[ - _send_to_comm_id(local_rank, - sent_array_to_send_node[sent_array])] + 1 - for sent_array in sent_array_dep_mapper(ary)), - default=nparts), - ) + 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( - start < end - for start, end - in materialized_ary_to_part_id_range.values()), \ + ( + 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: end - 1 - for ary, (_, end) - in materialized_ary_to_part_id_range.items()} + ary: materialized_ary_to_part_id_end[ary] - 1 + for ary in materialized_arrays} # }}} From 2ded7b795d59a4375cdce72044fa96951fdc690d Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 28 Nov 2022 16:17:10 -0600 Subject: [PATCH 9/9] exclude received arrays from materialized_arrays --- pytato/distributed/partition.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/pytato/distributed/partition.py b/pytato/distributed/partition.py index f5f499829..ee8ae6463 100644 --- a/pytato/distributed/partition.py +++ b/pytato/distributed/partition.py @@ -939,15 +939,16 @@ 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()) + # 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) @@ -990,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)