[distributed] InputArgumentBase should not become "extra placeholders"

[kaushikcfd]

The following program:

import pytato as pt
import numpy as np
from pytato.distributed import (staple_distributed_send, make_distributed_recv,
                                find_distributed_partition)

from pytato.partition import generate_code_for_partition

rank = 0
size = 2

img1 = pt.make_placeholder("img1", shape=(100, 3), dtype="float32")
scl_np = np.random.rand(3)
scl = pt.make_data_wrapper(scl_np, name="scl")

y1 = img1 * scl

img2 = staple_distributed_send(
    y1, dest_rank=(rank-1) % size,
    comm_tag=42,
    stapled_to=make_distributed_recv(shape=img1.shape, dtype=img1.dtype,
                                     src_rank=(rank+1) % size, comm_tag=42,
                                     axes=img1.axes))
out = img1 + img2*scl


distributed_parts = find_distributed_partition(
    pt.make_dict_of_named_arrays({"out": out}))
prg_per_partition = generate_code_for_partition(distributed_parts)

for _, prg in prg_per_partition.items():
    print(prg.program)

performs unnecessary copies (redundant copies have been annotated in the kernel)

---------------------------------------------------------------------------
KERNEL: _pt_kernel
---------------------------------------------------------------------------
ARGUMENTS:
_pt_dist_id_0: type: np:dtype('float64'), shape: (100, 3), dim_tags: (N1:stride:3, N0:stride:1) aspace: global
_pt_part_ph_id: type: np:dtype('float32'), shape: (100, 3), dim_tags: (N1:stride:3, N0:stride:1) aspace: global
_pt_part_ph_id_0: type: np:dtype('float64'), shape: (3), dim_tags: (N0:stride:1) aspace: global
img1: type: np:dtype('float32'), shape: (100, 3), dim_tags: (N1:stride:3, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
scl: type: np:dtype('float64'), shape: (3), dim_tags: (N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
---------------------------------------------------------------------------
DOMAINS:
{ [_pt_part_ph_id_dim0, _pt_part_ph_id_dim1] : 0 <= _pt_part_ph_id_dim0 <= 99 and 0 <= _pt_part_ph_id_dim1 <= 2 }
{ [_pt_part_ph_id_0_dim0] : 0 <= _pt_part_ph_id_0_dim0 <= 2 }
{ [_pt_dist_id_0_dim0, _pt_dist_id_0_dim1] : 0 <= _pt_dist_id_0_dim0 <= 99 and 0 <= _pt_dist_id_0_dim1 <= 2 }
---------------------------------------------------------------------------
INSTRUCTIONS:
   for _pt_part_ph_id_dim1, _pt_part_ph_id_dim0
        # WARNING: UNNECESSARY COPY!!!
↱      _pt_part_ph_id[_pt_part_ph_id_dim0, _pt_part_ph_id_dim1] = img1[_pt_part_ph_id_dim0, _pt_part_ph_id_dim1]  {id=_pt_part_ph_id_store}
│  end _pt_part_ph_id_dim1, _pt_part_ph_id_dim0
│  for _pt_part_ph_id_0_dim0
|     # WARNING UNNECESSARY COPY!!!
│↱   _pt_part_ph_id_0[_pt_part_ph_id_0_dim0] = scl[_pt_part_ph_id_0_dim0]  {id=_pt_part_ph_id_0_store}
││ end _pt_part_ph_id_0_dim0
││ for _pt_dist_id_0_dim0, _pt_dist_id_0_dim1
└└     _pt_dist_id_0[_pt_dist_id_0_dim0, _pt_dist_id_0_dim1] = _pt_part_ph_id[_pt_dist_id_0_dim0, _pt_dist_id_0_dim1]*_pt_part_ph_id_0[_pt_dist_id_0_dim1]  {id=_pt_dist_id_0_store}
   end _pt_dist_id_0_dim0, _pt_dist_id_0_dim1
---------------------------------------------------------------------------

[kaushikcfd]

In general, I would prefer if we don't materialize arrays that the user didn't ask for. I'm still going through the code to see how difficult that is i.e. making the partitions disjoint shouldn't be one of the aims.

[inducer]

This might be hard problem:tm:.

Partitioning currently is based on a few assumptions:

• We must be able to execute parts in any dependency-respecting order.
• Parts are compiled at partitioning time, so what inputs they take from memory vs. what they compute is decided at that time.
• No part may depend on its own outputs as inputs. This is harder to decide than it looks, since we don't have full knowledge of the computation graph. Some edges go off to other nodes and then come back.
  • As a first step towards making this last point tractable, we currently strengthen the requirement to create partition boundaries on every edge that goes between nodes that are/aren't a dependency of a receive or that feed/don't feed a send.

That last point is particularly relevant if you're looking to move partition boundaries.

For our DG use case, making fewer partitions goes a long ways towards mitigating this, but there's otherwise not a general golden solution that I can see. What makes DG happier is that, in general, we just have one batch of communication that needs to happen (or two for second-order operators), one per flux exchange really. We're currently much dumber than that, creating partitions based on scalar sends. I'd asked @matthiasdiener to look into joining communications into batches. I don't know how far he's gotten with it.

(I feel like I ought to copy-paste some of this into the docs. PR incoming. edit: #274)

[inducer]

I should add that what makes these batches easy to recognize programmatically and without user involvement is that, at the grudge level, they happen as flux exchanges on an array container.

[kaushikcfd]

As a first step towards making this last point tractable, we currently strengthen the requirement to create partition boundaries on every edge that goes between nodes that are/aren't a dependency of a receive or that feed/don't feed a send.

Just to establish the severity of this, on a single rank we were able to have a 2D CNS DG operator as 30 GPU kernel launches. But now we have 130 partitions of the DAG itself leading to ~150-200 GPU kernel launches. At this point loop fusion won't be of much use since the individual kernels themselves turn out to be lean.

[inducer]

Just to establish the severity of this

I'm aware! Like I said, I think the first mitigation of this should be creating fewer partitions. For CNS, we should be able to get away with ~10 at most.

[inducer]

Maybe we should just sit down together and hack this...

[inducer]

Went ahead and wrote the pytato part of this. #275

[kaushikcfd]

Thanks! Although this is helpful, I still think our partitioning strategy is way too automatic (and therefore a bit noisy, too many nodes that are converted to placeholders). I think having a partitioning implementation that generates the partitions assuming only the materialized nodes as inputs/outputs would keep that down. I'm trying to put up an implementation around this and should further assist the work in #275.

[inducer]

that generates the partitions assuming only the materialized nodes as inputs/outputs would keep that down

Ooooh. If you could get that to go, that'd be very sweet.

[inducer]

Should this be closed in view of #277?