bench_MGG_DSE_dist_ps.py

#!/usr/bin/env python3
import os
import sys

os.environ["PYTHONWARNINGS"] = "ignore"
os.environ["LD_LIBRARY_PATH"] += os.pathsep + 'local/openmpi-4.1.1/lib/'
os.environ["PATH"] += os.pathsep + 'local/openmpi-4.1.1/bin/'
os.environ["LD_LIBRARY_PATH"] += os.pathsep + 'local/cudnn-v8.2/lib64'
os.environ["NVSHMEM_SYMMETRIC_SIZE"] = '14763950080' 

hidden = 16
warpPerblock = 4 
interleaved_dist = int(sys.argv[1])
partSize = int(sys.argv[2])

num_GPUs = int(sys.argv[3])

dataset = [
        ( 'Reddit'                      , 602      	, 41),
]

GPU_avail = "CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 "
# GPU_avail = "CUDA_VISIBLE_DEVICES=4,5,6,7 "

pre_condit = GPU_avail + 'OMPI_MCA_plm_rsh_agent=sh\
              mpirun --allow-run-as-root -np {} '.format(num_GPUs)

command = "build/MGG_np_div_kernel "

for data, in_dim, out_dim in dataset:
        beg_file = "dataset/bin/{}_beg_pos.bin".format(data)
        csr_file = "dataset/bin/{}_csr.bin".format(data)
        weight_file = "dataset/bin/{}_weight.bin".format(data)
        os.system(pre_condit + "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}".
        format(command, beg_file, csr_file, weight_file,  
                num_GPUs, partSize, warpPerblock, interleaved_dist,
                in_dim, hidden, out_dim))