Added CUDA support for non-x86_64 platforms

Author: stemann

M/MLX/MLX_CUDA/build_tarballs.jl

@@ -19,6 +19,13 @@ for platform in platforms
 
     cuda_deps = CUDA.required_dependencies(platform; static_sdk=true)
 
+    # Download the CUDA nvcc redist for the host architecture (x86_64) for non-x86_64-linux-gnu platforms
+    platform_sources = BinaryBuilder.AbstractSource[sources...]
+    if !platforms_match(platform, Platform("x86_64", "linux"))
+        cuda_version = platform["cuda"]
+        push!(platform_sources, CUDA.cuda_nvcc_redist_source(cuda_version, "x86_64"))
+    end
+
     build_tarballs(ARGS, name, version, sources, script, [platform], products, [dependencies; cuda_deps];
                    lazy_artifacts=true,
                    julia_compat,

M/MLX/common.jl

@@ -50,6 +50,24 @@ if [[ $bb_full_target == *cuda* ]]; then
     )
 fi
 
+if $USE_CCACHE; then
+    export CMAKE_CUDA_COMPILER_LAUNCHER=ccache
+fi
+
+if [[ $bb_full_target == *cuda* && $target != x86_64-linux-gnu* ]]; then
+    # Add /usr/lib/csl-glibc-x86_64 to LD_LIBRARY_PATH to be able to run host `nvcc`/`ptxas`/`fatbinary`
+    # while keeping the default /usr/lib/csl-musl-x86_64,
+    export LD_LIBRARY_PATH=/usr/lib/csl-musl-x86_64:/usr/lib/csl-glibc-x86_64:$LD_LIBRARY_PATH
+
+    # Make sure the host CUDA executables are used by copying from the host (x86_64) nvcc redist
+    NVCC_DIR=(/workspace/srcdir/cuda_nvcc-linux-$(arch)-${cuda_version}*-archive)
+    rm -rfv $prefix/cuda/bin
+    cp -av ${NVCC_DIR}/bin $prefix/cuda/bin
+    
+    rm -rfv $prefix/cuda/nvvm/bin
+    cp -av ${NVCC_DIR}/nvvm/bin $prefix/cuda/nvvm/bin
+fi
+
 install_license LICENSE
 
 if [[ "$target" != aarch64-apple-darwin* ]]; then