pulsar build and CI changes

Summary:
Changes to CI and some minor fixes now that pulsar is part of pytorch3d. Most significantly, add CUB to CI builds.

Make CUB_HOME override the CUB already in cudatoolkit (important for cuda11.0 which uses cub 1.9.9 which pulsar doesn't work well with.
Make imageio available for testing.
Lint fixes.
Fix some test verbosity.
Avoid use of atomicAdd_block on older GPUs.

Reviewed By: nikhilaravi, classner

Differential Revision: D24773716

fbshipit-source-id: 2428356bb2e62735f2bc0c15cbe4cff35b1b24b8

packaging/build_conda.sh

@@ -17,4 +17,4 @@ setup_conda_pytorch_constraint
 setup_conda_cudatoolkit_constraint
 setup_visual_studio_constraint
 # shellcheck disable=SC2086
-conda build $CONDA_CHANNEL_FLAGS ${TEST_FLAG:-} -c defaults -c conda-forge --no-anaconda-upload -c fvcore --python "$PYTHON_VERSION" packaging/pytorch3d
+conda build $CONDA_CHANNEL_FLAGS ${TEST_FLAG:-} -c bottler -c defaults -c conda-forge --no-anaconda-upload -c fvcore --python "$PYTHON_VERSION" packaging/pytorch3d

packaging/build_wheel.sh

@@ -12,5 +12,6 @@ setup_env "$VERSION"
 setup_wheel_python
 pip_install numpy
 setup_pip_pytorch_version
+download_nvidiacub_if_needed
 python setup.py clean
 IS_WHEEL=1 python setup.py bdist_wheel

packaging/pkg_helpers.bash

@@ -251,24 +251,32 @@ setup_conda_pytorch_constraint() {
 # Translate CUDA_VERSION into CUDA_CUDATOOLKIT_CONSTRAINT
 setup_conda_cudatoolkit_constraint() {
   export CONDA_CPUONLY_FEATURE=""
+  export CONDA_CUB_CONSTRAINT=""
   if [[ "$(uname)" == Darwin ]]; then
     export CONDA_CUDATOOLKIT_CONSTRAINT=""
   else
     case "$CU_VERSION" in
       cu110)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.0,<11.1 # [not osx]"
+        # Even though cudatoolkit 11.0 provides CUB we need our own, to control the
+        # version, because the built-in 1.9.9 in the cudatoolkit causes problems.
+        export CONDA_CUB_CONSTRAINT="- nvidiacub"
         ;;
       cu102)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.2,<10.3 # [not osx]"
+        export CONDA_CUB_CONSTRAINT="- nvidiacub"
         ;;
       cu101)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.1,<10.2 # [not osx]"
+        export CONDA_CUB_CONSTRAINT="- nvidiacub"
         ;;
       cu100)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.0,<10.1 # [not osx]"
+        export CONDA_CUB_CONSTRAINT="- nvidiacub"
         ;;
       cu92)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=9.2,<9.3 # [not osx]"
+        export CONDA_CUB_CONSTRAINT="- nvidiacub"
         ;;
       cpu)
         export CONDA_CUDATOOLKIT_CONSTRAINT=""
@@ -292,3 +300,17 @@ setup_visual_studio_constraint() {
       cp packaging/$VSTOOLCHAIN_PACKAGE/conda_build_config.yaml packaging/pytorch3d/conda_build_config.yaml
   fi
 }
+
+download_nvidiacub_if_needed() {
+  case "$CU_VERSION" in
+    cu110|cu102|cu101|cu100|cu92)
+      echo "Downloading cub"
+      wget --no-verbose https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz
+      tar xzf 1.10.0.tar.gz
+      CUB_HOME=$(realpath ./cub-1.10.0)
+      export CUB_HOME
+      echo "CUB_HOME is now $CUB_HOME"
+      ;;
+  esac
+  # We don't need CUB for a cpu build or if cuda is 11.1 or higher
+}

packaging/pytorch3d/meta.yaml

@@ -8,6 +8,7 @@ source:
 requirements:
   build:
     - {{ compiler('c') }} # [win]
+    {{ environ.get('CONDA_CUB_CONSTRAINT') }}
 
   host:
     - python
@@ -31,6 +32,7 @@ build:
     - CUDA_HOME
     - FORCE_CUDA
     - NVCC_FLAGS
+    - MAX_JOBS
   features:
     {{ environ.get('CONDA_CPUONLY_FEATURE') }}
 
@@ -41,7 +43,7 @@ test:
     - tests
     - docs
   requires:
-    - ca-certificates
+    - imageio
   commands:
     #pytest .
     python -m unittest discover -v -s tests