Unify the backpressuring mechanism to handle both individual task launches and points from an index launch. A level of indirection here is likely all that is needed. The approach for doing the backpressuring on index launches is in #102.

It would save some time if we could already pack a tensor into a desired compressed format (like CSR) once while preprocessing the tensor in tensor_to_hdf5. Then, future operations can just load the tensor and avoid a somewhat expensive packing step, which can save some time while running experiments.

https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.55.6125&rep=rep1&type=pdf

The schedule just looks like

.distribute({i, j}, {in, jn}, {il, jl}, Grid(...))
.stagger(k, {in}, kos)
.communicate(a(i, j), in)
.communicate({b(i, k), c(k, j}, kos)
;

It's a hybrid between cannon and summa that only does shifting along one axis, and broadcasts along another.

A meta issue for me to track all of the Legion issues that I'm blocked on.

Legion

These are bugs blocking either implementation of benchmarks that I want to write, or evaluation of benchmarks that I already have.

GASNetEx

• Invalid gather from GPU device memory StanfordLegion/legion#1122 -- hard blocker. Shows up in MTTKRP, COSMA. Running with -gex:bindcuda 0 gets rid of the segfault, but comes at an unacceptable performance hit (basically puts us back at gasnet1 speeds).
• Crashes when running at high node counts StanfordLegion/legion#1118 (blocks large multi-GPU GEMM runs). Even with -gex:batch 0 I see some errors, and performance with -gex:batch 0 is noticeably worse.
• 2D Gathers with Gex StanfordLegion/legion#1138. I believe that multi-node adapt benchmark needs this. I think that this will also improve performance of the 2d algorithms on rectangular node counts.

Core Legion

• Assertion failure around futures / hangs with mttkrp + futures StanfordLegion/legion#1123. (TODO (rohany): When this is fixed, pull control replication to get the fix for memory corruption around instance reuse).
• Slow CPU fills StanfordLegion/legion#1113
• Memory corruption around instance reuse StanfordLegion/legion#1108 (sean has a patch for this in progress)
• Collective instances StanfordLegion/legion#546. These are necessary for algorithms which fully replicate a (small) single region among all GPUs. We have a hacky workaround for now, but it isn't great. Collective instances will also be needed for algorithms (like mttkrp) that replicate a portion of a region onto many nodes (StanfordLegion/legion#1125).
• Another hang around future maps: StanfordLegion/legion#1145

Feature Requests

• Virtual mappings for reduction instances StanfordLegion/legion#1119. Without this, we can't feasibly implement hierarchical distribution with algorithms that use reductions at the leaves (COSMA, Johnson, Solomonik). This does not look like it's going to be implemented by the PLDI deadline, so I should pivot things to performing flat decompositions over each GPU.

Legate

I'd like to benchmark against Legate, but there are some errors here blocking me from doing so.

• Sharding functor output error nv-legate/cunumeric#53
• Bad instance mapping leading to OOMs (I don't have an issue for this, but Manolis said he is looking into it).

For example, it could test whether GPUs are present and enable cuda, check how many sockets are present on the machine etc.

https://cmake.org/cmake/help/latest/module/FindBLAS.html

The lowerer currently doesn't handle a case where we do a position split, but some tensor dimensions are partitioned by an outer variable as well, so the constructed projection needs to handle dimensions partitioned from the outside. This can been seen in #97 .

Some fixes have gone into control replication that we need. Pull it in and ensure that nothing has broken.

the mapper likely needs to be changed when running GPU benchmarks with multiple OMP procs that have a heirarchical distribution. This line:

However, there's a chance that having a shard per node with control replication will cause this to not have any problems.

This script should handle the installation directories, and build legion, openblas, distal etc.

This was an annoying thing to rediscover...

Follow the example in #119

I still have to decide if / when I should do this, but it seems like something like this maybe possible? The idea is that we can create our partitions once up front, even for the hierarchical ones by explicitly holding onto all of the handles that sub-partitions tasks create. However, this leads to some extremely nested structures, where perhaps too many partitions are being created at once (consider cannon's GPU implementation which creates a partition of each tensor at each level of the computation for 4 levels -- node level, node level k loop, gpu level, gpu level k loop. This quickly balloons out to a large number of partition objects to manage at the top level. Before undertaking this, I would need to be more convinced that there are performance wins to gain here.

As soon as the collective instance branch is ready for testing, we need to move to it and introduce new concepts in the lowerer and mapper to correctly handle creating them. This is a three phase process.

1. Simple "replicated tensor" computations, and remove all of the hard-coded manual replication things. Codes that come to mind are

• SpMV weak scale
• TTV
• TTMC
• MTTKRP

2. More complicated launch patterns where subsets of launches need pieces of tensors

• Johnson's Algorithm
• COSMA

3. 2D matrix computations that do lock-step broadcast communcations, such as SUMMA. This step will be the hardest, as it requires changing alot of code. The problem with the current approach is that it does 1 2D launch, and then each launched sub-task launches a bunch more tasks. It's likely that we will need to convert this into a 3D launch with a projection functor that understands the ordering between tasks (also generated by DISTAL), and then chooses collectives to use for each row/column.

• PUMMA
• SUMMA
• 2.5D MatMul

The current strategy for naming partitions by their iteration space point identifier works great when all tensors are partitioned at the top level. However, when a tensor is partitioned only at lower levels of the computation, this becomes complicated as we need to ensure that all index spaces in the LegionTensor get a different color, and can also look up that color later. We avoid this problem for all tensors partitioned at the top level by letting the top level partition color be auto generated, giving all sub-partitions of the tensor a separate namespace.

In addition to a correctness motivation, there may be a performance benefit here. For example, consider the computation A(i, j) = B(i, k) * C(k, j), where we distribute i and communicate C under j. Here, we would want to create a single partition of C at the top level that all task loops over j operate on. The current strategy creates a partition of C for each subtask instead, which may not be as performant.

Hoisting these partitioning operations would simplify several aspects of the taco-legion interop (and avoid gotchas around index space reuse) as well as potentially allow for some performance improvements. The optimization is not deep and is just an application of LICM, but will be some work to implement.

Several small things to do:

• Be resilient to partial installations and allow for rebuilding certain components
• Control the number of threads to use with make
• Command line control over which builds are expected to succeed

We might be able to up front create the partition for the top level partitions, and then all recursive partitions of those can just be created and selected out with the projection functor.

• Inner product of two order 3 tensors ($$\sum_{ijk} = A_{ijk} * B_{ijk}$$) -- I don't think that the compiler can generate code that does this right now. It also seems like another case of something that requires collective instances for the broadcast reduction. I looked into this a bit more, and here's a plan of attack: 1) when doing a scalar reduction like this, we would want to return the result from intermediate tasks, and then use runtime->reduce_future_map to get the result, rather than replicating a dummy region and everyone reducing into that.

In order to avoid doing data dependent partitioning operations (or partitioning operations in general) when they are unneeded, we should expose a method to calculate all of the partitions needed (for at least the top level) of a distributed computation and returns these to the user. Then, the compute method should accept partitions for each of the tensors and not do any partitioning of its own in order to limit the runtime overhead of tracking all the extra metadata.

It's unclear how much of a benefit this will give us for dense computations, but it definitely will be useful for sparse computations where partitioning operations are data dependent. In such cases, avoiding a partitioning operation at each iteration of a solver will greatly improve performance.

openmp usages in fill and validate should be compiled as part of the library, not live in header files.

If it the cosma performance looks like what I saw with johnsons (where the computation is mainly bottlenecked on getting data then runs 1 task), add the steps back into the algorithm to get some pipelining going.

This makes it easier to track whats going on within DISTAL when things go wrong.

• Test that it works on Sapling in multi-node settings
• Lassen in multi-node settings
• stop printing output from the compilation command to stdout unless asked
• Implement GPU compilation support
• Move runtime library into the DISTAL directory
• Move runtime library within a namespace
• Move submodule dependencies into a folder called deps
• Transition compilation tests and generated code tests to use the new API, integrate with gtest
• Figure out installation details / exports of link and include directories to a target

Marking the whole task as untrack is a bit heavy handed, as this can result in the runtime collecting instances that we don't want collected. Instead, communicate should have a marker on it that says if a particular communication operation should allow for collection.

The precompute relation right now supports only when temporaries exist within a single machine. Based on discussions at PLDI, it would be useful to be able to factorize computations with precompute where the temporaries themselves are distributed across the machine, similarly to the input tensors.

This way people can set this stuff on construction of the LegionTensor, and don't need to be threaded through very far in the compilation pipeline

Support the basic subset of SpDISTAL's tensor distribution notation extensions that align naturally with Legion's partitioning operators + tensor layouts.