Added slurm support in 1D

.gitignore

@@ -25,3 +25,4 @@ gcn_distr_2d_gpu_gen.py
 gcn_distr_2d_gpu_floor_edit.py
 gcn_distr_2d_gpu_floor_old.py
 gcn_distr_2d_gpu_floor_gpu_gen.py
+*.pt

gcn_distr.py

@@ -682,6 +682,17 @@ def main():
         outputs = None
         if "OMPI_COMM_WORLD_RANK" in os.environ.keys():
             os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"]
+        # Initialize distributed environment with SLURM
+        if "SLURM_PROCID" in os.environ.keys():
+            os.environ["RANK"] = os.environ["SLURM_PROCID"]
+
+        if "SLURM_NTASKS" in os.environ.keys():
+            os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"]
+
+        if "MASTER_ADDR" not in os.environ.keys():
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+
+        os.environ["MASTER_PORT"] = "1234"
         dist.init_process_group(backend='nccl')
         rank = dist.get_rank()
         size = dist.get_world_size()

sparse-extension/build/temp.linux-x86_64-3.8/.ninja_log

@@ -0,0 +1,2 @@
+# ninja log v5
+0	18678	1619116367000000000	/global/cscratch1/sd/alokt/CAGNET/sparse-extension/build/temp.linux-x86_64-3.8/sparse_coo_tensor.o	73197118b8a3cab5

sparse-extension/build/temp.linux-x86_64-3.8/build.ninja

@@ -0,0 +1,20 @@
+ninja_required_version = 1.3
+cxx = c++
+
+cflags = -pthread -B /usr/common/software/pytorch/1.7.1-gpu/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/usr/common/software/sles15_cgpu/cuda/10.2.89/include -I/usr/common/software/pytorch/1.7.1-gpu/lib/python3.8/site-packages/torch/include -I/usr/common/software/pytorch/1.7.1-gpu/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/usr/common/software/pytorch/1.7.1-gpu/lib/python3.8/site-packages/torch/include/TH -I/usr/common/software/pytorch/1.7.1-gpu/lib/python3.8/site-packages/torch/include/THC -I/usr/common/software/pytorch/1.7.1-gpu/include/python3.8 -c
+post_cflags = -lcusparse -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1013"' -DTORCH_EXTENSION_NAME=sparse_coo_tensor_cpp -D_GLIBCXX_USE_CXX11_ABI=1 -std=c++14
+ldflags = 
+
+rule compile
+  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags
+  depfile = $out.d
+  deps = gcc
+
+
+
+build /global/cscratch1/sd/alokt/CAGNET/sparse-extension/build/temp.linux-x86_64-3.8/sparse_coo_tensor.o: compile /global/cscratch1/sd/alokt/CAGNET/sparse-extension/sparse_coo_tensor.cpp
+
+
+
+
+