#!/bin/bash

# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES
# SPDX-License-Identifier: BSD-3-Clause

# This script as an example of how a custom communication layer can be built for cuDSS
# so that distributed multi-node multi-GPU computations can be used.
# It requires an NVCC compiler and GPU-aware MPI or NCCL.
#
# Please check/set the following environment variables:
#  - $CUDA_PATH = Path to your CUDA installation directory.
#  - $CUDSS_PATH = Path to your CUDSS installation directory.
#  - $MPI_PATH  = Path to your MPI library installation directory.
#                 If your MPI library is installed in system
#                 directories as opposed to its own (root) directory,
#                 ${MPI_PATH}/include is expected to contain the mpi.h header.
#                 ${MPI_PATH}/lib64 or ${MPI_PATH}/lib is expected to contain libmpi.so.
#  - $NCCL_PATH = Path to your NCCL library installation directory.
#                 If your NCCL library is installed in system
#                 directories as opposed to its own (root) directory,
#                 ${NCCL_PATH}/include is expected to contain the mpi.h header.
#                 ${NCCL_PATH}/lib64 or ${NCCL_PATH}/lib is expected to contain libnccl.so.
#
# Run (inside this directory): source ./cudss_build_commlayer.sh

if [ "$1" == "-h" ] || [ "$1" == "--help" ]
  then
    echo "Usage example:"
    echo "chmod +x $0; CUDA_PATH=<CUDA CTK path> MPI_PATH=<MPI path with include/ and lib/ or lib64/> NCCL_PATH=<NCCL pat with include/ and lib/ or lib64/> ./$0.sh <openmpi|nccl>"
    exit 0
elif [ "$1" == "openmpi" ] || [ "$1" == "nccl" ]
  then
    echo "Building communication layer with ${1} backend"
else
    echo "Script input arguments (for help run with -h or --help):"
    echo "Args: $@"
    exit 1
fi


if [ -z "${CUDA_PATH}" ]
then
    echo "Environment variable CUDA_PATH is not set. Please set it to point to the CUDA root directory!"
    exit 2
fi

if [ -z "${CUDSS_PATH}" ]
then
    echo "Environment variable CUDSS_PATH is not set. Please set it to point to the CUDSS installation directory (for headers)!"
    exit 2
fi

backend=$1
#backend=openmpi
#backend=nccl

if [[ "$backend" == "openmpi" ]]; then
    COMM_LIB_PATH=${MPI_PATH}
    comm_lib_name="mpi"
    if [ -z "${MPI_PATH}" ]
    then
        echo "Environment variable MPI_PATH is not set. Plese set it to point to the MPI root directory!"
        echo "Note that MPI_PATH/include is expected to contain the mpi.h header."
        exit 3
    fi
elif [[ "$backend" == "nccl" ]]; then
    COMM_LIB_PATH=${NCCL_PATH}
    comm_lib_name="nccl"
    if [ -z "${NCCL_PATH}" ]
    then
        echo "Environment variable NCCL_PATH is not set. Plese set it to point to the NCCL root directory!"
        echo "Note that NCCL_PATH/include is expected to contain the nccl.h header."
        exit 3
    fi
fi

${CUDA_PATH}/bin/nvcc -forward-unknown-to-host-compiler -shared -fPIC \
    -Wall -Werror all-warnings \
    -I${CUDA_PATH}/include -I${CUDSS_PATH}/include -I${COMM_LIB_PATH}/include \
    cudss_commlayer_${backend}.cu \
    -L${COMM_LIB_PATH}/lib64 -L${COMM_LIB_PATH}/lib -l${comm_lib_name} \
    -o libcudss_commlayer_${backend}.so

# To enable usage of the communication layer in cuDSS routines, please set the environment
# variable below and follow the docs for the required source code changes
# export CUDSS_COMM_LIB=${PWD}/libcudss_commlayer_${backend}.so

ls libcudss_commlayer_${backend}.so
exit $?

