def visit(self, op, *op_args):
args = list()
for arg in op_args:
comm_pattern = get_comm_pattern(from_node=arg, to_node=op)
if comm_pattern:
pair = CommNodePair(from_node=arg,
to_node=op,
node_type=comm_pattern)
if pair.get_send_node():
self.send_nodes.add(pair.get_send_node())
if pair.get_recv_node():
recv_node = pair.get_recv_node()
if isinstance(recv_node, (dict)):
start_node = recv_node['start_node']
wait_node = recv_node['wait_node']
args.append(start_node)
op.add_control_dep(wait_node)
start_node.invalidate_property_cache('all_deps')
wait_node.invalidate_property_cache('all_deps')
else:
args.append(pair.get_recv_node())
else:
args.append(arg)
op._args = tuple(args)
# invalidate deps cache as op._args is updated
op.invalidate_property_cache('all_deps')
def visit(self, op):
args = list()
for arg in op.args:
comm_pattern = get_comm_pattern(from_node=arg, to_node=op)
if comm_pattern:
pair = CommNodePair(from_node=arg,
to_node=op,
node_type=comm_pattern)
self.send_nodes.add(pair.get_send_node())
args.append(pair.get_recv_node())
else:
args.append(arg)
op._args = tuple(args)
# invalidate deps cache as op._args is updated
op.invalidate_property_cache('all_deps')
def test_get_node_type(from_node, to_node, expected_type):
assert expected_type == get_comm_pattern(from_node, to_node)