Use parRSB.py to find halo info

3rd_party/gnn/dist-gnn/generate_halo_info.py

@@ -0,0 +1,38 @@
+import argparse
+from mpi4py import MPI
+from parrsb import Mesh
+import numpy as np
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Process command line arguments."
+    )
+    parser.add_argument(
+        "-p",
+        "--poly",
+        type=int,
+        required=True,
+        help="Specify the polynomial order",
+    )
+    parser.add_argument(
+        "-c",
+        "--case",
+        type=str,
+        required=True,
+        help="Specify the mesh file path (.re2 file, without the extension)",
+    )
+    parser.add_argument(
+        "-l",
+        "--log",
+        type=str,
+        default="info",
+        choices=["debug", "info", "warning", "error"],
+        required=False,
+        help="Specify the log level (default: info)",
+    )
+    args = parser.parse_args()
+
+    comm = MPI.COMM_WORLD
+    m = Mesh(args.case, comm)
+    partitions = m.partition()

scripts/ml/setup_case

@@ -245,7 +245,27 @@ function load_modules() {
   fi
 }
 
-function build_smartsim() {
+function install_parrsb() {
+  GSLIB_DIR=$(mktemp -d)
+  PARRSB_DIR=$(mktemp -d)
+
+  # Install gslib
+  git clone https://github.com/thilinarmtb/gslib.git -b general_graph ${GSLIB_DIR}
+  cmake -B ${GSLIB_DIR}/build -S ${GSLIB_DIR} -DCMAKE_INSTALL_PREFIX=${VENV}
+  cmake --build ${GSLIB_DIR}/build --target install
+
+  # Install parRSB
+  git clone https://github.com/thilinarmtb/parRSB.git ${PARRSB_DIR}
+  cmake -B ${PARRSB_DIR}/build -S ${PARRSB_DIR} -DCMAKE_INSTALL_PREFIX=${VENV} \
+    -Dgs_DIR=${VENV}/lib/cmake/gs
+  cmake --build ${PARRSB_DIR}/build --target install
+
+  # Install parRSB.py
+  pip install git+https://github.com/thilinarmtb/parRSB.py --prefix ${VENV} -vv \
+    -Ccmake.define.parRSB_DIR=${VENV}/lib/cmake/parRSB
+}
+
+function install_smartsim() {
   if [ "${SYSTEM}" == "polaris" ]; then
     export CC=cc
     export CXX=CC
@@ -272,7 +292,6 @@ function build_smartsim() {
 }
 
 function setup_venv() {
-  reuse=0
   if [ -d ${VENV_PATH} ]; then
     echo -e "\033[35mPython venv \"${VENV_PATH}\" already exists, reusing it ... \033[m"
     return 0
@@ -297,21 +316,21 @@ function setup_venv() {
 
   if [ ${MODEL} == "sr-gnn" ]; then
     if [ ${SYSTEM} == "aurora" ]; then
-      cwd=${PWD}
-      git clone https://github.com/rusty1s/pytorch_cluster.git ${VENV_PATH}/pytorch_cluster
-      cd ${VENV_PATH}/pytorch_cluster
-      # Need to force the install to not link with OpenMP, change lines 53-54 to if False:
-      sed -i 's/fopenmp/lgomp/' setup.py
-      CXX=$(which dpcpp) python setup.py install
-      cd ${cwd}
+      PYTCLUSTER_DIR=$(mktemp -d)
+      git clone https://github.com/rusty1s/pytorch_cluster.git $PYTCLUSTER_DIR
+      # Need to force the install to not link with OpenMP,
+      # change lines 53-54 to if False:
+      sed -i 's/fopenmp/lgomp/' ${PYTCLUSTER_DIR}/setup.py
+      CXX=$(which dpcpp) python ${PYTCLUSTER_DIR}/setup.py install --prefix ${VENV_PATH}
     fi
   fi
 
   pip install pymech
 
   if [ "${CLIENT}" == "smartredis" ]; then
-    build_smartsim
+    install_smartsim
   fi
+  install_parrsb
 
   deactivate
 }
@@ -331,3 +350,4 @@ setup_case
 load_modules
 setup_venv
 generate_script
+echo -e "\033[35mExecute \"run.sh\" script from a compute node to run the example.\033[m"