Add NS via cusolvermp to wheel build

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Author: ksivaman

build_tools/wheel_utils/Dockerfile.aarch

@@ -35,12 +35,23 @@ RUN dnf clean all
 RUN dnf -y install glog.aarch64 glog-devel.aarch64
 RUN dnf -y install libnccl libnccl-devel libnccl-static
 
+# expose system libs for TE CMake build.
+RUN dnf -y install \
+        libcusolvermp0-cuda-${CUDA_MAJOR} libcusolvermp0-devel-cuda-${CUDA_MAJOR} && \
+    dnf clean all
+RUN mkdir -p /opt/nvidia/cusolvermp && \
+    ln -s /usr/include/libcusolvermp/${CUDA_MAJOR} /opt/nvidia/cusolvermp/include && \
+    ln -s /usr/lib64/libcusolvermp/${CUDA_MAJOR}   /opt/nvidia/cusolvermp/lib && \
+    echo "/usr/lib64/libcusolvermp/${CUDA_MAJOR}" > /etc/ld.so.conf.d/999_nvidia_cusolvermp.conf && \
+    ldconfig
+
 ENV PATH="/usr/local/cuda/bin:${PATH}"
-ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
+ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:/opt/nvidia/cusolvermp/lib:${LD_LIBRARY_PATH}"
 ENV CUDA_HOME=/usr/local/cuda
 ENV CUDA_ROOT=/usr/local/cuda
 ENV CUDA_PATH=/usr/local/cuda
 ENV CUDADIR=/usr/local/cuda
+ENV CUSOLVERMP_HOME=/opt/nvidia/cusolvermp
 ENV NVTE_RELEASE_BUILD=1
 
 CMD ["/bin/bash", "-c", "bash /TransformerEngine/build_tools/wheel_utils/build_wheels.sh manylinux_2_28_aarch64 $BUILD_METAPACKAGE $BUILD_COMMON $BUILD_PYTORCH $BUILD_JAX $CUDA_MAJOR"]

build_tools/wheel_utils/Dockerfile.x86

@@ -35,12 +35,23 @@ RUN dnf clean all
 RUN dnf -y install glog.x86_64 glog-devel.x86_64
 RUN dnf -y install libnccl libnccl-devel libnccl-static
 
+# expose system libs for TE CMake build.
+RUN dnf -y install \
+        libcusolvermp0-cuda-${CUDA_MAJOR} libcusolvermp0-devel-cuda-${CUDA_MAJOR} && \
+    dnf clean all
+RUN mkdir -p /opt/nvidia/cusolvermp && \
+    ln -s /usr/include/libcusolvermp/${CUDA_MAJOR} /opt/nvidia/cusolvermp/include && \
+    ln -s /usr/lib64/libcusolvermp/${CUDA_MAJOR}   /opt/nvidia/cusolvermp/lib && \
+    echo "/usr/lib64/libcusolvermp/${CUDA_MAJOR}" > /etc/ld.so.conf.d/999_nvidia_cusolvermp.conf && \
+    ldconfig
+
 ENV PATH="/usr/local/cuda/bin:${PATH}"
-ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
+ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:/opt/nvidia/cusolvermp/lib:${LD_LIBRARY_PATH}"
 ENV CUDA_HOME=/usr/local/cuda
 ENV CUDA_ROOT=/usr/local/cuda
 ENV CUDA_PATH=/usr/local/cuda
 ENV CUDADIR=/usr/local/cuda
+ENV CUSOLVERMP_HOME=/opt/nvidia/cusolvermp
 ENV NVTE_RELEASE_BUILD=1
 
 CMD ["/bin/bash", "-c", "bash /TransformerEngine/build_tools/wheel_utils/build_wheels.sh manylinux_2_28_x86_64 $BUILD_METAPACKAGE $BUILD_COMMON $BUILD_PYTORCH $BUILD_JAX $CUDA_MAJOR"]

build_tools/wheel_utils/build_wheels.sh

@@ -25,6 +25,10 @@ git submodule update --init --recursive
 # Install deps
 /opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel
 
+# Enable optional build features. cuSolverMp is provided by the build image
+# (see Dockerfile.x86 / Dockerfile.aarch), which also sets CUSOLVERMP_HOME.
+export NVTE_WITH_CUSOLVERMP=1
+
 if $BUILD_METAPACKAGE ; then
         cd /TransformerEngine
         NVTE_BUILD_METAPACKAGE=1 /opt/python/cp310-cp310/bin/python setup.py bdist_wheel 2>&1 | tee /wheelhouse/logs/metapackage.txt