Edge placement update (#11)

Updated API for NCCL and NVSHMEM backends for DGraph with running OGB and Benchmark code. Incorporates caching to reduce overhead.

In detail:
* Adding Multi-partition function

* Adding preprocessing code

* Add padding functionality to NVSHMEM operations, enabling arbitrary
shaped inputs to the function

* Revamped the Distributed Graph implementation to simplify the code and
remove the old slicing based code

- Also adds documentation

* Updating Graphcast implementation with DGraph distributed

* Incorporate NCCL cache into the backend engine

* Update local tensor getter to use a placement tensor mask

* Add Graphcast distributed trainer

* Fix distributed graph object

* Save intermediate preprocessed file so it can be reused

* Remove unnecessary complexity of file passing and use a single torch tensor

* Graphcast update static graph generator with preprocessing code

* Updated mesh graph placement algorithm

* Fixed scatter test with new API

* Update distributed GCN with edge placement tensor

* Add GatherCacheGenerator and ScatterCacheGenerator

* Add Graphcast preprocessing
* Add static method for mesh partitioning
* Add grid vertex placement logic to MeshGraph

* Add OGBN-products update

* Overly complicated but correct graph data

* Add separated out benchmarking code for small tests

* Fix missing batch dim

* More general fix

* Disabled some incomplete NVSHMEM caching optimizations.  Added code to
set the default PyTorch device.

* Add NVSHMEM benchmark code
- Append backend type to output files
- Add sample plot generation code

* Adding torch distributed init with NVSHMEM communicator

* Apply suggestions from code review

* Fixed plotting script to grab the right log files.

* Fix the cached benchmarks

* Apply review suggestions, remove dead code

---------

Co-authored-by: Brian C. Van Essen <vanessen1@llnl.gov>

Author: szaman19

DGraph/Communicator.py

@@ -12,9 +12,9 @@
 #
 # SPDX-License-Identifier: (Apache-2.0)
 import torch
-from DGraph.distributed.mpi import MPIBackendEngine
+
 from DGraph.distributed.nccl import NCCLBackendEngine
-from DGraph.distributed.nvshmem import NVSHMEMBackendEngine
+
 from DGraph.CommunicatorBase import CommunicatorBase
 
 SUPPORTED_BACKENDS = ["nccl", "mpi", "nvshmem"]
@@ -38,8 +38,12 @@ def __init__(self, backend: str, **kwargs) -> None:
         if backend == "nccl":
             self.__backend_engine = NCCLBackendEngine()
         elif backend == "mpi":
+            from DGraph.distributed.mpi import MPIBackendEngine
+
             self.__backend_engine = MPIBackendEngine(**kwargs)
         elif backend == "nvshmem":
+            from DGraph.distributed.nvshmem import NVSHMEMBackendEngine
+
             self.__backend_engine = NVSHMEMBackendEngine()
         else:
             raise NotImplementedError(f"Backend {backend} not implemented")
@@ -65,6 +69,32 @@ def get_local_rank_slice(self, tensor: torch.Tensor, dim: int = -1) -> torch.Ten
         self.__check_init()
         return self.__backend_engine.get_local_rank_slice(tensor, dim)
 
+    def get_local_tensor(
+        self, tensor: torch.Tensor, placement_tensor: torch.Tensor, dim: int = -1
+    ) -> torch.Tensor:
+        """Returns the tensor corresponding to the current process based on the placement tensor.
+
+        Args:
+            tensor: The tensor to be sliced.
+            placement_tensor: A boolean tensor of the same shape as the tensor, where True indicates the process
+                that should receive the corresponding element.
+            dim: The dimension along which the tensor should be sliced.
+
+        Returns:
+            (torch.Tensor): The local tensor corresponding to the current process.
+        """
+        self.__check_init()
+        mask = (placement_tensor == self.get_rank()).bool()
+        mask_shape = [1] * tensor.ndim
+        mask_shape[dim] = mask.size(0)
+        mask_expanded = mask.view(mask_shape).expand_as(tensor)
+        masked_tensor = tensor[mask_expanded]
+        new_shape = list(tensor.shape)
+        new_shape[dim] = int(mask.sum().item())
+        masked_tensor = masked_tensor.view(new_shape)
+
+        return masked_tensor
+
     def scatter(self, *args, **kwargs) -> torch.Tensor:
         self.__check_init()
         return self.__backend_engine.scatter(*args, **kwargs)

DGraph/CommunicatorBase.py

@@ -15,6 +15,7 @@ class CommunicatorBase:
     _is_initialized = False
 
     def __init__(self):
+        self.backend = ""
         pass
 
     def init_process_group(self, backend: str, **kwargs):