build CUDA support on all Linux platforms

Author: leofang

recipe/build-mpi.sh

@@ -21,11 +21,11 @@ if [[ "$target_platform" == osx-* ]]; then
     fi
 fi
 
-if [[ -z $CUDA_HOME ]]; then
+if [[ -z $CUDA_PATH ]]; then
     build_with_cuda=""
 else
-    export CFLAGS="$CFLAGS -I$CUDA_HOME/include"
-    export CXXFLAGS="$CXXFLAGS -I$CUDA_HOME/include"
+    export CFLAGS="$CFLAGS -I$CUDA_PATH/include"
+    export CXXFLAGS="$CXXFLAGS -I$CUDA_PATH/include"
     build_with_cuda="--with-cuda --with-ucx=$PREFIX"
 fi
 

recipe/conda_build_config.yaml

@@ -1,10 +1,14 @@
 mpi_type:
-  - external  # [linux and x86_64]
+  - external  # [linux]
   - conda
-docker_image:  # [linux and x86_64]
-  - quay.io/condaforge/linux-anvil-cos7-cuda:10.2     # [linux and x86_64]
-cudatoolkit:  # [linux and x86_64]
-  - 10.2      # [linux and x86_64]
+docker_image:  # [linux]
+  - quay.io/condaforge/linux-anvil-cos7-cuda:10.2      # [linux64]
+  - quay.io/condaforge/linux-anvil-ppc64le-cuda:11.0   # [ppc64le]
+  - quay.io/condaforge/linux-anvil-aarch64-cuda:11.0   # [aarch64]
+cudatoolkit:  # [linux]
+  - 10.2      # [linux64]
+  - 11.0      # [ppc64le]
+  - 11.0      # [aarch64]
 enable_cuda:
-  - True      # [linux and x86_64]
-  - False     # [not (linux and x86_64)]
+  - True      # [linux]
+  - False     # [not linux]

recipe/meta.yaml

@@ -1,6 +1,6 @@
 {% set version = "4.1.4" %}
 {% set major = version.rpartition('.')[0] %}
-{% set build = 1 %}
+{% set build = 2 %}
 
 # give conda package a higher build number
 {% if mpi_type == 'conda' %}
@@ -30,8 +30,6 @@ outputs:
         - {{ pin_subpackage('openmpi', min_pin='x.x.x', max_pin='x') }}
       ignore_run_exports:
         - ucx  # [enable_cuda]
-      script_env:
-        - CUDA_HOME     # [enable_cuda]
     requirements:
       build:
         - {{ compiler('c') }}
@@ -50,7 +48,7 @@ outputs:
         - zlib
         - mpi 1.0 openmpi
       run_constrained:
-        - cudatoolkit  >= {{ cudatoolkit }}       # [enable_cuda]
+        - cudatoolkit  >= {{ cudatoolkit }}         # [enable_cuda]
         - {{ pin_compatible("ucx", max_pin="x") }}  # [enable_cuda]
     test:
       script: run_test.sh