diff --git a/.circleci/config.yml b/.circleci/config.yml index 4077f91d3a683..bcd2a5527b983 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1,7 +1,929 @@ +docker_config_defaults: &docker_config_defaults + user: jenkins + aws_auth: + # This IAM user only allows read-only access to ECR + aws_access_key_id: AKIAJ2J6FIG5OSZTQ3IA + aws_secret_access_key: ${CIRCLECI_AWS_SECRET_KEY_FOR_ECR_READ_ONLY} + +pytorch_linux_cpu_build_test_defaults: &pytorch_linux_cpu_build_test_defaults + resource_class: large + working_directory: /var/lib/jenkins/workspace + steps: + - checkout + - run: + name: Build + no_output_timeout: "10h" + command: | + export IN_CIRCLECI=1 + export SCCACHE_BUCKET=ossci-compiler-cache-circleci-v2 + export SCCACHE_MAX_JOBS=`expr $(nproc) - 1` + export MEMORY_LIMIT_MAX_JOBS=8 # the "large" resource class on CircleCI has 32 CPU cores, if we use all of them we'll OOM + export MAX_JOBS=$(( ${SCCACHE_MAX_JOBS} > ${MEMORY_LIMIT_MAX_JOBS} ? ${MEMORY_LIMIT_MAX_JOBS} : ${SCCACHE_MAX_JOBS} )) + # This IAM user allows write access to S3 bucket for sccache + export AWS_ACCESS_KEY_ID=AKIAJJZUW4G2ASX5W7KA + export AWS_SECRET_ACCESS_KEY=${CIRCLECI_AWS_SECRET_KEY_FOR_SCCACHE_S3_BUCKET} + git submodule update --init + .jenkins/pytorch/build.sh + .jenkins/pytorch/test.sh + +pytorch_linux_build_defaults: &pytorch_linux_build_defaults + resource_class: large + working_directory: /var/lib/jenkins/workspace + steps: + - checkout + - run: + name: Build + no_output_timeout: "10h" + command: | + export IN_CIRCLECI=1 + export SCCACHE_BUCKET=ossci-compiler-cache-circleci-v2 + if [ -n "${CUDA_VERSION}" ]; then + export TORCH_CUDA_ARCH_LIST=5.2 + fi + export SCCACHE_MAX_JOBS=`expr $(nproc) - 1` + export MEMORY_LIMIT_MAX_JOBS=8 # the "large" resource class on CircleCI has 32 CPU cores, if we use all of them we'll OOM + export MAX_JOBS=$(( ${SCCACHE_MAX_JOBS} > ${MEMORY_LIMIT_MAX_JOBS} ? ${MEMORY_LIMIT_MAX_JOBS} : ${SCCACHE_MAX_JOBS} )) + # This IAM user allows write access to S3 bucket for sccache + export AWS_ACCESS_KEY_ID=AKIAJJZUW4G2ASX5W7KA + export AWS_SECRET_ACCESS_KEY=${CIRCLECI_AWS_SECRET_KEY_FOR_SCCACHE_S3_BUCKET} + git submodule update --init || git submodule update --init || git submodule update --init + .jenkins/pytorch/build.sh + mkdir -p pytorch-ci-env/ + cp -r /opt/conda/lib/python${PYTHON_VERSION}/site-packages/torch pytorch-ci-env/torch + cp -r build/bin pytorch-ci-env/cpp_test_bin + if [ -d "../cpp-build" ]; then + cp -r ../cpp-build pytorch-ci-env/cpp-build + fi + - persist_to_workspace: + root: /var/lib/jenkins/workspace/pytorch-ci-env + paths: + - "*" + +pytorch_linux_test_defaults: &pytorch_linux_test_defaults + machine: + image: default + steps: + - checkout + - run: + name: Prepare workspace + command: | + sudo mkdir -p /opt/workspace + sudo chmod -R 777 /opt/workspace + - attach_workspace: + at: /opt/workspace + - run: + name: Build + no_output_timeout: "10h" + command: | + set -x + sudo pip install awscli + if [ -n "${CUDA_VERSION}" ]; then + curl -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - + echo "deb https://nvidia.github.io/libnvidia-container/ubuntu14.04/amd64 /" | sudo tee -a /etc/apt/sources.list.d/nvidia-docker.list + echo "deb https://nvidia.github.io/nvidia-container-runtime/ubuntu14.04/amd64 /" | sudo tee -a /etc/apt/sources.list.d/nvidia-docker.list + echo "deb https://nvidia.github.io/nvidia-docker/ubuntu14.04/amd64 /" | sudo tee -a /etc/apt/sources.list.d/nvidia-docker.list + fi + sudo apt-get update + sudo apt-get remove linux-image-generic linux-headers-generic linux-generic + sudo apt-get install linux-headers-$(uname -r) + sudo apt-get install linux-image-generic + if [ -n "${CUDA_VERSION}" ]; then + wget 'https://s3.amazonaws.com/ossci-linux/nvidia_driver/NVIDIA-Linux-x86_64-396.26.run' + sudo /bin/bash ./NVIDIA-Linux-x86_64-396.26.run -s --no-drm + sudo apt-get install -y nvidia-docker2 + fi + sudo pkill -SIGHUP dockerd + if [ -n "${CUDA_VERSION}" ]; then + nvidia-smi + fi + # This IAM user only allows read-only access to ECR + export AWS_ACCESS_KEY_ID=AKIAJ2J6FIG5OSZTQ3IA + export AWS_SECRET_ACCESS_KEY=${CIRCLECI_AWS_SECRET_KEY_FOR_ECR_READ_ONLY} + eval $(aws ecr get-login --region us-east-1 --no-include-email) + docker pull ${DOCKER_IMAGE} + if [ -n "${CUDA_VERSION}" ]; then + id=$(docker run --runtime=nvidia -t -d -w /var/lib/jenkins ${DOCKER_IMAGE}) + else + id=$(docker run -t -d -w /var/lib/jenkins ${DOCKER_IMAGE}) + fi + pwd + echo "declare -x IN_CIRCLECI=1" > /home/circleci/project/env + echo "declare -x PYTHON_VERSION=${PYTHON_VERSION}" >> /home/circleci/project/env + echo "declare -x SCCACHE_BUCKET=ossci-compiler-cache-circleci-v2" >> /home/circleci/project/env + # This IAM user allows write access to S3 bucket for sccache + echo "declare -x AWS_ACCESS_KEY_ID=AKIAJJZUW4G2ASX5W7KA" >> /home/circleci/project/env + echo "declare -x AWS_SECRET_ACCESS_KEY=${CIRCLECI_AWS_SECRET_KEY_FOR_SCCACHE_S3_BUCKET}" >> /home/circleci/project/env + mkdir -p /home/circleci/project/build + cp -r /opt/workspace/cpp_test_bin /home/circleci/project/build/bin + docker cp /home/circleci/project/. "$id:/var/lib/jenkins/workspace" + echo "mkdir -p /opt/conda/lib/python${PYTHON_VERSION}/site-packages" | docker exec -u jenkins -i "$id" bash + docker cp "/opt/workspace/torch" "$id:/opt/conda/lib/python${PYTHON_VERSION}/site-packages/torch" + if [ -d "/opt/workspace/cpp-build" ]; then + docker cp "/opt/workspace/cpp-build" "$id:/var/lib/jenkins/cpp-build" + fi + if [ -n "${MULTI_GPU}" ]; then + (echo "source ./workspace/env" && echo 'sudo chown -R jenkins workspace /opt/conda/lib/python${PYTHON_VERSION}/site-packages/torch && cd workspace && (git submodule update --init || git submodule update --init || git submodule update --init) && .jenkins/pytorch/multigpu-test.sh') | docker exec -u jenkins -i "$id" bash + else + (echo "source ./workspace/env" && echo 'sudo chown -R jenkins workspace /opt/conda/lib/python${PYTHON_VERSION}/site-packages/torch && cd workspace && (git submodule update --init || git submodule update --init || git submodule update --init) && .jenkins/pytorch/test.sh') | docker exec -u jenkins -i "$id" bash + fi + +caffe2_linux_build_defaults: &caffe2_linux_build_defaults + resource_class: large + working_directory: /var/lib/jenkins/workspace + steps: + - checkout + - run: + name: Build + no_output_timeout: "10h" + command: | + export IN_CIRCLECI=1 + export SCCACHE_BUCKET=ossci-compiler-cache-circleci-v2 + # This IAM user allows write access to S3 bucket for sccache + export AWS_ACCESS_KEY_ID=AKIAJJZUW4G2ASX5W7KA + export AWS_SECRET_ACCESS_KEY=${CIRCLECI_AWS_SECRET_KEY_FOR_SCCACHE_S3_BUCKET} + export SCCACHE_MAX_JOBS=`expr $(nproc) - 1` + export MEMORY_LIMIT_MAX_JOBS=8 # the "large" resource class on CircleCI has 32 CPU cores, if we use all of them we'll OOM + export MAX_JOBS=$(( ${SCCACHE_MAX_JOBS} > ${MEMORY_LIMIT_MAX_JOBS} ? ${MEMORY_LIMIT_MAX_JOBS} : ${SCCACHE_MAX_JOBS} )) + + set -ex + + # Need to checkout fetch PRs for onnxbot tracking PRs + git submodule update --init third_party/onnx || true + cd third_party/onnx && git fetch --tags --progress origin +refs/pull/*:refs/remotes/origin/pr/* && cd - + + # Reinitialize submodules + git submodule update --init --recursive + + # Ensure jenkins can write to the ccache root dir. + sudo chown jenkins:jenkins "${HOME}/.ccache" + + # Make ccache log to the workspace, so we can archive it after the build + mkdir -p build + ccache -o log_file=$PWD/build/ccache.log + + # Configure additional cmake arguments + cmake_args=() + cmake_args+=("$CMAKE_ARGS") + + if [[ $BUILD_ENVIRONMENT == *aten* ]]; then + cmake_args+=("-DBUILD_ATEN=ON") + fi + + # conda must be added to the path for Anaconda builds (this location must be + # the same as that in install_anaconda.sh used to build the docker image) + if [[ "${BUILD_ENVIRONMENT}" == conda* ]]; then + export PATH=/opt/conda/bin:$PATH + sudo chown -R jenkins:jenkins '/opt/conda' + fi + + # Build + if test -x ".jenkins/caffe2/build.sh"; then + ./.jenkins/caffe2/build.sh ${cmake_args[@]} + else + ./.jenkins/build.sh ${cmake_args[@]} + fi + + # Show sccache stats if it is running + if pgrep sccache > /dev/null; then + sccache --show-stats + fi + + # Copy all necessary binaries to shared workspace + mkdir -p caffe2-ci-env + cp -r third_party/onnx caffe2-ci-env/onnx + if [ -d "/usr/local/caffe2" ]; then + cp -r /usr/local/caffe2 caffe2-ci-env/caffe2 + fi + if [ -d "/opt/conda" ]; then + cp -r /opt/conda caffe2-ci-env/conda_env + fi + - persist_to_workspace: + root: /var/lib/jenkins/workspace/caffe2-ci-env + paths: + - "*" + +caffe2_linux_test_defaults: &caffe2_linux_test_defaults + machine: + image: default + steps: + - checkout + - run: + name: Prepare workspace + command: | + sudo mkdir -p /opt/workspace + sudo chmod -R 777 /opt/workspace + - attach_workspace: + at: /opt/workspace + - run: + name: Build + no_output_timeout: "10h" + command: | + set -x + sudo pip install awscli + if [ -n "${CUDA_VERSION}" ]; then + curl -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - + echo "deb https://nvidia.github.io/libnvidia-container/ubuntu14.04/amd64 /" | sudo tee -a /etc/apt/sources.list.d/nvidia-docker.list + echo "deb https://nvidia.github.io/nvidia-container-runtime/ubuntu14.04/amd64 /" | sudo tee -a /etc/apt/sources.list.d/nvidia-docker.list + echo "deb https://nvidia.github.io/nvidia-docker/ubuntu14.04/amd64 /" | sudo tee -a /etc/apt/sources.list.d/nvidia-docker.list + fi + sudo apt-get update + sudo apt-get remove linux-image-generic linux-headers-generic linux-generic + sudo apt-get install linux-headers-$(uname -r) + sudo apt-get install linux-image-generic + if [ -n "${CUDA_VERSION}" ]; then + wget 'https://s3.amazonaws.com/ossci-linux/nvidia_driver/NVIDIA-Linux-x86_64-396.26.run' + sudo /bin/bash ./NVIDIA-Linux-x86_64-396.26.run -s --no-drm + sudo apt-get install -y nvidia-docker2 + fi + sudo pkill -SIGHUP dockerd + if [ -n "${CUDA_VERSION}" ]; then + nvidia-smi + fi + # This IAM user only allows read-only access to ECR + export AWS_ACCESS_KEY_ID=AKIAJ2J6FIG5OSZTQ3IA + export AWS_SECRET_ACCESS_KEY=${CIRCLECI_AWS_SECRET_KEY_FOR_ECR_READ_ONLY} + eval $(aws ecr get-login --region us-east-1 --no-include-email) + docker pull ${DOCKER_IMAGE} + if [ -n "${CUDA_VERSION}" ]; then + id=$(docker run --runtime=nvidia -t -d -w /var/lib/jenkins ${DOCKER_IMAGE}) + else + id=$(docker run -t -d -w /var/lib/jenkins ${DOCKER_IMAGE}) + fi + pwd + echo "declare -x IN_CIRCLECI=1" > /home/circleci/project/env + echo "declare -x SCCACHE_BUCKET=ossci-compiler-cache-circleci-v2" >> /home/circleci/project/env + # This IAM user allows write access to S3 bucket for sccache + echo "declare -x AWS_ACCESS_KEY_ID=AKIAJJZUW4G2ASX5W7KA" >> /home/circleci/project/env + echo "declare -x AWS_SECRET_ACCESS_KEY=${CIRCLECI_AWS_SECRET_KEY_FOR_SCCACHE_S3_BUCKET}" >> /home/circleci/project/env + echo "declare -x BUILD_ENVIRONMENT=${BUILD_ENVIRONMENT}" >> /home/circleci/project/env + + # TODO: merge this into Caffe2 build.sh + cat >/home/circleci/project/ci_build_script.sh < /dev/null; then + printf "#!/bin/sh\nexec sccache $(which clang++) \$*" > "${SCCACHE_BIN}/clang++" + chmod a+x "${SCCACHE_BIN}/clang++" + + printf "#!/bin/sh\nexec sccache $(which clang) \$*" > "${SCCACHE_BIN}/clang" + chmod a+x "${SCCACHE_BIN}/clang" + + export PATH="${SCCACHE_BIN}:$PATH" + fi + + # Build + if [ "${BUILD_IOS:-0}" -eq 1 ]; then + scripts/build_ios.sh + elif [ -n "${CAFFE2_USE_ANACONDA}" ]; then + # All conda build logic should be in scripts/build_anaconda.sh + scripts/build_anaconda.sh + else + scripts/build_local.sh + fi + + # Show sccache stats if it is running + if which sccache > /dev/null; then + sccache --show-stats + fi + version: 2 jobs: - build: + pytorch_linux_trusty_py2_7_9_build_test: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-trusty-py2.7.9:238 + <<: *docker_config_defaults + <<: *pytorch_linux_cpu_build_test_defaults + + pytorch_linux_trusty_py2_7_build_test: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-trusty-py2.7:238 + <<: *docker_config_defaults + <<: *pytorch_linux_cpu_build_test_defaults + + pytorch_linux_trusty_py3_5_build_test: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-trusty-py3.5:238 + <<: *docker_config_defaults + <<: *pytorch_linux_cpu_build_test_defaults + + pytorch_linux_trusty_py3_6_gcc4_8_build_test: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-trusty-py3.6-gcc4.8:238 + <<: *docker_config_defaults + <<: *pytorch_linux_cpu_build_test_defaults + + pytorch_linux_trusty_py3_6_gcc5_4_build_test: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-trusty-py3.6-gcc5.4:238 + <<: *docker_config_defaults + <<: *pytorch_linux_cpu_build_test_defaults + + pytorch_linux_trusty_py3_6_gcc7_build_test: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-trusty-py3.6-gcc7:238 + <<: *docker_config_defaults + <<: *pytorch_linux_cpu_build_test_defaults + + pytorch_linux_trusty_pynightly_build_test: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-trusty-pynightly:238 + <<: *docker_config_defaults + <<: *pytorch_linux_cpu_build_test_defaults + + pytorch_linux_xenial_py3_clang5_asan_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3-clang5-asan:238 + <<: *docker_config_defaults + environment: + PYTHON_VERSION: "3.6" + <<: *pytorch_linux_build_defaults + + pytorch_linux_xenial_py3_clang5_asan_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3-clang5-asan:238" + PYTHON_VERSION: "3.6" + resource_class: large + <<: *pytorch_linux_test_defaults + + pytorch_linux_xenial_cuda8_cudnn6_py3_build: docker: - - image: circleci/python:3.7-node-browsers + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda8-cudnn6-py3:238 + <<: *docker_config_defaults + environment: + PYTHON_VERSION: "3.6" + CUDA_VERSION: "8" + <<: *pytorch_linux_build_defaults + + pytorch_linux_xenial_cuda8_cudnn6_py3_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda8-cudnn6-py3:238" + PYTHON_VERSION: "3.6" + CUDA_VERSION: "8" + resource_class: gpu.medium + <<: *pytorch_linux_test_defaults + + pytorch_linux_xenial_cuda8_cudnn6_py3_multigpu_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda8-cudnn6-py3:238" + PYTHON_VERSION: "3.6" + CUDA_VERSION: "8" + MULTI_GPU: "1" + resource_class: gpu.large + <<: *pytorch_linux_test_defaults + + pytorch_linux_xenial_cuda9_cudnn7_py2_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda9-cudnn7-py2:238 + <<: *docker_config_defaults + environment: + PYTHON_VERSION: "2.7" + CUDA_VERSION: "9" + <<: *pytorch_linux_build_defaults + + pytorch_linux_xenial_cuda9_cudnn7_py2_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda9-cudnn7-py2:238" + PYTHON_VERSION: "2.7" + CUDA_VERSION: "9" + resource_class: gpu.medium + <<: *pytorch_linux_test_defaults + + pytorch_linux_xenial_cuda9_cudnn7_py3_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda9-cudnn7-py3:238 + <<: *docker_config_defaults + environment: + PYTHON_VERSION: "3.6" + CUDA_VERSION: "9" + <<: *pytorch_linux_build_defaults + + pytorch_linux_xenial_cuda9_cudnn7_py3_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda9-cudnn7-py3:238" + PYTHON_VERSION: "3.6" + CUDA_VERSION: "9" + resource_class: gpu.medium + <<: *pytorch_linux_test_defaults + + pytorch_linux_xenial_cuda9_2_cudnn7_py3_gcc7_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda9.2-cudnn7-py3-gcc7:238 + <<: *docker_config_defaults + environment: + PYTHON_VERSION: "3.6" + CUDA_VERSION: "9.2" + <<: *pytorch_linux_build_defaults + + pytorch_linux_xenial_cuda9_2_cudnn7_py3_gcc7_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda9.2-cudnn7-py3-gcc7:238" + PYTHON_VERSION: "3.6" + CUDA_VERSION: "9.2" + resource_class: gpu.medium + <<: *pytorch_linux_test_defaults + + pytorch_macos_10_13_py3_build: + macos: + xcode: "9.0" + steps: + - checkout + - run: + name: Build + environment: + BUILD_ENVIRONMENT: pytorch-macos-10.13-py3 + no_output_timeout: "10h" + command: | + set -ex + + export IN_CIRCLECI=1 + + # Install sccache + sudo curl https://s3.amazonaws.com/ossci-macos/sccache --output /usr/local/bin/sccache + sudo chmod +x /usr/local/bin/sccache + + export SCCACHE_BUCKET=ossci-compiler-cache-circleci-v2 + # This IAM user allows write access to S3 bucket for sccache + export AWS_ACCESS_KEY_ID=AKIAJJZUW4G2ASX5W7KA + export AWS_SECRET_ACCESS_KEY=${CIRCLECI_AWS_SECRET_KEY_FOR_SCCACHE_S3_BUCKET} + + git submodule update --init + chmod a+x .jenkins/pytorch/macos-build.sh + .jenkins/pytorch/macos-build.sh + - persist_to_workspace: + root: /Users/distiller/pytorch-ci-env + paths: + - "*" + + pytorch_macos_10_13_py3_test: + macos: + xcode: "9.0" steps: - - run: echo "hello world" + - checkout + - run: + name: Prepare workspace + command: | + sudo mkdir -p /Users/distiller/pytorch-ci-env + sudo chmod -R 777 /Users/distiller/pytorch-ci-env + - attach_workspace: + at: /Users/distiller/pytorch-ci-env + - run: + name: Build + environment: + BUILD_ENVIRONMENT: pytorch-macos-10.13-py3 + no_output_timeout: "10h" + command: | + set -ex + export IN_CIRCLECI=1 + git submodule update --init + chmod a+x .jenkins/pytorch/macos-test.sh + .jenkins/pytorch/macos-test.sh + + pytorch_macos_10_13_cuda9_2_cudnn7_py3_build: + macos: + xcode: "9.0" + steps: + - checkout + - run: + name: Build + environment: + JOB_BASE_NAME: pytorch-macos-10.13-cuda9.2-cudnn7-py3-build + BUILD_ENVIRONMENT: pytorch-macos-10.13-cuda9.2-cudnn7-py3 + no_output_timeout: "10h" + command: | + set -ex + + export IN_CIRCLECI=1 + + # Install CUDA 9.2 + sudo rm -rf ~/cuda_9.2.64_mac_installer.app || true + curl https://s3.amazonaws.com/ossci-macos/cuda_9.2.64_mac_installer.zip -o ~/cuda_9.2.64_mac_installer.zip + unzip ~/cuda_9.2.64_mac_installer.zip -d ~/ + sudo ~/cuda_9.2.64_mac_installer.app/Contents/MacOS/CUDAMacOSXInstaller --accept-eula --no-window + sudo cp /usr/local/cuda/lib/libcuda.dylib /Developer/NVIDIA/CUDA-9.2/lib/libcuda.dylib + sudo rm -rf /usr/local/cuda || true + + # Install cuDNN 7.1 for CUDA 9.2 + curl https://s3.amazonaws.com/ossci-macos/cudnn-9.2-osx-x64-v7.1.tgz -o ~/cudnn-9.2-osx-x64-v7.1.tgz + rm -rf ~/cudnn-9.2-osx-x64-v7.1 && mkdir ~/cudnn-9.2-osx-x64-v7.1 + tar -xzvf ~/cudnn-9.2-osx-x64-v7.1.tgz -C ~/cudnn-9.2-osx-x64-v7.1 + sudo cp ~/cudnn-9.2-osx-x64-v7.1/cuda/include/cudnn.h /Developer/NVIDIA/CUDA-9.2/include/ + sudo cp ~/cudnn-9.2-osx-x64-v7.1/cuda/lib/libcudnn* /Developer/NVIDIA/CUDA-9.2/lib/ + sudo chmod a+r /Developer/NVIDIA/CUDA-9.2/include/cudnn.h /Developer/NVIDIA/CUDA-9.2/lib/libcudnn* + + # Install sccache + sudo curl https://s3.amazonaws.com/ossci-macos/sccache --output /usr/local/bin/sccache + sudo chmod +x /usr/local/bin/sccache + export SCCACHE_BUCKET=ossci-compiler-cache-circleci-v2 + # This IAM user allows write access to S3 bucket for sccache + export AWS_ACCESS_KEY_ID=AKIAJJZUW4G2ASX5W7KA + export AWS_SECRET_ACCESS_KEY=${CIRCLECI_AWS_SECRET_KEY_FOR_SCCACHE_S3_BUCKET} + + git submodule update --init + chmod a+x .jenkins/pytorch/macos-build.sh + .jenkins/pytorch/macos-build.sh + + caffe2_py2_cuda8_0_cudnn6_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda8.0-cudnn6-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + CUDA_VERSION: "8" + BUILD_ENVIRONMENT: "py2-cuda8.0-cudnn6-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_cuda8_0_cudnn6_ubuntu16_04_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda8.0-cudnn6-ubuntu16.04:190" + CUDA_VERSION: "8" + BUILD_ENVIRONMENT: "py2-cuda8.0-cudnn6-ubuntu16.04" + resource_class: gpu.medium + <<: *caffe2_linux_test_defaults + + caffe2_py2_cuda9_0_cudnn7_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda9.0-cudnn7-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + CUDA_VERSION: "9" + BUILD_ENVIRONMENT: "py2-cuda9.0-cudnn7-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_cuda9_0_cudnn7_ubuntu16_04_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda9.0-cudnn7-ubuntu16.04:190" + CUDA_VERSION: "9" + BUILD_ENVIRONMENT: "py2-cuda9.0-cudnn7-ubuntu16.04" + resource_class: gpu.medium + <<: *caffe2_linux_test_defaults + + caffe2_py2_cuda9_0_cudnn7_aten_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda9.0-cudnn7-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + CUDA_VERSION: "9" + BUILD_ENVIRONMENT: "py2-cuda9.0-cudnn7-aten-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_cuda9_0_cudnn7_aten_ubuntu16_04_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda9.0-cudnn7-ubuntu16.04:190" + CUDA_VERSION: "9" + BUILD_ENVIRONMENT: "py2-cuda9.0-cudnn7-aten-ubuntu16.04" + resource_class: gpu.medium + <<: *caffe2_linux_test_defaults + + caffe2_py2_cuda9_1_cudnn7_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda9.1-cudnn7-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + CUDA_VERSION: "9.1" + BUILD_ENVIRONMENT: "py2-cuda9.1-cudnn7-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_cuda9_1_cudnn7_ubuntu16_04_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda9.1-cudnn7-ubuntu16.04:190" + CUDA_VERSION: "9.1" + BUILD_ENVIRONMENT: "py2-cuda9.1-cudnn7-ubuntu16.04" + resource_class: gpu.medium + <<: *caffe2_linux_test_defaults + + caffe2_py2_mkl_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-mkl-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-mkl-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_mkl_ubuntu16_04_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-mkl-ubuntu16.04:190" + BUILD_ENVIRONMENT: "py2-mkl-ubuntu16.04" + resource_class: large + <<: *caffe2_linux_test_defaults + + caffe2_py2_gcc4_8_ubuntu14_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-gcc4.8-ubuntu14.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-gcc4.8-ubuntu14.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_gcc4_8_ubuntu14_04_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-gcc4.8-ubuntu14.04:190" + BUILD_ENVIRONMENT: "py2-gcc4.8-ubuntu14.04" + resource_class: large + <<: *caffe2_linux_test_defaults + + caffe2_onnx_py2_gcc5_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-gcc5-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "onnx-py2-gcc5-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_onnx_py2_gcc5_ubuntu16_04_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-gcc5-ubuntu16.04:190" + BUILD_ENVIRONMENT: "onnx-py2-gcc5-ubuntu16.04" + resource_class: large + <<: *caffe2_linux_test_defaults + + caffe2_conda2_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/conda2-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "conda2-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_conda2_ubuntu16_04_test: + environment: + DOCKER_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/conda2-ubuntu16.04:190" + BUILD_ENVIRONMENT: "conda2-ubuntu16.04" + resource_class: large + <<: *caffe2_linux_test_defaults + + caffe2_py2_cuda8_0_cudnn7_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda8.0-cudnn7-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-cuda8.0-cudnn7-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_gcc4_9_ubuntu14_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-gcc4.9-ubuntu14.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-gcc4.9-ubuntu14.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_clang3_8_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-clang3.8-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-clang3.8-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_clang3_9_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-clang3.9-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-clang3.9-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_gcc6_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-gcc6-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-gcc6-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_gcc7_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-gcc7-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-gcc7-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_cuda8_0_cudnn7_aten_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda8.0-cudnn7-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-cuda8.0-cudnn7-aten-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_android_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-android-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-android-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_conda3_cuda9_0_cudnn7_ubuntu16_04_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/conda3-cuda9.0-cudnn7-ubuntu16.04:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "conda3-cuda9.0-cudnn7-ubuntu16.04" + <<: *caffe2_linux_build_defaults + + caffe2_py2_cuda9_0_cudnn7_centos7_build: + docker: + - image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/caffe2/py2-cuda9.0-cudnn7-centos7:190 + <<: *docker_config_defaults + environment: + BUILD_ENVIRONMENT: "py2-cuda9.0-cudnn7-centos7" + <<: *caffe2_linux_build_defaults + + caffe2_py2_ios_macos10_13_build: + environment: + BUILD_IOS: "1" + PYTHON_INSTALLATION: "system" + PYTHON_VERSION: "2" + <<: *caffe2_macos_build_defaults + + caffe2_py2_system_macos10_13_build: + environment: + PYTHON_INSTALLATION: "system" + PYTHON_VERSION: "2" + <<: *caffe2_macos_build_defaults + +workflows: + version: 2 + build: + jobs: + - pytorch_linux_trusty_py2_7_9_build_test + - pytorch_linux_trusty_py2_7_build_test + - pytorch_linux_trusty_py3_5_build_test + - pytorch_linux_trusty_py3_6_gcc4_8_build_test + - pytorch_linux_trusty_py3_6_gcc5_4_build_test + - pytorch_linux_trusty_py3_6_gcc7_build_test + - pytorch_linux_trusty_pynightly_build_test + - pytorch_linux_xenial_py3_clang5_asan_build + - pytorch_linux_xenial_py3_clang5_asan_test: + requires: + - pytorch_linux_xenial_py3_clang5_asan_build + - pytorch_linux_xenial_cuda8_cudnn6_py3_build + - pytorch_linux_xenial_cuda8_cudnn6_py3_test: + requires: + - pytorch_linux_xenial_cuda8_cudnn6_py3_build + - pytorch_linux_xenial_cuda8_cudnn6_py3_multigpu_test: + requires: + - pytorch_linux_xenial_cuda8_cudnn6_py3_build + - pytorch_linux_xenial_cuda9_cudnn7_py2_build + - pytorch_linux_xenial_cuda9_cudnn7_py2_test: + requires: + - pytorch_linux_xenial_cuda9_cudnn7_py2_build + - pytorch_linux_xenial_cuda9_cudnn7_py3_build + - pytorch_linux_xenial_cuda9_cudnn7_py3_test: + requires: + - pytorch_linux_xenial_cuda9_cudnn7_py3_build + - pytorch_linux_xenial_cuda9_2_cudnn7_py3_gcc7_build + - pytorch_linux_xenial_cuda9_2_cudnn7_py3_gcc7_test: + requires: + - pytorch_linux_xenial_cuda9_2_cudnn7_py3_gcc7_build + + # - pytorch_macos_10_13_py3_build + # - pytorch_macos_10_13_py3_test: + # requires: + # - pytorch_macos_10_13_py3_build + # - pytorch_macos_10_13_cuda9_2_cudnn7_py3_build + + - caffe2_py2_cuda8_0_cudnn6_ubuntu16_04_build + - caffe2_py2_cuda8_0_cudnn6_ubuntu16_04_test: + requires: + - caffe2_py2_cuda8_0_cudnn6_ubuntu16_04_build + - caffe2_py2_cuda9_0_cudnn7_ubuntu16_04_build + - caffe2_py2_cuda9_0_cudnn7_ubuntu16_04_test: + requires: + - caffe2_py2_cuda9_0_cudnn7_ubuntu16_04_build + - caffe2_py2_cuda9_0_cudnn7_aten_ubuntu16_04_build + - caffe2_py2_cuda9_0_cudnn7_aten_ubuntu16_04_test: + requires: + - caffe2_py2_cuda9_0_cudnn7_aten_ubuntu16_04_build + - caffe2_py2_mkl_ubuntu16_04_build + - caffe2_py2_mkl_ubuntu16_04_test: + requires: + - caffe2_py2_mkl_ubuntu16_04_build + - caffe2_py2_cuda9_1_cudnn7_ubuntu16_04_build + - caffe2_py2_cuda9_1_cudnn7_ubuntu16_04_test: + requires: + - caffe2_py2_cuda9_1_cudnn7_ubuntu16_04_build + - caffe2_py2_gcc4_8_ubuntu14_04_build + - caffe2_py2_gcc4_8_ubuntu14_04_test: + requires: + - caffe2_py2_gcc4_8_ubuntu14_04_build + - caffe2_onnx_py2_gcc5_ubuntu16_04_build + - caffe2_onnx_py2_gcc5_ubuntu16_04_test: + requires: + - caffe2_onnx_py2_gcc5_ubuntu16_04_build + - caffe2_conda2_ubuntu16_04_build + - caffe2_conda2_ubuntu16_04_test: + requires: + - caffe2_conda2_ubuntu16_04_build + - caffe2_py2_cuda8_0_cudnn7_ubuntu16_04_build + - caffe2_py2_gcc4_9_ubuntu14_04_build + - caffe2_py2_clang3_8_ubuntu16_04_build + - caffe2_py2_clang3_9_ubuntu16_04_build + - caffe2_py2_gcc6_ubuntu16_04_build + - caffe2_py2_gcc7_ubuntu16_04_build + - caffe2_py2_cuda8_0_cudnn7_aten_ubuntu16_04_build + - caffe2_py2_android_ubuntu16_04_build + - caffe2_conda3_cuda9_0_cudnn7_ubuntu16_04_build + - caffe2_py2_cuda9_0_cudnn7_centos7_build + + # - caffe2_py2_ios_macos10_13_build + # - caffe2_py2_system_macos10_13_build diff --git a/.clang-tidy b/.clang-tidy index 3cb409f928607..aa2bf56185dc1 100644 --- a/.clang-tidy +++ b/.clang-tidy @@ -4,6 +4,7 @@ Checks: ' * ,clang-analyzer-* ,modernize-* + ,-cert-dcl21-cpp ,-cert-err58-cpp ,-cert-err60-cpp ,-clang-diagnostic-* @@ -12,10 +13,12 @@ Checks: ' ,-cppcoreguidelines-pro-bounds-constant-array-index ,-cppcoreguidelines-pro-type-member-init ,-cppcoreguidelines-pro-type-static-cast-downcast + ,-cppcoreguidelines-pro-type-union-access ,-cppcoreguidelines-pro-type-vararg ,-cppcoreguidelines-special-member-functions ,-fuchsia-* ,-google-build-using-namespace + ,-google-default-arguments ,-google-explicit-constructor ,-google-readability-braces-around-statements ,-google-readability-namespace-comments @@ -24,6 +27,7 @@ Checks: ' ,-google-runtime-references ,-hicpp-braces-around-statements ,-hicpp-explicit-conversions + ,-hicpp-member-init ,-hicpp-no-array-decay ,-hicpp-signed-bitwise ,-hicpp-special-member-functions diff --git a/.jenkins/caffe2/build.sh b/.jenkins/caffe2/build.sh index ac9235c8ed24b..4b4014ba42d37 100755 --- a/.jenkins/caffe2/build.sh +++ b/.jenkins/caffe2/build.sh @@ -64,6 +64,15 @@ if [ -z "${SCCACHE}" ] && which ccache > /dev/null; then export PATH="$CACHE_WRAPPER_DIR:$PATH" fi +# sccache will fail for CUDA builds if all cores are used for compiling +if [ -z "$MAX_JOBS" ]; then + if [[ "${BUILD_ENVIRONMENT}" == *-cuda* ]] && [ -n "${SCCACHE}" ]; then + MAX_JOBS=`expr $(nproc) - 1` + else + MAX_JOBS=$(nproc) + fi +fi + report_compile_cache_stats() { if [[ -n "${SCCACHE}" ]]; then "$SCCACHE" --show-stats @@ -184,13 +193,6 @@ if [[ -x "$(command -v cmake3)" ]]; then else CMAKE_BINARY=cmake fi -# sccache will fail for CUDA builds if all cores are used for compiling -if [[ "${BUILD_ENVIRONMENT}" == *-cuda* ]] && [ -n "${SCCACHE}" ]; then - MAX_JOBS=`expr $(nproc) - 1` -else - MAX_JOBS=$(nproc) -fi - ############################################################################### # Configure and make diff --git a/.jenkins/pytorch/build.sh b/.jenkins/pytorch/build.sh index 0f26005f74cb2..3ffed384b081b 100755 --- a/.jenkins/pytorch/build.sh +++ b/.jenkins/pytorch/build.sh @@ -8,13 +8,13 @@ if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-* ]]; then # TODO: move this to Docker sudo apt-get update - sudo apt-get install libnccl-dev=2.2.13-1+cuda9.0 libnccl2=2.2.13-1+cuda9.0 + sudo apt-get install -y --allow-downgrades --allow-change-held-packages libnccl-dev=2.2.13-1+cuda9.0 libnccl2=2.2.13-1+cuda9.0 fi if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda8-* ]] || [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-cudnn7-py2* ]]; then # TODO: move this to Docker sudo apt-get update - sudo apt-get install openmpi-bin libopenmpi-dev + sudo apt-get install -y --allow-downgrades --allow-change-held-packages openmpi-bin libopenmpi-dev sudo apt-get install -y --no-install-recommends openssh-client openssh-server sudo mkdir -p /var/run/sshd fi @@ -72,8 +72,10 @@ fi # sccache will fail for CUDA builds if all cores are used for compiling # gcc 7 with sccache seems to have intermittent OOM issue if all cores are used -if ([[ "$BUILD_ENVIRONMENT" == *cuda* ]] || [[ "$BUILD_ENVIRONMENT" == *gcc7* ]]) && which sccache > /dev/null; then - export MAX_JOBS=`expr $(nproc) - 1` +if [ -z "$MAX_JOBS" ]; then + if ([[ "$BUILD_ENVIRONMENT" == *cuda* ]] || [[ "$BUILD_ENVIRONMENT" == *gcc7* ]]) && which sccache > /dev/null; then + export MAX_JOBS=`expr $(nproc) - 1` + fi fi # Target only our CI GPU machine's CUDA arch to speed up the build diff --git a/.jenkins/pytorch/macos-build.sh b/.jenkins/pytorch/macos-build.sh index 41b272eae63a8..8439797024ad4 100755 --- a/.jenkins/pytorch/macos-build.sh +++ b/.jenkins/pytorch/macos-build.sh @@ -29,11 +29,15 @@ if [[ "${JOB_BASE_NAME}" == *cuda9.2* ]]; then export CUDA_HOME=/Developer/NVIDIA/CUDA-${CUDA_VERSION} export NO_CUDA=0 - # Eigen gives "explicit specialization of class must precede its first use" error - # when compiling with Xcode 9.1 toolchain, so we have to use Xcode 8.2 toolchain instead. - export DEVELOPER_DIR=/Library/Developer/CommandLineTools + if [ -z "${IN_CIRCLECI}" ]; then + # Eigen gives "explicit specialization of class must precede its first use" error + # when compiling with Xcode 9.1 toolchain, so we have to use Xcode 8.2 toolchain instead. + export DEVELOPER_DIR=/Library/Developer/CommandLineTools + fi else - export DEVELOPER_DIR=/Applications/Xcode9.app/Contents/Developer + if [ -z "${IN_CIRCLECI}" ]; then + export DEVELOPER_DIR=/Applications/Xcode9.app/Contents/Developer + fi fi export MACOSX_DEPLOYMENT_TARGET=10.9 @@ -62,5 +66,7 @@ export IMAGE_COMMIT_TAG=${BUILD_ENVIRONMENT}-${IMAGE_COMMIT_ID} python setup.py install # Upload torch binaries when the build job is finished -7z a ${IMAGE_COMMIT_TAG}.7z ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch* -aws s3 cp ${IMAGE_COMMIT_TAG}.7z s3://ossci-macos-build/pytorch/${IMAGE_COMMIT_TAG}.7z --acl public-read +if [ -z "${IN_CIRCLECI}" ]; then + 7z a ${IMAGE_COMMIT_TAG}.7z ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch* + aws s3 cp ${IMAGE_COMMIT_TAG}.7z s3://ossci-macos-build/pytorch/${IMAGE_COMMIT_TAG}.7z --acl public-read +fi diff --git a/.jenkins/pytorch/macos-test.sh b/.jenkins/pytorch/macos-test.sh index c6aa9429ed70d..87e0476e418ba 100755 --- a/.jenkins/pytorch/macos-test.sh +++ b/.jenkins/pytorch/macos-test.sh @@ -16,18 +16,22 @@ fi export PATH="${PYTORCH_ENV_DIR}/miniconda3/bin:$PATH" source ${PYTORCH_ENV_DIR}/miniconda3/bin/activate conda install -y mkl mkl-include numpy pyyaml setuptools cmake cffi ninja -rm -rf ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch* +if [ -z "${IN_CIRCLECI}" ]; then + rm -rf ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch* +fi git submodule update --init --recursive export CMAKE_PREFIX_PATH=${PYTORCH_ENV_DIR}/miniconda3/ # Test PyTorch -if [[ "${JOB_BASE_NAME}" == *cuda9.2* ]]; then - # Eigen gives "explicit specialization of class must precede its first use" error - # when compiling with Xcode 9.1 toolchain, so we have to use Xcode 8.2 toolchain instead. - export DEVELOPER_DIR=/Library/Developer/CommandLineTools -else - export DEVELOPER_DIR=/Applications/Xcode9.app/Contents/Developer +if [ -z "${IN_CIRCLECI}" ]; then + if [[ "${JOB_BASE_NAME}" == *cuda9.2* ]]; then + # Eigen gives "explicit specialization of class must precede its first use" error + # when compiling with Xcode 9.1 toolchain, so we have to use Xcode 8.2 toolchain instead. + export DEVELOPER_DIR=/Library/Developer/CommandLineTools + else + export DEVELOPER_DIR=/Applications/Xcode9.app/Contents/Developer + fi fi export MACOSX_DEPLOYMENT_TARGET=10.9 export CXX=clang++ @@ -38,9 +42,11 @@ export MAX_JOBS=2 export IMAGE_COMMIT_TAG=${BUILD_ENVIRONMENT}-${IMAGE_COMMIT_ID} # Download torch binaries in the test jobs -rm -rf ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch* -aws s3 cp s3://ossci-macos-build/pytorch/${IMAGE_COMMIT_TAG}.7z ${IMAGE_COMMIT_TAG}.7z -7z x ${IMAGE_COMMIT_TAG}.7z -o"${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages" +if [ -z "${IN_CIRCLECI}" ]; then + rm -rf ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch* + aws s3 cp s3://ossci-macos-build/pytorch/${IMAGE_COMMIT_TAG}.7z ${IMAGE_COMMIT_TAG}.7z + 7z x ${IMAGE_COMMIT_TAG}.7z -o"${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages" +fi test_python_all() { echo "Ninja version: $(ninja --version)" diff --git a/.jenkins/pytorch/multigpu-test.sh b/.jenkins/pytorch/multigpu-test.sh index ceee027f1a7d5..a8fc1c8baae69 100755 --- a/.jenkins/pytorch/multigpu-test.sh +++ b/.jenkins/pytorch/multigpu-test.sh @@ -8,4 +8,21 @@ COMPACT_JOB_NAME="${BUILD_ENVIRONMENT}-multigpu-test" source "$(dirname "${BASH_SOURCE[0]}")/common.sh" echo "Testing pytorch (distributed only)" + +if [ -n "${IN_CIRCLECI}" ]; then + if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-* ]]; then + # TODO: move this to Docker + sudo apt-get update + sudo apt-get install -y --allow-downgrades --allow-change-held-packages libnccl-dev=2.2.13-1+cuda9.0 libnccl2=2.2.13-1+cuda9.0 + fi + + if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda8-* ]] || [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-cudnn7-py2* ]]; then + # TODO: move this to Docker + sudo apt-get update + sudo apt-get install -y --allow-downgrades --allow-change-held-packages openmpi-bin libopenmpi-dev + sudo apt-get install -y --no-install-recommends openssh-client openssh-server + sudo mkdir -p /var/run/sshd + fi +fi + time python test/run_test.py --verbose -i distributed diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index a646537c3a074..05bd71602b978 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -9,6 +9,22 @@ source "$(dirname "${BASH_SOURCE[0]}")/common.sh" echo "Testing pytorch" +if [ -n "${IN_CIRCLECI}" ]; then + if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-* ]]; then + # TODO: move this to Docker + sudo apt-get update + sudo apt-get install -y --allow-downgrades --allow-change-held-packages libnccl-dev=2.2.13-1+cuda9.0 libnccl2=2.2.13-1+cuda9.0 + fi + + if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda8-* ]] || [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-cudnn7-py2* ]]; then + # TODO: move this to Docker + sudo apt-get update + sudo apt-get install -y --allow-downgrades --allow-change-held-packages openmpi-bin libopenmpi-dev + sudo apt-get install -y --no-install-recommends openssh-client openssh-server + sudo mkdir -p /var/run/sshd + fi +fi + # JIT C++ extensions require ninja. git clone https://github.com/ninja-build/ninja --quiet pushd ninja diff --git a/.travis.yml b/.travis.yml index bd6cec56fb897..be45e69f67cb2 100644 --- a/.travis.yml +++ b/.travis.yml @@ -26,3 +26,6 @@ matrix: python: "3.6" install: pip install mypy mypy-extensions script: mypy @mypy-files.txt + - env: CPP_DOC_CHECK + install: sudo apt-get install -y doxygen + script: cd docs/cpp && ./check-doxygen.sh diff --git a/CMakeLists.txt b/CMakeLists.txt index 5d354943aa223..9fdfe4aef7f94 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -82,14 +82,11 @@ cmake_dependent_option( option(USE_FFMPEG "Use ffmpeg" OFF) option(USE_GFLAGS "Use GFLAGS" ON) option(USE_GLOG "Use GLOG" ON) -option(USE_GLOO "Use Gloo" ON) -option(USE_GLOO_IBVERBS "Use Gloo IB verbs for distributed support" OFF) option(USE_LEVELDB "Use LEVELDB" ON) option(USE_LITE_PROTO "Use lite protobuf instead of full." OFF) option(USE_LMDB "Use LMDB" ON) option(USE_METAL "Use Metal for iOS build" ON) option(USE_MOBILE_OPENGL "Use OpenGL for mobile code" ON) -option(USE_MPI "Use MPI" ON) option(USE_NATIVE_ARCH "Use -march=native" OFF) option(USE_NCCL "Use NCCL" ON) option(USE_SYSTEM_NCCL "Use system-wide NCCL" OFF) @@ -116,7 +113,16 @@ option(USE_ZSTD "Use ZSTD" OFF) option(USE_MKLDNN "Use MKLDNN" OFF) option(USE_IDEEP "Use IDEEP interface in MKL BLAS" ON) option(USE_MKLML "Use MKLML interface in MKL BLAS" ON) -option(USE_DISTRIBUTED "Use THD (distributed)" OFF) +option(USE_DISTRIBUTED "Use distributed" ON) +cmake_dependent_option( + USE_MPI "Use MPI. Only available if USE_DISTRIBUTED is on." ON + "USE_DISTRIBUTED" OFF) +cmake_dependent_option( + USE_GLOO "Use Gloo. Only available if USE_DISTRIBUTED is on." ON + "USE_DISTRIBUTED" OFF) +cmake_dependent_option( + USE_GLOO_IBVERBS "Use Gloo IB verbs for distributed. Only available if USE_GLOO is on." OFF + "USE_GLOO" OFF) # Used when building Caffe2 through setup.py option(BUILDING_WITH_TORCH_LIBS "Tell cmake if Caffe2 is being built alongside torch libs" OFF) @@ -378,6 +384,7 @@ if (BUILD_SHARED_LIBS) ${PROJECT_SOURCE_DIR}/cmake/public/cuda.cmake ${PROJECT_SOURCE_DIR}/cmake/public/glog.cmake ${PROJECT_SOURCE_DIR}/cmake/public/gflags.cmake + ${PROJECT_SOURCE_DIR}/cmake/public/mkl.cmake ${PROJECT_SOURCE_DIR}/cmake/public/protobuf.cmake ${PROJECT_SOURCE_DIR}/cmake/public/threads.cmake ${PROJECT_SOURCE_DIR}/cmake/public/utils.cmake @@ -397,24 +404,16 @@ else() endif() # ---[ Modules -# TODO(orionr): Enable all of this for Windows DLL when we -# can figure out how to get it to build -if (NOT (MSVC AND BUILD_SHARED_LIBS)) add_subdirectory(modules) -endif() # ---[ Binaries # Binaries will be built after the Caffe2 main libraries and the modules # are built. For the binaries, they will be linked to the Caffe2 main # libraries, as well as all the modules that are built with Caffe2 (the ones # built in the previous Modules section above). -# TODO(orionr): Enable all of this for Windows DLL when we -# can figure out how to get it to build -if (NOT (MSVC AND BUILD_SHARED_LIBS)) if (BUILD_BINARY) add_subdirectory(binaries) endif() -endif() include(cmake/Summary.cmake) caffe2_print_configuration_summary() diff --git a/CODEOWNERS b/CODEOWNERS index 9a51f5d2a822e..113be035c9b99 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -4,6 +4,7 @@ /aten/ @apaszke @soumith @colesbury @gchanan @zdevito @ezyang /torch/ @apaszke @soumith @colesbury @gchanan @zdevito @ezyang /docs/source @apaszke @soumith @colesbury @gchanan @zdevito @ezyang @ssnl @zou3519 +/docs/cpp @goldsborough @ebetica @apaszke @soumith @colesbury @gchanan @zdevito @ezyang /test @apaszke @soumith @colesbury @gchanan @zdevito @ezyang /tools @apaszke @soumith @colesbury @gchanan @zdevito @ezyang /README.md @apaszke @soumith @colesbury @gchanan @zdevito @ezyang diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 11a219d6b32dd..0f46fa1cf62a7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -104,6 +104,18 @@ PyTorch uses [Google style](http://sphinxcontrib-napoleon.readthedocs.io/en/late for formatting docstrings. Length of line inside docstrings block must be limited to 80 characters to fit into Jupyter documentation popups. +For C++ documentation (https://pytorch.org/cppdocs), we use +[Doxygen](http://www.doxygen.nl/) and then convert it to +[Sphinx](http://www.sphinx-doc.org/) via +[Breathe](https://github.com/michaeljones/breathe) and +[Exhale](https://github.com/svenevs/exhale). Check the [Doxygen +reference](http://www.stack.nl/~dimitri/doxygen/manual/index.html) for more +information on the documentation syntax. To build the documentation locally, +`cd` into `docs/cpp` and then `make html`. + +We run Doxygen in CI (Travis) to verify that you do not use invalid Doxygen +commands. To run this check locally, run `./check-doxygen.sh` from inside +`docs/cpp`. ## Managing multiple build trees diff --git a/aten/src/ATen/ATen.h b/aten/src/ATen/ATen.h index d586b0740f9fb..cf074730bf072 100644 --- a/aten/src/ATen/ATen.h +++ b/aten/src/ATen/ATen.h @@ -16,7 +16,7 @@ #include "ATen/OptionsGuard.h" #include "ATen/core/Scalar.h" #include "ATen/ScalarOps.h" -#include "ATen/Storage.h" +#include "ATen/core/Storage.h" #include "ATen/Tensor.h" #include "ATen/TensorGeometry.h" #include "ATen/TensorMethods.h" diff --git a/aten/src/ATen/CPUApplyUtils.h b/aten/src/ATen/CPUApplyUtils.h index bba97c5e28ee7..530b470f6c8be 100644 --- a/aten/src/ATen/CPUApplyUtils.h +++ b/aten/src/ATen/CPUApplyUtils.h @@ -3,9 +3,99 @@ #include "ATen/Parallel.h" #include "ATen/TensorUtils.h" #include +#include +#include namespace at { +/* +[collapse dims] Updates sizes, and strides to reflect a "collapse" of +the info, possibly excluding the optional excludeDim. A "collapsed" version +of the info is the fewest dims that order the tensor's elements in the same +way as the original info. If excludeDim is specified, the collapse is the +fewest dims that order the tensor's elements as the original and preserve the +excluded dimension, unless the tensor collapses to a point. + +This function returns a pair of values. + +1) The (new) index of the preserved dimension if excludeDim is +specified. 0 if the tensor is collapsed to a point. -1 +otherwise. + +2) The new number of dimensions. +*/ +template +inline std::pair collapse_dims( + T* sizes, + T* strides, + int64_t dims, + const int excludeDim = -1) { + AT_CHECK( + excludeDim >= -1 && excludeDim < dims, + "expected excluded dim between -1 and dims - 1"); + + int64_t stopDim = (excludeDim == -1) ? dims : excludeDim; + int64_t newIndex = -1; + int64_t oldIndex = 0; + int64_t remappedExcludedDim = -1; + + while (oldIndex < dims) { + // Finds a dimension to collapse into + for (; oldIndex < stopDim; ++oldIndex) { + if (sizes[oldIndex] == 1) { + continue; + } + + ++newIndex; + sizes[newIndex] = sizes[oldIndex]; + strides[newIndex] = strides[oldIndex]; + ++oldIndex; + break; + } + + // Collapses dims + for (; oldIndex < stopDim; ++oldIndex) { + if (sizes[oldIndex] == 1) { + continue; + } + + if (strides[newIndex] == sizes[oldIndex] * strides[oldIndex]) { + sizes[newIndex] *= sizes[oldIndex]; + strides[newIndex] = strides[oldIndex]; + } else { + ++newIndex; + sizes[newIndex] = sizes[oldIndex]; + strides[newIndex] = strides[oldIndex]; + } + } + + // Handles excludeDim being set (oldIndex == excludeDim) + if (oldIndex != dims) { + // Preserves excluded dimension + ++newIndex; + sizes[newIndex] = sizes[oldIndex]; + strides[newIndex] = strides[oldIndex]; + remappedExcludedDim = newIndex; + + // Restarts iteration after excludeDim + ++oldIndex; + stopDim = dims; + } + } + + // Handles special case of all dims size 1 + if (newIndex == -1 || (newIndex == 0 && sizes[0] == 1)) { + dims = 1; + sizes[0] = 1; + strides[0] = 1; + + return std::pair(0, 1); + } + + dims = newIndex + 1; + return std::pair(remappedExcludedDim, dims); +} + /* * The basic strategy for apply is as follows: * @@ -50,27 +140,6 @@ inline Tensor sort_strides(Tensor& tensor_) { return tensor; } -template -inline void _setup_arrays(Tensor& tensor, Arg* iter) { - int64_t max_dim = tensor.ndimension(); - iter->dim_ = 0; - for (int64_t i = 0; i < max_dim; i++) { - int64_t size = tensor.size(i); - int64_t stride = tensor.stride(i); - while (stride > 0 && i + 1 < max_dim && - (tensor.size(i + 1) == 1 || - stride == tensor.size(i + 1) * tensor.stride(i + 1))) { - size = size * tensor.size(i + 1); - if (tensor.size(i + 1) != 1) - stride = tensor.stride(i + 1); - i++; - } - iter->sizes_[iter->dim_] = size; - iter->strides_[iter->dim_] = stride; - iter->dim_++; - } -} - template struct strided_tensor_iter_fixed { public: @@ -86,8 +155,14 @@ struct strided_tensor_iter_fixed { strided_tensor_iter_fixed(strided_tensor_iter_fixed&&) = default; strided_tensor_iter_fixed(Tensor& tensor, bool sort_strides = false) : data_(tensor.data()) { - memset(counter_, 0, sizeof(int64_t) * N); - _setup_arrays(tensor, this); + std::memset(counter_, 0, sizeof(int64_t) * N); + std::memcpy( + sizes_, tensor.sizes().data(), tensor.ndimension() * sizeof(int64_t)); + std::memcpy( + strides_, + tensor.strides().data(), + tensor.ndimension() * sizeof(int64_t)); + dim_ = std::get<1>(collapse_dims(sizes_, strides_, tensor.ndimension())); } }; @@ -111,7 +186,7 @@ struct strided_tensor_iter { counter_(dim_, 0), sizes_(tensor.sizes().vec()), strides_(tensor.strides().vec()) { - _setup_arrays(tensor, this); + dim_ = std::get<1>(collapse_dims(sizes_.data(), strides_.data(), dim_)); } }; diff --git a/aten/src/ATen/CPUTypeDefault.cpp b/aten/src/ATen/CPUTypeDefault.cpp new file mode 100644 index 0000000000000..99661d61d140f --- /dev/null +++ b/aten/src/ATen/CPUTypeDefault.cpp @@ -0,0 +1,20 @@ +#include + +#include +#include + +namespace at { + +Allocator* CPUTypeDefault::allocator() const { + return getCPUAllocator(); +} + +Device CPUTypeDefault::getDeviceFromPtr(void * data) const { + return DeviceType::CPU; +} + +std::unique_ptr CPUTypeDefault::generator() const { + return std::unique_ptr(new CPUGenerator(&at::globalContext())); +} + +} // namespace at diff --git a/aten/src/ATen/CPUTypeDefault.h b/aten/src/ATen/CPUTypeDefault.h new file mode 100644 index 0000000000000..c9776b7b0a2cc --- /dev/null +++ b/aten/src/ATen/CPUTypeDefault.h @@ -0,0 +1,14 @@ +#pragma once +#include + +namespace at { + +struct AT_API CPUTypeDefault : public TypeDefault { + CPUTypeDefault(TensorTypeId type_id, bool is_variable, bool is_undefined) + : TypeDefault(type_id, is_variable, is_undefined) {} + Allocator* allocator() const override; + Device getDeviceFromPtr(void * data) const override; + std::unique_ptr generator() const override; +}; + +} // namespace at diff --git a/aten/src/ATen/Config.h.in b/aten/src/ATen/Config.h.in index 8373c92404c34..2692da71e418d 100644 --- a/aten/src/ATen/Config.h.in +++ b/aten/src/ATen/Config.h.in @@ -8,3 +8,4 @@ #define AT_MKLDNN_ENABLED() @AT_MKLDNN_ENABLED@ #define AT_MKL_ENABLED() @AT_MKL_ENABLED@ +#define CAFFE2_STATIC_LINK_CUDA() @CAFFE2_STATIC_LINK_CUDA@ diff --git a/aten/src/ATen/Context.cpp b/aten/src/ATen/Context.cpp index 80c8e92647851..77953c9e18e60 100644 --- a/aten/src/ATen/Context.cpp +++ b/aten/src/ATen/Context.cpp @@ -118,4 +118,8 @@ Type& getMaybeVariableType(const TensorImpl* impl) { backend, impl->scalar_type(), impl->is_variable()); } +Allocator* getCPUAllocator() { + return getTHDefaultAllocator(); +} + } diff --git a/aten/src/ATen/Context.h b/aten/src/ATen/Context.h index 28bcee552d1a4..faf451ef72d5d 100644 --- a/aten/src/ATen/Context.h +++ b/aten/src/ATen/Context.h @@ -9,6 +9,7 @@ #include "ATen/core/Error.h" #include "ATen/detail/CUDAHooksInterface.h" #include "ATen/detail/VariableHooksInterface.h" +#include "ATen/detail/ComplexHooksInterface.h" // This is temporary #include "ATen/core/ATenCoreTest.h" @@ -26,7 +27,10 @@ class AT_API Context { return type_registry[static_cast(p)][static_cast(s)].get(); } Type * getNonVariableTypeOpt(Backend p, ScalarType s) { - if (p != Backend::Undefined) initCUDAIfNeeded(backendToDeviceType(p)); + if (p != Backend::Undefined) { + initCUDAIfNeeded(backendToDeviceType(p)); + initComplexIfNeeded(s); + } auto type = getNonVariableTypeRaw(p, s); if(!type) { @@ -54,6 +58,10 @@ class AT_API Context { return getNonVariableType(p, s); } } + void registerType(Backend b, ScalarType s, Type* t) { + type_registry[static_cast(b)][static_cast(s)].reset(t); + detail::getVariableHooks().registerVariableTypeFor(this, b, s); + } Generator & defaultGenerator(DeviceType device_type) { initCUDAIfNeeded(device_type); @@ -87,6 +95,11 @@ class AT_API Context { }); return thc_state.get(); } + void lazyInitComplex() { + std::call_once(complex_init_, [&] { + detail::getComplexHooks().registerComplexTypes(this); + }); + } THCState* getTHCState() { // AT_ASSERT(thc_state); @@ -124,7 +137,13 @@ class AT_API Context { lazyInitCUDA(); } } + void initComplexIfNeeded(ScalarType s) { + if (isComplexType(s)) { + lazyInitComplex(); + } + } std::once_flag thc_init; + std::once_flag complex_init_; bool enabled_cudnn = true; bool deterministic_cudnn = false; bool benchmark_cudnn = false; @@ -158,6 +177,8 @@ static inline Type& getNonVariableType(DeviceType p, ScalarType s) { AT_API Type& getMaybeVariableType(TensorOptions options); AT_API Type& getMaybeVariableType(const TensorImpl*); +AT_API Allocator* getCPUAllocator(); + static inline Type& CPU(ScalarType s) { return getNonVariableType(Backend::CPU, s); } diff --git a/aten/src/ATen/TensorImpl.cpp b/aten/src/ATen/TensorImpl.cpp index 798e0d7787386..a520f4a5457d6 100644 --- a/aten/src/ATen/TensorImpl.cpp +++ b/aten/src/ATen/TensorImpl.cpp @@ -5,11 +5,10 @@ #include #include #include +#include #include -#include - namespace at { Type& TensorImpl::type() const { diff --git a/aten/src/ATen/TensorImpl.h b/aten/src/ATen/TensorImpl.h index f114a4048fd7f..f41eb165279c0 100644 --- a/aten/src/ATen/TensorImpl.h +++ b/aten/src/ATen/TensorImpl.h @@ -3,7 +3,7 @@ #include #include -#include "ATen/Storage.h" +#include "ATen/core/Storage.h" #include "ATen/core/optional.h" #include "ATen/core/TensorTypeId.h" #include "ATen/core/TensorTypeIdRegistration.h" diff --git a/aten/src/ATen/UndefinedType.cpp b/aten/src/ATen/UndefinedType.cpp index 367db6fbb2868..bea9baf61892f 100644 --- a/aten/src/ATen/UndefinedType.cpp +++ b/aten/src/ATen/UndefinedType.cpp @@ -11,9 +11,14 @@ ScalarType UndefinedType::scalarType() const { Backend UndefinedType::backend() const { return Backend::Undefined; } -bool UndefinedType::is_cuda() const { return false; } -bool UndefinedType::is_sparse() const { return false; } -bool UndefinedType::is_distributed() const { return false; } + +Allocator* UndefinedType::allocator() const { + AT_ERROR("allocator not defined for UndefinedType"); +} + +Device UndefinedType::getDeviceFromPtr(void*) const { + AT_ERROR("getDeviceFromPtr not defined for UndefinedType"); +} Storage UndefinedType::storage(bool resizable) const { AT_ERROR("storage not defined for UndefinedType"); @@ -38,8 +43,9 @@ std::unique_ptr UndefinedType::generator() const { } const char * UndefinedType::toString() const { - return UndefinedType::typeString(); + return "UndefinedType"; } + TypeID UndefinedType::ID() const { return TypeID::Undefined; } @@ -61,10 +67,6 @@ Type & UndefinedType::toScalarType(ScalarType s) const { AT_ERROR("toScalarType not implemented for UndefinedType to non-UndefinedType"); } -const char * UndefinedType::typeString() { - return "UndefinedType"; -} - Tensor & UndefinedType::s_copy_(Tensor & self, const Tensor & src, bool non_blocking) const { AT_ERROR("s_copy not defined for UndefinedType"); } diff --git a/aten/src/ATen/UndefinedType.h b/aten/src/ATen/UndefinedType.h index c08b1a47156f5..594fb99e61dc0 100644 --- a/aten/src/ATen/UndefinedType.h +++ b/aten/src/ATen/UndefinedType.h @@ -15,9 +15,8 @@ struct UndefinedType final : public TypeDefault { explicit UndefinedType(); virtual ScalarType scalarType() const override; virtual Backend backend() const override; - virtual bool is_cuda() const override; - virtual bool is_sparse() const override; - virtual bool is_distributed() const override; + virtual Allocator* allocator() const override; + virtual Device getDeviceFromPtr(void* data) const override; virtual Storage storage(bool resizable = false) const override; virtual Storage storage(size_t size, bool resizable = false) const override; virtual Storage storageFromBlob(void * data, int64_t size, const std::function & deleter) const override; @@ -28,7 +27,6 @@ struct UndefinedType final : public TypeDefault { virtual Type & toBackend(Backend b) const override; virtual Type & toScalarType(ScalarType s) const override; virtual TypeID ID() const override; - static const char * typeString(); virtual Storage unsafeStorageFromTH(void * th_pointer, bool retain) const override; virtual Tensor unsafeTensorFromTH(void * th_pointer, bool retain) const override; diff --git a/aten/src/ATen/WrapDimUtils.h b/aten/src/ATen/WrapDimUtils.h index 09d4e2a53fae1..8e9db589c5267 100644 --- a/aten/src/ATen/WrapDimUtils.h +++ b/aten/src/ATen/WrapDimUtils.h @@ -1,30 +1,10 @@ #pragma once +#include "ATen/core/WrapDimMinimal.h" #include "ATen/TensorImpl.h" -#include namespace at { -static inline int64_t maybe_wrap_dim(int64_t dim, int64_t dim_post_expr, bool wrap_scalar=true) { - if (dim_post_expr <= 0) { - if (!wrap_scalar) { - std::ostringstream oss; - oss << "dimension specified as " << dim << " but tensor has no dimensions"; - throw std::runtime_error(oss.str()); - } - dim_post_expr = 1; // this will make range [-1, 0] - } - - int64_t min = -dim_post_expr; - int64_t max = dim_post_expr - 1; - AT_CHECK( - dim >= min && dim <= max, - "Dimension out of range (expected to be in range of [", - min, ", ", max, "], but got ", dim, ")"); - if (dim < 0) dim += dim_post_expr; - return dim; -} - static inline int64_t maybe_wrap_dim(int64_t dim, TensorImpl *tensor) { return maybe_wrap_dim(dim, tensor->dim()); } diff --git a/aten/src/ATen/core/Storage.h b/aten/src/ATen/core/Storage.h index a602acc2cb116..6f2a8fd68ee71 100644 --- a/aten/src/ATen/core/Storage.h +++ b/aten/src/ATen/core/Storage.h @@ -8,6 +8,8 @@ struct AT_API Storage { public: Storage() {} Storage(StorageImpl* storage_impl) : storage_impl_(c10::intrusive_ptr::reclaim(storage_impl)) {} + Storage(const c10::intrusive_ptr& ptr) : storage_impl_(ptr) {} + Storage(c10::intrusive_ptr&& ptr) : storage_impl_(std::move(ptr)) {} Storage( at::ScalarType, size_t size, diff --git a/aten/src/ATen/core/StorageImpl.h b/aten/src/ATen/core/StorageImpl.h index 4f73b44455aac..e80c11c6b0e21 100644 --- a/aten/src/ATen/core/StorageImpl.h +++ b/aten/src/ATen/core/StorageImpl.h @@ -6,11 +6,6 @@ #include -struct THFinalizer { - virtual void operator()() = 0; - virtual ~THFinalizer() {}; -}; - namespace at { struct Type; diff --git a/aten/src/ATen/core/WrapDimMinimal.h b/aten/src/ATen/core/WrapDimMinimal.h new file mode 100644 index 0000000000000..6971bac0b3f67 --- /dev/null +++ b/aten/src/ATen/core/WrapDimMinimal.h @@ -0,0 +1,23 @@ +#pragma once + +#include "ATen/core/Error.h" + +namespace at { + +static inline int64_t maybe_wrap_dim(int64_t dim, int64_t dim_post_expr, bool wrap_scalar=true) { + if (dim_post_expr <= 0) { + AT_CHECK(wrap_scalar, "dimension specified as ", dim, " but tensor has no dimensions"); + dim_post_expr = 1; // this will make range [-1, 0] + } + + int64_t min = -dim_post_expr; + int64_t max = dim_post_expr - 1; + AT_CHECK( + dim >= min && dim <= max, + "Dimension out of range (expected to be in range of [", + min, ", ", max, "], but got ", dim, ")"); + if (dim < 0) dim += dim_post_expr; + return dim; +} + +} diff --git a/aten/src/ATen/cuda/ATenCUDAGeneral.h b/aten/src/ATen/cuda/ATenCUDAGeneral.h index 366adf0f2396f..7b41f1fe3f723 100644 --- a/aten/src/ATen/cuda/ATenCUDAGeneral.h +++ b/aten/src/ATen/cuda/ATenCUDAGeneral.h @@ -1,7 +1,7 @@ #pragma once #ifdef _WIN32 -# if defined(ATen_cuda_EXPORTS) || defined(caffe2_gpu_EXPORTS) || defined(CAFFE2_BUILD_MAIN_LIB) +# if defined(ATen_cuda_EXPORTS) || defined(caffe2_gpu_EXPORTS) || defined(CAFFE2_CUDA_BUILD_MAIN_LIB) # define AT_CUDA_API __declspec(dllexport) # else # define AT_CUDA_API __declspec(dllimport) diff --git a/aten/src/ATen/cuda/CUDAContext.cpp b/aten/src/ATen/cuda/CUDAContext.cpp index 9547815b260fe..e30162d81eb21 100644 --- a/aten/src/ATen/cuda/CUDAContext.cpp +++ b/aten/src/ATen/cuda/CUDAContext.cpp @@ -1,5 +1,5 @@ #include "ATen/cuda/CUDAContext.h" -#include "THC/THCGeneral.h" +#include "THC/THCGeneral.hpp" namespace at { namespace cuda { @@ -45,6 +45,10 @@ void uncheckedSetCurrentCUDAStream(CUDAStream stream) { detail::CUDAStream_uncheckedSetStream(stream.internals()); } +Allocator* getCUDADeviceAllocator() { + return at::globalContext().getTHCState()->cudaDeviceAllocator; +} + /* Handles */ #ifndef __HIP_PLATFORM_HCC__ cusparseHandle_t getCurrentCUDASparseHandle() { @@ -54,4 +58,4 @@ void uncheckedSetCurrentCUDAStream(CUDAStream stream) { } // namespace cuda -} // namespace at \ No newline at end of file +} // namespace at diff --git a/aten/src/ATen/cuda/CUDAContext.h b/aten/src/ATen/cuda/CUDAContext.h index 4a04479dcbecb..fee250e6d935d 100644 --- a/aten/src/ATen/cuda/CUDAContext.h +++ b/aten/src/ATen/cuda/CUDAContext.h @@ -54,6 +54,8 @@ AT_API CUDAStream getCurrentCUDAStream(int64_t device = -1); AT_API void setCurrentCUDAStream(CUDAStream stream); AT_API void uncheckedSetCurrentCUDAStream(CUDAStream stream); +AT_API Allocator* getCUDADeviceAllocator(); + /* Handles */ #ifndef __HIP_PLATFORM_HCC__ AT_API cusparseHandle_t getCurrentCUDASparseHandle(); diff --git a/aten/src/ATen/cuda/CUDADevice.h b/aten/src/ATen/cuda/CUDADevice.h new file mode 100644 index 0000000000000..77dc3581a2130 --- /dev/null +++ b/aten/src/ATen/cuda/CUDADevice.h @@ -0,0 +1,16 @@ +#pragma once + +#include "ATen/cuda/Exceptions.h" + +#include "cuda.h" + +namespace at { +namespace cuda { + +inline Device getDeviceFromPtr(void* ptr) { + struct cudaPointerAttributes attr; + AT_CUDA_CHECK(cudaPointerGetAttributes(&attr, ptr)); + return {DeviceType::CUDA, attr.device}; +} + +}} // namespace at::cuda diff --git a/aten/src/ATen/cuda/CUDATypeDefault.cpp b/aten/src/ATen/cuda/CUDATypeDefault.cpp new file mode 100644 index 0000000000000..ea9ef19d1f6ec --- /dev/null +++ b/aten/src/ATen/cuda/CUDATypeDefault.cpp @@ -0,0 +1,19 @@ +#include + +#include +#include +#include + +namespace at { + +Allocator* CUDATypeDefault::allocator() const { + return cuda::getCUDADeviceAllocator(); +} +Device CUDATypeDefault::getDeviceFromPtr(void * data) const { + return cuda::getDeviceFromPtr(data); +} +std::unique_ptr CUDATypeDefault::generator() const { + return std::unique_ptr(new CUDAGenerator(&at::globalContext())); +} + +} // namespace at diff --git a/aten/src/ATen/cuda/CUDATypeDefault.h b/aten/src/ATen/cuda/CUDATypeDefault.h new file mode 100644 index 0000000000000..d1b0eadeb0c05 --- /dev/null +++ b/aten/src/ATen/cuda/CUDATypeDefault.h @@ -0,0 +1,16 @@ +#pragma once +#include +#include + +namespace at { + +struct AT_CUDA_API CUDATypeDefault : public TypeDefault { + CUDATypeDefault(TensorTypeId type_id, bool is_variable, bool is_undefined) + : TypeDefault(type_id, is_variable, is_undefined) {} + + Allocator* allocator() const override; + Device getDeviceFromPtr(void * data) const override; + std::unique_ptr generator() const override; +}; + +} // namespace at diff --git a/aten/src/ATen/cuda/_curand_mtgp32_host.h b/aten/src/ATen/cuda/_curand_mtgp32_host.h new file mode 100644 index 0000000000000..22c24084777b2 --- /dev/null +++ b/aten/src/ATen/cuda/_curand_mtgp32_host.h @@ -0,0 +1,34 @@ +#include + +// Forward declarations of functions that are defined in libcurand_static.a +// This is to avoid multiple-definitions of these when statically linking +// cudarand in both Caffe2 and ATen +#if CAFFE2_STATIC_LINK_CUDA() +curandStatus_t curandMakeMTGP32Constants( + const mtgp32_params_fast_t params[], + mtgp32_kernel_params_t * p); +void mtgp32_init_state( + unsigned int state[], + const mtgp32_params_fast_t *para, + unsigned int seed); +curandStatus_t CURANDAPI curandMakeMTGP32KernelState( + curandStateMtgp32_t *s, + mtgp32_params_fast_t params[], + mtgp32_kernel_params_t *k, + int n, + unsigned long long seed); +extern mtgp32_params_fast_t mtgp32dc_params_fast_11213[]; +int mtgp32_init_by_array( + unsigned int state[], + const mtgp32_params_fast_t *para, + unsigned int *array, int length); +int mtgp32_init_by_str( + unsigned int state[], + const mtgp32_params_fast_t *para, + unsigned char *array); +extern const int mtgpdc_params_11213_num; + +#else // CAFFE2_STATIC_LINK_CUDA +#include +#include +#endif // CAFFE2_STATIC_LINK_CUDA diff --git a/aten/src/ATen/cuda/detail/TensorInfo.cuh b/aten/src/ATen/cuda/detail/TensorInfo.cuh index 906f03d9c6efe..7dfa9051e1034 100644 --- a/aten/src/ATen/cuda/detail/TensorInfo.cuh +++ b/aten/src/ATen/cuda/detail/TensorInfo.cuh @@ -1,6 +1,7 @@ #pragma once -#include "ATen/ATen.h" +#include +#include namespace at { namespace cuda { @@ -22,18 +23,7 @@ struct TensorInfo { // slice) void reduceDim(int dim); - /* - Updates the TensorInfo's dims, sizes, and strides to reflect a "collapse" of - the info, possibly excluding the optional excludeDim. A "collapsed" version - of the info is the fewest dims that order the tensor's elements in the same - way as the original info. If excludeDim is specified, the collapse is the - fewest dims that order the tensor's elements as the original and preserve the - excluded dimension, unless the tensor collapses to a point. - - Returns the (new) index of the preserved dimension if excludeDim is - specified. Returns 0 if the tensor is collapsed to a point. Returns -1 - otherwise. - */ + // See note on [collapse dims]. int collapseDims(const int excludeDim = -1); // Contiguous tensors of more than one dimension are collapsed down @@ -79,71 +69,9 @@ TensorInfo::reduceDim(int dim) { template int TensorInfo::collapseDims(const int excludeDim) { - - AT_CHECK(excludeDim >= -1 && excludeDim < dims, - "expected excluded dim between -1 and dims - 1"); - - int stopDim = (excludeDim == -1) ? dims : excludeDim; - int newIndex = -1; - int oldIndex = 0; - int remappedExcludedDim = -1; - - while (oldIndex < dims) { - // Finds a dimension to collapse into - for (; oldIndex < stopDim; ++oldIndex) { - if (sizes[oldIndex] == 1) { - continue; - } - - ++newIndex; - sizes[newIndex] = sizes[oldIndex]; - strides[newIndex] = strides[oldIndex]; - ++oldIndex; - break; - } - - // Collapses dims - for (; oldIndex < stopDim; ++oldIndex) { - if (sizes[oldIndex] == 1) { - continue; - } - - if (strides[newIndex] == sizes[oldIndex] * strides[oldIndex]) { - sizes[newIndex] *= sizes[oldIndex]; - strides[newIndex] = strides[oldIndex]; - } else { - ++newIndex; - sizes[newIndex] = sizes[oldIndex]; - strides[newIndex] = strides[oldIndex]; - } - } - - // Handles excludeDim being set (oldIndex == excludeDim) - if (oldIndex != dims) { - - // Preserves excluded dimension - ++newIndex; - sizes[newIndex] = sizes[oldIndex]; - strides[newIndex] = strides[oldIndex]; - remappedExcludedDim = newIndex; - - // Restarts iteration after excludeDim - ++oldIndex; - stopDim = dims; - } - } - - // Handles special case of all dims size 1 - if (newIndex == -1 || (newIndex == 0 && sizes[0] == 1)) { - dims = 1; - sizes[0] = 1; - strides[0] = 1; - - return 0; - } - - dims = newIndex + 1; - return remappedExcludedDim; + auto result = at::collapse_dims(sizes, strides, dims, excludeDim); + dims = std::get<1>(result); + return std::get<0>(result); } // Translate a linear index for the apply to a T* offset; diff --git a/aten/src/ATen/cudnn/Descriptors.h b/aten/src/ATen/cudnn/Descriptors.h index b901313e9b070..d5b1c1fe54802 100644 --- a/aten/src/ATen/cudnn/Descriptors.h +++ b/aten/src/ATen/cudnn/Descriptors.h @@ -257,7 +257,7 @@ struct AT_CUDA_API DropoutDescriptor AT_CUDNN_CHECK(cudnnDropoutGetStatesSize(handle, &state_size)); AT_ASSERT(type.is_cuda()); AT_ASSERT(type.scalarType() == kByte); - state = at::empty({static_cast(state_size)}, type); + state = at::empty({static_cast(state_size)}, type.options()); AT_CUDNN_CHECK(cudnnSetDropoutDescriptor(mut_desc(), handle, dropout, state.data_ptr(), state_size, seed)); } @@ -291,7 +291,7 @@ struct AT_CUDA_API RNNDescriptor DropoutDescriptor dropout_desc_; void set(cudnnHandle_t handle, int hidden_size, int num_layers, DropoutDescriptor&& dropout_desc, cudnnRNNInputMode_t input_mode, cudnnDirectionMode_t bidirectional, - cudnnRNNMode_t mode, cudnnDataType_t datatype) { + cudnnRNNMode_t mode, cudnnDataType_t datatype, cudnnRNNAlgo_t algo) { dropout_desc_ = std::move(dropout_desc); AT_CUDNN_CHECK(cudnnSetRNNDescriptor_v6( handle, @@ -302,7 +302,7 @@ struct AT_CUDA_API RNNDescriptor input_mode, bidirectional, mode, - CUDNN_RNN_ALGO_STANDARD, + algo, datatype)); #if CUDNN_VERSION >= 7000 && CUDA_VERSION >= 9000 cudaDeviceProp* prop = at::cuda::getCurrentDeviceProperties(); diff --git a/aten/src/ATen/detail/ComplexHooksInterface.cpp b/aten/src/ATen/detail/ComplexHooksInterface.cpp new file mode 100644 index 0000000000000..9755e288ff5fe --- /dev/null +++ b/aten/src/ATen/detail/ComplexHooksInterface.cpp @@ -0,0 +1,25 @@ +#include + +namespace at { + +namespace detail { +const ComplexHooksInterface& getComplexHooks() { + static std::unique_ptr complex_hooks; + // NB: The once_flag here implies that if you try to call any Complex + // functionality before you load the complex library, you're toast. + // Same restriction as in getCUDAHooks() + static std::once_flag once; + std::call_once(once, [] { + complex_hooks = ComplexHooksRegistry()->Create("ComplexHooks", ComplexHooksArgs{}); + if (!complex_hooks) { + complex_hooks = + std::unique_ptr(new ComplexHooksInterface()); + } + }); + return *complex_hooks; +} +} // namespace detail + +AT_DEFINE_REGISTRY(ComplexHooksRegistry, ComplexHooksInterface, ComplexHooksArgs) + +} diff --git a/aten/src/ATen/detail/ComplexHooksInterface.h b/aten/src/ATen/detail/ComplexHooksInterface.h new file mode 100644 index 0000000000000..80ecfb6f26f83 --- /dev/null +++ b/aten/src/ATen/detail/ComplexHooksInterface.h @@ -0,0 +1,27 @@ +#pragma once + +#include +#include + +namespace at { + +class Context; + +struct AT_API ComplexHooksInterface { + virtual ~ComplexHooksInterface() {} + + virtual void registerComplexTypes(Context*) const { + AT_ERROR("Cannot register complex types without loading a library with complex support"); + } +}; + +struct AT_API ComplexHooksArgs {}; +AT_DECLARE_REGISTRY(ComplexHooksRegistry, ComplexHooksInterface, ComplexHooksArgs) +#define REGISTER_COMPLEX_HOOKS(clsname) \ + AT_REGISTER_CLASS(ComplexHooksRegistry, clsname, clsname) + +namespace detail { +AT_API const ComplexHooksInterface& getComplexHooks(); +} + +} diff --git a/aten/src/ATen/function_wrapper.py b/aten/src/ATen/function_wrapper.py index 9c6c52891ef4b..ac7b212199e35 100644 --- a/aten/src/ATen/function_wrapper.py +++ b/aten/src/ATen/function_wrapper.py @@ -93,7 +93,7 @@ def TypedDict(name, attrs, total=True): # type: ignore return at::native::${api_name}(${type_method_actuals}, options()); } """) -# 4. add virtual override to TypeDerived.h +# 4. add override to TypeDerived.h TYPE_DERIVED_DECLARATION = CodeTemplate("""\ ${return_type} ${method_prefix_derived}${api_name}(${type_method_formals}) const override; """) diff --git a/aten/src/ATen/gen.py b/aten/src/ATen/gen.py index d23831be4732b..762262c25ef19 100644 --- a/aten/src/ATen/gen.py +++ b/aten/src/ATen/gen.py @@ -256,6 +256,8 @@ def generate_storage_type_and_tensor(backend, density, scalar_type, declarations ] env['extra_cuda_headers'] = ['#include '] env['extra_cuda_headers'].append('#include ') + env['extra_cuda_headers'].append('#include ') + env['extra_cuda_headers'].append('#include ') sname = '' if scalar_name == "Float" else scalar_name env['THType'] = 'Cuda{}'.format(sname) env['THStorage'] = 'THCuda{}Storage'.format(sname) diff --git a/aten/src/ATen/native/Linear.cpp b/aten/src/ATen/native/Linear.cpp index 5137baa8683de..bdf9602fe9ae0 100644 --- a/aten/src/ATen/native/Linear.cpp +++ b/aten/src/ATen/native/Linear.cpp @@ -457,4 +457,62 @@ Tensor bilinear(const Tensor& input1, const Tensor& input2, const Tensor& weight return output; } +// implements tensordot, a matrix-multiplication-like contraction, but the dimensions given +// in the two dimension lists +Tensor tensordot(const Tensor& input1, const Tensor& input2, IntList dims1, IntList dims2) { + AT_CHECK(dims1.size() == dims2.size(), "both dimension lists should have same length"); + int64_t csize = 1; // total size of the contracted dimensions + Tensor t1 = input1; + Tensor t2 = input2; + for (size_t i = 0; i < dims1.size(); i++) { + int s1 = input1.size(dims1[i]); + int s2 = input2.size(dims2[i]); + if (s2 == 1) { // broadcasted dimensions can be summed right away + t1 = t1.sum(dims1[i], true); + } else if (s1 == 1) { + t2 = t2.sum(dims2[i], true); + } else { + AT_CHECK(s1 == s2, "contracted dimensions need to match, but first has size ", s1, " in dim ", dims1[i], + " and second has size ", s2, " in dim ", dims2[i]); + csize *= s1; + } + } + + auto cdims1 = dim_list_to_bitset(dims1, input1.dim()); + auto cdims2 = dim_list_to_bitset(dims2, input2.dim()); + std::vector p1, p2, rsizes; // p1, p2: input permutations, rsizes: sizes of the result + p1.reserve(input1.dim()); + p2.reserve(input2.dim()); + rsizes.reserve(input1.dim() + input2.dim() - (int64_t) dims1.size()); + int64_t size1 = 1; // number of non-contracted elements in input1 + int64_t size2 = 1; // number of non-contracted elements in input2 + + // fill the permutations and compute sizes + for (int64_t i = 0; i < input1.dim(); i++) { + if (! cdims1[i]) { + p1.emplace_back(i); + size1 *= t1.size(i); + rsizes.emplace_back(t1.size(i)); + } + } + for (size_t i = 0; i < dims1.size(); i++) { + p1.emplace_back(dims1[i]); + } + for (size_t i = 0; i < dims2.size(); i++) { + p2.emplace_back(dims2[i]); + } + for (int64_t i = 0; i < input2.dim(); i++) { + if (! cdims2[i]) { + p2.emplace_back(i); + size2 *= t2.size(i); + rsizes.emplace_back(t2.size(i)); + } + } + // permut and reshape for matrix multiplication + t1 = t1.permute(p1).reshape({size1, csize}); + t2 = t2.permute(p2).reshape({csize, size2}); + // multiply and reshape to target size + return at::mm(t1, t2).reshape(rsizes); +} + }} // namespace at::native diff --git a/aten/src/ATen/native/PackedSequence.cpp b/aten/src/ATen/native/PackedSequence.cpp new file mode 100644 index 0000000000000..9e660955d9224 --- /dev/null +++ b/aten/src/ATen/native/PackedSequence.cpp @@ -0,0 +1,166 @@ +#include "ATen/ATen.h" +#include "ATen/NativeFunctions.h" + +namespace at { namespace native { + +void checkLongTensor(const Tensor& tensor) { + auto & t = tensor.type(); + AT_CHECK(tensor.dim() == 1 && t.device_type() == at::kCPU && t.scalarType() == at::kLong, + "'lengths' argument should be a 1D CPU int64 tensor"); +} + +std::tuple _pack_padded_sequence(const Tensor& _input, const Tensor& _lengths, bool batch_first) { + auto input = batch_first ? _input.transpose(0, 1) : _input; + auto lengths_t = _lengths.contiguous(); + checkLongTensor(lengths_t); + + int64_t batch_size = input.size(1); + int64_t * lengths = lengths_t.data(); + AT_CHECK(lengths_t.size(0) == batch_size, + "Expected `len(lengths)` to be equal to batch_size, but got ", lengths_t.size(0), + " (batch_size=", batch_size, ")"); + AT_CHECK(lengths[batch_size - 1] > 0, + "Length of all samples has to be greater than 0, but found an element " + "in 'lengths' that is <= 0"); + + std::vector steps; + steps.reserve(batch_size); + at::Tensor batch_sizes_t = at::empty(lengths[0], _lengths.options()); + int64_t * batch_sizes = batch_sizes_t.data(); + + std::vector step_shape; // == [-1, *input.shape[2:]] + { + auto input_sizes = input.sizes(); + step_shape.reserve(input_sizes.size()); + auto s_input_sizes = input_sizes.slice(2); + step_shape.push_back(-1); + step_shape.insert(step_shape.end(), s_input_sizes.begin(), s_input_sizes.end()); + } + + // To understand what's going on in this loop imagine that the input is a padded 2D + // array that looks like this (x = valid entry, . = padding) + // + // 1 1 1 1 1 + // 2 2 2 . . + // 2 2 2 . . + // 4 . . . . + // 4 . . . . + // + // Where the vertical dimension corresponds to time, and horizontal dim to batch. + // In this example, the lengths array will be equal to [5, 3, 3, 1, 1], and we will + // iterate over them in reverse order (from the rightmost column to the left). + // We want to avoid eager slicing of the input at every time step, and wait for + // the moments where the length increases. In this example, that will happen at the + // first, second and fourth steps. Then, we slice out the whole block of the input + // that corresponds to this length, and hasn't been sliced yet (the steps at which each + // element is sliced are annotated in the array above). You can think of this as if we + // were scanning the sequences from the shortest one, and every time we realize there's + // more elements below in our column, we lower the counter (prev_l), and append the new + // block to the output. + int64_t prev_l = 0; + for (int64_t i = 0; i < batch_size; ++i) { + int64_t l = lengths[batch_size - 1 - i]; + if (l > prev_l) { + auto current_batch_size = batch_size - i; + steps.push_back(input.slice(0, prev_l, l).slice(1, 0, current_batch_size).contiguous().view(step_shape)); + for (int64_t j = 0; j < (l - prev_l); ++j) { + (*batch_sizes++) = current_batch_size; + } + prev_l = l; + } else if (prev_l > l) { + AT_ERROR("'lengths' array has to be sorted in decreasing order"); + } + } + + return std::make_tuple(at::cat(steps), batch_sizes_t); +} + +Tensor _pack_padded_sequence_backward(const Tensor& grad, at::IntList input_size, const Tensor& _batch_sizes, bool batch_first) { + std::vector input_size_after_t = input_size.vec(); + if (batch_first) { + AT_CHECK(input_size.size() >= 2); + std::swap(input_size_after_t[0], input_size_after_t[1]); + } + auto grad_input = at::zeros(input_size_after_t, grad.options()); + auto batch_sizes_t = _batch_sizes.contiguous(); + checkLongTensor(batch_sizes_t); + + int64_t offset = 0; + int64_t max_seq_len = batch_sizes_t.size(0); + int64_t * batch_sizes = batch_sizes_t.data(); + for (int64_t i = 0; i < max_seq_len; ++i) { + grad_input[i].slice(0, 0, batch_sizes[i]).copy_(grad.slice(0, offset, offset + batch_sizes[i])); + offset += batch_sizes[i]; + } + + if (batch_first) { + grad_input = grad_input.transpose(0, 1); + } + + return grad_input; +} + +std::tuple _pad_packed_sequence(const Tensor& data, const Tensor& _batch_sizes, bool batch_first, Scalar padding_value, int64_t total_length) { + auto batch_sizes_t = _batch_sizes.contiguous(); + checkLongTensor(batch_sizes_t); + + int64_t * batch_sizes = batch_sizes_t.data(); + int64_t max_batch_size = batch_sizes[0]; + int64_t max_real_seq_length = batch_sizes_t.size(0); + int64_t max_seq_length = max_real_seq_length; + if (total_length > 0) { + AT_CHECK(total_length >= max_seq_length, + "Expected total_length to be at least the length of the longest " + "sequence in input, but got total_length=", total_length, " and " + "max sequence length being ", max_seq_length); + max_seq_length = total_length; + } + + std::vector output_size; // == [max_seq_length, max_batch_size, *var_data.size()[1:]] + { + output_size.reserve(data.dim() + 1); + output_size.push_back(max_seq_length); + output_size.push_back(max_batch_size); + auto s_data_size = data.sizes().slice(1); + output_size.insert(output_size.end(), s_data_size.begin(), s_data_size.end()); + } + auto output = at::full(output_size, padding_value, data.options()); + + // This will be modified at every iteration, but we reserve memory for it now. + std::vector tmp_view_size = std::move(output_size); // == [-1, -1, *var_data.size()[1:]] + + at::Tensor lengths_t = at::empty(max_batch_size, batch_sizes_t.options()); + int64_t * lengths = lengths_t.data() + max_batch_size - 1; + int64_t data_offset = 0; + int64_t prev_batch_size = max_batch_size; + int64_t prev_i = 0; + for (int64_t i = 0; i <= max_real_seq_length; ++i) { + int64_t batch_size = i != max_real_seq_length ? batch_sizes[i] : 0; + if (batch_size != prev_batch_size) { + int64_t l = prev_batch_size * (i - prev_i); + // The lines below are equivalent to this: + // output[prev_i:i, :prev_batch_size] = tmp.view(i - prev_i, prev_batch_size, *input.shape[2:]) + auto tmp = data.slice(0, data_offset, data_offset + l); + tmp_view_size[0] = i - prev_i; + tmp_view_size[1] = prev_batch_size; + output.slice(0, prev_i, i).slice(1, 0, prev_batch_size).copy_(tmp.view(tmp_view_size)); + data_offset += l; + prev_i = i; + } + int64_t dec = prev_batch_size - batch_size; + if (dec > 0) { + for (int64_t j = 0; j < dec; ++j) { + (*lengths--) = i; + } + } + prev_batch_size = batch_size; + } + + if (batch_first) { + output = output.transpose(0, 1); + } + + return std::make_tuple(output, lengths_t); +} + +}} // namespace at::native diff --git a/aten/src/ATen/native/TensorCompare.cpp b/aten/src/ATen/native/TensorCompare.cpp index e2dafc43cb2b7..c774652478509 100644 --- a/aten/src/ATen/native/TensorCompare.cpp +++ b/aten/src/ATen/native/TensorCompare.cpp @@ -43,13 +43,15 @@ Tensor isclose(const Tensor& self, const Tensor& other, double rtol, double atol auto max_error = atol + rtol * other.abs(); auto close = actual_error <= max_error; - // Handle +/-inf - close.__ior__(self == other); - close.__iand__((self == INFINITY) == (other == INFINITY)); - close.__iand__((self == -INFINITY) == (other == -INFINITY)); - - if (equal_nan) { - close.__ior__((self != self).__and__((other != other))); + if (isFloatingType(self.type().scalarType()) && isFloatingType(other.type().scalarType())) { + // Handle +/-inf + close.__ior__(self == other); + close.__iand__((self == INFINITY) == (other == INFINITY)); + close.__iand__((self == -INFINITY) == (other == -INFINITY)); + + if (equal_nan) { + close.__ior__((self != self).__and__((other != other))); + } } return close; } diff --git a/aten/src/ATen/native/Unique.cpp b/aten/src/ATen/native/Unique.cpp index d5ff300c0dd9e..2faeaa2d82119 100644 --- a/aten/src/ATen/native/Unique.cpp +++ b/aten/src/ATen/native/Unique.cpp @@ -21,7 +21,7 @@ std::tuple _unique_cpu_template( const Tensor& input = self.contiguous(); const scalar_t* input_data = input.data(); std::unordered_set set(input_data, input_data + input.numel()); - Tensor output = at::empty({static_cast(set.size())}, input.type()); + Tensor output = at::empty({static_cast(set.size())}, input.options()); scalar_t* output_data = output.data(); if (sorted) { @@ -32,7 +32,7 @@ std::tuple _unique_cpu_template( std::copy(set.begin(), set.end(), output_data); } - Tensor inverse_indices = at::empty({0}, self.type().toScalarType(kLong)); + Tensor inverse_indices = at::empty({0}, self.options().dtype(kLong)); if (return_inverse) { inverse_indices.resize_(input.sizes()); int64_t* inverse_indices_data = inverse_indices.data(); @@ -103,12 +103,12 @@ std::tuple _unique_dim_cpu_template( return false; }); - Tensor input_sorted = at::empty(input_flat.sizes(), input_flat.type()); + Tensor input_sorted = at::empty(input_flat.sizes(), input_flat.options()); for (int i = 0; i < indices.size(); ++i) { input_sorted[i] = input_flat[indices[i]]; } - Tensor inverse_indices = at::empty(indices.size(), self.type().toScalarType(kLong)); + Tensor inverse_indices = at::empty(indices.size(), self.options().dtype(kLong)); std::vector input_unbind = at::unbind(input_sorted, 0); auto last = _unique_dim_cpu_impl( input_unbind.begin(), input_unbind.end(), indices, inverse_indices); diff --git a/aten/src/ATen/native/cuda/EmbeddingBag.cu b/aten/src/ATen/native/cuda/EmbeddingBag.cu index 4f989e3db5c56..853c04deb2215 100644 --- a/aten/src/ATen/native/cuda/EmbeddingBag.cu +++ b/aten/src/ATen/native/cuda/EmbeddingBag.cu @@ -175,15 +175,15 @@ Tensor embedding_bag_backward_cuda_sum_avg( Tensor &bag_size = const_cast(bag_size_); - auto grad_weight = at::zeros({num_weights, grad.size(1)}, grad.type()); + auto grad_weight = at::zeros({num_weights, grad.size(1)}, grad.options()); cudaStream_t stream = at::cuda::getCurrentCUDAStream(); ptrdiff_t numel = indices.numel(); int64_t stride = grad_weight.stride(0); - auto sorted_indices = indices.type().tensor(indices.sizes()); - auto orig_indices = indices.type().tensor(indices.sizes()); + auto sorted_indices = at::empty_like(indices); + auto orig_indices = at::empty_like(indices); using device_ptr = thrust::device_ptr; // Sort the inputs into sorted with the corresponding indices; we @@ -208,7 +208,7 @@ Tensor embedding_bag_backward_cuda_sum_avg( Tensor count; if (scale_grad_by_freq) { - count = indices.type().tensor(indices.sizes()); + count = at::empty_like(indices); auto allocator = THCThrustAllocator(globalContext().lazyInitCUDA()); auto policy = thrust::cuda::par(allocator).on(stream); @@ -278,7 +278,7 @@ Tensor embedding_bag_backward_cuda_max(const Tensor &grad, const Tensor &max_indices, int64_t num_weights) { - auto grad_weight = at::zeros({num_weights, grad.size(1)}, grad.type()); + auto grad_weight = at::zeros({num_weights, grad.size(1)}, grad.options()); int64_t stride = grad_weight.stride(0); diff --git a/aten/src/ATen/native/cuda/LossCTC.cu b/aten/src/ATen/native/cuda/LossCTC.cu index 681d3a48115fb..16d7935f3d49f 100644 --- a/aten/src/ATen/native/cuda/LossCTC.cu +++ b/aten/src/ATen/native/cuda/LossCTC.cu @@ -185,7 +185,7 @@ std::tuple ctc_loss_gpu_template(const Tensor& log_probs, const int64_t tg_target_stride; int64_t max_target_length; - auto tg_batch_offsets = at::empty({batch_size}, TensorOptions(at::CPU(kLong))); + auto tg_batch_offsets = at::empty({batch_size}, at::device(at::kCPU).dtype(at::kLong)); auto tg_batch_offsets_data = tg_batch_offsets.data(); if (targets.dim() == 1) { // concatenated targets int64_t pos = 0; @@ -219,8 +219,8 @@ std::tuple ctc_loss_gpu_template(const Tensor& log_probs, const " (while checking arguments for ", c, ")"); } - auto target_lengths_t = at::tensor(target_lengths, targets.options().device(at::Device(at::Device::Type::CPU)).dtype(kLong)).toType(targets.type().toScalarType(kLong)); - auto input_lengths_t = at::tensor(input_lengths, targets.options().device(at::Device(at::Device::Type::CPU)).dtype(kLong)).toType(targets.type().toScalarType(kLong)); + auto target_lengths_t = at::tensor(target_lengths, targets.options().dtype(kLong)); + auto input_lengths_t = at::tensor(input_lengths, targets.options().dtype(kLong)); tg_batch_offsets = tg_batch_offsets.toType(targets.type().toScalarType(kLong)); Tensor log_alpha = at::empty({batch_size, log_probs.size(0), 2*max_target_length+1}, log_probs.options()); diff --git a/aten/src/ATen/native/cudnn/RNN.cpp b/aten/src/ATen/native/cudnn/RNN.cpp index e1427609a43f7..8fc896afe23a1 100644 --- a/aten/src/ATen/native/cudnn/RNN.cpp +++ b/aten/src/ATen/native/cudnn/RNN.cpp @@ -98,7 +98,7 @@ namespace { cudnnDirectionMode_t bidirectional; cudnnRNNMode_t mode; cudnnDataType_t datatype; - + cudnnRNNAlgo_t algo = CUDNN_RNN_ALGO_STANDARD; cudnnRNNInputMode_t input_mode = CUDNN_LINEAR_INPUT; int64_t num_directions() const { @@ -131,6 +131,10 @@ namespace { void set_bidirectional(bool fn_bidirectional) { bidirectional = fn_bidirectional ? CUDNN_BIDIRECTIONAL : CUDNN_UNIDIRECTIONAL; } + + void set_algo(cudnnRNNAlgo_t algo){ + this->algo = algo; + } void set(int64_t mode, int64_t hidden_size, int64_t num_layers, bool bidirectional, cudnnDataType_t datatype) { this->set_mode(mode); @@ -143,7 +147,7 @@ namespace { RNNDescriptor descriptor(cudnnHandle_t handle, DropoutDescriptor&& dropout_desc) const { RNNDescriptor rnn_desc; - rnn_desc.set(handle, hidden_size, num_layers, std::move(dropout_desc), input_mode, bidirectional, mode, datatype); + rnn_desc.set(handle, hidden_size, num_layers, std::move(dropout_desc), input_mode, bidirectional, mode, datatype, algo); return rnn_desc; } @@ -557,6 +561,30 @@ namespace { } } + cudnnRNNAlgo_t get_algo(const RNNDescriptorParams& rnn, const TensorDescriptorListParams& tensors){ +#if CUDNN_VERSION < 7200 || CUDA_VERSION < 9010 + return CUDNN_RNN_ALGO_STANDARD; +#else + cudaDeviceProp* prop = at::cuda::getCurrentDeviceProperties(); + const int64_t bsize = tensors.mini_batch; + if (prop->major == 7 && rnn.datatype == CUDNN_DATA_HALF && !tensors.is_input_packed()) { + if (rnn.num_layers == 1 && rnn.hidden_size <= 1024 && tensors.input_size <=1024 && rnn.num_directions() == 1 && + rnn.hidden_size % 128 == 0 && tensors.input_size % 128 == 0){ + //technically, batch size should be multiple of 8, but there are quite a few multiple-of-8 batchsizes that give bad perf, + //weed them out + if ((bsize % 16 == 0 && bsize != 80 && bsize !=112) || bsize == 8){ + if ((tensors.seq_length >=40 && bsize <=128) || + (tensors.seq_length >=20 && bsize <=96) || + (tensors.seq_length >=10 && bsize <=32)) { + return CUDNN_RNN_ALGO_PERSIST_STATIC; + } + } + } + } + return CUDNN_RNN_ALGO_STANDARD; +#endif + } + } // anonymous namespace // NB: does inplace update into TensorList @@ -676,6 +704,8 @@ std::tuple _cudnn_rnn( auto y = output; auto handle = getCudnnHandle(); + cudnnRNNAlgo_t algo = get_algo(fn.rnn, fn.tensors); + fn.rnn.set_algo(algo); RNNDescriptors descs(fn, handle, x, y, hx, cx); FilterDescriptor w_desc; @@ -858,6 +888,8 @@ std::tuple _cudnn_rnn_backward_input( throw std::runtime_error("Gradients aren't CUDA tensors"); } + cudnnRNNAlgo_t algo = get_algo(fn.rnn, fn.tensors); + fn.rnn.set_algo(algo); RNNDescriptors descs(fn, handle, x, y, hx, cx); FilterDescriptor w_desc; @@ -968,6 +1000,8 @@ std::vector _cudnn_rnn_backward_weight( const auto& y = output; auto dw = weight_buf.type().tensor(weight_buf.sizes()).zero_(); + cudnnRNNAlgo_t algo = get_algo(fn.rnn, fn.tensors); + fn.rnn.set_algo(algo); RNNDescriptors descs(fn, handle, x, y, hx, cx); FilterDescriptor w_desc; diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 80e4a94e0633e..4b5e758492f76 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -1721,6 +1721,9 @@ CPU: _tanh_out_cpu CUDA: _tanh_out_cuda +- func: tensordot(Tensor self, Tensor other, IntList dims_self, IntList dims_other) -> Tensor + variants: function + - func: transpose(Tensor self, int64_t dim0, int64_t dim1) -> Tensor - func: transpose_(Tensor self, int64_t dim0, int64_t dim1) -> Tensor @@ -2208,3 +2211,13 @@ - func: rnn_relu_cell(Tensor input, Tensor hx, Tensor w_ih, Tensor w_hh, Tensor? b_ih={}, Tensor? b_hh={}) -> Tensor variants: function + +# PackedSequence utilities +- func: _pack_padded_sequence(Tensor input, Tensor lengths, bool batch_first) -> (Tensor, Tensor) + variants: function + +- func: _pack_padded_sequence_backward(Tensor grad, IntList input_size, Tensor batch_sizes, bool batch_first) -> Tensor + variants: function + +- func: _pad_packed_sequence(Tensor data, Tensor batch_sizes, bool batch_first, Scalar padding_value, int64_t total_length) -> (Tensor, Tensor) + variants: function diff --git a/aten/src/ATen/native/sparse/SparseTensor.cpp b/aten/src/ATen/native/sparse/SparseTensor.cpp index 091ec2ce431e6..c96e577ba9d47 100644 --- a/aten/src/ATen/native/sparse/SparseTensor.cpp +++ b/aten/src/ATen/native/sparse/SparseTensor.cpp @@ -125,7 +125,7 @@ SparseTensor new_with_dims_and_size_sparse(const SparseType& dtype, int64_t spar SparseTensor self = new_sparse(dtype); AT_CHECK(size.size() != 0, "cannot construct sparse tensor with 0 dimensions and no values; you must specify at least 1 dimension if you want to create a sparse tensor with no elements, \ -or you must provide a single-element `values` tensor (e.g. x=torch.sparse_coo_tensor(torch.zeros(0,1), 12.3, [])) if you want to create a scalar sparse tensor"); +or you must provide a single-element `values` tensor (e.g. x = torch.sparse_coo_tensor(torch.zeros(0, 1), 12.3, [])) if you want to create a scalar sparse tensor"); _get_sparse_impl(self)->resize_and_clear_(sparseDims, denseDims, size); return self; } @@ -173,17 +173,24 @@ SparseTensor new_with_tensor_and_size_sparse(const LongTensor& indices, const Te // Check to make sure all indices are within the boundaries of `sizes` if (indices.numel() > 0) { + LongTensor min_indices = std::get(indices.min(/* dim */ 1, /* keepdim */ false)); LongTensor max_indices = std::get(indices.max(/* dim */ 1, /* keepdim */ false)); - LongTensor cpu_max_indices; - if (max_indices.is_cuda()) { + LongTensor cpu_min_indices, cpu_max_indices; + if (indices.is_cuda()) { + cpu_min_indices = at::CPU(kLong).copy(min_indices); cpu_max_indices = at::CPU(kLong).copy(max_indices); } else { + cpu_min_indices = min_indices; cpu_max_indices = max_indices; } + auto cpu_min_indices_accessor = cpu_min_indices.accessor(); auto cpu_max_indices_accessor = cpu_max_indices.accessor(); for (int64_t d = 0; d < sparseDims; d++) { // NB: This used to sync ndim times to access each entry; now we copy // everything to CPU first and then access it. + int64_t min_index_in_dim = cpu_min_indices_accessor[d]; + AT_CHECK(min_index_in_dim >= 0, + "found negative index ", min_index_in_dim, " for dim ", d); int64_t max_index_in_dim = cpu_max_indices_accessor[d]; int64_t dim_size = sizes[static_cast(d)]; AT_CHECK(max_index_in_dim < dim_size, diff --git a/aten/src/ATen/native/sparse/SparseUtils.h b/aten/src/ATen/native/sparse/SparseUtils.h index d48b837bab7c7..3ce0eee53353e 100644 --- a/aten/src/ATen/native/sparse/SparseUtils.h +++ b/aten/src/ATen/native/sparse/SparseUtils.h @@ -108,16 +108,9 @@ inline LongTensor _newFlattenedIndices(const SparseTensor& self, bool forceClone // TODO: Expose this for real in ATen, some day? // NB: Doesn't preserve data. inline Tensor _new_values_with_size_of(const Tensor& values, int64_t nnz) { - if (values.numel() == 0) { // values tensor uninitialized - // TODO: This logic looks bogus; if we have an uninitialized - // values tensor, why should we believe that denseDims == 0? - // That's the assumption this code makes. - return values.type().tensor({nnz}); - } else { - std::vector size = values.sizes().vec(); - size[0] = nnz; - return values.type().tensor(size); - } + std::vector size = values.sizes().vec(); + size[0] = nnz; + return values.type().tensor(size); } diff --git a/aten/src/ATen/native/sparse/cuda/SparseCUDATensorMath.cu b/aten/src/ATen/native/sparse/cuda/SparseCUDATensorMath.cu index 2715f81f5f6c6..d7701479e9b24 100644 --- a/aten/src/ATen/native/sparse/cuda/SparseCUDATensorMath.cu +++ b/aten/src/ATen/native/sparse/cuda/SparseCUDATensorMath.cu @@ -78,7 +78,7 @@ Tensor& s_addmm_out_sparse_dense_cuda(Tensor& r_, const Tensor& t, const SparseT LongTensor rowIndices = indices.select(0, 0); LongTensor colIndices = indices.select(0, 1); IntTensor csr = _to_csr_int(rowIndices, m, nnz); - IntTensor colIndicesInt = at::empty({colIndices.size(0)}, indices.type().toScalarType(kInt)); + IntTensor colIndicesInt = at::empty({colIndices.size(0)}, indices.options().dtype(kInt)); colIndicesInt.copy_(colIndices); // No half support, so we don't have to use CUDATypeConversion @@ -153,7 +153,7 @@ Tensor s_addmm_sparse_dense_cuda( Scalar beta, Scalar alpha ) { - Tensor r = t.type().tensor(); + Tensor r = at::empty({0}, t.options()); s_addmm_out_sparse_dense_cuda(r, t, sparse, dense, beta, alpha); return r; } @@ -208,7 +208,7 @@ SparseTensor& hspmm_out_sparse_cuda(SparseTensor& r_, const SparseTensor& sparse LongTensor indices = at::empty({1, nnz}, CUDA(kLong)); // create values in column-major format to avoid copying in spaddmm - Tensor values = at::empty({n, nnz}, dense.type()); + Tensor values = at::empty({n, nnz}, dense.options()); values.transpose_(0, 1); // why does sparse need to be cloned? If this is really necessary maybe we @@ -434,7 +434,7 @@ SparseTensor& mul_out_sparse_cuda(SparseTensor& r_, const SparseTensor& t_, cons Tensor t_values_ = t._values(); LongTensor s_indices_ = src._indices(); Tensor s_values_ = src._values(); - LongTensor r_indices_ = t_indices_.type().tensor({sparseDims, max_nnz}); + LongTensor r_indices_ = at::empty({sparseDims, max_nnz}, t_indices_.options()); Tensor r_values_ = _new_values_with_size_of(t_values_, max_nnz).zero_(); r_.resize_as_(src); _get_sparse_impl(r_)->set_indices_and_values_unsafe(r_indices_, r_values_); diff --git a/aten/src/ATen/templates/Functions.h b/aten/src/ATen/templates/Functions.h index 0afb07418e782..b4a2e05e759ea 100644 --- a/aten/src/ATen/templates/Functions.h +++ b/aten/src/ATen/templates/Functions.h @@ -5,7 +5,7 @@ #include "ATen/core/Scalar.h" #include "ATen/Type.h" #include "ATen/Tensor.h" -#include "ATen/Storage.h" +#include "ATen/core/Storage.h" #include "ATen/core/Generator.h" #include "ATen/core/Deprecated.h" #include "ATen/NativeFunctions.h" diff --git a/aten/src/ATen/templates/SparseTypeDerived.cpp b/aten/src/ATen/templates/SparseTypeDerived.cpp index 8985e7d373b76..064da2bc2186e 100644 --- a/aten/src/ATen/templates/SparseTypeDerived.cpp +++ b/aten/src/ATen/templates/SparseTypeDerived.cpp @@ -28,42 +28,18 @@ namespace at { ${Type}::${Type}() - : TypeDefault(${Backend}TensorId(), /*is_variable=*/false, /*is_undefined=*/false) {} + : ${DenseBackend}TypeDefault(${Backend}TensorId(), /*is_variable=*/false, /*is_undefined=*/false) {} ScalarType ${Type}::scalarType() const { return ScalarType::${ScalarName}; } Backend ${Type}::backend() const { return Backend::${Backend}; } -bool ${Type}::is_cuda() const { return backend() == Backend::CUDA || backend() == Backend::SparseCUDA; } -bool ${Type}::is_sparse() const { return backend() == Backend::SparseCPU || backend() == Backend::SparseCUDA; } -bool ${Type}::is_distributed() const { return false; } - -Storage ${Type}::storage(bool resizable) const { - AT_ERROR("storage not supported on sparse"); -} -Storage ${Type}::storage(size_t size, bool resizable) const { - AT_ERROR("storage not supported on sparse"); -} -Storage ${Type}::storageFromBlob(void * data, int64_t size, const std::function & deleter) const { - AT_ERROR("storage not supported on sparse"); -} -Storage ${Type}::storageWithAllocator(int64_t size, Allocator* allocator) const { - AT_ERROR("storage not supported on sparse"); -} -Tensor ${Type}::unsafeTensorFromTH(void * th_pointer, bool retain) const { - AT_ERROR("unsafeTensorFromTH not supported on sparse"); -} -Storage ${Type}::unsafeStorageFromTH(void * th_pointer, bool retain) const { - AT_ERROR("unsafeTensorFromTH not supported on sparse"); -} -std::unique_ptr ${Type}::generator() const { - return std::unique_ptr(new ${Generator}(&at::globalContext())); -} const char * ${Type}::toString() const { - return ${Type}::typeString(); + return "${Type}"; } + TypeID ${Type}::ID() const { return ${TypeID}; } @@ -72,10 +48,6 @@ size_t ${Type}::elementSizeInBytes() const { return sizeof(${ScalarType}); } -const char * ${Type}::typeString() { - return "${Type}"; -} - ${type_derived_method_definitions} } diff --git a/aten/src/ATen/templates/Tensor.h b/aten/src/ATen/templates/Tensor.h index 0228bcd0b7da1..967a123b983d9 100644 --- a/aten/src/ATen/templates/Tensor.h +++ b/aten/src/ATen/templates/Tensor.h @@ -52,9 +52,9 @@ struct AT_API Tensor { } } Tensor(const c10::intrusive_ptr& ptr) - : tensor_impl_(std::move(ptr)) {} - Tensor(c10::intrusive_ptr&& ptr) : tensor_impl_(ptr) {} + Tensor(c10::intrusive_ptr&& ptr) + : tensor_impl_(std::move(ptr)) {} Tensor(const Tensor&) = default; Tensor(Tensor&&) = default; diff --git a/aten/src/ATen/templates/Type.h b/aten/src/ATen/templates/Type.h index f605b1b3d756f..df1425f589da9 100644 --- a/aten/src/ATen/templates/Type.h +++ b/aten/src/ATen/templates/Type.h @@ -41,6 +41,10 @@ static inline void noop_deleter(void*) {} enum class TypeID { ${type_ids} + CPUComplexFloat, + CPUComplexDouble, + CUDAComplexFloat, + CUDAComplexDouble, Undefined, NumOptions }; @@ -58,6 +62,8 @@ struct AT_API Type { virtual bool is_distributed() const = 0; bool is_variable() const noexcept { return is_variable_; } bool is_undefined() const noexcept { return is_undefined_; } + virtual Allocator * allocator() const = 0; + virtual Device getDeviceFromPtr(void * data) const = 0; virtual Storage storage(bool resizable = false) const = 0; virtual Storage storage(size_t size, bool resizable = false) const = 0; virtual Storage storageFromBlob(void * data, int64_t size, const std::function & deleter=noop_deleter) const = 0; diff --git a/aten/src/ATen/templates/TypeDefault.cpp b/aten/src/ATen/templates/TypeDefault.cpp index 376fcd4348931..6fd208b5914b0 100644 --- a/aten/src/ATen/templates/TypeDefault.cpp +++ b/aten/src/ATen/templates/TypeDefault.cpp @@ -6,7 +6,7 @@ #include "ATen/NativeFunctions.h" #include "ATen/core/Scalar.h" #include "ATen/core/SparseTensorRef.h" -#include "ATen/Storage.h" +#include "ATen/core/Storage.h" #include "ATen/Tensor.h" #include "ATen/core/TensorOptions.h" #include "ATen/DeviceGuard.h" @@ -78,6 +78,34 @@ Tensor TypeDefault::tensorWithAllocator(IntList sizes, IntList strides, Allocato auto storage = storageWithAllocator(computeStorageSize(sizes, strides), std::move(allocator)); return tensor(storage, 0, sizes, strides); } + +Storage TypeDefault::storage(bool resizable) const { + return Storage(scalarType(), 0, allocator(), resizable); +} +Storage TypeDefault::storage(size_t size, bool resizable) const { + return Storage(scalarType(), size, allocator(), resizable); +} +Storage TypeDefault::storageFromBlob(void * data, int64_t size, const std::function & deleter) const { + return Storage( + scalarType(), + InefficientStdFunctionContext::makeDataPtr(data, deleter, getDeviceFromPtr(data)), + size, + deleter); +} +Storage TypeDefault::storageWithAllocator(int64_t size, Allocator* allocator) const { + return Storage(scalarType(), size, allocator); +} +Tensor TypeDefault::unsafeTensorFromTH(void * th_pointer, bool retain) const { + return Tensor(static_cast(th_pointer), retain); +} +Storage TypeDefault::unsafeStorageFromTH(void * th_pointer, bool retain) const { + if (retain && th_pointer) { + c10::raw::intrusive_ptr::incref(static_cast(th_pointer)); + } + return Storage(static_cast(th_pointer)); +} + + Tensor TypeDefault::scalarTensor(Scalar s) const { return tensor({}).fill_(s); } diff --git a/aten/src/ATen/templates/TypeDefault.h b/aten/src/ATen/templates/TypeDefault.h index 69731c717478e..af9f37b704da6 100644 --- a/aten/src/ATen/templates/TypeDefault.h +++ b/aten/src/ATen/templates/TypeDefault.h @@ -12,9 +12,15 @@ struct AT_API TypeDefault : public Type { // Make sure overload resolution considers the nullary virtual method. // (A single argument overload is generated in the list.) - bool is_cuda() const override = 0; - bool is_sparse() const override = 0; - bool is_distributed() const override = 0; + bool is_cuda() const override { + return backend() == Backend::CUDA || backend() == Backend::SparseCUDA; + } + bool is_sparse() const override { + return backend() == Backend::SparseCPU || backend() == Backend::SparseCUDA; + } + bool is_distributed() const override { + return false; + } Type & toBackend(Backend b) const override; Type & toScalarType(ScalarType s) const override; @@ -28,6 +34,13 @@ struct AT_API TypeDefault : public Type { Tensor tensorWithAllocator(IntList sizes, IntList strides, Allocator* allocator) const override; Tensor scalarTensor(Scalar s) const override; + Storage storage(bool resizable = false) const override; + Storage storage(size_t size, bool resizable = false) const override; + Storage storageFromBlob(void * data, int64_t size, const std::function & deleter) const override; + Storage storageWithAllocator(int64_t size, Allocator* allocator) const override; + Storage unsafeStorageFromTH(void * th_pointer, bool retain) const override; + Tensor unsafeTensorFromTH(void * th_pointer, bool retain) const override; + // example // virtual Tensor * add(Tensor & a, Tensor & b) = 0; ${type_method_declarations} diff --git a/aten/src/ATen/templates/TypeDerived.cpp b/aten/src/ATen/templates/TypeDerived.cpp index c24f7faf963a4..5b26400ea52f1 100644 --- a/aten/src/ATen/templates/TypeDerived.cpp +++ b/aten/src/ATen/templates/TypeDerived.cpp @@ -30,81 +30,19 @@ namespace at { -#if ${isCUDA} -static int getPointerDevice(void* ptr) { - struct cudaPointerAttributes attr; - THCudaCheck(cudaPointerGetAttributes(&attr, ptr)); - return attr.device; -} -#endif - ${Type}::${Type}() - : TypeDefault(${Backend}TensorId(), /*is_variable=*/false, /*is_undefined=*/false) {} + : ${DenseBackend}TypeDefault(${Backend}TensorId(), /*is_variable=*/false, /*is_undefined=*/false) {} ScalarType ${Type}::scalarType() const { return ScalarType::${ScalarName}; } Backend ${Type}::backend() const { return Backend::${Backend}; } -bool ${Type}::is_cuda() const { return backend() == Backend::CUDA || backend() == Backend::SparseCUDA; } -bool ${Type}::is_sparse() const { return backend() == Backend::SparseCPU || backend() == Backend::SparseCUDA; } -bool ${Type}::is_distributed() const { return false; } - -Storage ${Type}::storage(bool resizable) const { - return Storage( - ScalarType::${ScalarName}, - 0, -#if ${isCUDA} - globalContext().getTHCState()->cudaDeviceAllocator, -#else - getTHDefaultAllocator(), -#endif - resizable - ); -} -Storage ${Type}::storage(size_t size, bool resizable) const { - return Storage( - ScalarType::${ScalarName}, - size, -#if ${isCUDA} - globalContext().getTHCState()->cudaDeviceAllocator, -#else - getTHDefaultAllocator(), -#endif - resizable - ); -} -Storage ${Type}::storageFromBlob(void * data, int64_t size, const std::function & deleter) const { - return Storage( - ScalarType::${ScalarName}, - InefficientStdFunctionContext::makeDataPtr(data, deleter, -#if ${isCUDA} - Device(DeviceType::CUDA, getPointerDevice(data)) -#else - DeviceType::CPU -#endif - ), - size, - deleter); -} -Storage ${Type}::storageWithAllocator(int64_t size, Allocator* allocator) const { - return Storage(ScalarType::${ScalarName}, size, allocator); -} -Tensor ${Type}::unsafeTensorFromTH(void * th_pointer, bool retain) const { - return Tensor(static_cast(th_pointer), retain); -} -Storage ${Type}::unsafeStorageFromTH(void * th_pointer, bool retain) const { - if (retain) - ${THStorage}_retain(${state,} (${THStorage}*) th_pointer); - return Storage((${THStorage}*) th_pointer); -} -std::unique_ptr ${Type}::generator() const { - return std::unique_ptr(new ${Generator}(&at::globalContext())); -} const char * ${Type}::toString() const { - return ${Type}::typeString(); + return "${Type}"; } + TypeID ${Type}::ID() const { return ${TypeID}; } @@ -113,10 +51,6 @@ size_t ${Type}::elementSizeInBytes() const { return sizeof(${ScalarType}); } -const char * ${Type}::typeString() { - return "${Type}"; -} - /* example Tensor * ${Type}::add(Tensor & a, Tensor & b) { std::cout << "add ${Tensor}\n"; diff --git a/aten/src/ATen/templates/TypeDerived.h b/aten/src/ATen/templates/TypeDerived.h index 663a913da172c..be26dfc09a9b1 100644 --- a/aten/src/ATen/templates/TypeDerived.h +++ b/aten/src/ATen/templates/TypeDerived.h @@ -2,11 +2,13 @@ // ${generated_comment} -#include "ATen/TypeDefault.h" +#include "ATen/CPUTypeDefault.h" #include "ATen/Context.h" #include "ATen/TensorMethods.h" #include "ATen/CheckGenerator.h" +$extra_cuda_headers + #ifdef _MSC_VER #ifdef Type #undef Type @@ -15,24 +17,13 @@ namespace at { -struct ${Type} final : public TypeDefault { +struct ${Type} final : public ${DenseBackend}TypeDefault { explicit ${Type}(); virtual ScalarType scalarType() const override; virtual Backend backend() const override; - virtual bool is_cuda() const override; - virtual bool is_sparse() const override; - virtual bool is_distributed() const override; - virtual Storage storage(bool resizable = false) const override; - virtual Storage storage(size_t size, bool resizable = false) const override; - virtual Storage storageFromBlob(void * data, int64_t size, const std::function & deleter) const override; - virtual Storage storageWithAllocator(int64_t size, Allocator* allocator) const override; - virtual std::unique_ptr generator() const override; virtual const char * toString() const override; virtual size_t elementSizeInBytes() const override; virtual TypeID ID() const override; - static const char * typeString(); - virtual Storage unsafeStorageFromTH(void * th_pointer, bool retain) const override; - virtual Tensor unsafeTensorFromTH(void * th_pointer, bool retain) const override; // example // virtual Tensor * add(Tensor & a, Tensor & b) override; diff --git a/aten/src/ATen/test/native_test.cpp b/aten/src/ATen/test/native_test.cpp index fac85da04aa60..473c9ff12b52e 100644 --- a/aten/src/ATen/test/native_test.cpp +++ b/aten/src/ATen/test/native_test.cpp @@ -6,6 +6,8 @@ using namespace at; +using Catch::Matchers::StartsWith; + #define REQUIRE_EQUAL(t1, t2) \ REQUIRE(t1.equal(t2)); @@ -73,10 +75,10 @@ void test(Type & T, Type & AccT) { SECTION( "size / stride" ) { auto scalar = randn({}, T); - REQUIRE_THROWS_WITH(scalar.size(0), "dimension specified as 0 but tensor has no dimensions"); - REQUIRE_THROWS_WITH(scalar.size(-1), "dimension specified as -1 but tensor has no dimensions"); - REQUIRE_THROWS_WITH(scalar.stride(0), "dimension specified as 0 but tensor has no dimensions"); - REQUIRE_THROWS_WITH(scalar.stride(-1), "dimension specified as -1 but tensor has no dimensions"); + REQUIRE_THROWS_WITH(scalar.size(0), StartsWith("dimension specified as 0 but tensor has no dimensions")); + REQUIRE_THROWS_WITH(scalar.size(-1), StartsWith("dimension specified as -1 but tensor has no dimensions")); + REQUIRE_THROWS_WITH(scalar.stride(0), StartsWith("dimension specified as 0 but tensor has no dimensions")); + REQUIRE_THROWS_WITH(scalar.stride(-1), StartsWith("dimension specified as -1 but tensor has no dimensions")); auto empty = randn({0}, T); REQUIRE(empty.size(0) == 0); diff --git a/aten/src/ATen/test/verify_api_visibility.cpp b/aten/src/ATen/test/verify_api_visibility.cpp index ed296ce2ffe64..3384023acb74c 100644 --- a/aten/src/ATen/test/verify_api_visibility.cpp +++ b/aten/src/ATen/test/verify_api_visibility.cpp @@ -12,4 +12,8 @@ #error "AT_MKLDNN_ENABLED should not be visible in public headers" #endif +#ifdef CAFFE2_STATIC_LINK_CUDA +#error "CAFFE2_STATIC_LINK_CUDA should not be visible in public headers" +#endif + auto main() -> int {} diff --git a/aten/src/THC/THCGeneral.h.in b/aten/src/THC/THCGeneral.h.in index adbd5906d2165..d949e87f8967c 100644 --- a/aten/src/THC/THCGeneral.h.in +++ b/aten/src/THC/THCGeneral.h.in @@ -22,7 +22,7 @@ #endif #ifdef _WIN32 -# if defined(ATen_cuda_EXPORTS) || defined(caffe2_gpu_EXPORTS) || defined(CAFFE2_BUILD_MAIN_LIB) +# if defined(ATen_cuda_EXPORTS) || defined(caffe2_gpu_EXPORTS) || defined(CAFFE2_CUDA_BUILD_MAIN_LIB) # define THC_API THC_EXTERNC __declspec(dllexport) # define THC_CLASS __declspec(dllexport) # else diff --git a/aten/src/THC/THCSortUtils.cuh b/aten/src/THC/THCSortUtils.cuh index 518063a229c72..67b667e5afb06 100644 --- a/aten/src/THC/THCSortUtils.cuh +++ b/aten/src/THC/THCSortUtils.cuh @@ -142,7 +142,7 @@ bitonicSortKVInPlace(TensorInfo keys, IndexType keySliceStride, TensorInfo values, IndexType valueSliceStride, - const Comparator& comp) { + Comparator comp) { // Find the slice of the tensor that we are sorting const IndexType linearIndex = getLinearBlockId(); // Tiling the slices could have us be out of bounds, if there are a diff --git a/aten/src/THC/THCTensorRandom.cu b/aten/src/THC/THCTensorRandom.cu index c5ac9f1ab66a9..8eb580169cc95 100644 --- a/aten/src/THC/THCTensorRandom.cu +++ b/aten/src/THC/THCTensorRandom.cu @@ -6,12 +6,13 @@ #include "THCReduceApplyUtils.cuh" #include "THCTensorRandom.cuh" #include "THCGenerator.hpp" +#include "ATen/Config.h" + +#include "ATen/cuda/_curand_mtgp32_host.h" #include #include #include -#include -#include #define MAX_NUM_BLOCKS 200 #define BLOCK_SIZE 256 diff --git a/binaries/benchmark_helper.cc b/binaries/benchmark_helper.cc index c5becbb2b4021..255fa8ce4b2aa 100644 --- a/binaries/benchmark_helper.cc +++ b/binaries/benchmark_helper.cc @@ -68,7 +68,7 @@ bool backendCudaSet(const string& backend) { void setDeviceType(caffe2::NetDef* net_def, caffe2::DeviceType& run_dev) { for (int j = 0; j < net_def->op_size(); j++) { caffe2::OperatorDef* op = net_def->mutable_op(j); - op->mutable_device_option()->set_device_type(run_dev); + op->mutable_device_option()->set_device_type(caffe2::TypeToProto(run_dev)); } } diff --git a/binaries/core_overhead_benchmark_gpu.cc b/binaries/core_overhead_benchmark_gpu.cc index 5cb0a62797553..018880432d4b7 100644 --- a/binaries/core_overhead_benchmark_gpu.cc +++ b/binaries/core_overhead_benchmark_gpu.cc @@ -167,7 +167,7 @@ static void BM_OperatorCreationCPU(benchmark::State& state) { OperatorDef def; Workspace ws; def.set_type("DummyEmpty"); - def.mutable_device_option()->set_device_type(CPU); + def.mutable_device_option()->set_device_type(PROTO_CPU); while (state.KeepRunning()) { op = CreateOperator(def, &ws); } @@ -180,7 +180,7 @@ static void BM_OperatorCreationCUDA(benchmark::State& state) { OperatorDef def; Workspace ws; def.set_type("DummyEmpty"); - def.mutable_device_option()->set_device_type(CUDA); + def.mutable_device_option()->set_device_type(PROTO_CUDA); while (state.KeepRunning()) { op = CreateOperator(def, &ws); } diff --git a/binaries/print_registered_core_operators.cc b/binaries/print_registered_core_operators.cc index c76ea3eacac95..412ce88c44922 100644 --- a/binaries/print_registered_core_operators.cc +++ b/binaries/print_registered_core_operators.cc @@ -52,8 +52,8 @@ int main(int argc, char** argv) { for (const auto& pair : *caffe2::gDeviceTypeRegistry()) { std::cout << "Device type " << pair.first #ifndef CAFFE2_USE_LITE_PROTO - << " (" << caffe2::DeviceType_Name( - static_cast(pair.first)) + << " (" + << at::DeviceTypeName(static_cast(pair.first)) << ")" #endif << std::endl; diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index dea1a8e812a9b..2a338b85dff4d 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -65,11 +65,6 @@ endif() # ---[ Caffe2 build # Note: the folders that are being commented out have not been properly # addressed yet. -# TODO(orionr): Enable all of this for Windows DLL when we -# can figure out how to get it to build -if (MSVC AND BUILD_SHARED_LIBS) -add_subdirectory(utils) -else() add_subdirectory(proto) add_subdirectory(contrib) add_subdirectory(core) @@ -101,7 +96,6 @@ add_subdirectory(sgd) add_subdirectory(share) # add_subdirectory(test) # todo: use caffe2_gtest_main instead of gtest_main because we will need to call GlobalInit add_subdirectory(transforms) -endif() # Advanced: if we have white list specified, we will do intersections for all # main lib srcs. @@ -166,16 +160,8 @@ if (FALSE) endif() # ---[ List of libraries to link with -# TODO(orionr): Enable all of this for Windows DLL when we -# can figure out how to get it to build -if (NOT (MSVC AND BUILD_SHARED_LIBS)) add_library(caffe2_protos STATIC $) add_dependencies(caffe2_protos Caffe2_PROTO) -else() -# Do not include caffe2 or caffe protos, but rather have it only be -# a library to attach local protobuf. -add_library(caffe2_protos STATIC utils/dummy.cpp) -endif() # If we are going to link protobuf locally inside caffe2 libraries, what we will do is # to create a helper static library that always contains libprotobuf source files, and # link the caffe2 related dependent libraries to it. @@ -341,7 +327,7 @@ if(USE_CUDA) # NB: This must be target_compile_definitions, not target_compile_options, # as the latter is not respected by nvcc if (MSVC) - target_compile_definitions(caffe2_gpu PRIVATE "-DCAFFE2_BUILD_MAIN_LIB") + target_compile_definitions(caffe2_gpu PRIVATE "-DCAFFE2_CUDA_BUILD_MAIN_LIB") endif() # Set standard properties on the target @@ -401,9 +387,6 @@ if ($ENV{WERROR}) endif() # ---[ Test binaries. -# TODO(orionr): Enable all of this for Windows DLL when we -# can figure out how to get it to build -if (NOT (MSVC AND BUILD_SHARED_LIBS)) if (BUILD_TEST) set(Caffe2_ALL_TEST_SRCS ${Caffe2_CPU_TEST_SRCS}) if (USE_CUDA) @@ -414,6 +397,8 @@ if (BUILD_TEST) get_filename_component(test_name ${test_src} NAME_WE) add_executable(${test_name} "${test_src}") target_link_libraries(${test_name} ${Caffe2_MAIN_LIBS} gtest_main) + target_include_directories(${test_name} PRIVATE $) + target_include_directories(${test_name} PRIVATE ${Caffe2_CPU_INCLUDE}) if (${CMAKE_MAJOR_VERSION}.${CMAKE_MINOR_VERSION} GREATER 3.0) target_compile_features(${test_name} PRIVATE cxx_range_for) endif() @@ -429,6 +414,8 @@ if (BUILD_TEST) get_filename_component(test_name ${test_src} NAME_WE) hip_add_executable(${test_name} "${test_src}") target_link_libraries(${test_name} ${Caffe2_MAIN_LIBS} gtest_main) + target_include_directories(${test_name} PRIVATE $) + target_include_directories(${test_name} PRIVATE ${Caffe2_CPU_INCLUDE}) if (${CMAKE_MAJOR_VERSION}.${CMAKE_MINOR_VERSION} GREATER 3.0) target_compile_features(${test_name} PRIVATE cxx_range_for) endif() @@ -439,13 +426,13 @@ if (BUILD_TEST) endforeach() endif() endif() -endif() set(__aten_test_dir "test/aten") if (NOT USE_ROCM) foreach(test_src ${ATen_CPU_TEST_SRCS}) get_filename_component(test_name ${test_src} NAME_WE) add_executable(${test_name} "${test_src}") + target_include_directories(${test_name} PRIVATE $) target_include_directories(${test_name} PRIVATE ${Caffe2_CPU_INCLUDE}) target_include_directories(${test_name} SYSTEM PRIVATE ${Caffe2_DEPENDENCY_INCLUDE}) target_link_libraries(${test_name} ${Caffe2_MAIN_LIBS}) @@ -457,6 +444,7 @@ if (NOT USE_ROCM) foreach(test_src ${ATen_CUDA_TEST_SRCS}) get_filename_component(test_name ${test_src} NAME_WE) torch_cuda_based_add_executable(${test_name} "${test_src}") + target_include_directories(${test_name} PRIVATE $) target_include_directories(${test_name} PRIVATE ${Caffe2_CPU_INCLUDE}) target_include_directories(${test_name} SYSTEM PRIVATE ${Caffe2_DEPENDENCY_INCLUDE}) target_link_libraries(${test_name} ${Caffe2_MAIN_LIBS}) @@ -466,9 +454,6 @@ if (NOT USE_ROCM) endif() endif() -# TODO(orionr): Enable all of this for Windows DLL when we -# can figure out how to get it to build -if (NOT (MSVC AND BUILD_SHARED_LIBS)) if (BUILD_PYTHON) # Python site-packages # Get canonical directory for python site packages (relative to install @@ -525,6 +510,9 @@ if (BUILD_PYTHON) if (APPLE) set_target_properties(caffe2_pybind11_state PROPERTIES LINK_FLAGS "-undefined dynamic_lookup") endif() + target_include_directories(caffe2_pybind11_state PRIVATE $) + target_include_directories(caffe2_pybind11_state PRIVATE ${Caffe2_CPU_INCLUDE}) + target_link_libraries( caffe2_pybind11_state caffe2_library) if (WIN32) @@ -547,6 +535,8 @@ if (BUILD_PYTHON) if (APPLE) set_target_properties(caffe2_pybind11_state_gpu PROPERTIES LINK_FLAGS "-undefined dynamic_lookup") endif() + target_include_directories(caffe2_pybind11_state_gpu PRIVATE $) + target_include_directories(caffe2_pybind11_state_gpu PRIVATE ${Caffe2_CPU_INCLUDE}) target_link_libraries( caffe2_pybind11_state_gpu caffe2_library caffe2_gpu_library) if (WIN32) @@ -570,6 +560,8 @@ if (BUILD_PYTHON) if (APPLE) set_target_properties(caffe2_pybind11_state_hip PROPERTIES LINK_FLAGS "-undefined dynamic_lookup") endif() + target_include_directories(caffe2_pybind11_state_hip PRIVATE $) + target_include_directories(caffe2_pybind11_state_hip PRIVATE ${Caffe2_CPU_INCLUDE}) target_link_libraries( caffe2_pybind11_state_hip caffe2_library caffe2_hip_library) if (WIN32) @@ -657,7 +649,6 @@ if (BUILD_PYTHON) install(DIRECTORY ${CMAKE_BINARY_DIR}/caffe2 DESTINATION ${PYTHON_LIB_REL_PATH} FILES_MATCHING PATTERN "*.py") endif() -endif() # Finally, set the Caffe2_MAIN_LIBS variable in the parent scope. set(Caffe2_MAIN_LIBS ${Caffe2_MAIN_LIBS} PARENT_SCOPE) diff --git a/caffe2/contrib/nccl/cuda_nccl_op_gpu.cc b/caffe2/contrib/nccl/cuda_nccl_op_gpu.cc index 9722d5891334d..27be681f4f482 100644 --- a/caffe2/contrib/nccl/cuda_nccl_op_gpu.cc +++ b/caffe2/contrib/nccl/cuda_nccl_op_gpu.cc @@ -68,6 +68,47 @@ class NCCLAllreduceOp final : public Operator { } } + static std::vector ShapeInference( + const OperatorDef& def, + const std::vector& in) { + auto n_outputs = def.output_size(); + CAFFE_ENFORCE( + n_outputs == 1 || n_outputs == in.size(), + "NCCLAllreduce only supports N-1 or N-N reductions"); + + for (auto i = 0; i < in.size(); i++) { + CAFFE_ENFORCE( + in[0].dims_size() == in[i].dims_size(), + "NCCLAllreduce requires inputs of same dimension"); + for (auto j = 0; j < in[0].dims_size(); j++) { + CAFFE_ENFORCE( + in[0].dims(j) == in[i].dims(j), + "NCCLAllreduce requires inputs to be of same shape"); + } + } + + std::vector out(n_outputs); + for (auto i = 0; i < out.size(); i++) { + out[i] = in[0]; + } + return out; + } + + static struct OpSchema::Cost CostInference( + const OperatorDef& def, + const vector& inputs) { + CAFFE_ENFORCE_GE(inputs.size(), 1, "Conv requires at least 1 input"); + const TensorShape X0 = inputs[0]; + const auto nElem = nElemFromDim(inputs[0]); + + struct OpSchema::Cost c; + c.flops = (inputs.size() - 1) * nElem; + c.bytes_read = inputs.size() * nElem; + c.bytes_written = def.output_size() * nElem; + c.params_bytes = 0; + return c; + } + protected: }; @@ -173,6 +214,8 @@ REGISTER_CUDA_OPERATOR(NCCLAllreduce, NCCLAllreduceOp); OPERATOR_SCHEMA(NCCLAllreduce) .NumInputs(1, CAFFE2_COMPILE_TIME_MAX_GPUS) .NumOutputs(1, CAFFE2_COMPILE_TIME_MAX_GPUS) + .CostInferenceFunction(NCCLAllreduceOp::CostInference) + .TensorInferenceFunction(NCCLAllreduceOp::ShapeInference) .IdenticalTypeAndShape() .InputsCanCrossDevices() .AllowOneToOneInplace() diff --git a/caffe2/contrib/nervana/nervana_fc_op_gpu_test.cc b/caffe2/contrib/nervana/nervana_fc_op_gpu_test.cc index 012eea69c9dc6..972d9231dcf9c 100644 --- a/caffe2/contrib/nervana/nervana_fc_op_gpu_test.cc +++ b/caffe2/contrib/nervana/nervana_fc_op_gpu_test.cc @@ -19,7 +19,7 @@ namespace { static void AddConstInput(const std::vector& shape, const float value, const string& name, Workspace* ws) { DeviceOption option; - option.set_device_type(CUDA); + option.set_device_type(PROTO_CUDA); CUDAContext context(option); Blob* blob = ws->CreateBlob(name); auto* tensor = blob->GetMutableTensor(CUDA); @@ -43,7 +43,7 @@ TEST(NervanaFullyConnectedTest, Test) { def.add_input("W"); def.add_input("B"); def.add_output("Y"); - def.mutable_device_option()->set_device_type(CUDA); + def.mutable_device_option()->set_device_type(PROTO_CUDA); def.set_engine("NERVANA"); AddConstInput(std::vector{5, 10}, 1., "X", &ws); AddConstInput(std::vector{6, 10}, 1., "W", &ws); diff --git a/caffe2/contrib/opencl/context.h b/caffe2/contrib/opencl/context.h index 64426b1be80d5..ce788a39a7cdc 100644 --- a/caffe2/contrib/opencl/context.h +++ b/caffe2/contrib/opencl/context.h @@ -36,7 +36,7 @@ class OpenCLContext final { public: explicit OpenCLContext(); explicit OpenCLContext(const DeviceOption& option) { - DCHECK_EQ(option.device_type(), OPENCL); + DCHECK_EQ(option.device_type(), PROTO_OPENCL); OpenCLContext(); } ~OpenCLContext() {} diff --git a/caffe2/core/blob_gpu_test.cc b/caffe2/core/blob_gpu_test.cc index 24ee596b5238b..0f0fcb54906dc 100644 --- a/caffe2/core/blob_gpu_test.cc +++ b/caffe2/core/blob_gpu_test.cc @@ -193,7 +193,7 @@ TEST(TensorTest, TensorSerializationMultiDevices) { EXPECT_EQ(tensor_proto.float_data(i), i); } EXPECT_TRUE(tensor_proto.has_device_detail()); - EXPECT_EQ(tensor_proto.device_detail().device_type(), CUDA); + EXPECT_EQ(tensor_proto.device_detail().device_type(), PROTO_CUDA); EXPECT_EQ(tensor_proto.device_detail().cuda_gpu_id(), gpu_id); // Test if the restored blob is still of the same device. blob.Reset(); diff --git a/caffe2/core/blob_test.cc b/caffe2/core/blob_test.cc index 024ec836a0df7..97d17c6e5924f 100644 --- a/caffe2/core/blob_test.cc +++ b/caffe2/core/blob_test.cc @@ -883,7 +883,7 @@ TYPED_TEST(TypedTensorTest, BigTensorSerialization) { { DeviceOption option; - option.set_device_type(CPU); + option.set_device_type(PROTO_CPU); Argument db_type_arg = MakeArgument("db_type", "vector_db"); Argument absolute_path_arg = MakeArgument("absolute_path", true); Argument db_source_arg = MakeArgument("db", db_source); @@ -996,7 +996,7 @@ TEST(ContentChunks, Serialization) { { DeviceOption option; - option.set_device_type(CPU); + option.set_device_type(PROTO_CPU); Argument db_type_arg = MakeArgument("db_type", "vector_db"); Argument absolute_path_arg = MakeArgument("absolute_path", true); Argument db_source_arg = MakeArgument("db", db_source); diff --git a/caffe2/core/common_gpu.h b/caffe2/core/common_gpu.h index 35ba71854fc7c..4658cd0d75609 100644 --- a/caffe2/core/common_gpu.h +++ b/caffe2/core/common_gpu.h @@ -25,6 +25,42 @@ #include "caffe2/core/logging.h" #include "caffe2/core/common.h" +// Defines CAFFE2_CUDA_EXPORT and CAFFE2_CUDA_IMPORT. On Windows, this corresponds to +// different declarations (dllexport and dllimport). On Linux/Mac, it just +// resolves to the same "default visibility" setting. +#if defined(_MSC_VER) +#if defined(CAFFE2_BUILD_SHARED_LIBS) +#define CAFFE2_CUDA_EXPORT __declspec(dllexport) +#define CAFFE2_CUDA_IMPORT __declspec(dllimport) +#else +#define CAFFE2_CUDA_EXPORT +#define CAFFE2_CUDA_IMPORT +#endif +#else +#if defined(__GNUC__) +#define CAFFE2_CUDA_EXPORT __attribute__((__visibility__("default"))) +#else +#define CAFFE2_CUDA_EXPORT +#endif +#define CAFFE2_CUDA_IMPORT CAFFE2_CUDA_EXPORT +#endif + +// CAFFE2_CUDA_API is a macro that, depends on whether you are building the +// main caffe2 library or not, resolves to either CAFFE2_CUDA_EXPORT or +// CAFFE2_CUDA_IMPORT. +// +// This is used in e.g. Caffe2's protobuf files: when building the main library, +// it is defined as CAFFE2_CUDA_EXPORT to fix a Windows global-variable-in-dll +// issue, and for anyone dependent on Caffe2 it will be defined as +// CAFFE2_CUDA_IMPORT. + +#ifdef CAFFE2_CUDA_BUILD_MAIN_LIB +#define CAFFE2_CUDA_API CAFFE2_CUDA_EXPORT +#else +#define CAFFE2_CUDA_API CAFFE2_CUDA_IMPORT +#endif + + // This is a macro defined for cuda fp16 support. In default, cuda fp16 is // supported by NVCC 7.5, but it is also included in the Tegra X1 platform with // a (custom?) NVCC 7.0. As a result, we would normally just check the cuda diff --git a/caffe2/core/context.h b/caffe2/core/context.h index 657ad3db98dca..b12dc13e88657 100644 --- a/caffe2/core/context.h +++ b/caffe2/core/context.h @@ -48,7 +48,7 @@ class CAFFE2_API CPUContext final : public BaseContext { : random_seed_( option.has_random_seed() ? option.random_seed() : RandomNumberSeed()) { - CAFFE_ENFORCE_EQ(option.device_type(), CPU); + CAFFE_ENFORCE_EQ(option.device_type(), PROTO_CPU); } ~CPUContext() noexcept override {} diff --git a/caffe2/core/context_base.cc b/caffe2/core/context_base.cc index b6a68336e01f2..388a390d5aaa7 100644 --- a/caffe2/core/context_base.cc +++ b/caffe2/core/context_base.cc @@ -3,20 +3,20 @@ namespace caffe2 { // TODO: rename context.h -> context_cpu.h & context_base.h -> context.h -std::array& -GetStaticContexts() { - static std::array - static_contexts; +StaticContextMap& GetStaticContexts() { + static StaticContextMap static_contexts; return static_contexts; } -void set_static_context(int d, BaseStaticContext* ptr) { +void set_static_context(DeviceType t, BaseStaticContext* ptr) { auto& static_contexts = GetStaticContexts(); - static_contexts[d] = ptr; + static_contexts[t] = ptr; } -BaseStaticContext* get_static_context(int d) { - return GetStaticContexts()[d]; +BaseStaticContext* get_static_context(DeviceType t) { + auto* ptr = GetStaticContexts()[t]; + CAFFE_ENFORCE(ptr, "StaticContext is not registered yet."); + return ptr; } } // namespace caffe2 diff --git a/caffe2/core/context_base.h b/caffe2/core/context_base.h index e8e4ea24bf449..196bcb30ec79f 100644 --- a/caffe2/core/context_base.h +++ b/caffe2/core/context_base.h @@ -37,7 +37,7 @@ class CAFFE2_API BaseStaticContext { * current context and the a data pointer */ virtual void ExtractDeviceOption(DeviceOption* device, const void* /*data*/) { - device->set_device_type(GetDeviceType()); + device->set_device_type(TypeToProto(GetDeviceType())); } }; @@ -169,22 +169,21 @@ class CAFFE2_API BaseContext { } }; -CAFFE2_API std::array& -GetStaticContexts(); -CAFFE2_API void set_static_context(int d, BaseStaticContext* ptr); -CAFFE2_API BaseStaticContext* get_static_context(int d); +using StaticContextMap = CaffeMap; +CAFFE2_API StaticContextMap& GetStaticContexts(); +CAFFE2_API void set_static_context(DeviceType t, BaseStaticContext* ptr); +CAFFE2_API BaseStaticContext* get_static_context(DeviceType t); -template +template struct StaticContextFunctionRegisterer { explicit StaticContextFunctionRegisterer(BaseStaticContext* ptr) { - static_assert(d < COMPILE_TIME_MAX_DEVICE_TYPES, ""); - set_static_context(d, ptr); + set_static_context(t, ptr); } }; -#define REGISTER_STATIC_CONTEXT(d, f) \ +#define REGISTER_STATIC_CONTEXT(t, f) \ namespace { \ - static StaticContextFunctionRegisterer g_static_context_##d(f); \ + static StaticContextFunctionRegisterer g_static_context_##d(f); \ } } // namespace caffe2 diff --git a/caffe2/core/context_gpu.cu b/caffe2/core/context_gpu.cu index 81b16ac2bac00..1eaa579ee0cdb 100644 --- a/caffe2/core/context_gpu.cu +++ b/caffe2/core/context_gpu.cu @@ -257,7 +257,7 @@ CUDAContext::CUDAContext(const DeviceOption& option) option.has_random_seed() ? option.random_seed() : RandomNumberSeed()) { static Caffe2CudaInitializerHelper g_cuda_initializer_; - DCHECK_EQ(option.device_type(), CUDA); + DCHECK_EQ(option.device_type(), PROTO_CUDA); } // shared mutex to lock out alloc / free during NCCL launches diff --git a/caffe2/core/context_gpu.h b/caffe2/core/context_gpu.h index a8380fda5239a..3090ca57aedc3 100644 --- a/caffe2/core/context_gpu.h +++ b/caffe2/core/context_gpu.h @@ -33,7 +33,7 @@ enum class CudaMemoryPoolType { * * The memory pool is set up during caffe2's global initialization time. */ -CAFFE2_API CudaMemoryPoolType GetCudaMemoryPoolType(); +CAFFE2_CUDA_API CudaMemoryPoolType GetCudaMemoryPoolType(); /** * A struct to host thread-local cuda objects. @@ -44,7 +44,7 @@ CAFFE2_API CudaMemoryPoolType GetCudaMemoryPoolType(); * and deallocating these objects at the thread scope. This class is solely * used inside CUDAContext and should not be used externally. */ -class CAFFE2_API ThreadLocalCUDAObjects { +class CAFFE2_CUDA_API ThreadLocalCUDAObjects { friend class CUDAContext; private: @@ -135,9 +135,9 @@ class CAFFE2_API ThreadLocalCUDAObjects { #endif // CAFFE2_USE_CUDNN }; -CAFFE2_API BaseStaticContext* GetCUDAStaticContext(); +CAFFE2_CUDA_API BaseStaticContext* GetCUDAStaticContext(); -class CAFFE2_API CUDAContext final : public BaseContext { +class CAFFE2_CUDA_API CUDAContext final : public BaseContext { public: // The default cuda context constructor. explicit CUDAContext(const int gpu_id = -1); @@ -332,7 +332,7 @@ inline void CPUContext::CopyBytes( * GPU present during runtime, at global initialization time we will set * the CPU memory allocator to allocate pinned memory. */ -struct CAFFE2_API PinnedCPUAllocator final : CPUAllocator { +struct CAFFE2_CUDA_API PinnedCPUAllocator final : CPUAllocator { PinnedCPUAllocator() {} ~PinnedCPUAllocator() override {} std::pair New(size_t nbytes) override { @@ -381,7 +381,7 @@ struct CAFFE2_API PinnedCPUAllocator final : CPUAllocator { DefaultCPUAllocator baseAllocator_; }; -class CAFFE2_API CUDAStaticContext final : public BaseStaticContext { +class CAFFE2_CUDA_API CUDAStaticContext final : public BaseStaticContext { public: std::pair New(size_t nbytes) const override; @@ -403,7 +403,7 @@ class CAFFE2_API CUDAStaticContext final : public BaseStaticContext { } void ExtractDeviceOption(DeviceOption* device, const void* data) override { - device->set_device_type(GetDeviceType()); + device->set_device_type(TypeToProto(GetDeviceType())); device->set_cuda_gpu_id(GetGPUIDForPointer(data)); } diff --git a/caffe2/core/event.h b/caffe2/core/event.h index 5137d7e868033..9d725ad755cf2 100644 --- a/caffe2/core/event.h +++ b/caffe2/core/event.h @@ -1,13 +1,15 @@ #ifndef CAFFE2_CORE_EVENT_H_ #define CAFFE2_CORE_EVENT_H_ +#include #include "caffe2/core/common.h" #include "caffe2/core/logging.h" #include "caffe2/proto/caffe2_pb.h" namespace caffe2 { -constexpr int MaxDeviceTypes = DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES; +constexpr int MaxDeviceTypes = + DeviceTypeProto::PROTO_COMPILE_TIME_MAX_DEVICE_TYPES; class Event; enum EventStatus { @@ -65,20 +67,22 @@ class CAFFE2_API Event { ~Event() {} void Record( - int recorder_type, + DeviceType recorder_type, const void* context, const char* err_msg = nullptr) { + auto recorder_index = TypeToProto(recorder_type); CAFFE_ENFORCE_EQ( - recorder_type, + recorder_index, type_, "You are trying to record with a wrong device type."); - CAFFE_ENFORCE(event_recorder_[recorder_type]); - event_recorder_[recorder_type](this, context, err_msg); + CAFFE_ENFORCE(event_recorder_[recorder_index]); + event_recorder_[recorder_index](this, context, err_msg); } - void Wait(int waiter_type, void* context) const { - CAFFE_ENFORCE(event_waiter_[waiter_type][type_]); - event_waiter_[waiter_type][type_](this, context); + void Wait(DeviceType waiter_type, void* context) const { + auto waiter_index = TypeToProto(waiter_type); + CAFFE_ENFORCE(event_waiter_[waiter_index][type_]); + event_waiter_[waiter_index][type_](this, context); } void Finish() const { @@ -185,57 +189,57 @@ class CAFFE2_API Event { static EventSetCallbackFunction event_callback_setter_[MaxDeviceTypes]; - template + template friend struct EventCreateFunctionRegisterer; - template + template friend struct EventRecordFunctionRegisterer; - template + template friend struct EventWaitFunctionRegisterer; - template + template friend struct EventFinishFunctionRegisterer; - template + template friend struct EventQueryFunctionRegisterer; - template + template friend struct EventErrorMessageFunctionRegisterer; - template + template friend struct EventSetFinishedFunctionRegisterer; - template + template friend struct EventSetCallbackFunctionRegisterer; - template + template friend struct EventResetFunctionRegisterer; }; -template +template struct EventCreateFunctionRegisterer { explicit EventCreateFunctionRegisterer(EventCreateFunction f) { - static_assert(d < MaxDeviceTypes, ""); + auto d = TypeToProto(t); Event::event_creator_[d] = f; } }; -#define REGISTER_EVENT_CREATE_FUNCTION(d, f) \ +#define REGISTER_EVENT_CREATE_FUNCTION(t, f) \ namespace { \ - static EventCreateFunctionRegisterer g_event_create_##d(f); \ + static EventCreateFunctionRegisterer g_event_create_##d(f); \ } -template +template struct EventRecordFunctionRegisterer { explicit EventRecordFunctionRegisterer(EventRecordFunction f) { - static_assert(d < MaxDeviceTypes, ""); + auto d = TypeToProto(t); Event::event_recorder_[d] = f; } }; -#define REGISTER_EVENT_RECORD_FUNCTION(d, f) \ +#define REGISTER_EVENT_RECORD_FUNCTION(t, f) \ namespace { \ - static EventRecordFunctionRegisterer g_event_record_##d(f); \ + static EventRecordFunctionRegisterer g_event_record_##d(f); \ } -template +template struct EventWaitFunctionRegisterer { explicit EventWaitFunctionRegisterer(EventWaitFunction f) { - static_assert(waiter_type < MaxDeviceTypes, ""); - static_assert(event_type < MaxDeviceTypes, ""); - Event::event_waiter_[waiter_type][event_type] = f; + auto waiter_index = TypeToProto(waiter_type); + auto event_index = TypeToProto(event_type); + Event::event_waiter_[waiter_index][event_index] = f; } }; #define REGISTER_EVENT_WAIT_FUNCTION(w, d, f) \ @@ -243,76 +247,76 @@ struct EventWaitFunctionRegisterer { static EventWaitFunctionRegisterer g_event_wait_##w##_##d(f); \ } -template +template struct EventQueryFunctionRegisterer { explicit EventQueryFunctionRegisterer(EventQueryFunction f) { - static_assert(d < MaxDeviceTypes, ""); + auto d = TypeToProto(t); Event::event_querier_[d] = f; } }; -#define REGISTER_EVENT_QUERY_FUNCTION(d, f) \ +#define REGISTER_EVENT_QUERY_FUNCTION(t, f) \ namespace { \ - static EventQueryFunctionRegisterer g_event_query_##d(f); \ + static EventQueryFunctionRegisterer g_event_query_##d(f); \ } -template +template struct EventErrorMessageFunctionRegisterer { explicit EventErrorMessageFunctionRegisterer(EventErrorMessageFunction f) { - static_assert(d < MaxDeviceTypes, ""); + auto d = TypeToProto(t); Event::event_err_msg_getter_[d] = f; } }; -#define REGISTER_EVENT_ERROR_MESSAGE_FUNCTION(d, f) \ +#define REGISTER_EVENT_ERROR_MESSAGE_FUNCTION(t, f) \ namespace { \ - static EventErrorMessageFunctionRegisterer g_event_err_msg_##d(f); \ + static EventErrorMessageFunctionRegisterer g_event_err_msg_##d(f); \ } -template +template struct EventSetFinishedFunctionRegisterer { explicit EventSetFinishedFunctionRegisterer(EventSetFinishedFunction f) { - static_assert(d < MaxDeviceTypes, ""); + auto d = TypeToProto(t); Event::event_finished_setter_[d] = f; } }; -#define REGISTER_EVENT_SET_FINISHED_FUNCTION(d, f) \ +#define REGISTER_EVENT_SET_FINISHED_FUNCTION(t, f) \ namespace { \ - static EventSetFinishedFunctionRegisterer g_event_set_finished_##d(f); \ + static EventSetFinishedFunctionRegisterer g_event_set_finished_##d(f); \ } -template +template struct EventSetCallbackFunctionRegisterer { explicit EventSetCallbackFunctionRegisterer(EventSetCallbackFunction f) { - static_assert(d < MaxDeviceTypes, ""); + auto d = TypeToProto(t); Event::event_callback_setter_[d] = f; } }; -#define REGISTER_EVENT_SET_CALLBACK_FUNCTION(d, f) \ +#define REGISTER_EVENT_SET_CALLBACK_FUNCTION(t, f) \ namespace { \ - static EventSetCallbackFunctionRegisterer g_event_set_callback_##d(f); \ + static EventSetCallbackFunctionRegisterer g_event_set_callback_##d(f); \ } -template +template struct EventFinishFunctionRegisterer { explicit EventFinishFunctionRegisterer(EventFinishFunction f) { - static_assert(d < MaxDeviceTypes, ""); + auto d = TypeToProto(t); Event::event_finisher_[d] = f; } }; -#define REGISTER_EVENT_FINISH_FUNCTION(d, f) \ +#define REGISTER_EVENT_FINISH_FUNCTION(t, f) \ namespace { \ - static EventFinishFunctionRegisterer g_event_finish_##d(f); \ + static EventFinishFunctionRegisterer g_event_finish_##d(f); \ } -template +template struct EventResetFunctionRegisterer { explicit EventResetFunctionRegisterer(EventResetFunction f) { - static_assert(d < MaxDeviceTypes, ""); + auto d = TypeToProto(t); Event::event_resetter_[d] = f; } }; -#define REGISTER_EVENT_RESET_FUNCTION(d, f) \ +#define REGISTER_EVENT_RESET_FUNCTION(t, f) \ namespace { \ - static EventResetFunctionRegisterer g_event_reset_##d(f); \ + static EventResetFunctionRegisterer g_event_reset_##d(f); \ } } // namespace caffe2 diff --git a/caffe2/core/event_cpu.h b/caffe2/core/event_cpu.h index 130841c10e72d..95143524635f1 100644 --- a/caffe2/core/event_cpu.h +++ b/caffe2/core/event_cpu.h @@ -9,8 +9,9 @@ struct CPUEventWrapper { explicit CPUEventWrapper(const DeviceOption& option) : status_(EventStatus::EVENT_INITIALIZED) { CAFFE_ENFORCE( - option.device_type() == CPU || option.device_type() == MKLDNN || - option.device_type() == IDEEP, + option.device_type() == PROTO_CPU || + option.device_type() == PROTO_MKLDNN || + option.device_type() == PROTO_IDEEP, "Expected CPU/MKLDNN/IDEEP device type"); } ~CPUEventWrapper() {} diff --git a/caffe2/core/event_gpu.cc b/caffe2/core/event_gpu.cc index f61fac45e0c55..6253ca19c9ab7 100644 --- a/caffe2/core/event_gpu.cc +++ b/caffe2/core/event_gpu.cc @@ -11,7 +11,7 @@ struct CudaEventWrapper { : cuda_stream_(nullptr), cuda_gpu_id_(option.cuda_gpu_id()), status_(EventStatus::EVENT_INITIALIZED) { - CAFFE_ENFORCE(option.device_type(), CUDA); + CAFFE_ENFORCE(option.device_type(), PROTO_CUDA); DeviceGuard g(cuda_gpu_id_); CUDA_ENFORCE(cudaEventCreate( &cuda_event_, cudaEventDefault | cudaEventDisableTiming)); diff --git a/caffe2/core/event_gpu_test.cc b/caffe2/core/event_gpu_test.cc index ec83d9818e2e5..18fe152198e23 100644 --- a/caffe2/core/event_gpu_test.cc +++ b/caffe2/core/event_gpu_test.cc @@ -9,9 +9,9 @@ TEST(EventCUDATest, EventBasics) { if (!HasCudaGPU()) return; DeviceOption device_cpu; - device_cpu.set_device_type(CPU); + device_cpu.set_device_type(PROTO_CPU); DeviceOption device_cuda; - device_cuda.set_device_type(CUDA); + device_cuda.set_device_type(PROTO_CUDA); CPUContext context_cpu(device_cpu); CUDAContext context_cuda(device_cuda); diff --git a/caffe2/core/event_test.cc b/caffe2/core/event_test.cc index 02994ec6c59b7..2c18693017aa3 100644 --- a/caffe2/core/event_test.cc +++ b/caffe2/core/event_test.cc @@ -6,7 +6,7 @@ namespace caffe2 { TEST(EventCPUTest, EventBasics) { DeviceOption device_option; - device_option.set_device_type(CPU); + device_option.set_device_type(PROTO_CPU); Event event(device_option); CPUContext context; diff --git a/caffe2/core/hip/context_hip.cc b/caffe2/core/hip/context_hip.cc index 45446d75aef57..0fabb20a642c9 100644 --- a/caffe2/core/hip/context_hip.cc +++ b/caffe2/core/hip/context_hip.cc @@ -263,7 +263,7 @@ HIPContext::HIPContext(const DeviceOption& option) option.has_random_seed() ? option.random_seed() : RandomNumberSeed()) { static Caffe2HipInitializerHelper g_hip_initializer_; - DCHECK_EQ(option.device_type(), HIP); + DCHECK_EQ(option.device_type(), PROTO_HIP); } // shared mutex to lock out alloc / free during NCCL launches diff --git a/caffe2/core/hip/context_hip.h b/caffe2/core/hip/context_hip.h index e62c26b86b04d..cd309e6473be4 100644 --- a/caffe2/core/hip/context_hip.h +++ b/caffe2/core/hip/context_hip.h @@ -392,7 +392,7 @@ class HIPStaticContext final : public BaseStaticContext { } void ExtractDeviceOption(DeviceOption* device, const void* data) override { - device->set_device_type(GetDeviceType()); + device->set_device_type(TypeToProto(GetDeviceType())); device->set_hip_gpu_id(GetGPUIDForPointer(data)); } diff --git a/caffe2/core/hip/event_hip.cc b/caffe2/core/hip/event_hip.cc index 4d2ace0029e7a..6f0db4642ddbb 100644 --- a/caffe2/core/hip/event_hip.cc +++ b/caffe2/core/hip/event_hip.cc @@ -13,9 +13,10 @@ struct HipEventWrapper hip_gpu_id_(option.hip_gpu_id()), status_(EventStatus::EVENT_INITIALIZED) { - CAFFE_ENFORCE(option.device_type(), HIP); - DeviceGuard g(hip_gpu_id_); - HIP_ENFORCE(hipEventCreate(&hip_event_ /*, hipEventDefault | hipEventDisableTiming*/)); + CAFFE_ENFORCE(option.device_type(), PROTO_HIP); + DeviceGuard g(hip_gpu_id_); + HIP_ENFORCE(hipEventCreate( + &hip_event_ /*, hipEventDefault | hipEventDisableTiming*/)); } ~HipEventWrapper() { diff --git a/caffe2/core/hip/net_async_dag_hip.cc b/caffe2/core/hip/net_async_dag_hip.cc index be50c4078c84a..fa35b2a8c2161 100644 --- a/caffe2/core/hip/net_async_dag_hip.cc +++ b/caffe2/core/hip/net_async_dag_hip.cc @@ -88,20 +88,18 @@ AsyncDAGNet::AsyncDAGNet(const std::shared_ptr& net_def, Workspace int AsyncDAGNet::stream(const DeviceOption& device_option) { int stream_id = 0; - if(device_option.device_type() == HIP) - { - int gpu_id = device_option.hip_gpu_id(); - CAFFE_ENFORCE_GE(gpu_id, 0, "Invalid gpu id: " + caffe2::to_string(gpu_id)); - if((unsigned)gpu_id >= stream_counters_.size()) - { - stream_counters_.resize(gpu_id + 1, 0); - } - do - { - stream_id = stream_counters_[gpu_id]++; - stream_counters_[gpu_id] %= FLAGS_caffe2_streams_per_gpu; - } while(FLAGS_caffe2_net_async_check_stream_status && - !HIPContext::IsStreamFree(device_option, stream_id)); + if (device_option.device_type() == PROTO_HIP) { + int gpu_id = device_option.hip_gpu_id(); + CAFFE_ENFORCE_GE( + gpu_id, 0, "Invalid gpu id: " + caffe2::to_string(gpu_id)); + if ((unsigned)gpu_id >= stream_counters_.size()) { + stream_counters_.resize(gpu_id + 1, 0); + } + do { + stream_id = stream_counters_[gpu_id]++; + stream_counters_[gpu_id] %= FLAGS_caffe2_streams_per_gpu; + } while (FLAGS_caffe2_net_async_check_stream_status && + !HIPContext::IsStreamFree(device_option, stream_id)); } return stream_id; } diff --git a/caffe2/core/net_async_base.cc b/caffe2/core/net_async_base.cc index 164cb69076247..b40a8fa33778a 100644 --- a/caffe2/core/net_async_base.cc +++ b/caffe2/core/net_async_base.cc @@ -136,13 +136,13 @@ TaskThreadPool* AsyncNetBase::pool_getter( TaskThreadPool* AsyncNetBase::pool(const DeviceOption& device_option) { if (use_single_pool_) { - return pool_getter(cpu_pools_, CPU, -1, num_workers_); + return pool_getter(cpu_pools_, PROTO_CPU, -1, num_workers_); } static const std::unordered_set cpu_types{ - CPU, - MKLDNN, - IDEEP, - ONLY_FOR_TEST, + PROTO_CPU, + PROTO_MKLDNN, + PROTO_IDEEP, + PROTO_ONLY_FOR_TEST, }; if (cpu_types.find(device_option.device_type()) != cpu_types.end()) { auto numa_node_id = -1; @@ -155,13 +155,13 @@ TaskThreadPool* AsyncNetBase::pool(const DeviceOption& device_option) { FLAGS_caffe2_net_async_max_numa_nodes, "Invalid NUMA node id: ", numa_node_id); - return pool_getter(cpu_pools_, CPU, numa_node_id, num_workers_); - } else if (device_option.device_type() == CUDA) { + return pool_getter(cpu_pools_, PROTO_CPU, numa_node_id, num_workers_); + } else if (device_option.device_type() == PROTO_CUDA) { auto gpu_id = device_option.cuda_gpu_id(); CAFFE_ENFORCE( gpu_id >= 0 && gpu_id < FLAGS_caffe2_net_async_max_gpus, "Invalid GPU id: " + caffe2::to_string(gpu_id)); - return pool_getter(gpu_pools_, CUDA, gpu_id, num_workers_); + return pool_getter(gpu_pools_, PROTO_CUDA, gpu_id, num_workers_); } else { CAFFE_THROW( "Unsupported device type " + @@ -172,7 +172,7 @@ TaskThreadPool* AsyncNetBase::pool(const DeviceOption& device_option) { int AsyncNetBase::stream(int task_id) { const auto& device_option = event(task_id).GetDeviceOption(); int stream_id = 0; - if (device_option.device_type() == CUDA) { + if (device_option.device_type() == PROTO_CUDA) { int gpu_id = device_option.cuda_gpu_id(); CAFFE_ENFORCE_GE(gpu_id, 0, "Invalid gpu id: " + caffe2::to_string(gpu_id)); if ((unsigned)gpu_id >= getStreamCounters().size()) { diff --git a/caffe2/core/net_async_dag_gpu.cc b/caffe2/core/net_async_dag_gpu.cc index a6e6ad79a362c..225337d1452b9 100644 --- a/caffe2/core/net_async_dag_gpu.cc +++ b/caffe2/core/net_async_dag_gpu.cc @@ -111,7 +111,7 @@ AsyncDAGNet::AsyncDAGNet( int AsyncDAGNet::stream(const DeviceOption& device_option) { int stream_id = 0; - if (device_option.device_type() == CUDA) { + if (device_option.device_type() == PROTO_CUDA) { int gpu_id = device_option.cuda_gpu_id(); CAFFE_ENFORCE_GE(gpu_id, 0, "Invalid gpu id: " + caffe2::to_string(gpu_id)); if ((unsigned)gpu_id >= stream_counters_.size()) { diff --git a/caffe2/core/net_async_polling.cc b/caffe2/core/net_async_polling.cc index 67d2c6422d6ae..94d28edd97ef5 100644 --- a/caffe2/core/net_async_polling.cc +++ b/caffe2/core/net_async_polling.cc @@ -16,10 +16,9 @@ AsyncPollingNet::AsyncPollingNet( task_timers_[task_id] = caffe2::make_unique(); } - stats_.reserve(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES); - for (auto device_idx = 0; - device_idx < DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES; - ++device_idx) { + auto MAX_DEVICE_TYPES = DeviceTypeProto::PROTO_COMPILE_TIME_MAX_DEVICE_TYPES; + stats_.reserve(MAX_DEVICE_TYPES); + for (auto device_idx = 0; device_idx < MAX_DEVICE_TYPES; ++device_idx) { stats_.emplace_back( "async_net/stats/" + net_def->name() + "/" + caffe2::DeviceTypeName(device_idx)); @@ -38,7 +37,7 @@ bool AsyncPollingNet::DoRunAsync() { Timer timer; bool success = pollAndSchedule(); if (FLAGS_caffe2_dag_net_collect_stats) { - CAFFE_EVENT(stats_[CPU], poll_time_ms, timer.MilliSeconds()); + CAFFE_EVENT(stats_[PROTO_CPU], poll_time_ms, timer.MilliSeconds()); } if (!success) { finalizeEvents(); @@ -133,7 +132,7 @@ bool AsyncPollingNet::pollAndSchedule() { } if (FLAGS_caffe2_dag_net_collect_stats) { CAFFE_EVENT( - stats_[CPU], poll_status_update_time_us, timer.MicroSeconds()); + stats_[PROTO_CPU], poll_status_update_time_us, timer.MicroSeconds()); } std::unordered_set visited_children; diff --git a/caffe2/core/net_dag.cc b/caffe2/core/net_dag.cc index 6234edd026446..e7d079db7bc5e 100644 --- a/caffe2/core/net_dag.cc +++ b/caffe2/core/net_dag.cc @@ -72,10 +72,10 @@ DAGNetBase::DAGNetBase( task_timers_[idx] = caffe2::make_unique(); } } - stats_.reserve(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES); - for (auto device_idx = 0; - device_idx < DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES; - ++device_idx) { + constexpr auto MAX_DEVICE_TYPES = + DeviceTypeProto::PROTO_COMPILE_TIME_MAX_DEVICE_TYPES; + stats_.reserve(MAX_DEVICE_TYPES); + for (auto device_idx = 0; device_idx < MAX_DEVICE_TYPES; ++device_idx) { stats_.emplace_back( "dag_net/stats/" + net_def->name() + "/" + caffe2::DeviceTypeName(device_idx)); diff --git a/caffe2/core/net_gpu_test.cc b/caffe2/core/net_gpu_test.cc index 8bc82b53f5b54..eaea9377f9bca 100644 --- a/caffe2/core/net_gpu_test.cc +++ b/caffe2/core/net_gpu_test.cc @@ -34,11 +34,11 @@ class NetTestDummyOp final : public OperatorBase { // Simulate CUDA operator behavior bool HasAsyncPart() const override { - return debug_def().device_option().device_type() == CUDA; + return debug_def().device_option().device_type() == PROTO_CUDA; } bool SupportsAsyncScheduling() const override { - return debug_def().device_option().device_type() == CUDA; + return debug_def().device_option().device_type() == PROTO_CUDA; } protected: diff --git a/caffe2/core/net_test.cc b/caffe2/core/net_test.cc index 1b9397038e9be..43eb6562a4769 100644 --- a/caffe2/core/net_test.cc +++ b/caffe2/core/net_test.cc @@ -36,11 +36,11 @@ class NetTestDummyOp final : public OperatorBase { // Simulate CUDA operator behavior bool HasAsyncPart() const override { - return debug_def().device_option().device_type() == CUDA; + return debug_def().device_option().device_type() == PROTO_CUDA; } bool SupportsAsyncScheduling() const override { - return debug_def().device_option().device_type() == CUDA; + return debug_def().device_option().device_type() == PROTO_CUDA; } protected: diff --git a/caffe2/core/nomnigraph/CMakeLists.txt b/caffe2/core/nomnigraph/CMakeLists.txt index fcde4b7f0f092..8ac0a52871314 100644 --- a/caffe2/core/nomnigraph/CMakeLists.txt +++ b/caffe2/core/nomnigraph/CMakeLists.txt @@ -1,31 +1,17 @@ # ---[ CPU files. file(GLOB_RECURSE NOMNI_SRCS *.cc) -file(GLOB_RECURSE NOMNI_TEST_SRCS *test.cc) +file(GLOB_RECURSE NOMNI_TEST_SRCS *Test.cc) exclude(NOMNI_SRCS "${NOMNI_SRCS}" "${NOMNI_TEST_SRCS}") -# TODO(orionr): The nomnigraph source should likely just be included -# in the Caffe2 source list, since this won't live separately -add_library(nomnigraph STATIC "${NOMNI_SRCS}") -target_compile_options(nomnigraph PRIVATE "-DCAFFE2_BUILD_MAIN_LIB") -add_dependencies(nomnigraph Caffe2_PROTO) +list(APPEND Caffe2_CPU_SRCS ${NOMNI_SRCS}) +list(APPEND Caffe2_CPU_INCLUDE ${CMAKE_CURRENT_SOURCE_DIR}/include) +list(APPEND Caffe2_CPU_TEST_SRCS ${NOMNI_TEST_SRCS}) -target_include_directories(nomnigraph PUBLIC - $ - $) -list(APPEND Caffe2_PUBLIC_DEPENDENCY_LIBS nomnigraph) -set(Caffe2_PUBLIC_DEPENDENCY_LIBS ${Caffe2_PUBLIC_DEPENDENCY_LIBS} PARENT_SCOPE) - -install(TARGETS nomnigraph EXPORT Caffe2Targets DESTINATION lib) -install(DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/include +install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/include DESTINATION include FILES_MATCHING PATTERN "*.h") -if (BUILD_TEST) - foreach(test_src ${NOMNI_TEST_SRCS}) - get_filename_component(test_name ${test_src} NAME_WE) - add_executable(${test_name} "${test_src}") - target_link_libraries(${test_name} nomnigraph gtest_main) - add_test(NAME ${test_name} COMMAND $) - install(TARGETS ${test_name} DESTINATION test) - endforeach() -endif() +# ---[ Send the lists to the parent scope. +set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} PARENT_SCOPE) +set(Caffe2_CPU_INCLUDE ${Caffe2_CPU_INCLUDE} PARENT_SCOPE) +set(Caffe2_CPU_TEST_SRCS ${Caffe2_CPU_TEST_SRCS} PARENT_SCOPE) diff --git a/caffe2/core/nomnigraph/include/nomnigraph/Generated/OpClasses.h b/caffe2/core/nomnigraph/include/nomnigraph/Generated/OpClasses.h index a6362066f15f4..6516722fc26ec 100644 --- a/caffe2/core/nomnigraph/include/nomnigraph/Generated/OpClasses.h +++ b/caffe2/core/nomnigraph/include/nomnigraph/Generated/OpClasses.h @@ -18,62 +18,62 @@ class Conv : public NeuralNetOperator { int group = 1, vector dilations = {1, 1}) : NeuralNetOperator(NNKind::Conv), - KernelShape(kernelShape), - Pads(pads), - Strides(strides), - Group(group), - Dilations(dilations) {} + kernelShape_(kernelShape), + pads_(pads), + strides_(strides), + group_(group), + dilations_(dilations) {} ~Conv() {} NOMNIGRAPH_DEFINE_NN_RTTI(Conv); vector getKernelShape() const { - return KernelShape; + return kernelShape_; } vector getPads() const { - return Pads; + return pads_; } vector getStrides() const { - return Strides; + return strides_; } int getGroup() const { - return Group; + return group_; } vector getDilations() const { - return Dilations; + return dilations_; } void setKernelShape(vector kernelShape) { - KernelShape = kernelShape; + kernelShape_ = kernelShape; } void setPads(vector pads) { - Pads = pads; + pads_ = pads; } void setStrides(vector strides) { - Strides = strides; + strides_ = strides; } void setGroup(int group) { - Group = group; + group_ = group; } void setDilations(vector dilations) { - Dilations = dilations; + dilations_ = dilations; } private: - vector KernelShape; - vector Pads; - vector Strides; - int Group; - vector Dilations; + vector kernelShape_; + vector pads_; + vector strides_; + int group_; + vector dilations_; }; class ConvRelu : public NeuralNetOperator { @@ -85,70 +85,70 @@ class ConvRelu : public NeuralNetOperator { int group = 1, vector dilations = {1, 1}) : NeuralNetOperator(NNKind::ConvRelu), - KernelShape(kernelShape), - Pads(pads), - Strides(strides), - Group(group), - Dilations(dilations) {} + kernelShape_(kernelShape), + pads_(pads), + strides_(strides), + group_(group), + dilations_(dilations) {} ConvRelu(const Conv& conv) : NeuralNetOperator(NNKind::ConvRelu), - KernelShape(conv.getKernelShape()), - Pads(conv.getPads()), - Strides(conv.getStrides()), - Group(conv.getGroup()), - Dilations(conv.getDilations()) {} + kernelShape_(conv.getKernelShape()), + pads_(conv.getPads()), + strides_(conv.getStrides()), + group_(conv.getGroup()), + dilations_(conv.getDilations()) {} ~ConvRelu() {} NOMNIGRAPH_DEFINE_NN_RTTI(ConvRelu); vector getKernelShape() const { - return KernelShape; + return kernelShape_; } vector getPads() const { - return Pads; + return pads_; } vector getStrides() const { - return Strides; + return strides_; } int getGroup() const { - return Group; + return group_; } vector getDilations() const { - return Dilations; + return dilations_; } void setKernelShape(vector kernelShape) { - KernelShape = kernelShape; + kernelShape_ = kernelShape; } void setPads(vector pads) { - Pads = pads; + pads_ = pads; } void setStrides(vector strides) { - Strides = strides; + strides_ = strides; } void setGroup(int group) { - Group = group; + group_ = group; } void setDilations(vector dilations) { - Dilations = dilations; + dilations_ = dilations; } private: - vector KernelShape; - vector Pads; - vector Strides; - int Group; - vector Dilations; + vector kernelShape_; + vector pads_; + vector strides_; + int group_; + vector dilations_; }; class ConvTranspose : public NeuralNetOperator { @@ -160,62 +160,62 @@ class ConvTranspose : public NeuralNetOperator { int group = 1, vector dilations = {1, 1}) : NeuralNetOperator(NNKind::ConvTranspose), - KernelShape(kernelShape), - Pads(pads), - Strides(strides), - Group(group), - Dilations(dilations) {} + kernelShape_(kernelShape), + pads_(pads), + strides_(strides), + group_(group), + dilations_(dilations) {} ~ConvTranspose() {} NOMNIGRAPH_DEFINE_NN_RTTI(ConvTranspose); vector getKernelShape() const { - return KernelShape; + return kernelShape_; } vector getPads() const { - return Pads; + return pads_; } vector getStrides() const { - return Strides; + return strides_; } int getGroup() const { - return Group; + return group_; } vector getDilations() const { - return Dilations; + return dilations_; } void setKernelShape(vector kernelShape) { - KernelShape = kernelShape; + kernelShape_ = kernelShape; } void setPads(vector pads) { - Pads = pads; + pads_ = pads; } void setStrides(vector strides) { - Strides = strides; + strides_ = strides; } void setGroup(int group) { - Group = group; + group_ = group; } void setDilations(vector dilations) { - Dilations = dilations; + dilations_ = dilations; } private: - vector KernelShape; - vector Pads; - vector Strides; - int Group; - vector Dilations; + vector kernelShape_; + vector pads_; + vector strides_; + int group_; + vector dilations_; }; class AveragePool : public NeuralNetOperator { @@ -225,42 +225,42 @@ class AveragePool : public NeuralNetOperator { vector pads = {0, 0}, vector strides = {1, 1}) : NeuralNetOperator(NNKind::AveragePool), - KernelShape(kernelShape), - Pads(pads), - Strides(strides) {} + kernelShape_(kernelShape), + pads_(pads), + strides_(strides) {} ~AveragePool() {} NOMNIGRAPH_DEFINE_NN_RTTI(AveragePool); vector getKernelShape() const { - return KernelShape; + return kernelShape_; } vector getPads() const { - return Pads; + return pads_; } vector getStrides() const { - return Strides; + return strides_; } void setKernelShape(vector kernelShape) { - KernelShape = kernelShape; + kernelShape_ = kernelShape; } void setPads(vector pads) { - Pads = pads; + pads_ = pads; } void setStrides(vector strides) { - Strides = strides; + strides_ = strides; } private: - vector KernelShape; - vector Pads; - vector Strides; + vector kernelShape_; + vector pads_; + vector strides_; }; class AveragePoolRelu : public NeuralNetOperator { @@ -270,48 +270,48 @@ class AveragePoolRelu : public NeuralNetOperator { vector pads = {0, 0}, vector strides = {1, 1}) : NeuralNetOperator(NNKind::AveragePoolRelu), - KernelShape(kernelShape), - Pads(pads), - Strides(strides) {} + kernelShape_(kernelShape), + pads_(pads), + strides_(strides) {} AveragePoolRelu(const AveragePool& averagePool) : NeuralNetOperator(NNKind::AveragePoolRelu), - KernelShape(averagePool.getKernelShape()), - Pads(averagePool.getPads()), - Strides(averagePool.getStrides()) {} + kernelShape_(averagePool.getKernelShape()), + pads_(averagePool.getPads()), + strides_(averagePool.getStrides()) {} ~AveragePoolRelu() {} NOMNIGRAPH_DEFINE_NN_RTTI(AveragePoolRelu); vector getKernelShape() const { - return KernelShape; + return kernelShape_; } vector getPads() const { - return Pads; + return pads_; } vector getStrides() const { - return Strides; + return strides_; } void setKernelShape(vector kernelShape) { - KernelShape = kernelShape; + kernelShape_ = kernelShape; } void setPads(vector pads) { - Pads = pads; + pads_ = pads; } void setStrides(vector strides) { - Strides = strides; + strides_ = strides; } private: - vector KernelShape; - vector Pads; - vector Strides; + vector kernelShape_; + vector pads_; + vector strides_; }; class MaxPool : public NeuralNetOperator { @@ -321,42 +321,42 @@ class MaxPool : public NeuralNetOperator { vector pads = {0, 0}, vector strides = {1, 1}) : NeuralNetOperator(NNKind::MaxPool), - KernelShape(kernelShape), - Pads(pads), - Strides(strides) {} + kernelShape_(kernelShape), + pads_(pads), + strides_(strides) {} ~MaxPool() {} NOMNIGRAPH_DEFINE_NN_RTTI(MaxPool); vector getKernelShape() const { - return KernelShape; + return kernelShape_; } vector getPads() const { - return Pads; + return pads_; } vector getStrides() const { - return Strides; + return strides_; } void setKernelShape(vector kernelShape) { - KernelShape = kernelShape; + kernelShape_ = kernelShape; } void setPads(vector pads) { - Pads = pads; + pads_ = pads; } void setStrides(vector strides) { - Strides = strides; + strides_ = strides; } private: - vector KernelShape; - vector Pads; - vector Strides; + vector kernelShape_; + vector pads_; + vector strides_; }; class MaxPoolRelu : public NeuralNetOperator { @@ -366,48 +366,48 @@ class MaxPoolRelu : public NeuralNetOperator { vector pads = {0, 0}, vector strides = {1, 1}) : NeuralNetOperator(NNKind::MaxPoolRelu), - KernelShape(kernelShape), - Pads(pads), - Strides(strides) {} + kernelShape_(kernelShape), + pads_(pads), + strides_(strides) {} MaxPoolRelu(const MaxPool& maxPool) : NeuralNetOperator(NNKind::MaxPoolRelu), - KernelShape(maxPool.getKernelShape()), - Pads(maxPool.getPads()), - Strides(maxPool.getStrides()) {} + kernelShape_(maxPool.getKernelShape()), + pads_(maxPool.getPads()), + strides_(maxPool.getStrides()) {} ~MaxPoolRelu() {} NOMNIGRAPH_DEFINE_NN_RTTI(MaxPoolRelu); vector getKernelShape() const { - return KernelShape; + return kernelShape_; } vector getPads() const { - return Pads; + return pads_; } vector getStrides() const { - return Strides; + return strides_; } void setKernelShape(vector kernelShape) { - KernelShape = kernelShape; + kernelShape_ = kernelShape; } void setPads(vector pads) { - Pads = pads; + pads_ = pads; } void setStrides(vector strides) { - Strides = strides; + strides_ = strides; } private: - vector KernelShape; - vector Pads; - vector Strides; + vector kernelShape_; + vector pads_; + vector strides_; }; class Sum : public NeuralNetOperator { @@ -437,42 +437,43 @@ class SumRelu : public NeuralNetOperator { class Send : public NeuralNetOperator { public: Send(string destination) - : NeuralNetOperator(NNKind::Send), Destination(destination) {} + : NeuralNetOperator(NNKind::Send), destination_(destination) {} ~Send() {} NOMNIGRAPH_DEFINE_NN_RTTI(Send); string getDestination() const { - return Destination; + return destination_; } void setDestination(string destination) { - Destination = destination; + destination_ = destination; } private: - string Destination; + string destination_; }; class Receive : public NeuralNetOperator { public: - Receive(string source) : NeuralNetOperator(NNKind::Receive), Source(source) {} + Receive(string source) + : NeuralNetOperator(NNKind::Receive), source_(source) {} ~Receive() {} NOMNIGRAPH_DEFINE_NN_RTTI(Receive); string getSource() const { - return Source; + return source_; } void setSource(string source) { - Source = source; + source_ = source; } private: - string Source; + string source_; }; class BatchNormalization : public NeuralNetOperator { @@ -483,82 +484,82 @@ class BatchNormalization : public NeuralNetOperator { bool spatial = true, bool isTest = false) : NeuralNetOperator(NNKind::BatchNormalization), - Epsilon(epsilon), - Momentum(momentum), - Spatial(spatial), - IsTest(isTest) {} + epsilon_(epsilon), + momentum_(momentum), + spatial_(spatial), + isTest_(isTest) {} ~BatchNormalization() {} NOMNIGRAPH_DEFINE_NN_RTTI(BatchNormalization); float getEpsilon() const { - return Epsilon; + return epsilon_; } float getMomentum() const { - return Momentum; + return momentum_; } bool getSpatial() const { - return Spatial; + return spatial_; } bool getIsTest() const { - return IsTest; + return isTest_; } void setEpsilon(float epsilon) { - Epsilon = epsilon; + epsilon_ = epsilon; } void setMomentum(float momentum) { - Momentum = momentum; + momentum_ = momentum; } void setSpatial(bool spatial) { - Spatial = spatial; + spatial_ = spatial; } void setIsTest(bool isTest) { - IsTest = isTest; + isTest_ = isTest; } private: - float Epsilon; - float Momentum; - bool Spatial; - bool IsTest; + float epsilon_; + float momentum_; + bool spatial_; + bool isTest_; }; class Clip : public NeuralNetOperator { public: Clip(float min, float max) - : NeuralNetOperator(NNKind::Clip), Min(min), Max(max) {} + : NeuralNetOperator(NNKind::Clip), min_(min), max_(max) {} ~Clip() {} NOMNIGRAPH_DEFINE_NN_RTTI(Clip); float getMin() const { - return Min; + return min_; } float getMax() const { - return Max; + return max_; } void setMin(float min) { - Min = min; + min_ = min; } void setMax(float max) { - Max = max; + max_ = max; } private: - float Min; - float Max; + float min_; + float max_; }; class FC : public NeuralNetOperator { @@ -586,31 +587,31 @@ class GivenTensorFill : public NeuralNetOperator { class Concat : public NeuralNetOperator { public: Concat(int axis = -1, bool addAxis = false) - : NeuralNetOperator(NNKind::Concat), Axis(axis), AddAxis(addAxis) {} + : NeuralNetOperator(NNKind::Concat), axis_(axis), addAxis_(addAxis) {} ~Concat() {} NOMNIGRAPH_DEFINE_NN_RTTI(Concat); int getAxis() const { - return Axis; + return axis_; } bool getAddAxis() const { - return AddAxis; + return addAxis_; } void setAxis(int axis) { - Axis = axis; + axis_ = axis; } void setAddAxis(bool addAxis) { - AddAxis = addAxis; + addAxis_ = addAxis; } private: - int Axis; - bool AddAxis; + int axis_; + bool addAxis_; }; class Softmax : public NeuralNetOperator { diff --git a/caffe2/core/nomnigraph/include/nomnigraph/Graph/Graph.h b/caffe2/core/nomnigraph/include/nomnigraph/Graph/Graph.h index fcaba7a4e8ac6..c6b10f0912eca 100644 --- a/caffe2/core/nomnigraph/include/nomnigraph/Graph/Graph.h +++ b/caffe2/core/nomnigraph/include/nomnigraph/Graph/Graph.h @@ -1,9 +1,5 @@ //===- nomnigraph/Graph/Graph.h - Basic graph implementation ----*- C++ -*-===// // -// TODO Licensing. -// -//===----------------------------------------------------------------------===// -// // This file defines a basic graph API for generic and flexible use with // graph algorithms. // @@ -43,7 +39,7 @@ class Node; // \brief Edge within a Graph. template -class Edge : public StorageType { +class CAFFE2_API Edge : public StorageType { public: using NodeRef = typename Graph::NodeRef; Edge(NodeRef tail, NodeRef head, U... args) @@ -77,7 +73,7 @@ class Edge : public StorageType { // \brief Node within a Graph. template -class Node : public StorageType, public Notifier> { +class CAFFE2_API Node : public StorageType, public Notifier> { public: using NodeRef = typename Graph::NodeRef; using EdgeRef = typename Graph::EdgeRef; @@ -156,7 +152,7 @@ class Node : public StorageType, public Notifier> { /// for example. /// template -class Subgraph { +class CAFFE2_API Subgraph { public: Subgraph() { DEBUG_PRINT("Creating instance of Subgraph: %p\n", this); @@ -223,7 +219,7 @@ class Subgraph { /// Everything is owned by the graph to simplify storage concerns. /// template -class Graph { +class CAFFE2_API Graph { public: using SubgraphType = Subgraph; using NodeRef = Node*; @@ -248,27 +244,8 @@ class Graph { return createNodeInternal(Node()); } - void importNode(NodeRef node, Graph& otherGraph) { - for (auto it = nodes_.begin(); it != nodes_.end(); ++it) { - if (&(*it) == node) { - std::list>& otherNodes = otherGraph.nodes_; - otherNodes.splice(otherNodes.end(), nodes_, it, ++it); - otherGraph.nodeRefs_.insert(node); - break; - } - } - } - - void importEdge(EdgeRef edge, Graph& otherGraph) { - std::list>& otherEdges = otherGraph.edges_; - for (auto it = edges_.begin(); it != edges_.end(); ++it) { - if (&(*it) == edge) { - otherEdges.splice(otherEdges.end(), edges_, it, ++it); - break; - } - } - } - + // Swap two nodes. + // Any edge V -> N1 becomes V -> N2, and N1 -> V becomes N2 -> V. void swapNodes(NodeRef n1, NodeRef n2) { // First rectify the edges for (auto& inEdge : n1->getInEdges()) { @@ -295,34 +272,14 @@ class Graph { n2->setInEdges(n1InEdges); } - /// \brief Replace a node in the graph with a generic - /// set of nodes. + /// \brief Replace a node in the graph with another node. /// \note The node replaced simply has its edges cut, but it not /// deleted from the graph. Call Graph::deleteNode to delete it. - /// \p old A node to be replaced in the graph. - /// \p newTail The node that inherit the old node's in-edges - /// \p newHead (optional) The node that inherit the old node's out-edges - void replaceNode( - const NodeRef& old, - const NodeRef& newTail, - const NodeRef& newHead_ = nullptr) { - // If no newHead is specified, make the tail the head as well. - // We are effectively replacing the node with one node in this case. - const NodeRef newHead = newHead_ ? newHead_ : newTail; - const auto inEdges = old->getInEdges(); - const auto outEdges = old->getOutEdges(); - - for (const auto& inEdge : inEdges) { - inEdge->setHead(newTail); - old->removeInEdge(inEdge); - newTail->addInEdge(inEdge); - } - - for (const auto& outEdge : outEdges) { - outEdge->setTail(newHead); - old->removeOutEdge(outEdge); - newTail->addOutEdge(outEdge); - } + /// \p oldNode A node to be replaced in the graph. + /// \p newNode The node that inherit the old node's in-edges and out-edges. + void replaceNode(const NodeRef& oldNode, const NodeRef& newNode) { + replaceInEdges(oldNode, newNode); + replaceOutEdges(oldNode, newNode); } // All out-edges oldNode -> V will be replaced with newNode -> V @@ -361,36 +318,46 @@ class Graph { return e; } - /// \brief Get a reference to the edge between two nodes if it exists. - /// note: will fail assertion if the edge does not exist. - EdgeRef getEdge(NodeRef tail, NodeRef head) const { + /// \brief Get a reference to the edge between two nodes if it exists. Returns + /// nullptr if the edge does not exist. + EdgeRef getEdgeIfExists(NodeRef tail, NodeRef head) const { for (auto& inEdge : head->getInEdges()) { if (inEdge->tail() == tail) { return inEdge; } } - assert(0 && "Edge doesn't exist."); return nullptr; } + /// \brief Returns true if there is an edge between the given two nodes. + bool hasEdge(NodeRef tail, NodeRef head) const { + return getEdgeIfExists(tail, head); + } + + /// \brief Get a reference to the edge between two nodes if it exists. + /// note: will fail assertion if the edge does not exist. + EdgeRef getEdge(NodeRef tail, NodeRef head) const { + auto result = getEdgeIfExists(tail, head); + assert(result && "Edge doesn't exist."); + return result; + } + /// \brief Deletes a node from the graph. /// \param n A reference to the node. - /// \param deleteEdges (optional) Whether or not to delete the edges - /// related to the node. - void deleteNode(NodeRef n, bool deleteEdges = true) { + void deleteNode(NodeRef n) { if (!hasNode(n)) { return; } - if (deleteEdges) { - auto inEdges = n->inEdges_; - for (auto& edge : inEdges) { - deleteEdge(edge); - } - auto outEdges = n->outEdges_; - for (auto& edge : outEdges) { - deleteEdge(edge); - } + + auto inEdges = n->inEdges_; + for (auto& edge : inEdges) { + deleteEdge(edge); + } + auto outEdges = n->outEdges_; + for (auto& edge : outEdges) { + deleteEdge(edge); } + for (auto i = nodes_.begin(); i != nodes_.end(); ++i) { if (&*i == n) { nodeRefs_.erase(n); @@ -401,11 +368,9 @@ class Graph { } // Delete all nodes in the set. - void deleteNodes( - const std::unordered_set& nodes, - bool deleteEdges = true) { + void deleteNodes(const std::unordered_set& nodes) { for (auto node : nodes) { - deleteNode(node, deleteEdges); + deleteNode(node); } } @@ -415,11 +380,9 @@ class Graph { /// \brief Deletes a edge from the graph. /// \p e A reference to the edge. - void deleteEdge(EdgeRef e, bool removeRef = true) { - if (removeRef) { - e->tail_->removeOutEdge(e); - e->head_->removeInEdge(e); - } + void deleteEdge(EdgeRef e) { + e->tail_->removeOutEdge(e); + e->head_->removeInEdge(e); for (auto i = edges_.begin(); i != edges_.end(); ++i) { if (&*i == e) { edges_.erase(i); diff --git a/caffe2/core/nomnigraph/include/nomnigraph/Representations/Compiler.h b/caffe2/core/nomnigraph/include/nomnigraph/Representations/Compiler.h index 0b6aefa3d4c4a..d8e9c1090b3ed 100644 --- a/caffe2/core/nomnigraph/include/nomnigraph/Representations/Compiler.h +++ b/caffe2/core/nomnigraph/include/nomnigraph/Representations/Compiler.h @@ -8,7 +8,7 @@ namespace nom { namespace repr { -class Value { +class CAFFE2_API Value { public: enum class ValueKind { Value, Instruction, Data }; Value(ValueKind K) : kind_(K) {} @@ -22,7 +22,7 @@ class Value { const ValueKind kind_; }; -class Data : public Value { +class CAFFE2_API Data : public Value { public: Data() : Value(ValueKind::Data) {} static bool classof(const Value* V) { @@ -41,7 +41,7 @@ class Data : public Value { size_t version_ = 0; }; -class Instruction : public Value { +class CAFFE2_API Instruction : public Value { public: /// \brief All the different types of execution. enum class Opcode { @@ -66,7 +66,7 @@ class Instruction : public Value { Opcode op_; }; -class Terminator : public Instruction { +class CAFFE2_API Terminator : public Instruction { public: Terminator(Instruction::Opcode op) : Instruction(op) {} @@ -80,17 +80,17 @@ class Terminator : public Instruction { } }; -class Branch : public Terminator { +class CAFFE2_API Branch : public Terminator { public: Branch() : Terminator(Instruction::Opcode::Branch) {} }; -class Return : public Terminator { +class CAFFE2_API Return : public Terminator { public: Return() : Terminator(Instruction::Opcode::Return) {} }; -class Phi : public Instruction { +class CAFFE2_API Phi : public Instruction { public: Phi() : Instruction(Instruction::Opcode::Phi) {} }; diff --git a/caffe2/core/nomnigraph/include/nomnigraph/Representations/NeuralNet.h b/caffe2/core/nomnigraph/include/nomnigraph/Representations/NeuralNet.h index d5836836d9df2..f4c3940acd071 100644 --- a/caffe2/core/nomnigraph/include/nomnigraph/Representations/NeuralNet.h +++ b/caffe2/core/nomnigraph/include/nomnigraph/Representations/NeuralNet.h @@ -41,7 +41,7 @@ class NeuralNetData; /// a saved void* pointer for external use. Derived classes /// add richer semantics to the annotation and it is encouraged /// to use them. -class Annotation { +class CAFFE2_API Annotation { public: enum class AnnotationKind { Generic, Caffe2 }; @@ -60,7 +60,7 @@ class Annotation { const AnnotationKind kind_; }; -class NeuralNetOperator : public Instruction { +class CAFFE2_API NeuralNetOperator : public Instruction { public: /// Discriminator for LLVM-style RTTI (isa<>) enum class NNKind { @@ -135,7 +135,7 @@ class NeuralNetOperator : public Instruction { std::unique_ptr extraAnnotation_; }; -class NeuralNetData : public Data { +class CAFFE2_API NeuralNetData : public Data { public: /// Discriminator for LLVM-style RTTI (isa<>) enum class NNDataKind { Generic, Tensor }; @@ -159,7 +159,7 @@ class NeuralNetData : public Data { size_t version_ = 0; }; -class Tensor : public NeuralNetData { +class CAFFE2_API Tensor : public NeuralNetData { public: enum class DataType { Generic, Float, Half, Int8 }; enum class Layout { Generic, NCHW, NHWC }; @@ -207,21 +207,21 @@ class Tensor : public NeuralNetData { #include "nomnigraph/Generated/OpClasses.h" -class While : public NeuralNetOperator { +class CAFFE2_API While : public NeuralNetOperator { public: While() : NeuralNetOperator(NNKind::While, Opcode::Branch) {} NOMNIGRAPH_DEFINE_NN_RTTI(While); ~While() {} }; -class NNPhi : public NeuralNetOperator { +class CAFFE2_API NNPhi : public NeuralNetOperator { public: NNPhi() : NeuralNetOperator(NNKind::NNPhi, Opcode::Phi) {} NOMNIGRAPH_DEFINE_NN_RTTI(NNPhi); ~NNPhi() {} }; -class GenericOperator : public NeuralNetOperator { +class CAFFE2_API GenericOperator : public NeuralNetOperator { public: GenericOperator() : NeuralNetOperator(NNKind::GenericOperator) {} GenericOperator(std::string name) @@ -243,7 +243,7 @@ using NNGraph = nom::Graph>; using NNSubgraph = nom::Subgraph>; using NNCFGraph = nom::repr::ControlFlowGraph; -struct NNModule { +struct CAFFE2_API NNModule { NNGraph dataFlow; NNCFGraph controlFlow; std::unordered_set inputs; @@ -262,7 +262,7 @@ template using enable_if_t = typename std::enable_if::type; template -struct inheritedFrom { +struct CAFFE2_EXPORT inheritedFrom { static constexpr bool value = std::is_base_of::value && !std::is_same::value; }; @@ -270,14 +270,14 @@ struct inheritedFrom { // This is just a way to fix issues when the isa<> implementation // can't automatically downcast. template -struct is_impl { +struct CAFFE2_EXPORT is_impl { inline static bool impl(N n) { return isa(n->data()); } }; template -struct is_impl::value>> { +struct CAFFE2_EXPORT is_impl::value>> { inline static bool impl(N n) { if (!isa(n->data().get())) { return false; @@ -288,7 +288,7 @@ struct is_impl::value>> { }; template -struct is_impl::value>> { +struct CAFFE2_EXPORT is_impl::value>> { inline static bool impl(N n) { if (!isa(n->data().get())) { return false; @@ -306,14 +306,14 @@ inline bool is(N n) { // This is just a way to fix issues when the dyn_cast<> implementation // can't automatically downcast. template -struct get_impl { +struct CAFFE2_EXPORT get_impl { inline static T* impl(N n) { return dyn_cast(n->data().get()); } }; template -struct get_impl::value>> { +struct CAFFE2_EXPORT get_impl::value>> { inline static T* impl(N n) { if (!is(n)) { assert(0 && "Cannot get type from node"); @@ -325,7 +325,7 @@ struct get_impl::value>> { }; template -struct get_impl::value>> { +struct CAFFE2_EXPORT get_impl::value>> { inline static T* impl(N n) { if (!is(n)) { assert(0 && "Cannot get type from node"); @@ -425,9 +425,9 @@ CAFFE2_API std::vector getOutputs(NNGraph::NodeRef n); CAFFE2_API void coalesceInsertedDataDependencies(repr::NNModule* m); template -struct NodeHelper {}; +struct CAFFE2_EXPORT NodeHelper {}; -struct NNNodeMatchCriteria { +struct CAFFE2_API NNNodeMatchCriteria { std::function predicate; std::string debugString; @@ -452,7 +452,7 @@ struct NNNodeMatchCriteria { } }; -std::ostream& operator<<( +CAFFE2_API std::ostream& operator<<( std::ostream& oss, const NNNodeMatchCriteria& criteria); @@ -481,7 +481,7 @@ NNNodeMatchCriteria matchOp( const std::function predicate, const std::string& debugString = "matchOpWithPredicate") { return NNNodeMatchCriteria( - [&predicate](NNGraph::NodeRef nodeRef) { + [predicate](NNGraph::NodeRef nodeRef) { NOM_REQUIRE_OR_RET_FALSE(is(nodeRef)); NodeType* node = get(nodeRef); return predicate(*node); @@ -489,7 +489,7 @@ NNNodeMatchCriteria matchOp( debugString); }; -struct NNNodeMatch { +struct CAFFE2_API NNNodeMatch { static bool isMatch( const NNGraph::NodeRef& node, const NNNodeMatchCriteria& criteria) { @@ -503,7 +503,7 @@ using NNSubgraphMatcher = // This helper method makes it easy to create matching criteria in NNGraph. // For example, operatorSubgraph(opMatch, ...) will refer to a tree like this: // ... -> opMatch -> opMatch_Output -NNMatchGraph::NodeRef operatorSubgraph( +CAFFE2_API NNMatchGraph::NodeRef operatorSubgraph( NNMatchGraph& g, const NNNodeMatchCriteria& root, const std::vector& childrenCriteria = {}, diff --git a/caffe2/core/nomnigraph/op_gen.py b/caffe2/core/nomnigraph/op_gen.py index 2d1125f5762ad..49cd2abb2cef7 100755 --- a/caffe2/core/nomnigraph/op_gen.py +++ b/caffe2/core/nomnigraph/op_gen.py @@ -99,30 +99,33 @@ def gen_class(op, op_def): attribute_setters = [] for attr in attributes: lower_name = attr[0][0].lower() + attr[0][1:] + private_name = lower_name + "_" default_arg = "" if len(attr) < 3 else " = {}".format(attr[2]) name = attr[0] t = attr[1] attr_arg = "{type} {lower_name}".format( type=t, lower_name=lower_name + default_arg ) - attr_init = "{name}({lower_name})".format(name=name, lower_name=lower_name) - attr_declare = "{type} {name};".format(type=t, name=name) + attr_init = "{private_name}({lower_name})".format( + private_name=private_name, lower_name=lower_name) + attr_declare = "{type} {private_name};".format( + type=t, private_name=private_name) attr_get = dedent( """ {type} get{name}() const {{ - return {name}; + return {private_name}; }} """.format( - type=t, name=name + type=t, name=name, private_name=private_name ) ) attr_set = dedent( """ void set{name}({type} {lower_name}) {{ - {name} = {lower_name}; + {private_name} = {lower_name}; }} """.format( - type=t, name=name, lower_name=lower_name + type=t, name=name, private_name=private_name, lower_name=lower_name ) ) attribute_args.append(attr_arg) @@ -137,9 +140,11 @@ def gen_class(op, op_def): lower_other_op = other_op[0].lower() + other_op[1:] other_init = [default_init] for attr in attributes: + lower_name = attr[0][0].lower() + attr[0][1:] + private_name = lower_name + "_" other_init.append( - "{name}({other_op}.get{name}())".format( - name=attr[0], other_op=lower_other_op + "{private_name}({other_op}.get{name}())".format( + name=attr[0], private_name=private_name, other_op=lower_other_op ) ) init = dedent( diff --git a/caffe2/core/nomnigraph/tests/dominator_tree_test.cc b/caffe2/core/nomnigraph/tests/AlgorithmsTest.cc similarity index 100% rename from caffe2/core/nomnigraph/tests/dominator_tree_test.cc rename to caffe2/core/nomnigraph/tests/AlgorithmsTest.cc diff --git a/caffe2/core/nomnigraph/tests/binary_match_test.cc b/caffe2/core/nomnigraph/tests/BinaryMatchImplTest.cc similarity index 100% rename from caffe2/core/nomnigraph/tests/binary_match_test.cc rename to caffe2/core/nomnigraph/tests/BinaryMatchImplTest.cc diff --git a/caffe2/core/nomnigraph/tests/GraphTest.cc b/caffe2/core/nomnigraph/tests/GraphTest.cc new file mode 100644 index 0000000000000..0138a9b5759d3 --- /dev/null +++ b/caffe2/core/nomnigraph/tests/GraphTest.cc @@ -0,0 +1,120 @@ +#include "test_util.h" + +#include "nomnigraph/Converters/Dot.h" +#include "nomnigraph/Graph/Algorithms.h" +#include "nomnigraph/Graph/Graph.h" +#include "nomnigraph/Support/Casting.h" + +#include + +using TestGraph = nom::Graph; +TestGraph::NodeRef createTestNode(TestGraph& g) { + return g.createNode(TestClass()); +} + +TEST(Basic, CreateNodeAndEdge) { + TestGraph g; + auto n1 = createTestNode(g); + auto n2 = createTestNode(g); + g.createEdge(n1, n2); + EXPECT_TRUE(g.hasNode(n1)); + EXPECT_TRUE(g.hasNode(n2)); +} + +TEST(Basic, DeleteNode) { + TestGraph g; + auto n1 = createTestNode(g); + auto n2 = createTestNode(g); + g.createEdge(n1, n2); + EXPECT_TRUE(g.hasNode(n1)); + g.deleteNode(n1); + EXPECT_FALSE(g.hasNode(n1)); +} + +TEST(Basic, DeleteEdge) { + TestGraph g; + auto n1 = createTestNode(g); + auto n2 = createTestNode(g); + auto e = g.createEdge(n1, n2); + g.deleteEdge(e); +} + +TEST(Basic, ReplaceEdges) { + TestGraph g; + auto n1 = createTestNode(g); + auto n2 = createTestNode(g); + auto n3 = createTestNode(g); + auto n4 = createTestNode(g); + auto n5 = createTestNode(g); + + g.createEdge(n1, n3); + g.createEdge(n2, n3); + g.createEdge(n3, n4); + /* + 1 2 5 + | + 3 + | + 4 + */ + + EXPECT_FALSE(g.hasEdge(n1, n5)); + EXPECT_FALSE(g.hasEdge(n2, n5)); + g.replaceInEdges(n3, n5); + /* + 1 2 3 + | | + 5 4 + */ + EXPECT_TRUE(g.hasEdge(n1, n5)); + EXPECT_TRUE(g.hasEdge(n2, n5)); + + EXPECT_FALSE(g.hasEdge(n5, n4)); + g.replaceOutEdges(n3, n5); + /* + 1 2 3 + | + 5 + | + 4 + */ + EXPECT_TRUE(g.hasEdge(n5, n4)); + + g.replaceNode(n5, n3); + // Back to the original graph. + /* + 1 2 5 + | + 3 + | + 4 + */ + EXPECT_TRUE(g.hasEdge(n1, n3)); + EXPECT_TRUE(g.hasEdge(n2, n3)); +} + +TEST(Basic, HasNode) { + TestGraph g; + auto n1 = createTestNode(g); + auto n2 = createTestNode(g); + auto n3 = createTestNode(g); + g.createEdge(n1, n2); + g.createEdge(n1, n3); + // Current graph: 1 -> 2 -> 3 + EXPECT_TRUE(g.hasNode(n1)); + EXPECT_TRUE(g.hasNode(n2)); + EXPECT_TRUE(g.hasNode(n3)); + g.swapNodes(n1, n3); + // Current graph: 3 -> 2 -> 1 + EXPECT_TRUE(g.hasNode(n1)); + EXPECT_TRUE(g.hasNode(n3)); + g.deleteNode(n1); + // Current graph: 3 -> 2 + EXPECT_FALSE(g.hasNode(n1)); + auto n4 = createTestNode(g); + EXPECT_TRUE(g.hasNode(n4)); + g.replaceNode(n2, n4); + // Current graph: 3 -> 4 , 2 + // replaceNode doesn't delete n2. + EXPECT_TRUE(g.hasNode(n2)); +} diff --git a/caffe2/core/nomnigraph/tests/match_test.cc b/caffe2/core/nomnigraph/tests/MatchTest.cc similarity index 100% rename from caffe2/core/nomnigraph/tests/match_test.cc rename to caffe2/core/nomnigraph/tests/MatchTest.cc diff --git a/caffe2/core/nomnigraph/tests/neural_net_test.cc b/caffe2/core/nomnigraph/tests/NeuralNetTest.cc similarity index 100% rename from caffe2/core/nomnigraph/tests/neural_net_test.cc rename to caffe2/core/nomnigraph/tests/NeuralNetTest.cc diff --git a/caffe2/core/nomnigraph/tests/subgraph_matcher_test.cc b/caffe2/core/nomnigraph/tests/SubgraphMatcherTest.cc similarity index 100% rename from caffe2/core/nomnigraph/tests/subgraph_matcher_test.cc rename to caffe2/core/nomnigraph/tests/SubgraphMatcherTest.cc diff --git a/caffe2/core/nomnigraph/tests/tarjans_test.cc b/caffe2/core/nomnigraph/tests/TarjansImplTest.cc similarity index 100% rename from caffe2/core/nomnigraph/tests/tarjans_test.cc rename to caffe2/core/nomnigraph/tests/TarjansImplTest.cc diff --git a/caffe2/core/nomnigraph/tests/basic_test.cc b/caffe2/core/nomnigraph/tests/basic_test.cc deleted file mode 100644 index bb1474eb1af91..0000000000000 --- a/caffe2/core/nomnigraph/tests/basic_test.cc +++ /dev/null @@ -1,83 +0,0 @@ -#include "test_util.h" - -#include "nomnigraph/Converters/Dot.h" -#include "nomnigraph/Graph/Algorithms.h" -#include "nomnigraph/Graph/Graph.h" -#include "nomnigraph/Transformations/Match.h" -#include "nomnigraph/Support/Casting.h" - -#include - -TEST(Basic, CreateNodeAndEdge) { - TestClass t1; - TestClass t2; - nom::Graph g; - nom::Graph::NodeRef n1 = g.createNode(std::move(t1)); - nom::Graph::NodeRef n2 = g.createNode(std::move(t2)); - g.createEdge(n1, n2); - EXPECT_TRUE(g.hasNode(n1)); - EXPECT_TRUE(g.hasNode(n2)); -} - -TEST(Basic, DeleteNode) { - TestClass t1; - TestClass t2; - nom::Graph g; - nom::Graph::NodeRef n1 = g.createNode(std::move(t1)); - nom::Graph::NodeRef n2 = g.createNode(std::move(t2)); - g.createEdge(n1, n2); - EXPECT_TRUE(g.hasNode(n1)); - g.deleteNode(n1); - EXPECT_FALSE(g.hasNode(n1)); -} - -TEST(Basic, DeleteEdge) { - TestClass t1; - TestClass t2; - nom::Graph g; - nom::Graph::NodeRef n1 = g.createNode(std::move(t1)); - nom::Graph::NodeRef n2 = g.createNode(std::move(t2)); - nom::Graph::EdgeRef e = g.createEdge(n1, n2); - g.deleteEdge(e); -} - -TEST(Basic, HasNode) { - TestClass t1; - TestClass t2; - TestClass t3; - nom::Graph g; - nom::Graph::NodeRef n1 = g.createNode(std::move(t1)); - nom::Graph::NodeRef n2 = g.createNode(std::move(t2)); - nom::Graph::NodeRef n3 = g.createNode(std::move(t3)); - g.createEdge(n1, n2); - g.createEdge(n1, n3); - // Current graph: 1 -> 2 -> 3 - EXPECT_TRUE(g.hasNode(n1)); - EXPECT_TRUE(g.hasNode(n2)); - EXPECT_TRUE(g.hasNode(n3)); - g.swapNodes(n1, n3); - // Current graph: 3 -> 2 -> 1 - EXPECT_TRUE(g.hasNode(n1)); - EXPECT_TRUE(g.hasNode(n3)); - g.deleteNode(n1); - // Current graph: 3 -> 2 - EXPECT_FALSE(g.hasNode(n1)); - TestClass t4; - nom::Graph::NodeRef n4 = g.createNode(std::move(t4)); - EXPECT_TRUE(g.hasNode(n4)); - g.replaceNode(n2, n4); - // Current graph: 3 -> 4 , 2 - // replaceNode doesn't delete n2. - EXPECT_TRUE(g.hasNode(n2)); - - // Create a second graph g2, and import the nodes from g2 to g. - TestClass t5; - nom::Graph g2; - nom::Graph::NodeRef n5 = g2.createNode(std::move(t5)); - EXPECT_TRUE(g2.hasNode(n5)); - - EXPECT_FALSE(g.hasNode(n5)); - g2.importNode(n5, g); - // Current graph (g1): 3 -> 4, 2, 5 - EXPECT_TRUE(g.hasNode(n5)); -} diff --git a/caffe2/core/nomnigraph/tests/test_util.h b/caffe2/core/nomnigraph/tests/test_util.h index f5693c03d36a2..ae842f7cb8f91 100644 --- a/caffe2/core/nomnigraph/tests/test_util.h +++ b/caffe2/core/nomnigraph/tests/test_util.h @@ -1,6 +1,7 @@ #ifndef NOM_TESTS_TEST_UTIL_H #define NOM_TESTS_TEST_UTIL_H +#include "caffe2/core/common.h" #include "nomnigraph/Graph/Graph.h" #include "nomnigraph/Graph/Algorithms.h" #include "nomnigraph/Representations/NeuralNet.h" @@ -101,9 +102,9 @@ class TestRandom { * return labelMap; * }); */ -nom::Graph createGraph(); +CAFFE2_API nom::Graph createGraph(); -nom::Graph createGraphWithCycle(); +CAFFE2_API nom::Graph createGraphWithCycle(); std::map BBPrinter(typename nom::repr::NNCFGraph::NodeRef node); diff --git a/caffe2/core/operator.cc b/caffe2/core/operator.cc index 06e05ee3dca13..38cc117141b1c 100644 --- a/caffe2/core/operator.cc +++ b/caffe2/core/operator.cc @@ -72,8 +72,8 @@ PerOpEnginePrefType& g_per_op_engine_pref() { } GlobalEnginePrefType& g_global_engine_pref() { - static auto* g_global_engine_pref_ = new GlobalEnginePrefType{ - {DeviceType::CUDA, {"CUDNN"}}, {DeviceType::HIP, {"MIOPEN"}}}; + static auto* g_global_engine_pref_ = + new GlobalEnginePrefType{{CUDA, {"CUDNN"}}, {HIP, {"MIOPEN"}}}; return *g_global_engine_pref_; } @@ -81,7 +81,8 @@ unique_ptr TryCreateC2Operator( const string& key, const OperatorDef& operator_def, Workspace* ws) { - const auto& type = operator_def.device_option().device_type(); + const auto& type_proto = operator_def.device_option().device_type(); + const auto& type = ProtoToType(static_cast(type_proto)); CAFFE_ENFORCE( gDeviceTypeRegistry()->count(type), "Device type ", @@ -123,7 +124,9 @@ unique_ptr _CreateOperator( Workspace* ws) { static StaticLinkingProtector g_protector; const auto& op_type = operator_def.type(); - const auto& device_type = operator_def.device_option().device_type(); + const auto& device_type_proto = operator_def.device_option().device_type(); + const auto& device_type = + ProtoToType(static_cast(device_type_proto)); #ifndef CAFFE2_NO_OPERATOR_SCHEMA // first, check with OpSchema if the operator is legal. @@ -267,9 +270,11 @@ void SetEnginePref( void SetOpEnginePref( const std::string& op_type, - const CaffeMap& op_pref) { + const CaffeMap& op_pref) { for (const auto& device_pref_pair : op_pref) { - const auto& device_type = device_pref_pair.first; + const auto& device_type_proto = device_pref_pair.first; + const auto& device_type = + ProtoToType(static_cast(device_type_proto)); CAFFE_ENFORCE( gDeviceTypeRegistry()->count(device_type), "Device type ", @@ -306,8 +311,8 @@ unique_ptr CreateOperator( } } -std::map* gDeviceTypeRegistry() { - static std::map g_device_type_registry; +std::map* gDeviceTypeRegistry() { + static std::map g_device_type_registry; return &g_device_type_registry; } @@ -316,21 +321,21 @@ CAFFE_DEFINE_REGISTRY( OperatorBase, const OperatorDef&, Workspace*); -CAFFE_REGISTER_DEVICE_TYPE(DeviceType::CPU, CPUOperatorRegistry); +CAFFE_REGISTER_DEVICE_TYPE(CPU, CPUOperatorRegistry); CAFFE_DEFINE_REGISTRY( CUDAOperatorRegistry, OperatorBase, const OperatorDef&, Workspace*); -CAFFE_REGISTER_DEVICE_TYPE(DeviceType::CUDA, CUDAOperatorRegistry); +CAFFE_REGISTER_DEVICE_TYPE(CUDA, CUDAOperatorRegistry); CAFFE_DEFINE_REGISTRY( HIPOperatorRegistry, OperatorBase, const OperatorDef&, Workspace*); -CAFFE_REGISTER_DEVICE_TYPE(DeviceType::HIP, HIPOperatorRegistry); +CAFFE_REGISTER_DEVICE_TYPE(HIP, HIPOperatorRegistry); CAFFE_DEFINE_REGISTRY( GradientRegistry, @@ -642,11 +647,11 @@ std::map> ValidateTensorDevices( &_capacity, &blob_device); - if (blob_device.device_type() == CUDA && + if (blob_device.device_type() == PROTO_CUDA && blob_device.cuda_gpu_id() != op_device.cuda_gpu_id()) { mismatches[blob_name] = std::make_pair(op_device, blob_device); - } - else if (blob_device.device_type() == HIP && + } else if ( + blob_device.device_type() == PROTO_HIP && blob_device.hip_gpu_id() != op_device.hip_gpu_id()) { mismatches[blob_name] = std::make_pair(op_device, blob_device); } diff --git a/caffe2/core/operator.h b/caffe2/core/operator.h index 3e901b205c9bd..2def93f0b51d0 100644 --- a/caffe2/core/operator.h +++ b/caffe2/core/operator.h @@ -96,9 +96,6 @@ class CAFFE2_API OperatorBase : public Observable { // TODO(jerryzh): We'll need to check device type in Get() later // Get() -> Get(type) const auto& tensor = inputs_.at(idx)->template Get(); - CAFFE_ENFORCE( - tensor.is_contiguous(), - "Tensor must be contiguous for caffe2 operators"); return tensor; } catch (::caffe2::EnforceNotMet& enf) { if (has_debug_def()) { @@ -125,11 +122,7 @@ class CAFFE2_API OperatorBase : public Observable { static_assert( std::is_same::value, "Output(int, DeviceType) is only available for Tensor"); - auto* tensor = outputs_.at(idx)->GetMutableTensor(type); - CAFFE_ENFORCE( - tensor->is_contiguous(), - "Tensor must be contiguous for caffe2 operators"); - return tensor; + return outputs_.at(idx)->GetMutableTensor(type); } template @@ -377,7 +370,7 @@ class CAFFE2_API OperatorBase : public Observable { } public: - static constexpr int kNoNetPositionSet = -1; + static const int kNoNetPositionSet = -1; private: Workspace* operator_ws_; @@ -454,7 +447,7 @@ class CAFFE2_API OperatorBase : public Observable { // run on different devices. You should then implement the RunOnDevice() // function. template -class CAFFE2_API Operator : public OperatorBase { +class Operator : public OperatorBase { public: explicit Operator(const OperatorDef& operator_def, Workspace* ws) : OperatorBase(operator_def, ws), context_(operator_def.device_option()) { @@ -806,12 +799,12 @@ typedef Registry< std::unique_ptr, const OperatorDef&, Workspace*>* (*RegistryFunction)(); -CAFFE2_API std::map* gDeviceTypeRegistry(); +CAFFE2_API std::map* gDeviceTypeRegistry(); struct CAFFE2_API DeviceTypeRegisterer { - explicit DeviceTypeRegisterer(int32_t type, RegistryFunction func) { + explicit DeviceTypeRegisterer(DeviceType type, RegistryFunction func) { if (gDeviceTypeRegistry()->count(type)) { - std::cerr << "Device type " << type + std::cerr << "Device type " << DeviceTypeName(type) << "registered twice. This should not happen. Did you have " "duplicated numbers assigned to different devices?"; std::exit(1); @@ -842,7 +835,7 @@ CAFFE_DECLARE_REGISTRY( #define REGISTER_CPU_OPERATOR_CREATOR(key, ...) \ CAFFE_REGISTER_CREATOR(CPUOperatorRegistry, key, __VA_ARGS__) #define REGISTER_CPU_OPERATOR(name, ...) \ - extern void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(); \ + CAFFE2_IMPORT void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name();\ static void CAFFE2_UNUSED CAFFE_ANONYMOUS_VARIABLE_CPU##name() { \ CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(); \ } \ @@ -861,7 +854,7 @@ CAFFE_DECLARE_REGISTRY( #define REGISTER_CUDA_OPERATOR_CREATOR(key, ...) \ CAFFE_REGISTER_CREATOR(CUDAOperatorRegistry, key, __VA_ARGS__) #define REGISTER_CUDA_OPERATOR(name, ...) \ - extern void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(); \ + CAFFE2_IMPORT void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(); \ static void CAFFE2_UNUSED CAFFE_ANONYMOUS_VARIABLE_CUDA##name() { \ CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(); \ } \ @@ -886,7 +879,7 @@ CAFFE_DECLARE_REGISTRY( #define REGISTER_HIP_OPERATOR_CREATOR(key, ...) \ CAFFE_REGISTER_CREATOR(HIPOperatorRegistry, key, __VA_ARGS__) #define REGISTER_HIP_OPERATOR(name, ...) \ - extern void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(); \ + CAFFE2_IMPORT void CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(); \ static void CAFFE2_UNUSED CAFFE_ANONYMOUS_VARIABLE_HIP##name() { \ CAFFE2_PLEASE_ADD_OPERATOR_SCHEMA_FOR_##name(); \ } \ @@ -967,9 +960,9 @@ CAFFE2_API const std::string OpRegistryKey( using EnginePrefType = std::vector; // {device_type -> {operator_name -> EnginePrefType}} using PerOpEnginePrefType = - CaffeMap>; + CaffeMap>; // {device_type -> EnginePrefType} -using GlobalEnginePrefType = CaffeMap; +using GlobalEnginePrefType = CaffeMap; CAFFE2_API void SetPerOpEnginePref(const PerOpEnginePrefType& per_op_engine_pref); CAFFE2_API void SetGlobalEnginePref(const GlobalEnginePrefType& global_engine_pref); CAFFE2_API void SetEnginePref( @@ -977,7 +970,7 @@ CAFFE2_API void SetEnginePref( const GlobalEnginePrefType& global_engine_pref); CAFFE2_API void SetOpEnginePref( const std::string& op_type, - const CaffeMap& op_pref); + const CaffeMap& op_pref); CAFFE2_API TensorShape GetTensorShapeOfBlob(const Blob* b); diff --git a/caffe2/core/operator_gpu_test.cc b/caffe2/core/operator_gpu_test.cc index 103370abaa1f0..80c58b7a3c752 100644 --- a/caffe2/core/operator_gpu_test.cc +++ b/caffe2/core/operator_gpu_test.cc @@ -48,7 +48,7 @@ TEST(EnginePrefTest, GPUDeviceDefaultPreferredEngines) { return; OperatorDef op_def; Workspace ws; - op_def.mutable_device_option()->set_device_type(CUDA); + op_def.mutable_device_option()->set_device_type(PROTO_CUDA); op_def.set_type("JustTest"); { diff --git a/caffe2/core/operator_schema.h b/caffe2/core/operator_schema.h index a87c1851b029a..0653de28c68b0 100644 --- a/caffe2/core/operator_schema.h +++ b/caffe2/core/operator_schema.h @@ -190,7 +190,7 @@ class CAFFE2_API OpSchema { uint64_t flops{0}; // Floating point operations. uint64_t bytes_read{0}; // Total memory read. uint64_t bytes_written{0}; // Total memory written. - uint64_t params_bytes{0}; // Memory footprint of parameters + uint64_t params_bytes{0}; // Memory read for parameters. }; /** * @brief Registers a function that takes in an OperatorDef @@ -518,6 +518,8 @@ inline vector GetDimsVector(const TensorShape& shape) { // Helper function inline uint64_t nElemFromDim(const TensorShape& X, int dim = 0) { + CAFFE_ENFORCE_GE(dim, 0, "Invalid maximum index specified"); + uint64_t nElem = 1; for (int i = dim; i < X.dims_size(); ++i) { nElem *= X.dims(i); @@ -525,6 +527,18 @@ inline uint64_t nElemFromDim(const TensorShape& X, int dim = 0) { return nElem; } +// Helper function +inline uint64_t nElemBetweenDim(const TensorShape& X, int start, int stop) { + CAFFE_ENFORCE_GE(start, 0, "Invalid maximum index specified"); + CAFFE_ENFORCE_LE(stop, X.dims_size(), "Invalid maximum index specified"); + + uint64_t nElem = 1; + for (int i = start; i < stop; ++i) { + nElem *= X.dims(i); + } + return nElem; +} + // Helper function for infer op inputs and outputs device information. inline std::pair, std::vector> InferOpInputOutputDevice(const OperatorDef& op) { diff --git a/caffe2/core/operator_test.cc b/caffe2/core/operator_test.cc index 93d85c97df3ef..1ce881dae172e 100644 --- a/caffe2/core/operator_test.cc +++ b/caffe2/core/operator_test.cc @@ -70,7 +70,7 @@ REGISTER_CUDA_OPERATOR(JustTest, JustTest); REGISTER_CPU_OPERATOR(JustTestWithSomeOutput, JustTestWithSomeOutput); TEST(OperatorTest, DeviceTypeRegistryWorks) { - EXPECT_EQ(gDeviceTypeRegistry()->count(DeviceType::CPU), 1); + EXPECT_EQ(gDeviceTypeRegistry()->count(CPU), 1); } TEST(OperatorTest, RegistryWorks) { @@ -83,7 +83,7 @@ TEST(OperatorTest, RegistryWorks) { // as it needs to instantiate an Event object with CUDAContext. Thus we will // guard this test below. if (HasCudaRuntime()) { - op_def.mutable_device_option()->set_device_type(CUDA); + op_def.mutable_device_option()->set_device_type(PROTO_CUDA); op = CreateOperator(op_def, &ws); EXPECT_NE(nullptr, op.get()); } @@ -93,7 +93,7 @@ TEST(OperatorTest, RegistryWrongDevice) { OperatorDef op_def; Workspace ws; op_def.set_type("JustTypeCPUOnly"); - op_def.mutable_device_option()->set_device_type(CUDA); + op_def.mutable_device_option()->set_device_type(PROTO_CUDA); try { CreateOperator(op_def, &ws); LOG(FATAL) << "No exception was thrown"; @@ -333,7 +333,7 @@ REGISTER_GRADIENT(Foo, GetFooGradient); TEST(OperatorGradientRegistryTest, GradientSimple) { Argument arg = MakeArgument("arg", 1); DeviceOption option; - option.set_device_type(CPU); + option.set_device_type(PROTO_CPU); OperatorDef def = CreateOperatorDef( "Foo", "", std::vector{"in"}, std::vector{"out"}, std::vector{arg}, option, "DUMMY_ENGINE"); @@ -351,7 +351,7 @@ TEST(OperatorGradientRegistryTest, GradientSimple) { EXPECT_EQ(grad_op.output(0), "in_grad"); // Checks the engine, device option and arguments. EXPECT_EQ(grad_op.engine(), "DUMMY_ENGINE"); - EXPECT_EQ(grad_op.device_option().device_type(), CPU); + EXPECT_EQ(grad_op.device_option().device_type(), PROTO_CPU); EXPECT_EQ(grad_op.arg_size(), 1); EXPECT_EQ(grad_op.arg(0).SerializeAsString(), MakeArgument("arg", 1).SerializeAsString()); @@ -366,7 +366,7 @@ TEST(EnginePrefTest, PerOpEnginePref) { Workspace ws; op_def.set_type("JustTest"); - SetPerOpEnginePref({{DeviceType::CPU, {{"JustTest", {"BAR"}}}}}); + SetPerOpEnginePref({{CPU, {{"JustTest", {"BAR"}}}}}); { const auto op = CreateOperator(op_def, &ws); EXPECT_NE(nullptr, op.get()); @@ -377,8 +377,7 @@ TEST(EnginePrefTest, PerOpEnginePref) { // Invalid operator type ASSERT_THROW( - SetPerOpEnginePref({{DeviceType::CPU, {{"NO_EXIST", {"BAR"}}}}}), - EnforceNotMet); + SetPerOpEnginePref({{CPU, {{"NO_EXIST", {"BAR"}}}}}), EnforceNotMet); } TEST(EnginePrefTest, GlobalEnginePref) { @@ -386,7 +385,7 @@ TEST(EnginePrefTest, GlobalEnginePref) { Workspace ws; op_def.set_type("JustTest"); - SetGlobalEnginePref({{DeviceType::CPU, {"FOO", "BAR"}}}); + SetGlobalEnginePref({{CPU, {"FOO", "BAR"}}}); { const auto op = CreateOperator(op_def, &ws); EXPECT_NE(nullptr, op.get()); @@ -395,7 +394,7 @@ TEST(EnginePrefTest, GlobalEnginePref) { // clear SetGlobalEnginePref({}); - SetGlobalEnginePref({{DeviceType::CPU, {"FOO"}}}); + SetGlobalEnginePref({{CPU, {"FOO"}}}); { const auto op = CreateOperator(op_def, &ws); EXPECT_NE(nullptr, op.get()); @@ -405,7 +404,8 @@ TEST(EnginePrefTest, GlobalEnginePref) { SetGlobalEnginePref({}); // Invalid device type - ASSERT_THROW(SetGlobalEnginePref({{8888, {"FOO"}}}), EnforceNotMet); + // This check is no longer necessary with the enum class + // ASSERT_THROW(SetGlobalEnginePref({{8888, {"FOO"}}}), EnforceNotMet); } TEST(EnginePrefTest, GlobalEnginePrefAndPerOpEnginePref) { @@ -413,8 +413,8 @@ TEST(EnginePrefTest, GlobalEnginePrefAndPerOpEnginePref) { Workspace ws; op_def.set_type("JustTest"); - SetPerOpEnginePref({{DeviceType::CPU, {{"JustTest", {"BAR"}}}}}); - SetGlobalEnginePref({{DeviceType::CPU, {"BAZ"}}}); + SetPerOpEnginePref({{CPU, {{"JustTest", {"BAR"}}}}}); + SetGlobalEnginePref({{CPU, {"BAZ"}}}); { const auto op = CreateOperator(op_def, &ws); EXPECT_NE(nullptr, op.get()); @@ -432,8 +432,8 @@ TEST(EnginePrefTest, GlobalEnginePrefAndPerOpEnginePrefAndOpDef) { op_def.set_type("JustTest"); op_def.set_engine("BAR"); - SetPerOpEnginePref({{DeviceType::CPU, {{"JustTest", {"BAZ"}}}}}); - SetGlobalEnginePref({{DeviceType::CPU, {"BAZ"}}}); + SetPerOpEnginePref({{CPU, {{"JustTest", {"BAZ"}}}}}); + SetGlobalEnginePref({{CPU, {"BAZ"}}}); { const auto op = CreateOperator(op_def, &ws); EXPECT_NE(nullptr, op.get()); @@ -450,8 +450,8 @@ TEST(EnginePrefTest, SetOpEnginePref) { Workspace ws; op_def.set_type("JustTest"); - SetPerOpEnginePref({{DeviceType::CPU, {{"JustTest", {"BAZ"}}}}}); - SetOpEnginePref("JustTest", {{DeviceType::CPU, {"BAR"}}}); + SetPerOpEnginePref({{CPU, {{"JustTest", {"BAZ"}}}}}); + SetOpEnginePref("JustTest", {{CPU, {"BAR"}}}); { const auto op = CreateOperator(op_def, &ws); EXPECT_NE(nullptr, op.get()); @@ -468,8 +468,8 @@ TEST(EnginePrefTest, SetDefaultEngine) { Workspace ws; op_def.set_type("JustTest"); - SetPerOpEnginePref({{DeviceType::CPU, {{"JustTest", {"DEFAULT"}}}}}); - SetGlobalEnginePref({{DeviceType::CPU, {"BAR"}}}); + SetPerOpEnginePref({{CPU, {{"JustTest", {"DEFAULT"}}}}}); + SetGlobalEnginePref({{CPU, {"BAR"}}}); { const auto op = CreateOperator(op_def, &ws); EXPECT_NE(nullptr, op.get()); diff --git a/caffe2/core/storage.h b/caffe2/core/storage.h index 5d833ae370c72..973b07ee63064 100644 --- a/caffe2/core/storage.h +++ b/caffe2/core/storage.h @@ -16,14 +16,15 @@ #include "caffe2/core/logging.h" #include "caffe2/core/typeid.h" +#include +#include #include namespace caffe2 { using DataType = TypeMeta; -// TODO: changed to DataPtr in Aten when shared folder -// is ready using DataPtr = std::shared_ptr; +using at::DeviceType; class CAFFE2_API StorageImpl : public c10::intrusive_ptr_target { public: @@ -150,7 +151,6 @@ class CAFFE2_API StorageImpl : public c10::intrusive_ptr_target { // allocator_ takes precedence over StaticContext from device_type_ // Allocator* allocator_; DeviceType device_type_ = CPU; - }; class CAFFE2_API Storage { diff --git a/caffe2/core/tensor.h b/caffe2/core/tensor.h index 8c4bed094f27c..392cb523c21b4 100644 --- a/caffe2/core/tensor.h +++ b/caffe2/core/tensor.h @@ -17,8 +17,6 @@ CAFFE2_DECLARE_int64(caffe2_max_keep_on_shrink_memory); namespace caffe2 { -using DimVector = std::vector; - /** * A utility function to convert vector to vector. */ @@ -97,12 +95,12 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { // and immediately discard it in Resize() since // reset_tensor will be true and FreeMemory will be called, // we might want to avoid creating Storage twice? - explicit TensorImpl(const vector& dims, DeviceType device_type) + explicit TensorImpl(const vector& dims, at::DeviceType device_type) : storage_(device_type) { Resize(dims); } - explicit TensorImpl(const vector& dims, DeviceType device_type) + explicit TensorImpl(const vector& dims, at::DeviceType device_type) : storage_(device_type) { Resize(dims); } @@ -113,16 +111,16 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { TensorImpl( const TensorImpl& src, BaseContext* context_for_copy, - DeviceType device_type) + at::DeviceType device_type) : storage_(device_type) { CopyFrom(src, context_for_copy); } /** - * @brief: Create a Tensor of DeviceType `type` and initialize it with + * @brief: Create a Tensor of at::DeviceType `type` and initialize it with * src Tensor */ - TensorImpl(const TensorImpl& src, DeviceType device_type) + TensorImpl(const TensorImpl& src, at::DeviceType device_type) : storage_(device_type) { CopyFrom(src); } @@ -192,7 +190,7 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { return GetStaticContext()->CreateContext(); } - DeviceType GetDeviceType() const { + at::DeviceType GetDeviceType() const { return storage_.device_type(); } @@ -204,16 +202,11 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { if ((void*)&src == (void*)this) { return; } - CAFFE_ENFORCE_WITH_CALLER( - src.is_contiguous(), - "Source Tensor must be contiguous in order to be copied."); if (storage_.dtype() != src.meta()) { storage_ = Storage(GetDeviceType(), src.meta()); } if (src.size() == -1) { dims_.clear(); - strides_.clear(); - is_contiguous_ = true; numel_ = -1; storage_.reset(); return; @@ -275,8 +268,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { * complexity. */ void Extend(TIndex num, float growthPct, BaseContext* context) { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, "Tensor must be contiguous in order to call Extend."); CAFFE_ENFORCE_GE_WITH_CALLER(dims_.size(), 1); CAFFE_ENFORCE_GE_WITH_CALLER( num, 0, "`num` must be non-negative for Extend"); @@ -320,8 +311,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { * that the extra capacity after the end of the shurnk tensor is maintained. */ void ShrinkTo(TIndex outer_dim) { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, "Tensor must be contiguous in order to call ShrinkTo."); CAFFE_ENFORCE_WITH_CALLER(dims_.size() >= 1, "Tensor must be at least 1D"); CAFFE_ENFORCE_WITH_CALLER( outer_dim <= dims_[0], @@ -345,9 +334,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { */ template void ReserveSpace(const T& outer_dim) { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, - "Tensor must be contiguous in order to call ReserveSpace."); CAFFE_ENFORCE( numel_ != -1, "size should be initialized before calling ReserveSpace"); CAFFE_ENFORCE( @@ -417,9 +403,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { * sugar wrapper that essentially calls Resize(src_tensor.dims()). */ inline void ResizeLike(const TensorImpl& src_tensor) { - CAFFE_ENFORCE_WITH_CALLER( - src_tensor.is_contiguous(), - "Tensor must be contiguous in order to call ResizeLike."); // Note: need casting for different context types. if (static_cast(this) != static_cast(&src_tensor)) { Resize(src_tensor.dims()); @@ -431,8 +414,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { * This requires the total size of the tensor to remains constant. */ inline void Reshape(const vector& dims) { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, "Tensor must be contiguous in order to call Reshape."); TIndex new_size = 1; for (auto d : dims) { CAFFE_ENFORCE_GE_WITH_CALLER(d, 0); @@ -482,8 +463,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { void swap(TensorImpl& other) noexcept { std::swap(dims_, other.dims_); - std::swap(strides_, other.strides_); - std::swap(is_contiguous_, other.is_contiguous_); std::swap(numel_, other.numel_); std::swap(storage_, other.storage_); } @@ -541,9 +520,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { const TypeMeta& data_type, size_t capacity = 0, Deleter d = nullptr) { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, - "Tensor must be contiguous in order to call ShareExternalPointer."); CAFFE_ENFORCE_WITH_CALLER( data_type.id() != TypeIdentifier::uninitialized(), "To share with a raw external pointer you need to pass in an " @@ -567,9 +543,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { * or raw_mutable_data() must have been called prior to this function call. */ inline const void* raw_data() const { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, - "Tensor must be contiguous in order to call raw_data()"); CAFFE_ENFORCE_WITH_CALLER(storage_.data() || numel_ == 0); return storage_.data(); } @@ -582,8 +555,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { */ template inline const T* data() const { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, "Tensor must be contiguous in order to call data()"); CAFFE_ENFORCE_WITH_CALLER( storage_.data() || numel_ == 0, "The tensor is of non-zero shape, but its data is not allocated yet. " @@ -610,9 +581,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { * and a new storage will be created. */ inline void* raw_mutable_data(const TypeMeta& meta) { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, - "Tensor must be contiguous in order to call raw_mutable_data()"); // For 0-size tensors it's fine to return any pointer (including nullptr) if (storage_.dtype() == meta && (storage_.data() || numel_ == 0)) { return storage_.data(); @@ -677,9 +645,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { * and a new storage will be created. */ inline void* raw_mutable_data() { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, - "Tensor must be contiguous in order to call raw_mutable_data()"); CAFFE_ENFORCE_WITH_CALLER( storage_.dtype().id() != TypeIdentifier::uninitialized(), "Calling raw_mutable_data() without meta, but the current meta is " @@ -695,9 +660,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { */ template inline T* mutable_data() { - CAFFE_ENFORCE_WITH_CALLER( - is_contiguous_, - "Tensor must be contiguous in order to call mutable_data()"); if ((numel_ == 0 || storage_.data()) && IsType()) { return static_cast(storage_.data()); } @@ -774,25 +736,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { return canonical_axis_index_(axis_index, ndim()); } - inline int64_t stride(int64_t dim) const { -#ifndef NDEBUG - // TODO: dim wrapping? - CAFFE_ENFORCE_LT_WITH_CALLER(dim, strides_.size(), "Exceeding ndim limit"); - CAFFE_ENFORCE_GE_WITH_CALLER( - dim, 0, "Cannot have negative dimension index"); -#endif - return strides_[dim]; - } - - // TODO: Change to ArrayRef later - inline DimVector strides() { - return strides_; - } - - inline bool is_contiguous() const { - return is_contiguous_; - } - /** * Checks if the tensor content is of the given data type. */ @@ -849,10 +792,9 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { } protected: + using DimVector = std::vector; DimVector dims_; // sizes_ - DimVector strides_; TIndex numel_ = -1; // numel_ - bool is_contiguous_ = true; // we decide to keep reserved_ and it will // live in Tensor after the split // The logic is that if Extend() or ReserveSpace() were ever called, @@ -873,7 +815,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { new_numel *= src[i]; dims_[i] = src[i]; } - update_strides(); numel_ = new_numel; return numel_ != old_numel; } @@ -881,7 +822,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { bool SetDims() { auto old_numel = numel_; dims_.resize(0); - update_strides(); numel_ = 1; return numel_ != old_numel; } @@ -893,7 +833,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { auto old_numel = numel_; dims_.resize(1); dims_[0] = d0; - update_strides(); numel_ = d0; return numel_ != old_numel; } @@ -903,7 +842,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { dims_.resize(2); dims_[0] = d0; dims_[1] = d1; - update_strides(); numel_ = d0 * d1; return numel_ != old_numel; } @@ -914,7 +852,6 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { dims_[0] = d0; dims_[1] = d1; dims_[2] = d2; - update_strides(); numel_ = d0 * d1 * d2; return numel_ != old_numel; } @@ -927,21 +864,9 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target { dims_[1] = d1; dims_[2] = d2; dims_[3] = d3; - update_strides(); numel_ = d0 * d1 * d2 * d3; return numel_ != old_numel; } - inline void update_strides() { - strides_.resize(dims_.size()); - for (auto i = ndim() - 1; i >= 0; --i) { - if (i == ndim() - 1) { - strides_[i] = 1; - } else { - strides_[i] = strides_[i + 1] * std::max(dims_[i + 1], 1); - } - } - is_contiguous_ = true; - } }; class CAFFE2_API UndefinedTensorImpl final : public TensorImpl { @@ -1178,18 +1103,6 @@ class CAFFE2_API Tensor final { return impl_.get()->canonical_axis_index(axis_index); } - inline int64_t stride(int64_t dim) const { - return impl_.get()->stride(dim); - } - - inline DimVector strides() { - return impl_.get()->strides(); - } - - inline bool is_contiguous() const { - return impl_.get()->is_contiguous(); - } - template inline bool IsType() const { return impl_.get()->IsType(); diff --git a/caffe2/distributed/store_ops.h b/caffe2/distributed/store_ops.h index 20edde426891d..b0097a393e3d1 100644 --- a/caffe2/distributed/store_ops.h +++ b/caffe2/distributed/store_ops.h @@ -6,7 +6,7 @@ namespace caffe2 { -class CAFFE2_API StoreSetOp final : public Operator { +class StoreSetOp final : public Operator { public: StoreSetOp(const OperatorDef& operator_def, Workspace* ws); bool RunOnDevice() override; @@ -17,7 +17,7 @@ class CAFFE2_API StoreSetOp final : public Operator { INPUT_TAGS(HANDLER, DATA); }; -class CAFFE2_API StoreGetOp final : public Operator { +class StoreGetOp final : public Operator { public: StoreGetOp(const OperatorDef& operator_def, Workspace* ws); bool RunOnDevice() override; @@ -29,7 +29,7 @@ class CAFFE2_API StoreGetOp final : public Operator { OUTPUT_TAGS(DATA); }; -class CAFFE2_API StoreAddOp final : public Operator { +class StoreAddOp final : public Operator { public: StoreAddOp(const OperatorDef& operator_def, Workspace* ws); bool RunOnDevice() override; @@ -42,7 +42,7 @@ class CAFFE2_API StoreAddOp final : public Operator { OUTPUT_TAGS(VALUE); }; -class CAFFE2_API StoreWaitOp final : public Operator { +class StoreWaitOp final : public Operator { public: StoreWaitOp(const OperatorDef& operator_def, Workspace* ws); bool RunOnDevice() override; diff --git a/caffe2/ideep/operators/operator_fallback_ideep.h b/caffe2/ideep/operators/operator_fallback_ideep.h index 31df729a21785..9ae2323442120 100644 --- a/caffe2/ideep/operators/operator_fallback_ideep.h +++ b/caffe2/ideep/operators/operator_fallback_ideep.h @@ -43,12 +43,12 @@ class IDEEPFallbackOp final : public IDEEPOperator { IDEEPFallbackOp(const OperatorDef& def, Workspace* ws) : IDEEPOperator(def, ws) { - CAFFE_ENFORCE_EQ(def.device_option().device_type(), IDEEP); + CAFFE_ENFORCE_EQ(def.device_option().device_type(), PROTO_IDEEP); base_def_.CopyFrom(def); // base_def_ runs on CPU, so we will set its device option to CPU. // Copy to allow random_seed to be correctly propagated. base_def_.mutable_device_option()->CopyFrom(def.device_option()); - base_def_.mutable_device_option()->set_device_type(CPU); + base_def_.mutable_device_option()->set_device_type(PROTO_CPU); // Create output blobs in parent workspace, // then forward output blobs to local workspace. std::unordered_map forwarded_output_blobs; diff --git a/caffe2/ideep/utils/ideep_context.h b/caffe2/ideep/utils/ideep_context.h index c7215e0ed28b3..d84e45b1facb8 100644 --- a/caffe2/ideep/utils/ideep_context.h +++ b/caffe2/ideep/utils/ideep_context.h @@ -18,7 +18,7 @@ class IDEEPContext final : public BaseContext { : random_seed_( option.has_random_seed() ? option.random_seed() : RandomNumberSeed()) { - CAFFE_ENFORCE_EQ(option.device_type(), IDEEP); + CAFFE_ENFORCE_EQ(option.device_type(), PROTO_IDEEP); } ~IDEEPContext() noexcept override {} diff --git a/caffe2/mkl/mklmemory_serialization.cc b/caffe2/mkl/mklmemory_serialization.cc index 35f787b8f6f61..e59a9a15f4226 100644 --- a/caffe2/mkl/mklmemory_serialization.cc +++ b/caffe2/mkl/mklmemory_serialization.cc @@ -26,7 +26,7 @@ class MKLMemorySerializer : public BlobSerializerBase { blob_proto.set_type(kTensorBlobType); TensorProto* proto = blob_proto.mutable_tensor(); auto* device_detail = proto->mutable_device_detail(); - device_detail->set_device_type(MKLDNN); + device_detail->set_device_type(PROTO_MKLDNN); proto->set_name(name); if (blob.IsType>()) { const MKLMemory& src = blob.Get>(); diff --git a/caffe2/mkl/operators/operator_fallback_mkl.h b/caffe2/mkl/operators/operator_fallback_mkl.h index 4a7cf3845fee5..ce4c85d2c231e 100644 --- a/caffe2/mkl/operators/operator_fallback_mkl.h +++ b/caffe2/mkl/operators/operator_fallback_mkl.h @@ -44,12 +44,12 @@ class MKLFallbackOp final : public Operator { USE_OPERATOR_FUNCTIONS(MKLContext); MKLFallbackOp(const OperatorDef& def, Workspace* ws) : Operator(def, ws) { - CAFFE_ENFORCE_EQ(def.device_option().device_type(), MKLDNN); + CAFFE_ENFORCE_EQ(def.device_option().device_type(), PROTO_MKLDNN); OperatorDef base_def_(def); // base_def_ runs on CPU, so we will set its device option to CPU. // Copy to allow random_seed to be correctly propagated. base_def_.mutable_device_option()->CopyFrom(def.device_option()); - base_def_.mutable_device_option()->set_device_type(CPU); + base_def_.mutable_device_option()->set_device_type(PROTO_CPU); // Set up the symbols for the local workspace. for (const string& name : def.input()) { local_input_blobs_.push_back(local_ws_.CreateBlob(name)); diff --git a/caffe2/mkl/operators/spatial_batch_norm_op.cc b/caffe2/mkl/operators/spatial_batch_norm_op.cc index 41ee61585d71f..7d9856f1adfde 100644 --- a/caffe2/mkl/operators/spatial_batch_norm_op.cc +++ b/caffe2/mkl/operators/spatial_batch_norm_op.cc @@ -9,10 +9,16 @@ namespace caffe2 { namespace mkl { template -class MKLBNOp final : public SpatialBNOp { +class MKLBNOp final : public Operator { public: MKLBNOp(const OperatorDef& operator_def, Workspace* ws) - : SpatialBNOp(operator_def, ws) { + : Operator(operator_def, ws), + OP_SINGLE_ARG(int, OpSchema::Arg_IsTest, is_test_, 0), + OP_SINGLE_ARG(float, "epsilon", epsilon_, 1e-5f), + OP_SINGLE_ARG(float, "momentum", momentum_, 0.9f), + order_(StringToStorageOrder( + this->template GetSingleArgument("order", "NCHW"))), + OP_SINGLE_ARG(int, "num_batches", num_batches_, 1) { OPERATOR_NEEDS_FEATURE( order_ == StorageOrder::NCHW, "Only NCHW order supported."); OPERATOR_NEEDS_FEATURE( @@ -20,7 +26,7 @@ class MKLBNOp final : public SpatialBNOp { "Inplace BN not supported"); } - bool RunOnDevice() { + bool RunOnDevice() override { auto& X = OperatorBase::Input>(INPUT); auto& scale = OperatorBase::Input>(SCALE); auto& bias = OperatorBase::Input>(BIAS); @@ -34,7 +40,7 @@ class MKLBNOp final : public SpatialBNOp { MKLMemory* saved_mean = OperatorBase::Output>(SAVED_MEAN & (is_test_ - 1)); MKLMemory* saved_var = - OperatorBase::Output>(SAVED_INV_VAR & (is_test_ - 1)); + OperatorBase::Output>(SAVED_INV_STD & (is_test_ - 1)); // current code supports only NCHW - // have to look for MKL related changes for NHWC later @@ -134,6 +140,12 @@ class MKLBNOp final : public SpatialBNOp { } private: + const bool is_test_; + float epsilon_; + const float momentum_; + const StorageOrder order_; + const int num_batches_; + vector cached_input_dims_; LayoutWrapper scale_bias_layout_; LayoutWrapper saved_mean_layout_; @@ -150,6 +162,16 @@ class MKLBNOp final : public SpatialBNOp { PrimitiveWrapper primitive_; MKLMemory buffer_; void* resources_[dnnResourceNumber] = {0}; + + INPUT_TAGS( + INPUT, + SCALE, + BIAS, + EST_MEAN, + EST_VAR, + BATCH_MEAN_SUM, + BATCH_VAR_SUM); + OUTPUT_TAGS(OUTPUT, RUNNING_MEAN, RUNNING_VAR, SAVED_MEAN, SAVED_INV_STD); }; } // namespace mkl diff --git a/caffe2/mkl/utils/mkl_context.h b/caffe2/mkl/utils/mkl_context.h index 20936fc643a2d..18c492694db99 100644 --- a/caffe2/mkl/utils/mkl_context.h +++ b/caffe2/mkl/utils/mkl_context.h @@ -27,7 +27,7 @@ class MKLContext : public BaseContext { : random_seed_( option.has_random_seed() ? option.random_seed() : RandomNumberSeed()) { - CAFFE_ENFORCE_EQ(option.device_type(), MKLDNN); + CAFFE_ENFORCE_EQ(option.device_type(), PROTO_MKLDNN); } ~MKLContext() override {} diff --git a/caffe2/mkl/utils/mkl_memory.cc b/caffe2/mkl/utils/mkl_memory.cc index 97163d5d483aa..26e423220e3f0 100644 --- a/caffe2/mkl/utils/mkl_memory.cc +++ b/caffe2/mkl/utils/mkl_memory.cc @@ -25,7 +25,7 @@ static vector GetMKLTensorInfo( DeviceOption* device) { const mkl::MKLMemory* tc = static_cast*>(c); *capacity = tc->size() * sizeof(T); - device->set_device_type(MKLDNN); + device->set_device_type(PROTO_MKLDNN); device->set_cuda_gpu_id(0); return tc->dims(); } diff --git a/caffe2/mobile/contrib/arm-compute/core/context.h b/caffe2/mobile/contrib/arm-compute/core/context.h index 1869eb8d56078..4085e4983cc81 100644 --- a/caffe2/mobile/contrib/arm-compute/core/context.h +++ b/caffe2/mobile/contrib/arm-compute/core/context.h @@ -44,7 +44,7 @@ class GLContext final { static bool initialized; explicit GLContext(); explicit GLContext(const DeviceOption &option) { - DCHECK_EQ(option.device_type(), OPENGL); + DCHECK_EQ(option.device_type(), PROTO_OPENGL); GLContext(); } ~GLContext() {} diff --git a/caffe2/mobile/contrib/arm-compute/core/rewrite_net.cc b/caffe2/mobile/contrib/arm-compute/core/rewrite_net.cc index 764b44460a723..c714731acc38e 100644 --- a/caffe2/mobile/contrib/arm-compute/core/rewrite_net.cc +++ b/caffe2/mobile/contrib/arm-compute/core/rewrite_net.cc @@ -230,7 +230,7 @@ NetDef rewritePredictNetForOpenGL(const NetDef& predictNet, bool runFusion, std: for (auto i = 0; i < net.op().size(); ++i) { auto op = net.mutable_op(i); if (std::find(cpuOps.begin(), cpuOps.end(), op->type()) == cpuOps.end()) { - op->mutable_device_option()->set_device_type(OPENGL); + op->mutable_device_option()->set_device_type(PROTO_OPENGL); } } diff --git a/caffe2/mobile/contrib/arm-compute/test/gl_model_test.h b/caffe2/mobile/contrib/arm-compute/test/gl_model_test.h index 27e543dd53a49..e7f389abf8bfd 100644 --- a/caffe2/mobile/contrib/arm-compute/test/gl_model_test.h +++ b/caffe2/mobile/contrib/arm-compute/test/gl_model_test.h @@ -51,7 +51,7 @@ namespace caffe2 { for (auto i = 0; i < predict_net_def.op().size(); ++i) { auto op = predict_net_def.mutable_op(i); if (std::find(cpu_ops.begin(), cpu_ops.end(), op->type()) == cpu_ops.end()) { - op->mutable_device_option()->set_device_type(CPU); + op->mutable_device_option()->set_device_type(PROTO_CPU); } } predict_net_def.set_type("simple"); diff --git a/caffe2/mobile/contrib/arm-compute/test/gl_operator_test.h b/caffe2/mobile/contrib/arm-compute/test/gl_operator_test.h index 50b457c7ba86d..daa7ef008fc7b 100644 --- a/caffe2/mobile/contrib/arm-compute/test/gl_operator_test.h +++ b/caffe2/mobile/contrib/arm-compute/test/gl_operator_test.h @@ -7,12 +7,12 @@ namespace caffe2 { -#define DECLARE_OPENGL_OPERATOR(_name) \ - OperatorDef _name; \ - _name.mutable_device_option()->set_device_type(OPENGL); +#define DECLARE_OPENGL_OPERATOR(_name) \ + OperatorDef _name; \ + _name.mutable_device_option()->set_device_type(PROTO_OPENGL); -#define MAKE_OPENGL_OPERATOR(_op) \ - _op->mutable_device_option()->set_device_type(OPENGL); +#define MAKE_OPENGL_OPERATOR(_op) \ + _op->mutable_device_option()->set_device_type(PROTO_OPENGL); #define ADD_ARG(_op, _name, _type, _val) \ { \ diff --git a/caffe2/mpi/mpi_ops_gpu.cc b/caffe2/mpi/mpi_ops_gpu.cc index b28a31f00cda3..5a16bfa2014b2 100644 --- a/caffe2/mpi/mpi_ops_gpu.cc +++ b/caffe2/mpi/mpi_ops_gpu.cc @@ -60,19 +60,19 @@ REGISTER_CUDA_OPERATOR(MPIAllgather, MPIAllgatherOp); REGISTER_CUDA_OPERATOR(MPISendTensor, MPISendTensorOp); REGISTER_CUDA_OPERATOR(MPIReceiveTensor, MPIReceiveTensorOp); #else -REGISTER_CUDA_OPERATOR(MPIBroadcast, GPUFallbackOp>); +REGISTER_CUDA_OPERATOR(MPIBroadcast, GPUFallbackOp); REGISTER_CUDA_OPERATOR( MPIReduce, - GPUFallbackOp>); + GPUFallbackOp); REGISTER_CUDA_OPERATOR( MPIAllgather, - GPUFallbackOp>); + GPUFallbackOp); REGISTER_CUDA_OPERATOR( MPISendTensor, - GPUFallbackOp>); + GPUFallbackOp); REGISTER_CUDA_OPERATOR( MPIReceiveTensor, - GPUFallbackOp, SkipIndices<1, 2>>); + GPUFallbackOpEx>); #endif #if CAFFE2_HAS_CUDA_MPI_ALLREDUCE @@ -80,7 +80,7 @@ REGISTER_CUDA_OPERATOR(MPIAllreduce, MPIAllreduceOp); #else REGISTER_CUDA_OPERATOR( MPIAllreduce, - GPUFallbackOp>); + GPUFallbackOp); #endif } // namespace caffe2 diff --git a/caffe2/operators/activation_ops_cudnn.h b/caffe2/operators/activation_ops_cudnn.h index cc682ced0c0d0..5b1ebec13bd2c 100644 --- a/caffe2/operators/activation_ops_cudnn.h +++ b/caffe2/operators/activation_ops_cudnn.h @@ -19,7 +19,7 @@ class CuDNNActivationOpBase : public Operator { CUDNN_ENFORCE(cudnnCreateActivationDescriptor(&act_desc_)); } - ~CuDNNActivationOpBase() { + virtual ~CuDNNActivationOpBase() { CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(data_desc_)); CUDNN_ENFORCE(cudnnDestroyActivationDescriptor(act_desc_)); } diff --git a/caffe2/operators/affine_channel_op.h b/caffe2/operators/affine_channel_op.h index 8e4e009f08e7e..e271a04581045 100644 --- a/caffe2/operators/affine_channel_op.h +++ b/caffe2/operators/affine_channel_op.h @@ -43,24 +43,13 @@ class AffineChannelOp final : public Operator { const int N = X.dim32(0); const int C = X.dim32(1); const int HxW = X.size() / (N * C); - const std::array X_dims = {N, C, HxW}; - const std::array scale_dims = {1, C, 1}; Y->ResizeLike(X); - math::Mul( - 3, - X_dims.data(), - 3, - scale_dims.data(), + math::AffineChannel( + N, + C, + HxW, X.template data(), scale.template data(), - Y->template mutable_data(), - &context_); - math::Add( - 3, - X_dims.data(), - 3, - scale_dims.data(), - Y->template data(), bias.template data(), Y->template mutable_data(), &context_); @@ -80,21 +69,16 @@ class AffineChannelOp final : public Operator { "is_learnable = true."); } const int ndim = X.ndim(); + const int N = X.dim32(0); const int C = X.dim32(ndim - 1); - const int rows = X.size() / C; - const int cols = C; + const int HxW = X.size() / (N * C); Y->ResizeLike(X); - math::RowwiseMul( - rows, - cols, + math::AffineChannel( + N, + C, + HxW, X.template data(), scale.template data(), - Y->template mutable_data(), - &context_); - math::RowwiseAdd( - rows, - cols, - Y->template data(), bias.template data(), Y->template mutable_data(), &context_); diff --git a/caffe2/operators/batch_matmul_op_gpu_test.cc b/caffe2/operators/batch_matmul_op_gpu_test.cc index 33a5363b0afc3..57a09e3e60c73 100644 --- a/caffe2/operators/batch_matmul_op_gpu_test.cc +++ b/caffe2/operators/batch_matmul_op_gpu_test.cc @@ -15,14 +15,14 @@ class BatchMatMulOpGPUTest : public testing::Test { if (!HasCudaGPU()) { return; } - option_.set_device_type(CUDA); + option_.set_device_type(PROTO_CUDA); cuda_context_ = make_unique(option_); def_.set_name("test"); def_.set_type("BatchMatMul"); def_.add_input("A"); def_.add_input("B"); def_.add_output("Y"); - def_.mutable_device_option()->set_device_type(CUDA); + def_.mutable_device_option()->set_device_type(PROTO_CUDA); } void AddConstInput( diff --git a/caffe2/operators/conv_gradient_op.cc b/caffe2/operators/conv_gradient_op.cc index 8e357e1224c74..e50e5407ccc3a 100644 --- a/caffe2/operators/conv_gradient_op.cc +++ b/caffe2/operators/conv_gradient_op.cc @@ -4,8 +4,89 @@ namespace caffe2 { +std::vector TensorInferenceForConvGradient( + const OperatorDef& def, + const std::vector& in) { + CAFFE_ENFORCE_EQ(in.size(), 3, "ConvGradient requires 3 inputs"); + + if (in[0].unknown_shape()) { + std::vector out(1); + out[0].set_unknown_shape(true); + return out; + } + ArgumentHelper helper(def); + const auto no_bias = helper.GetSingleArgument("no_bias", 0); + const auto n_outputs = def.output_size(); + vector out(n_outputs); + + // FILTER_GRAD has the same shape as FILTER + out[0] = in[1]; + if (!no_bias) { + vector bias_shape = {in[1].dims(0)}; + out[1] = CreateTensorShape(bias_shape, in[1].data_type()); + } + + if (n_outputs == 3 || (no_bias && n_outputs == 2)) { + // INPUT_GRAD has the same shape as INPUT + out[out.size() - 1] = in[0]; + } + + return out; +} + +OpSchema::Cost CostInferenceForConvGradient( + const OperatorDef& def, + const vector& inputs) { + CAFFE_ENFORCE_EQ(inputs.size(), 3, "ConvGradient requires 3 inputs"); + ArgumentHelper helper(def); + const auto order = + StringToStorageOrder(helper.GetSingleArgument("order", "NCHW")); + const auto no_bias = helper.GetSingleArgument("no_bias", 0); + const auto n_outputs = def.output_size(); + + const auto& outputs = TensorInferenceForConvGradient(def, inputs); + const auto& X = inputs[0]; + const auto& filter = inputs[1]; + const auto& dY = inputs[2]; + const auto N = X.dims(0); + const auto M = filter.dims(0); + const auto C = + (order == StorageOrder::NCHW ? X.dims(1) : X.dims(X.dims_size() - 1)); + const auto output_image_size = + (order == StorageOrder::NCHW + ? nElemFromDim(dY, 2) + : nElemBetweenDim(dY, 1, dY.dims_size() - 1)); + auto kernel_elem = + (order == StorageOrder::NCHW + ? nElemFromDim(filter, 2) + : nElemBetweenDim(filter, 1, filter.dims_size() - 1)); + + struct OpSchema::Cost c; + c.flops = N * 2 * M * kernel_elem * C * output_image_size; + if (!no_bias) { + c.flops += N * (M * output_image_size); + } + if (n_outputs == 3 || (no_bias && n_outputs == 2)) { + c.flops += N * 2 * M * kernel_elem * C * output_image_size; + } + + c.bytes_read = (nElemFromDim(X) + nElemFromDim(filter) + nElemFromDim(dY)) * + sizeof(float); + + for (auto i = 0; i < n_outputs; i++) { + c.bytes_written += nElemFromDim(outputs[i]) * sizeof(float); + } + c.params_bytes = nElemFromDim(filter) * sizeof(float); + + return c; +} + REGISTER_CPU_OPERATOR(ConvGradient, ConvGradientOp); -OPERATOR_SCHEMA(ConvGradient).NumInputs(2, 3).NumOutputs(1, 3); +OPERATOR_SCHEMA(ConvGradient) + .NumInputs(2, 3) + .NumOutputs(1, 3) + .TensorInferenceFunction(TensorInferenceForConvGradient) + .CostInferenceFunction(CostInferenceForConvGradient); REGISTER_CPU_OPERATOR(Conv1DGradient, ConvGradientOp); OPERATOR_SCHEMA(Conv1DGradient).NumInputs(2, 3).NumOutputs(1, 3); diff --git a/caffe2/operators/cross_entropy_op.cu b/caffe2/operators/cross_entropy_op.cu index cab3c9692a42f..7030e846b714e 100644 --- a/caffe2/operators/cross_entropy_op.cu +++ b/caffe2/operators/cross_entropy_op.cu @@ -454,9 +454,7 @@ REGISTER_CUDA_OPERATOR(MakeTwoClassGradient, MakeTwoClassGradientOp); //TODO(surya) Add full GPU/CUDA support for the CrossEntropyOp -REGISTER_CUDA_OPERATOR(CrossEntropy, - GPUFallbackOp>); -REGISTER_CUDA_OPERATOR(CrossEntropyGradient, - GPUFallbackOp>); +REGISTER_CUDA_OPERATOR(CrossEntropy, GPUFallbackOp); +REGISTER_CUDA_OPERATOR(CrossEntropyGradient, GPUFallbackOp); } // namespace caffe2 diff --git a/caffe2/operators/elementwise_op_gpu_test.cc b/caffe2/operators/elementwise_op_gpu_test.cc index e652a9e551f40..d3c56af57db10 100644 --- a/caffe2/operators/elementwise_op_gpu_test.cc +++ b/caffe2/operators/elementwise_op_gpu_test.cc @@ -13,7 +13,7 @@ void CopyVector(const int N, const bool* x, bool* y) { template <> caffe2::OperatorDef CreateOperatorDef() { caffe2::OperatorDef def; - def.mutable_device_option()->set_device_type(caffe2::CUDA); + def.mutable_device_option()->set_device_type(caffe2::PROTO_CUDA); return def; } diff --git a/caffe2/operators/filler_op.cc b/caffe2/operators/filler_op.cc index c5a121e3a222d..9402e93560df3 100644 --- a/caffe2/operators/filler_op.cc +++ b/caffe2/operators/filler_op.cc @@ -349,7 +349,7 @@ OPERATOR_SCHEMA(UniformIntFill) .NumInputs({0, 1, 3}) .NumOutputs(1) .AllowInplace({{0, 0}}) - .TensorInferenceFunction(FillerTensorInference<>) + .TensorInferenceFunction(FillerTensorInference) .SetDoc(R"DOC( Fill the output tensor with int32 samples from uniform distribution [`min`, `max`]. diff --git a/caffe2/operators/filler_op.cu b/caffe2/operators/filler_op.cu index 65918cc04b5d5..a754d361442e3 100644 --- a/caffe2/operators/filler_op.cu +++ b/caffe2/operators/filler_op.cu @@ -64,8 +64,6 @@ REGISTER_CUDA_OPERATOR(GaussianFill, GaussianFillOp); REGISTER_CUDA_OPERATOR(XavierFill, XavierFillOp); REGISTER_CUDA_OPERATOR(MSRAFill, MSRAFillOp); REGISTER_CUDA_OPERATOR(RangeFill, RangeFillOp); -REGISTER_CUDA_OPERATOR( - LengthsRangeFill, - GPUFallbackOp>); +REGISTER_CUDA_OPERATOR(LengthsRangeFill, GPUFallbackOp); } // namespace caffe2 diff --git a/caffe2/operators/group_norm_op.cc b/caffe2/operators/group_norm_op.cc index 9bc6316e8ec58..4cb6b7a1ae830 100644 --- a/caffe2/operators/group_norm_op.cc +++ b/caffe2/operators/group_norm_op.cc @@ -156,23 +156,13 @@ bool GroupNormOp::RunOnDeviceImpl( T* Y_data, T* mu_data, T* rsig_data) { - const std::array dims = order_ == StorageOrder::NCHW - ? std::array{N, G, D, HxW} - : std::array{N, HxW, G, D}; - const std::array axes = order_ == StorageOrder::NCHW - ? std::array{2, 3} - : std::array{1, 3}; - - // Computes mean and variance. - math::Moments( - 4, dims.data(), 2, axes.data(), X_data, mu_data, rsig_data, &context_); - - // Uses rsqrt to computes 1 / std which is much faster than computes std. - EigenArrayMap(rsig_data, G, N) += epsilon_; - math::Rsqrt(N * G, rsig_data, rsig_data, &context_); - - // Computes Y = gamma * (X - mu) * rsig + beta. if (order_ == StorageOrder::NCHW) { + const std::array dims = {N * G, D * HxW}; + const int axis = 1; + math::Moments( + 2, dims.data(), 1, &axis, X_data, mu_data, rsig_data, &context_); + math::InvStd( + N * G, static_cast(epsilon_), rsig_data, rsig_data, &context_); GroupNormForwardNCHW( N, G, @@ -185,6 +175,12 @@ bool GroupNormOp::RunOnDeviceImpl( beta_data, Y_data); } else { + const std::array dims = {N, HxW, G, D}; + const std::array axes = {1, 3}; + math::Moments( + 4, dims.data(), 2, axes.data(), X_data, mu_data, rsig_data, &context_); + math::InvStd( + N * G, static_cast(epsilon_), rsig_data, rsig_data, &context_); GroupNormForwardNHWC( N, G, diff --git a/caffe2/operators/group_norm_op.cu b/caffe2/operators/group_norm_op.cu index cfdd308bceb05..161a914b53d1f 100644 --- a/caffe2/operators/group_norm_op.cu +++ b/caffe2/operators/group_norm_op.cu @@ -23,20 +23,6 @@ namespace { template using BlockReduce = cub::BlockReduce; -__global__ void InvStdCUDAKernel( - const int size, - const float epsilon, - const float* var, - float* rsig) { - CUDA_1D_KERNEL_LOOP(i, size) { -#if __CUDA_ARCH__ >= 350 - rsig[i] = rsqrtf(__ldg(var + i) + epsilon); -#else - rsig[i] = rsqrtf(var[i] + epsilon); -#endif - } -} - template __global__ void GroupNormForwardCUDAKernel( const int size, @@ -212,27 +198,14 @@ bool GroupNormOp::RunOnDeviceImpl( float* Y_data, float* mu_data, float* rsig_data) { - const std::array dims = order_ == StorageOrder::NCHW - ? std::array{N, G, D, HxW} - : std::array{N, HxW, G, D}; - const std::array axes = order_ == StorageOrder::NCHW - ? std::array{2, 3} - : std::array{1, 3}; - - // Computes mean and variance. - math::Moments( - 4, dims.data(), 2, axes.data(), X_data, mu_data, rsig_data, &context_); - - // Uses rsqrt to computes 1 / std which is much faster than computes std. - InvStdCUDAKernel<<< - CAFFE_GET_BLOCKS(N * G), - CAFFE_CUDA_NUM_THREADS, - 0, - context_.cuda_stream()>>>(N * G, epsilon_, rsig_data, rsig_data); - - // Computes Y = gamma * (X - mu) * rsig + beta. const int size = N * G * D * HxW; if (order_ == StorageOrder::NCHW) { + const std::array dims = {N * G, D * HxW}; + const int axis = 1; + math::Moments( + 2, dims.data(), 1, &axis, X_data, mu_data, rsig_data, &context_); + math::InvStd( + N * G, epsilon_, rsig_data, rsig_data, &context_); GroupNormForwardCUDAKernel <<::RunOnDeviceImpl( beta_data, Y_data); } else { + const std::array dims = {N, HxW, G, D}; + const std::array axes = {1, 3}; + math::Moments( + 4, dims.data(), 2, axes.data(), X_data, mu_data, rsig_data, &context_); + math::InvStd( + N * G, epsilon_, rsig_data, rsig_data, &context_); GroupNormForwardCUDAKernel << { vector shape_; }; +class Float16UniformFillOp : public Operator { + public: + Float16UniformFillOp(const OperatorDef& operator_def, Workspace* ws) + : Operator(operator_def, ws), + shape_(this->template GetRepeatedArgument("shape")), + min_(this->template GetSingleArgument("min", 0)), + max_(this->template GetSingleArgument("max", 1)) { + if (InputSize() == 3) { + CAFFE_ENFORCE( + !this->template HasSingleArgumentOfType("min"), + "Cannot set both min arg and min input blob"); + CAFFE_ENFORCE( + !this->template HasSingleArgumentOfType("max"), + "Cannot set both max arg and max input blob"); + } else { + CAFFE_ENFORCE_LT( + min_, max_, "Max value should be bigger than min value."); + } + } + + USE_OPERATOR_FUNCTIONS(CPUContext); + virtual ~Float16UniformFillOp() {} + + bool RunOnDevice() override; + + private: + vector shape_; + float min_; + float max_; +}; + inline std::vector Float16FillerTensorInference( const OperatorDef& def, const vector& in) { diff --git a/caffe2/operators/hip/spatial_batch_norm_op_miopen.cc b/caffe2/operators/hip/spatial_batch_norm_op_miopen.cc index 48d1ebd550728..3115e9e30cb6e 100644 --- a/caffe2/operators/hip/spatial_batch_norm_op_miopen.cc +++ b/caffe2/operators/hip/spatial_batch_norm_op_miopen.cc @@ -211,7 +211,7 @@ bool MIOpenSpatialBNOp::DoRunWithType() { } // Save the mean and inv var results. auto* save_mean = Output(SAVED_MEAN); - auto* save_var = Output(SAVED_INV_VAR); + auto* save_var = Output(SAVED_INV_STD); save_mean->Resize(C); save_var->Resize(C); void* save_mean_data = save_mean->template mutable_data(); @@ -294,7 +294,7 @@ bool MIOpenSpatialBNGradientOp::DoRunWithType() { auto* dBias_data = dBias->template mutable_data(); const auto& saved_mean = Input(SAVED_MEAN); - const auto& saved_var = Input(SAVED_INV_VAR); + const auto& saved_var = Input(SAVED_INV_STD); const void* saved_mean_data = saved_mean.template data(); const void* saved_var_data = saved_var.template data(); @@ -336,10 +336,6 @@ bool MIOpenSpatialBNGradientOp::RunOnDevice() { return true; } -// Since there is no default implementation for spatial batch normalization, -// we will register the miopen version as the default as well. -REGISTER_HIP_OPERATOR(SpatialBN, MIOpenSpatialBNOp); -REGISTER_HIP_OPERATOR(SpatialBNGradient, MIOpenSpatialBNGradientOp); REGISTER_MIOPEN_OPERATOR(SpatialBN, MIOpenSpatialBNOp); REGISTER_MIOPEN_OPERATOR(SpatialBNGradient, MIOpenSpatialBNGradientOp); diff --git a/caffe2/operators/load_save_op.cc b/caffe2/operators/load_save_op.cc index a77aa76b27dec..ffef2f8b39fb5 100644 --- a/caffe2/operators/load_save_op.cc +++ b/caffe2/operators/load_save_op.cc @@ -5,7 +5,8 @@ namespace caffe2 { template <> void LoadOp::SetCurrentDevice(BlobProto* proto) { if (proto->has_tensor()) { - proto->mutable_tensor()->mutable_device_detail()->set_device_type(CPU); + proto->mutable_tensor()->mutable_device_detail()->set_device_type( + PROTO_CPU); } } diff --git a/caffe2/operators/load_save_op_gpu.cc b/caffe2/operators/load_save_op_gpu.cc index 54a495926c145..cd70e9c2b5df2 100644 --- a/caffe2/operators/load_save_op_gpu.cc +++ b/caffe2/operators/load_save_op_gpu.cc @@ -7,7 +7,7 @@ template <> void LoadOp::SetCurrentDevice(BlobProto* proto) { if (proto->has_tensor()) { auto* device_detail = proto->mutable_tensor()->mutable_device_detail(); - device_detail->set_device_type(CUDA); + device_detail->set_device_type(PROTO_CUDA); device_detail->set_cuda_gpu_id(CaffeCudaGetDevice()); } } diff --git a/caffe2/operators/operator_fallback_gpu.h b/caffe2/operators/operator_fallback_gpu.h index 4af9bcde1ad9d..b0ee9611a6904 100644 --- a/caffe2/operators/operator_fallback_gpu.h +++ b/caffe2/operators/operator_fallback_gpu.h @@ -27,32 +27,34 @@ namespace caffe2 { * to register the CPU side, you can create its corresponding GPU operator * (with performance hits of course) via * REGISTER_CUDA_OPERATOR(MyMagic, - * GPUFallbackOp); + * GPUFallbackOp); + * Note that you will need to make sure that the operators actually share the + * same name. * * Advanced usage: if you want to have some specific outputs never copied, you * can use the SkipOutputCopy template argument to do that. For example, if * MyMagic produces two outputs and the first output is always going to live on * the CPU, you can do * REGISTER_CUDA_OPERATOR(MyMagic, - * GPUFallbackOp>); + * GPUFallbackOpEx>); */ -template > -class GPUFallbackOp final : public Operator { +template +class GPUFallbackOpEx final : public Operator { public: USE_OPERATOR_FUNCTIONS(CUDAContext); - GPUFallbackOp(const OperatorDef& def, Workspace* ws) + GPUFallbackOpEx(const OperatorDef& def, Workspace* ws) : Operator(def, ws) { - CAFFE_ENFORCE_EQ(def.device_option().device_type(), CUDA); + CAFFE_ENFORCE_EQ(def.device_option().device_type(), PROTO_CUDA); OperatorDef base_def_(def); // base_def_ runs on CPU, so we will set its device option to CPU. base_def_.clear_device_option(); - base_def_.mutable_device_option()->set_device_type(CPU); + base_def_.mutable_device_option()->set_device_type(PROTO_CPU); // Set up the symbols for the local workspace. for (const string& name : def.input()) { local_input_blobs_.push_back(local_ws_.CreateBlob(name)); CHECK_NOTNULL(local_input_blobs_.back()); } - base_op_.reset(new CPUOp(base_def_, &local_ws_)); + base_op_ = CreateOperator(base_def_, &local_ws_); for (const string& name : def.output()) { local_output_blobs_.push_back(local_ws_.GetBlob(name)); CHECK_NOTNULL(local_output_blobs_.back()); @@ -105,9 +107,11 @@ class GPUFallbackOp final : public Operator { Workspace local_ws_; vector local_input_blobs_; vector local_output_blobs_; - std::unique_ptr base_op_; + unique_ptr base_op_; }; +using GPUFallbackOp = GPUFallbackOpEx>; + } // namespace caffe2 #endif // CAFFE2_OPERATORS_OPERATOR_FALLBACK_H_ diff --git a/caffe2/operators/operator_fallback_gpu_test.cc b/caffe2/operators/operator_fallback_gpu_test.cc index e562858c073ec..59d765d586049 100644 --- a/caffe2/operators/operator_fallback_gpu_test.cc +++ b/caffe2/operators/operator_fallback_gpu_test.cc @@ -29,8 +29,7 @@ OPERATOR_SCHEMA(IncrementByOne) .NumInputs(1).NumOutputs(1).AllowInplace({{0, 0}}); REGISTER_CPU_OPERATOR(IncrementByOne, IncrementByOneOp); -REGISTER_CUDA_OPERATOR(IncrementByOne, - GPUFallbackOp); +REGISTER_CUDA_OPERATOR(IncrementByOne, GPUFallbackOp); TEST(OperatorFallbackTest, IncrementByOneOp) { OperatorDef op_def = CreateOperatorDef( @@ -59,7 +58,7 @@ TEST(OperatorFallbackTest, GPUIncrementByOneOp) { OperatorDef op_def = CreateOperatorDef( "IncrementByOne", "", vector{"X"}, vector{"X"}); - op_def.mutable_device_option()->set_device_type(CUDA); + op_def.mutable_device_option()->set_device_type(PROTO_CUDA); Workspace ws; Tensor source_tensor(vector{2, 3}, CPU); for (int i = 0; i < 6; ++i) { diff --git a/caffe2/operators/pack_segments.cu b/caffe2/operators/pack_segments.cu index ae573adde3db4..a76cac999e2a8 100644 --- a/caffe2/operators/pack_segments.cu +++ b/caffe2/operators/pack_segments.cu @@ -299,7 +299,7 @@ bool UnpackSegmentsOp::DoRunWithType2() { Data_T* out_ptr = static_cast(out->raw_mutable_data(data.meta())); // Return empty out (with the proper shape) if any of the dimensions is 0. - if (!(data.dim(0) * data.dim(1))) { + if (data.dim(0) == 0 || data.dim(1) == 0) { return true; } diff --git a/caffe2/operators/pool_op.cu b/caffe2/operators/pool_op.cu index 4af3be93a6b8f..1efd3ac9e189c 100644 --- a/caffe2/operators/pool_op.cu +++ b/caffe2/operators/pool_op.cu @@ -1667,6 +1667,7 @@ bool PoolGradientOp::RunOnDeviceWithOrderNHWC() { stride_h(), pad_t(), dX->template mutable_data()); + break; case 2: MaxPool2DBackwardNHWC <<CreateBlob(name); auto* tensor = blob->GetMutableTensor(CUDA); @@ -38,7 +38,7 @@ TEST(ReshapeOpGPUTest, testReshapeWithScalar) { def.add_output("XNew"); def.add_output("OldShape"); def.add_arg()->CopyFrom(MakeArgument("shape", vector{1})); - def.mutable_device_option()->set_device_type(CUDA); + def.mutable_device_option()->set_device_type(PROTO_CUDA); AddConstInput(vector(), 3.14, "X", &ws); // execute the op unique_ptr op(CreateOperator(def, &ws)); diff --git a/caffe2/operators/rnn/recurrent_network_executor_gpu.cc b/caffe2/operators/rnn/recurrent_network_executor_gpu.cc index 56e51ab342407..e16e2073f7fd1 100644 --- a/caffe2/operators/rnn/recurrent_network_executor_gpu.cc +++ b/caffe2/operators/rnn/recurrent_network_executor_gpu.cc @@ -69,7 +69,9 @@ void CUDARecurrentNetworkExecutor::_ExecRange(int from, int to) { continue; } - if (gpu_id == -1 && rnn_op.op->device_option().device_type() == DeviceType::CUDA) { + if (gpu_id == -1 && + rnn_op.op->device_option().device_type() == + DeviceTypeProto::PROTO_CUDA) { gpu_id = rnn_op.op->device_option().cuda_gpu_id(); } else { CAFFE_ENFORCE( diff --git a/caffe2/operators/roi_align_op_gpu_test.cc b/caffe2/operators/roi_align_op_gpu_test.cc index 92eafefcb65eb..eca60a6006a27 100644 --- a/caffe2/operators/roi_align_op_gpu_test.cc +++ b/caffe2/operators/roi_align_op_gpu_test.cc @@ -136,7 +136,8 @@ void CreateAndRun( def.add_input("X"); def.add_input("R"); def.add_output("Y"); - def.mutable_device_option()->set_device_type(GetDeviceType()); + def.mutable_device_option()->set_device_type( + TypeToProto(GetDeviceType())); def.add_arg()->CopyFrom(MakeArgument("spatial_scale", 1.0f / 16.0f)); def.add_arg()->CopyFrom(MakeArgument("pooled_h", 6)); def.add_arg()->CopyFrom(MakeArgument("pooled_w", 8)); @@ -151,7 +152,7 @@ void CreateAndRun( def_roialign.add_input("R"); def_roialign.add_output("Y_NHWC"); def_roialign.mutable_device_option()->set_device_type( - GetDeviceType()); + TypeToProto(GetDeviceType())); def_roialign.add_arg()->CopyFrom( MakeArgument("spatial_scale", 1.0f / 16.0f)); def_roialign.add_arg()->CopyFrom(MakeArgument("pooled_h", 6)); @@ -164,14 +165,16 @@ void CreateAndRun( def_x.set_type("NCHW2NHWC"); def_x.add_input("X"); def_x.add_output("X_NHWC"); - def_x.mutable_device_option()->set_device_type(GetDeviceType()); + def_x.mutable_device_option()->set_device_type( + TypeToProto(GetDeviceType())); OperatorDef def_y; def_y.set_name("test_y"); def_y.set_type("NHWC2NCHW"); def_y.add_input("Y_NHWC"); def_y.add_output("Y"); - def_y.mutable_device_option()->set_device_type(GetDeviceType()); + def_y.mutable_device_option()->set_device_type( + TypeToProto(GetDeviceType())); ops.push_back(CreateOperator(def_x, &ws)); ops.push_back(CreateOperator(def_roialign, &ws)); diff --git a/caffe2/operators/sparse_normalize_op_gpu.cu b/caffe2/operators/sparse_normalize_op_gpu.cu index d5e6cc73d1910..d71deb6a3d8cd 100644 --- a/caffe2/operators/sparse_normalize_op_gpu.cu +++ b/caffe2/operators/sparse_normalize_op_gpu.cu @@ -5,5 +5,5 @@ namespace caffe2 { REGISTER_CUDA_OPERATOR( SparseNormalize, - GPUFallbackOp>); + GPUFallbackOp); } diff --git a/caffe2/operators/spatial_batch_norm_gradient_op.cc b/caffe2/operators/spatial_batch_norm_gradient_op.cc index 5a9d55341f27e..49d07f68cc2eb 100644 --- a/caffe2/operators/spatial_batch_norm_gradient_op.cc +++ b/caffe2/operators/spatial_batch_norm_gradient_op.cc @@ -1,125 +1,140 @@ #include "caffe2/operators/spatial_batch_norm_op.h" +#include + #include "caffe2/utils/eigen_utils.h" namespace caffe2 { template <> -bool SpatialBNGradientOp::RunOnDevice() { - const auto& X = Input(INPUT); - const auto& dY = Input(OUTPUT_GRAD); - const auto& scale = Input(SCALE); - - CAFFE_ENFORCE(X.ndim() >= 3 && X.ndim() <= 5); - const int N = X.dim32(0); - const int C = - (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(X.ndim() - 1)); - const int H = (order_ == StorageOrder::NCHW ? X.dim32(2) : X.dim32(1)); - const int W = X.ndim() > 3 - ? (order_ == StorageOrder::NCHW ? X.dim32(3) : X.dim32(2)) - : 1; - const int D = X.ndim() > 4 - ? (order_ == StorageOrder::NCHW ? X.dim32(4) : X.dim32(3)) - : 1; - - const int sample_size = H * W * D; - - CAFFE_ENFORCE_EQ(scale.ndim(), 1); - CAFFE_ENFORCE_EQ(scale.dim32(0), C); - - ConstEigenVectorArrayMap scale_arr(scale.data(), C); - ConstEigenVectorArrayMap mean_arr(Input(SAVED_MEAN).data(), C); - ConstEigenVectorArrayMap inv_var_arr( - Input(SAVED_INV_VAR).data(), C); - - auto* dX = Output(INPUT_GRAD); - dX->ResizeLike(X); - - auto* dScale = Output(SCALE_GRAD); - auto* dBias = Output(BIAS_GRAD); - - if (num_batches_ == 1) { - dScale->ResizeLike(scale); - dBias->ResizeLike(scale); - } - - // dBias = np.sum(dY, axis=0) - // dScale = np.sum((X - mean) / inv_std * dy, axis=0) - // dX = (1. / N) * scale * inv_var * (N * dY - np.sum(dY, axis=0) - (X - mean) - // * inv_var * inv_var * np.sum(dY * (X - mean), axis=0)) - - EigenVectorArrayMap dBias_arr( - dBias->template mutable_data(), C); - EigenVectorArrayMap dScale_arr( - dScale->template mutable_data(), C); - - if (num_batches_ == 1) { - dBias_arr.setZero(); - dScale_arr.setZero(); - } - - const auto scaleInvVarNHW = scale_arr * inv_var_arr / (N * sample_size); +template +void SpatialBNGradientOp:: + ComputeMultiBatchScaleBiasGradientsAndFusedParams( + const int N, + const int C, + const int HxW, + const T* scale, + const T* mean, + const T* rstd, + const T* dscale_sum, + const T* dbias_sum, + T* dscale, + T* dbias, + T* alpha, + T* beta, + T* gamma) { + ConstEigenVectorArrayMap scale_arr(scale, C); + ConstEigenVectorArrayMap mean_arr(mean, C); + ConstEigenVectorArrayMap rstd_arr(rstd, C); + EigenVectorArrayMap dscale_arr(dscale, C); + EigenVectorArrayMap dbias_arr(dbias, C); + EigenVectorArrayMap alpha_arr(alpha, C); + EigenVectorArrayMap beta_arr(beta, C); + EigenVectorArrayMap gamma_arr(gamma, C); + const T inv_num_batches = T(1) / static_cast(num_batches_); + math::Scale( + C, inv_num_batches, dscale_sum, dscale, &context_); + math::Scale( + C, inv_num_batches, dbias_sum, dbias, &context_); + const T inv_nhw = T(1) / static_cast(N * HxW); + alpha_arr = scale_arr * rstd_arr; + beta_arr = dscale_arr * rstd_arr; + gamma_arr = alpha_arr * (mean_arr * beta_arr - dbias_arr) * inv_nhw; + beta_arr *= -alpha_arr * inv_nhw; +} - switch (order_) { - case StorageOrder::NCHW: { - ConstEigenArrayMap X_arr(X.data(), sample_size, N * C); - ConstEigenArrayMap dY_arr(dY.data(), sample_size, N * C); - EigenArrayMap dX_arr( - dX->template mutable_data(), sample_size, N * C); - dX_arr.setZero(); - if (N == 0) { - return true; - } - if (num_batches_ == 1) { - for (int nc = 0; nc < N * C; ++nc) { - int c = nc % C; - dBias_arr(c) += dY_arr.col(nc).sum(); - dScale_arr(c) += - ((X_arr.col(nc) - mean_arr(c)) * inv_var_arr(c) * dY_arr.col(nc)) - .sum(); - } - } else { - for (int c = 0; c < C; ++c) { - dBias_arr(c) /= num_batches_; - dScale_arr(c) /= num_batches_; - } - } - for (int nc = 0; nc < N * C; ++nc) { - int c = nc % C; - dX_arr.col(nc) += scaleInvVarNHW(c) * - (dY_arr.col(nc) * N * sample_size - dBias_arr(c) - - (X_arr.col(nc) - mean_arr[c]) * dScale_arr(c) * inv_var_arr(c)); +template <> +template +void SpatialBNGradientOp::ComputeScaleBiasGradientsAndFusedParams( + const int N, + const int C, + const int HxW, + const T* dY, + const T* X, + const T* scale, + const T* mean, + const T* rstd, + T* dscale, + T* dbias, + T* alpha, + T* beta, + T* gamma) { + ConstEigenVectorArrayMap scale_arr(scale, C); + ConstEigenVectorArrayMap mean_arr(mean, C); + ConstEigenVectorArrayMap rstd_arr(rstd, C); + EigenVectorArrayMap dscale_arr(dscale, C); + EigenVectorArrayMap dbias_arr(dbias, C); + EigenVectorArrayMap alpha_arr(alpha, C); + EigenVectorArrayMap beta_arr(beta, C); + EigenVectorArrayMap gamma_arr(gamma, C); + math::Set(C, T(0), dscale, &context_); + math::Set(C, T(0), dbias, &context_); + if (order_ == StorageOrder::NCHW) { + ConstEigenArrayMap dY_arr(dY, HxW, N * C); + ConstEigenArrayMap X_arr(X, HxW, N * C); + for (int i = 0; i < N; ++i) { + for (int j = 0; j < C; ++j) { + const int c = i * C + j; + dscale_arr(j) += + (dY_arr.col(c) * (X_arr.col(c) - mean_arr(j)) * rstd_arr(j)).sum(); + dbias_arr(j) += dY_arr.col(c).sum(); } - break; } - case StorageOrder::NHWC: { - ConstEigenArrayMap X_arr(X.data(), C, N * sample_size); - ConstEigenArrayMap dY_arr(dY.data(), C, N * sample_size); - EigenArrayMap dX_arr( - dX->template mutable_data(), C, N * sample_size); - dX_arr.setZero(); - if (N == 0) { - return true; - } + } else { + const int outer_size = N * HxW; + ConstEigenArrayMap dY_arr(dY, C, outer_size); + ConstEigenArrayMap X_arr(X, C, outer_size); + for (int i = 0; i < outer_size; ++i) { + dscale_arr += dY_arr.col(i) * (X_arr.col(i) - mean_arr) * rstd_arr; + dbias_arr += dY_arr.col(i); + } + } + const T inv_nhw = T(1) / static_cast(N * HxW); + alpha_arr = scale_arr * rstd_arr; + beta_arr = dscale_arr * rstd_arr; + gamma_arr = alpha_arr * (mean_arr * beta_arr - dbias_arr) * inv_nhw; + beta_arr *= -alpha_arr * inv_nhw; +} - const auto dYRowSum = dY_arr.rowwise().sum(); - const auto XMinusMean = X_arr.colwise() - mean_arr; - const auto dYMulXMinusMeanRowSum = (dY_arr * XMinusMean).rowwise().sum(); - const auto invVarSqr = inv_var_arr * inv_var_arr; - for (int nhw = 0; nhw < N * sample_size; ++nhw) { - dBias_arr += dY_arr.col(nhw); - dScale_arr += - (X_arr.col(nhw) - mean_arr) * inv_var_arr * dY_arr.col(nhw); - dX_arr.col(nhw) += scaleInvVarNHW * - (dY_arr.col(nhw) * N * sample_size - dYRowSum - - XMinusMean.col(nhw) * invVarSqr * dYMulXMinusMeanRowSum); - } - break; +template <> +template +void SpatialBNGradientOp::ComputeXGradient( + const int N, + const int C, + const int HxW, + const T* dY, + const T* X, + const T* alpha, + const T* beta, + const T* gamma, + T* dX) { + ConstEigenVectorArrayMap alpha_arr(alpha, C); + ConstEigenVectorArrayMap beta_arr(beta, C); + ConstEigenVectorArrayMap gamma_arr(gamma, C); + if (order_ == NCHW) { + const int stride = C * HxW; + const T* dY_ptr = dY; + const T* X_ptr = X; + T* dX_ptr = dX; + for (int i = 0; i < N; ++i) { + EigenArrayMap(dX_ptr, HxW, C) = + (ConstEigenArrayMap(dY_ptr, HxW, C).rowwise() * + alpha_arr.transpose() + + ConstEigenArrayMap(X_ptr, HxW, C).rowwise() * + beta_arr.transpose()) + .rowwise() + + gamma_arr.transpose(); + dY_ptr += stride; + X_ptr += stride; + dX_ptr += stride; } - default: - CAFFE_THROW("Unknown storage order: ", order_); + } else { + EigenArrayMap(dX, C, N * HxW) = + (ConstEigenArrayMap(dY, C, N * HxW).colwise() * alpha_arr + + ConstEigenArrayMap(X, C, N * HxW).colwise() * beta_arr) + .colwise() + + gamma_arr; } - return true; } REGISTER_CPU_OPERATOR(SpatialBNGradient, SpatialBNGradientOp); @@ -131,17 +146,20 @@ OPERATOR_SCHEMA(SpatialBNGradient) .NumOutputs(3) .AllowInplace({{5, 1}, {6, 2}}); +namespace { + // Spatial batch normalization's gradient, depending on the various input sizes, // is a bit more complex than usual gradient operators. class GetSpatialBNGradient : public GradientMakerBase { using GradientMakerBase::GradientMakerBase; - vector GetGradientDefs() override { + std::vector GetGradientDefs() override { // Check if we are in training or testing mode. - bool is_test = + const bool is_test = ArgumentHelper::GetSingleArgument(def_, OpSchema::Arg_IsTest, 0); - int num_batches = ArgumentHelper::GetSingleArgument(def_, "num_batches", 1); - vector grad_outputs{GI(0), GI(1), GI(2)}; - vector grad_inputs; + const int num_batches = + ArgumentHelper::GetSingleArgument(def_, "num_batches", 1); + const std::vector grad_outputs = {GI(0), GI(1), GI(2)}; + std::vector grad_inputs; if (is_test) { // This is in testing mode. The operator should have five inputs: // X, scale, bias, estimated_mean, estimated_variance @@ -149,19 +167,24 @@ class GetSpatialBNGradient : public GradientMakerBase { // X, scale, dY, estimated_mean, estimated_variance CAFFE_ENFORCE_EQ(def_.input_size(), 5); CAFFE_ENFORCE_EQ(def_.output_size(), 1); - grad_inputs = vector{I(0), I(1), GO(0), I(3), I(4)}; + grad_inputs = std::vector{I(0), I(1), GO(0), I(3), I(4)}; } else if (num_batches > 1) { CAFFE_ENFORCE_EQ(def_.input_size(), 7); CAFFE_ENFORCE_EQ(def_.output_size(), 5); - grad_inputs = vector{I(0), I(1), GO(0), O(3), O(4), GI(1), GI(2)}; + grad_inputs = + std::vector{I(0), I(1), GO(0), O(3), O(4), GI(1), GI(2)}; } else { CAFFE_ENFORCE_EQ(def_.input_size(), 5); CAFFE_ENFORCE_EQ(def_.output_size(), 5); - grad_inputs = vector{I(0), I(1), GO(0), O(3), O(4)}; + grad_inputs = std::vector{I(0), I(1), GO(0), O(3), O(4)}; } return SingleGradientDef( "SpatialBNGradient", "", grad_inputs, grad_outputs); } }; + +} // namespace + REGISTER_GRADIENT(SpatialBN, GetSpatialBNGradient); -} + +} // namespace caffe2 diff --git a/caffe2/operators/spatial_batch_norm_op.cc b/caffe2/operators/spatial_batch_norm_op.cc index 09f2b04fd9f25..dc064cd951acf 100644 --- a/caffe2/operators/spatial_batch_norm_op.cc +++ b/caffe2/operators/spatial_batch_norm_op.cc @@ -1,176 +1,76 @@ #include "caffe2/operators/spatial_batch_norm_op.h" +#include + #include "caffe2/utils/eigen_utils.h" namespace caffe2 { template <> -bool SpatialBNOp::RunOnDevice() { - const auto& X = Input(INPUT); - const auto& scale = Input(SCALE); - const auto& bias = Input(BIAS); - - CAFFE_ENFORCE(X.ndim() >= 3 && X.ndim() <= 5); - const int N = X.dim32(0); - const int C = - (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(X.ndim() - 1)); - const int H = (order_ == StorageOrder::NCHW ? X.dim32(2) : X.dim32(1)); - const int W = X.ndim() > 3 - ? (order_ == StorageOrder::NCHW ? X.dim32(3) : X.dim32(2)) - : 1; - const int D = X.ndim() > 4 - ? (order_ == StorageOrder::NCHW ? X.dim32(4) : X.dim32(3)) - : 1; - - const int sample_size = H * W * D; - CAFFE_ENFORCE_EQ(scale.ndim(), 1); - CAFFE_ENFORCE_EQ(bias.ndim(), 1); - CAFFE_ENFORCE_EQ(scale.dim32(0), C); - CAFFE_ENFORCE_EQ(bias.dim32(0), C); - - ConstEigenVectorArrayMap scale_arr(scale.data(), C); - ConstEigenVectorArrayMap bias_arr(bias.data(), C); - - auto* Y = Output(OUTPUT); - Y->ResizeLike(X); - auto* Y_data = Y->mutable_data(); - - if (!is_test_) { - // training mode - // Get the mean and variance. - // Note that, to be consistent with cudnn, we will output saved inverse - // std as output 5, but we will still use the same storage place to - // compute var as well. The inverse is going to be carried out at the end - // of the op. - Output(SAVED_MEAN)->Resize(C); - Output(SAVED_INV_VAR)->Resize(C); - EigenVectorArrayMap mean( - Output(SAVED_MEAN)->template mutable_data(), C); - EigenVectorArrayMap var( - Output(SAVED_INV_VAR)->mutable_data(), C); - if (N > 0) { - if (num_batches_ > 1) { - ConstEigenVectorArrayMap sums(Input(SUMS).data(), C); - ConstEigenVectorArrayMap sumsq(Input(SUMSQ).data(), C); - const auto multi_batch_size = N * num_batches_ * sample_size; - mean = sums / multi_batch_size; - var = (sumsq - (sums * sums) / multi_batch_size) / multi_batch_size; - } else { - mean.setZero(); - var.setZero(); - switch (order_) { - case StorageOrder::NCHW: { - ConstEigenArrayMap X_arr( - X.data(), sample_size, N * C); - for (int nc = 0; nc < N * C; ++nc) { - mean(nc % C) += X_arr.col(nc).sum(); - } - mean /= N * sample_size; - for (int nc = 0; nc < N * C; ++nc) { - var(nc % C) += - (X_arr.col(nc) - mean(nc % C)).matrix().squaredNorm(); - } - var /= N * sample_size; - break; - } - case StorageOrder::NHWC: { - ConstEigenArrayMap X_arr( - X.data(), C, N * sample_size); - for (int i = 0; i < N * sample_size; ++i) { - mean += X_arr.col(i); - } - mean /= N * sample_size; - for (int i = 0; i < N * sample_size; ++i) { - var += (X_arr.col(i) - mean) * (X_arr.col(i) - mean); - } - var /= N * sample_size; - break; - } - default: - CAFFE_THROW("Unknown storage order: ", order_); - } - } +template +void SpatialBNOp::ComputeFusedParam( + const int C, + const T* scale, + const T* bias, + const T* mean, + const T* var, + T* alpha, + T* beta) { + EigenVectorArrayMap alpha_arr(alpha, C); + EigenVectorArrayMap beta_arr(beta, C); + alpha_arr = ConstEigenVectorArrayMap(scale, C) * + (ConstEigenVectorArrayMap(var, C) + static_cast(epsilon_)).rsqrt(); + beta_arr = ConstEigenVectorArrayMap(bias, C) - + alpha_arr * ConstEigenVectorArrayMap(mean, C); +} - // Compute the running mean and running inv variance. - auto* running_mean = Output(RUNNING_MEAN); - auto* running_var = Output(RUNNING_VAR); - // Check if they are initialized - if (!running_mean->size()) { - running_mean->Resize(C); - EigenVectorArrayMap running_mean_map( - running_mean->mutable_data(), C); - running_mean_map.setZero(); - } - if (!running_var->size()) { - running_var->Resize(C); - EigenVectorArrayMap running_var_map( - running_var->mutable_data(), C); - running_var_map.setZero(); - } - EigenVectorArrayMap running_mean_arr( - running_mean->mutable_data(), C); - EigenVectorArrayMap running_var_arr( - running_var->mutable_data(), C); - running_mean_arr = running_mean_arr * momentum_ + mean * (1. - momentum_); - running_var_arr = running_var_arr * momentum_ + var * (1. - momentum_); - } else { - // set empty batch's mean / var output to zeros - mean.setZero(); - var.setZero(); - return true; - } - } +template <> +template +void SpatialBNOp::ComputeBatchMoments( + const int N, + const int C, + const int HxW, + const T* batch_mean_sum, + const T* batch_var_sum, + T* mean, + T* var) { + const T scale = T(1) / static_cast(num_batches_ * N * HxW); + EigenVectorArrayMap mean_arr(mean, C); + EigenVectorArrayMap var_arr(var, C); + mean_arr = ConstEigenVectorArrayMap(batch_mean_sum, C) * scale; + var_arr = + ConstEigenVectorArrayMap(batch_var_sum, C) * scale - mean_arr.square(); +} - // Regardless of training or testing, we will apply the estimated mean - // and standard deviation to the input. For testing, they are - // specified directly by the input, and for training, they are computed - // by the op. - Eigen::Array inv_std(C); - if (is_test_) { - ConstEigenVectorArrayMap var_arr(Input(EST_VAR).data(), C); - inv_std = (var_arr + epsilon_).sqrt().inverse(); - } else { - EigenVectorArrayMap saved_inv_std( - Output(SAVED_INV_VAR)->template mutable_data(), C); - saved_inv_std = (saved_inv_std + epsilon_).inverse().sqrt(); - inv_std = saved_inv_std; - } - ConstEigenVectorArrayMap mean_arr( - is_test_ ? Input(EST_MEAN).data() - : Output(SAVED_MEAN)->data(), - C); - // We can fuse the output computation as follows: - // ((x - est_mean) * (inv_var) * scale + bias - // to - // (x * inv_var * scale) + (bias - est_mean * inv_var * scale) - Eigen::Array new_scale = inv_std * scale_arr; - Eigen::Array new_bias = - bias_arr - mean_arr * inv_std * scale_arr; - switch (order_) { - case StorageOrder::NHWC: { - EigenArrayMap(Y_data, C, N * sample_size) = - (ConstEigenArrayMap(X.data(), C, N * sample_size) - .colwise() * - new_scale) - .colwise() + - new_bias; - break; - } - case StorageOrder::NCHW: { - EigenArrayMap Y_arr(Y_data, sample_size, N * C); - ConstEigenArrayMap X_arr(X.data(), sample_size, N * C); - for (int nc = 0; nc < N * C; ++nc) { - Y_arr.col(nc) = X_arr.col(nc) * new_scale(nc % C) + new_bias(nc % C); - } - break; - } - default: - CAFFE_THROW("Unknown storage order: ", order_); - } - return true; +template <> +template +void SpatialBNOp::ComputeRunningMomentsAndFusedParam( + const int C, + const T* scale, + const T* bias, + const T* mean, + const T* var, + T* running_mean, + T* running_var, + T* rstd, + T* alpha, + T* beta) { + const T a = T(1) - static_cast(momentum_); + const T b = static_cast(momentum_); + math::Axpby(C, a, mean, b, running_mean, &context_); + math::Axpby(C, a, var, b, running_var, &context_); + math::InvStd( + C, static_cast(epsilon_), var, rstd, &context_); + EigenVectorArrayMap alpha_arr(alpha, C); + EigenVectorArrayMap beta_arr(beta, C); + alpha_arr = ConstEigenVectorArrayMap(scale, C) * + ConstEigenVectorArrayMap(rstd, C); + beta_arr = ConstEigenVectorArrayMap(bias, C) - + alpha_arr * ConstEigenVectorArrayMap(mean, C); } namespace { + OpSchema::Cost CostInferenceForSpatialBN( const OperatorDef& def, const vector& in) { @@ -184,6 +84,7 @@ OpSchema::Cost CostInferenceForSpatialBN( cost.params_bytes = 2 * C * sizeof(float); return cost; } + } // namespace REGISTER_CPU_OPERATOR(SpatialBN, SpatialBNOp); @@ -191,35 +92,35 @@ REGISTER_CPU_OPERATOR(SpatialBN, SpatialBNOp); OPERATOR_SCHEMA(SpatialBN) .NumInputs({5, 7}) .NumOutputs({1, 5}) - .AllowInplace({{0, 0}}) - .CostInferenceFunction(CostInferenceForSpatialBN) + .AllowInplace({{0, 0}, {5, 3}, {6, 4}}) .EnforceInplace({{3, 1}, {4, 2}}) - .TensorInferenceFunction( - [](const OperatorDef& def, const vector& in) { - ArgumentHelper helper(def); - bool is_test = helper.GetSingleArgument(OpSchema::Arg_IsTest, 0); - - if (!is_test) { - vector out; - StorageOrder order = StringToStorageOrder( - helper.GetSingleArgument("order", "NCHW")); - const TensorShape& X = in[0]; - const int C = - (order == StorageOrder::NCHW ? X.dims(1) - : X.dims(X.dims_size() - 1)); - - out.push_back(in[0]); - TensorShape meanvar_tp = - CreateTensorShape(vector{C}, TensorProto::FLOAT); - out.push_back(meanvar_tp); // RUNNING_MEAN - out.push_back(meanvar_tp); // RUNNING_MEAN - out.push_back(meanvar_tp); // SAVED_MEAN - out.push_back(meanvar_tp); // SAVED_VAR - return out; - } else { - return vector{in[0]}; - } - }) + .CostInferenceFunction(CostInferenceForSpatialBN) + .TensorInferenceFunction([](const OperatorDef& def, + const vector& in) { + ArgumentHelper helper(def); + bool is_test = helper.GetSingleArgument(OpSchema::Arg_IsTest, 0); + + if (!is_test) { + vector out; + StorageOrder order = StringToStorageOrder( + helper.GetSingleArgument("order", "NCHW")); + const TensorShape& X = in[0]; + const int C = + (order == StorageOrder::NCHW ? X.dims(1) + : X.dims(X.dims_size() - 1)); + + out.push_back(in[0]); + TensorShape meanvar_tp = + CreateTensorShape(vector{C}, TensorProto::FLOAT); + out.push_back(meanvar_tp); // RUNNING_MEAN + out.push_back(meanvar_tp); // RUNNING_MEAN + out.push_back(meanvar_tp); // SAVED_MEAN + out.push_back(meanvar_tp); // SAVED_VAR + return out; + } else { + return vector{in[0]}; + } + }) .SetDoc(R"DOC( Applies spatial batch normalization to the input tensor as described in the original paper, [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://arxiv.org/abs/1502.03167). Be aware, this operator has two different output sets, depending on the value of *is_test*. According to the paper, the primary operation of spatial batch normalization is: @@ -240,11 +141,20 @@ Github Links: - https://github.com/pytorch/pytorch/blob/master/caffe2/operators/spatial_batch_norm_op.h )DOC") - .ArgIsTest("*(type: int; default: 0)* If set to nonzero, run spatial batch normalization in test mode.") - .Arg("epsilon", "*(type: float; default: 1e-5)* The epsilon value to use to avoid division by zero.") - .Arg("order", "*(type: string; default: \"NCHW\")* Specifies the order of the input data blob, where $N$ is batch size, $C$ is number of channels, $H$ is spatial height, and $W$ is spatial width. The only other valid option is \"NHWC\".") - .Arg("momentum", "*(type: float; default: 0.9)* Factor used in computing the running mean and variance. e.g., running_mean = running_mean x momentum + mean x (1 - momentum)") - .Arg("num_batches", "*(type: int; default: 1)* Specifies the number of batches to apply normalization on. Requires specifying the optional sums and sumsq inputs that provide statistics across multiple batches from which mean and variance can be determined.") + .ArgIsTest( + "*(type: int; default: 0)* If set to nonzero, run spatial batch normalization in test mode.") + .Arg( + "epsilon", + "*(type: float; default: 1e-5)* The epsilon value to use to avoid division by zero.") + .Arg( + "order", + "*(type: string; default: \"NCHW\")* Specifies the order of the input data blob, where $N$ is batch size, $C$ is number of channels, $H$ is spatial height, and $W$ is spatial width. The only other valid option is \"NHWC\".") + .Arg( + "momentum", + "*(type: float; default: 0.9)* Factor used in computing the running mean and variance. e.g., running_mean = running_mean x momentum + mean x (1 - momentum)") + .Arg( + "num_batches", + "*(type: int; default: 1)* Specifies the number of batches to apply normalization on. Requires specifying the optional sums and sumsq inputs that provide statistics across multiple batches from which mean and variance can be determined.") .Input( 0, "X", diff --git a/caffe2/operators/spatial_batch_norm_op.cu b/caffe2/operators/spatial_batch_norm_op.cu new file mode 100644 index 0000000000000..ecfc2001e5efd --- /dev/null +++ b/caffe2/operators/spatial_batch_norm_op.cu @@ -0,0 +1,10 @@ +#include "caffe2/operators/spatial_batch_norm_op.h" + +#include "caffe2/operators/spatial_batch_norm_op_gpu_impl.cuh" + +namespace caffe2 { + +REGISTER_CUDA_OPERATOR(SpatialBN, SpatialBNOp); +REGISTER_CUDA_OPERATOR(SpatialBNGradient, SpatialBNGradientOp); + +} // namespace caffe2 diff --git a/caffe2/operators/spatial_batch_norm_op.h b/caffe2/operators/spatial_batch_norm_op.h index 4a85ef7b3fd1a..754ac50f3725f 100644 --- a/caffe2/operators/spatial_batch_norm_op.h +++ b/caffe2/operators/spatial_batch_norm_op.h @@ -1,6 +1,12 @@ #ifndef CAFFE2_OPERATORS_SPATIAL_BATCH_NORM_OP_H_ #define CAFFE2_OPERATORS_SPATIAL_BATCH_NORM_OP_H_ +#include +#include +#include +#include +#include + #include "caffe2/core/context.h" #include "caffe2/core/operator.h" #include "caffe2/utils/math.h" @@ -11,69 +17,411 @@ template class SpatialBNOp : public Operator { public: USE_OPERATOR_CONTEXT_FUNCTIONS; + SpatialBNOp(const OperatorDef& operator_def, Workspace* ws) : Operator(operator_def, ws), - is_test_(this->template GetSingleArgument(OpSchema::Arg_IsTest, 0)), - epsilon_(this->template GetSingleArgument("epsilon", 1e-5f)), - momentum_(this->template GetSingleArgument("momentum", 0.9f)), + OP_SINGLE_ARG(bool, OpSchema::Arg_IsTest, is_test_, false), + OP_SINGLE_ARG(double, "epsilon", epsilon_, 1e-5), + OP_SINGLE_ARG(float, "momentum", momentum_, 0.9f), order_(StringToStorageOrder( - this->template GetSingleArgument("order", "NCHW"))), - num_batches_(this->template GetSingleArgument("num_batches", 1)) { - // TODO(jiayq): update the input and output size checks. + this->template GetSingleArgument("order", "NCHW"))), + OP_SINGLE_ARG(int, "num_batches", num_batches_, 1) { + CAFFE_ENFORCE_NE( + order_, + StorageOrder::UNKNOWN, + "order should be either \"NCHW\" or \"NHWC\"."); CAFFE_ENFORCE( (is_test_ && OutputSize() == 1) || (!is_test_ && OutputSize() == 5)); CAFFE_ENFORCE_GT(epsilon_, 0); CAFFE_ENFORCE_GE(momentum_, 0); CAFFE_ENFORCE_LE(momentum_, 1); } - ~SpatialBNOp() {} + + virtual ~SpatialBNOp() = default; bool RunOnDevice() override { + return DispatchHelper>::call(this, Input(0)); + } + + template + bool DoRunWithType() { + const auto& X = Input(INPUT); + const auto& scale = Input(SCALE); + const auto& bias = Input(BIAS); + auto* Y = Output(OUTPUT); + + const int ndim = X.ndim(); + CAFFE_ENFORCE_GE(ndim, 3); + const int N = X.dim32(0); + const int C = + (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(ndim - 1)); + const std::vector X_dims(X.dims().cbegin(), X.dims().cend()); + const int HxW = + std::accumulate( + X_dims.cbegin() + 1, X_dims.cend(), 1, std::multiplies()) / + C; + CAFFE_ENFORCE_EQ(scale.size(), C); + CAFFE_ENFORCE_EQ(bias.size(), C); + + Y->ResizeLike(X); + const T* X_data = X.template data(); + const T* scale_data = scale.template data(); + const T* bias_data = bias.template data(); + T* Y_data = Y->template mutable_data(); + alpha_.Resize(C); + beta_.Resize(C); + T* alpha_data = alpha_.template mutable_data(); + T* beta_data = beta_.template mutable_data(); + if (is_test_) { + if (N == 0) { + return true; + } + const auto& mean = Input(EST_MEAN); + const auto& var = Input(EST_VAR); + CAFFE_ENFORCE_EQ(mean.size(), C); + CAFFE_ENFORCE_EQ(var.size(), C); + ComputeFusedParam( + C, + scale_data, + bias_data, + mean.template data(), + var.template data(), + alpha_data, + beta_data); + } else { + auto* saved_mean = Output(SAVED_MEAN); + auto* saved_rstd = Output(SAVED_INV_STD); + if (num_batches_ == 1) { + saved_mean->Resize(C); + saved_rstd->Resize(C); + } else { + const auto& batch_mean_sum = Input(BATCH_MEAN_SUM); + const auto& batch_var_sum = Input(BATCH_VAR_SUM); + if (saved_mean != &batch_mean_sum) { + saved_mean->Resize(C); + } + if (saved_rstd != &batch_var_sum) { + saved_rstd->Resize(C); + } + } + T* saved_mean_data = saved_mean->template mutable_data(); + T* saved_rstd_data = saved_rstd->template mutable_data(); + auto* running_mean = Output(RUNNING_MEAN); + auto* running_var = Output(RUNNING_VAR); + if (running_mean->size() != C) { + running_mean->Resize(C); + math::Set( + C, T(0), running_mean->template mutable_data(), &context_); + } + if (running_var->size() != C) { + running_var->Resize(C); + math::Set( + C, T(0), running_var->template mutable_data(), &context_); + } + T* running_mean_data = running_mean->template mutable_data(); + T* running_var_data = running_var->template mutable_data(); + if (N == 0) { + math::Set(C, T(0), saved_mean_data, &context_); + math::Set(C, T(0), saved_rstd_data, &context_); + return true; + } + if (num_batches_ > 1) { + const auto& batch_mean_sum = Input(BATCH_MEAN_SUM); + const auto& batch_var_sum = Input(BATCH_VAR_SUM); + CAFFE_ENFORCE_EQ(batch_mean_sum.size(), C); + CAFFE_ENFORCE_EQ(batch_var_sum.size(), C); + ComputeBatchMoments( + N, + C, + HxW, + batch_mean_sum.template data(), + batch_var_sum.template data(), + saved_mean_data, + saved_rstd_data); + } else { + if (order_ == StorageOrder::NCHW) { + const std::array dims = {N, C, HxW}; + const std::array axes = {0, 2}; + math::Moments( + 3, + dims.data(), + 2, + axes.data(), + X_data, + saved_mean_data, + saved_rstd_data, + &context_); + } else { + const std::array dims = {N * HxW, C}; + const int axis = 0; + math::Moments( + 2, + dims.data(), + 1, + &axis, + X_data, + saved_mean_data, + saved_rstd_data, + &context_); + } + } + ComputeRunningMomentsAndFusedParam( + C, + scale_data, + bias_data, + saved_mean_data, + saved_rstd_data, + running_mean_data, + running_var_data, + saved_rstd_data, + alpha_data, + beta_data); + } + if (order_ == StorageOrder::NCHW) { + math::AffineChannel( + N, C, HxW, X_data, alpha_data, beta_data, Y_data, &context_); + } else { + math::AffineChannel( + N, C, HxW, X_data, alpha_data, beta_data, Y_data, &context_); + } + return true; } protected: - bool is_test_; + template + void ComputeFusedParam( + const int C, + const T* scale, + const T* bias, + const T* mean, + const T* var, + T* alpha, + T* beta); + + template + void ComputeBatchMoments( + const int N, + const int C, + const int HxW, + const T* batch_mean_sum, + const T* batch_var_sum, + T* mean, + T* var); + + template + void ComputeRunningMomentsAndFusedParam( + const int C, + const T* scale, + const T* bias, + const T* mean, + const T* var, + T* running_mean, + T* running_var, + T* rstd, + T* alpha, + T* beta); + + const bool is_test_; double epsilon_; - double momentum_; - StorageOrder order_; - int num_batches_; - INPUT_TAGS(INPUT, SCALE, BIAS, EST_MEAN, EST_VAR, SUMS, SUMSQ); - OUTPUT_TAGS(OUTPUT, RUNNING_MEAN, RUNNING_VAR, SAVED_MEAN, SAVED_INV_VAR); + const float momentum_; + const StorageOrder order_; + const int num_batches_; + + Tensor alpha_{Context::GetDeviceType()}; + Tensor beta_{Context::GetDeviceType()}; + + INPUT_TAGS( + INPUT, + SCALE, + BIAS, + EST_MEAN, + EST_VAR, + BATCH_MEAN_SUM, + BATCH_VAR_SUM); + OUTPUT_TAGS(OUTPUT, RUNNING_MEAN, RUNNING_VAR, SAVED_MEAN, SAVED_INV_STD); }; template class SpatialBNGradientOp : public Operator { public: USE_OPERATOR_CONTEXT_FUNCTIONS; + SpatialBNGradientOp(const OperatorDef& operator_def, Workspace* ws) : Operator(operator_def, ws), - is_test_(this->template GetSingleArgument(OpSchema::Arg_IsTest, 0)), - epsilon_(this->template GetSingleArgument("epsilon", 1e-5f)), + OP_SINGLE_ARG(double, "epsilon", epsilon_, 1e-5), order_(StringToStorageOrder( this->template GetSingleArgument("order", "NCHW"))), - num_batches_(this->template GetSingleArgument("num_batches", 1)) { + OP_SINGLE_ARG(int, "num_batches", num_batches_, 1) { + CAFFE_ENFORCE_NE( + order_, + StorageOrder::UNKNOWN, + "order should be either \"NCHW\" or \"NHWC\"."); CAFFE_ENFORCE(InputSize() == 5 || InputSize() == 7); - CAFFE_ENFORCE(OutputSize() == 3); + CAFFE_ENFORCE_EQ(OutputSize(), 3); } - ~SpatialBNGradientOp() {} + + virtual ~SpatialBNGradientOp() = default; bool RunOnDevice() override { + return DispatchHelper>::call(this, Input(0)); + } + + template + bool DoRunWithType() { + const auto& X = Input(INPUT); + const auto& dY = Input(OUTPUT_GRAD); + const auto& scale = Input(SCALE); + const auto& mean = Input(SAVED_MEAN); + const auto& rstd = Input(SAVED_INV_STD); + const int ndim = X.ndim(); + CAFFE_ENFORCE_GE(ndim, 3); + const int N = X.dim32(0); + const int C = + (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(ndim - 1)); + const std::vector X_dims(X.dims().cbegin(), X.dims().cend()); + const int HxW = + std::accumulate( + X_dims.cbegin() + 1, X_dims.cend(), 1, std::multiplies()) / + C; + CAFFE_ENFORCE_EQ(scale.size(), C); + CAFFE_ENFORCE_EQ(mean.size(), C); + CAFFE_ENFORCE_EQ(rstd.size(), C); + auto* dX = Output(INPUT_GRAD); + auto* dscale = Output(SCALE_GRAD); + auto* dbias = Output(BIAS_GRAD); + dX->ResizeLike(X); + if (num_batches_ == 1) { + dscale->ResizeLike(scale); + dbias->ResizeLike(scale); + } else { + const auto& dscale_sum = Input(AGGREGATE_SCALE_GRAD); + const auto& dbias_sum = Input(AGGREGATE_BIAS_GRAD); + if (dscale != &dscale_sum) { + dscale->ResizeLike(dscale_sum); + } + if (dbias != &dbias_sum) { + dbias->ResizeLike(dbias_sum); + } + } + const T* X_data = X.template data(); + const T* dY_data = dY.template data(); + const T* scale_data = scale.template data(); + const T* mean_data = mean.template data(); + const T* rstd_data = rstd.template data(); + T* dX_data = dX->template mutable_data(); + T* dscale_data = dscale->template mutable_data(); + T* dbias_data = dbias->template mutable_data(); + + if (N == 0) { + math::Set(C, T(0), dscale_data, &context_); + math::Set(C, T(0), dbias_data, &context_); + return true; + } + alpha_.Resize(C); + beta_.Resize(C); + gamma_.Resize(C); + T* alpha_data = alpha_.template mutable_data(); + T* beta_data = beta_.template mutable_data(); + T* gamma_data = gamma_.template mutable_data(); + if (num_batches_ > 1) { + const auto& dscale_sum = Input(AGGREGATE_SCALE_GRAD); + const auto& dbias_sum = Input(AGGREGATE_BIAS_GRAD); + ComputeMultiBatchScaleBiasGradientsAndFusedParams( + N, + C, + HxW, + scale_data, + mean_data, + rstd_data, + dscale_sum.template data(), + dbias_sum.template data(), + dscale_data, + dbias_data, + alpha_data, + beta_data, + gamma_data); + } else { + ComputeScaleBiasGradientsAndFusedParams( + N, + C, + HxW, + dY_data, + X_data, + scale_data, + mean_data, + rstd_data, + dscale_data, + dbias_data, + alpha_data, + beta_data, + gamma_data); + } + ComputeXGradient( + N, C, HxW, dY_data, X_data, alpha_data, beta_data, gamma_data, dX_data); + return true; } protected: - bool is_test_; + template + void ComputeMultiBatchScaleBiasGradientsAndFusedParams( + const int N, + const int C, + const int HxW, + const T* scale, + const T* mean, + const T* rstd, + const T* dscale_sum, + const T* dbias_sum, + T* dscale, + T* dbias, + T* alpha, + T* beta, + T* gamma); + + template + void ComputeScaleBiasGradientsAndFusedParams( + const int N, + const int C, + const int HxW, + const T* dY, + const T* X, + const T* scale, + const T* mean, + const T* rstd, + T* dscale, + T* dbias, + T* alpha, + T* beta, + T* gamma); + + template + void ComputeXGradient( + const int N, + const int C, + const int HxW, + const T* dY, + const T* X, + const T* alpha, + const T* beta, + const T* gamma, + T* dX); + double epsilon_; - StorageOrder order_; - int num_batches_; + const StorageOrder order_; + const int num_batches_; + + Tensor alpha_{Context::GetDeviceType()}; + Tensor beta_{Context::GetDeviceType()}; + Tensor gamma_{Context::GetDeviceType()}; INPUT_TAGS( INPUT, SCALE, OUTPUT_GRAD, SAVED_MEAN, - SAVED_INV_VAR, + SAVED_INV_STD, AGGREGATE_SCALE_GRAD, AGGREGATE_BIAS_GRAD); OUTPUT_TAGS(INPUT_GRAD, SCALE_GRAD, BIAS_GRAD); diff --git a/caffe2/operators/spatial_batch_norm_op_cudnn.cc b/caffe2/operators/spatial_batch_norm_op_cudnn.cc deleted file mode 100644 index 486485623ef7f..0000000000000 --- a/caffe2/operators/spatial_batch_norm_op_cudnn.cc +++ /dev/null @@ -1,376 +0,0 @@ -#include - -#include "caffe2/core/context_gpu.h" -#include "caffe2/core/cudnn_wrappers.h" -#include "caffe2/operators/spatial_batch_norm_op.h" -#include "caffe2/utils/math.h" - -// Note: Instead of directly failing, we will choose to not build this operator -// if cudnn version is not high enough. -static_assert(CUDNN_VERSION >= 5000, - "CudnnSpatialBN requires cudnn version 5.0 or above."); - -namespace caffe2 { - -class CudnnSpatialBNOp final : public SpatialBNOp { - public: - USE_OPERATOR_FUNCTIONS(CUDAContext); - CudnnSpatialBNOp(const OperatorDef& operator_def, Workspace* ws) - : SpatialBNOp(operator_def, ws), cudnn_wrapper_(&context_) { - CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&data_desc_)); - CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&bn_param_desc_)); - if (epsilon_ <= CUDNN_BN_MIN_EPSILON - FLT_EPSILON) { - LOG(ERROR) << "Provided epsilon is smaller than " - << "CUDNN_BN_MIN_EPSILON. Setting it to " - << "CUDNN_BN_MIN_EPSILON instead."; - } - epsilon_ = std::max(epsilon_, CUDNN_BN_MIN_EPSILON); -#if CUDNN_VERSION_MIN(7,0,0) - mode_ = CUDNN_BATCHNORM_SPATIAL_PERSISTENT; -#else - mode_ = CUDNN_BATCHNORM_SPATIAL; -#endif - } - - ~CudnnSpatialBNOp() { - CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(data_desc_)); - CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(bn_param_desc_)); - } - - template - bool DoRunWithType(); - bool RunOnDevice() override; - - protected: - CuDNNWrapper cudnn_wrapper_; - cudnnTensorDescriptor_t data_desc_; - cudnnTensorDescriptor_t bn_param_desc_; - vector cudnn_input_dims_; - - cudnnBatchNormMode_t mode_; -}; - -class CudnnSpatialBNGradientOp final : public SpatialBNGradientOp { - public: - USE_OPERATOR_FUNCTIONS(CUDAContext); - CudnnSpatialBNGradientOp(const OperatorDef& operator_def, Workspace* ws) - : SpatialBNGradientOp(operator_def, ws), - cudnn_wrapper_(&context_) { - CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&data_desc_)); - CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&bn_param_desc_)); - if (epsilon_ <= CUDNN_BN_MIN_EPSILON - FLT_EPSILON) { - LOG(ERROR) << "Provided epsilon is smaller than " - << "CUDNN_BN_MIN_EPSILON. Setting it to " - << "CUDNN_BN_MIN_EPSILON instead."; - } - epsilon_ = std::max(epsilon_, CUDNN_BN_MIN_EPSILON); -#if CUDNN_VERSION_MIN(7,0,0) - mode_ = CUDNN_BATCHNORM_SPATIAL_PERSISTENT; -#else - mode_ = CUDNN_BATCHNORM_SPATIAL; -#endif - } - - ~CudnnSpatialBNGradientOp() { - CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(data_desc_)); - CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(bn_param_desc_)); - } - - template - bool DoRunWithType(); - - bool RunOnDevice() override; - - protected: - CuDNNWrapper cudnn_wrapper_; - cudnnTensorDescriptor_t data_desc_; - cudnnTensorDescriptor_t bn_param_desc_; - vector cudnn_input_dims_; - - cudnnBatchNormMode_t mode_; -}; - - -//////////////////////////////////////////////////////////////////////////////// -// Implementations -//////////////////////////////////////////////////////////////////////////////// - -template -bool CudnnSpatialBNOp::DoRunWithType() { - - // QoL - typedef typename cudnnTypeWrapper::BNParamType BNParamType; - - const auto& X = Input(INPUT); - const auto& scale = Input(SCALE); - const auto& bias = Input(BIAS); - - CAFFE_ENFORCE_GE(X.ndim(), 3); - const int N = X.dim32(0); - const int C = X.ndim() > 3 - ? (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(X.ndim() - 1)) - : (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(2)); - const int H = (order_ == StorageOrder::NCHW ? X.dim32(2) : X.dim32(1)); - const int W = X.ndim() > 3 - ? (order_ == StorageOrder::NCHW ? X.dim32(3) : X.dim32(2)) - : 1; - const int D = X.ndim() > 4 - ? (order_ == StorageOrder::NCHW ? X.dim32(4) : X.dim32(3)) - : 1; - CAFFE_ENFORCE_EQ(scale.ndim(), 1); - CAFFE_ENFORCE_EQ(bias.ndim(), 1); - CAFFE_ENFORCE_EQ(scale.dim32(0), C); - CAFFE_ENFORCE_EQ(bias.dim32(0), C); - // See if we need to reshape. - if (N > 0 && X.dims() != cudnn_input_dims_) { - VLOG(1) << "Setting descriptors."; - cudnn_input_dims_ = X.dims(); - if (order_ == StorageOrder::NCHW) { - vector dims = {N, C, H, W, D}; - vector strides = {C * H * W * D, H * W * D, W * D, D, 1}; - CUDNN_ENFORCE(cudnnSetTensorNdDescriptor( - data_desc_, - cudnnTypeWrapper::type, - X.ndim() > 3 ? X.ndim() : 4, - dims.data(), - strides.data())); - } else { - vector dims = {N, C, H, W, D}; - vector strides = {H * W * D * C, 1, W * D * C, D * C, C}; - CUDNN_ENFORCE(cudnnSetTensorNdDescriptor( - data_desc_, - cudnnTypeWrapper::type, - X.ndim() > 3 ? X.ndim() : 4, - dims.data(), - strides.data())); - } - CUDNN_ENFORCE(cudnnDeriveBNTensorDescriptor( - bn_param_desc_, data_desc_, mode_)); - } - - // Now, depending on whether we are running test or not, we have two paths. - if (is_test_) { - // Run inference mode. - const auto& est_mean = Input(EST_MEAN); - const auto& est_var = Input(EST_VAR); - CAFFE_ENFORCE_EQ(est_mean.ndim(), 1); - CAFFE_ENFORCE_EQ(est_var.ndim(), 1); - CAFFE_ENFORCE_EQ(est_mean.dim32(0), C); - CAFFE_ENFORCE_EQ(est_var.dim32(0), C); - - auto* Y = Output(OUTPUT); - Y->ResizeLike(X); - T* Y_data = Y->template mutable_data(); - if (N == 0) { - return true; - } - CUDNN_ENFORCE(cudnnBatchNormalizationForwardInference( - cudnn_wrapper_.inline_cudnn_handle(), - // Note: PERSISTENT not implemented for inference - CUDNN_BATCHNORM_SPATIAL, - cudnnTypeWrapper::kOne(), - cudnnTypeWrapper::kZero(), - data_desc_, - X.template data(), - data_desc_, - Y_data, - bn_param_desc_, - scale.template data(), - bias.template data(), - est_mean.template data(), - est_var.template data(), - epsilon_)); - } else { - // Run training mode. - auto* Y = Output(OUTPUT); - Y->ResizeLike(X); - T* Y_data = Y->template mutable_data(); - // obtain running mean and running inv var, and see if we need to - // initialize them. - auto* running_mean = Output(RUNNING_MEAN); - auto* running_var = Output(RUNNING_VAR); - double this_factor = 1. - momentum_; - BNParamType* running_mean_data = nullptr; - BNParamType* running_var_data = nullptr; - if (!running_mean->size()) { - // If the input mean and var are not initialized yet, this is the first - // run and we will initialize the storage. - VLOG(1) << "Initializing running mean and var."; - // Need to do initialization - running_mean->Resize(C); - running_var->Resize(C); - running_mean_data = running_mean->template mutable_data(); - running_var_data = running_var->template mutable_data(); - // In principle, setting this_momentum to 1 will wipe existing data. - // This has a caveat that if cudnn does not deal with 0*NaN cases we - // will be having an issue. Thus we choose a safe path by explicitly - // setting zero. - math::Set(C, 0, running_mean_data, &context_); - math::Set(C, 0, running_var_data, &context_); - } else { - // Does not need to do initialization. - CAFFE_ENFORCE_EQ(running_mean->ndim(), 1); - CAFFE_ENFORCE_EQ(running_var->ndim(), 1); - CAFFE_ENFORCE_EQ(running_mean->dim32(0), C); - CAFFE_ENFORCE_EQ(running_var->dim32(0), C); - running_mean_data = running_mean->template mutable_data(); - running_var_data = running_var->template mutable_data(); - } - // Save the mean and inv var results. - auto* save_mean = Output(SAVED_MEAN); - auto* save_var = Output(SAVED_INV_VAR); - save_mean->Resize(C); - save_var->Resize(C); - void* save_mean_data = save_mean->template mutable_data(); - void* save_var_data = save_var->template mutable_data(); - if (N == 0) { - // set empty batch's mean / var output to zeros - math::Set( - C, 0, (BNParamType*)save_mean_data, &context_); - math::Set( - C, 0, (BNParamType*)save_var_data, &context_); - return true; - } - CUDNN_ENFORCE(cudnnBatchNormalizationForwardTraining( - cudnn_wrapper_.inline_cudnn_handle(), - mode_, - cudnnTypeWrapper::kOne(), - cudnnTypeWrapper::kZero(), - data_desc_, - X.template data(), - data_desc_, - Y_data, - bn_param_desc_, - scale.template data(), - bias.template data(), - this_factor, - running_mean_data, - running_var_data, - epsilon_, - save_mean_data, - save_var_data)); - } - return true; -} - -bool CudnnSpatialBNOp::RunOnDevice() { - if (Input(0).IsType()) { - return DoRunWithType(); - } else if (Input(0).IsType()) { - return DoRunWithType(); - } else { - LOG(FATAL) << "Unsupported input types"; - } - return true; -} - -template -bool CudnnSpatialBNGradientOp::DoRunWithType() { - // QoL - typedef typename cudnnTypeWrapper::BNParamType BNParamType; - - const auto& X = Input(INPUT); - const auto& scale = Input(SCALE); - const auto& dY = Input(OUTPUT_GRAD); - - CAFFE_ENFORCE_GE(X.ndim(), 3); - const int N = X.dim32(0); - const int C = X.ndim() > 3 - ? (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(X.ndim() - 1)) - : (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(2)); - const int H = (order_ == StorageOrder::NCHW ? X.dim32(2) : X.dim32(1)); - const int W = X.ndim() > 3 - ? (order_ == StorageOrder::NCHW ? X.dim32(3) : X.dim32(2)) - : 1; - const int D = X.ndim() > 4 - ? (order_ == StorageOrder::NCHW ? X.dim32(4) : X.dim32(3)) - : 1; - CAFFE_ENFORCE_EQ(scale.ndim(), 1); - CAFFE_ENFORCE_EQ(scale.dim32(0), C); - // See if we need to reshape. - if (N > 0 && X.dims() != cudnn_input_dims_) { - if (order_ == StorageOrder::NCHW) { - vector dims = {N, C, H, W, D}; - vector strides = {C * H * W * D, H * W * D, W * D, D, 1}; - CUDNN_ENFORCE(cudnnSetTensorNdDescriptor( - data_desc_, - cudnnTypeWrapper::type, - X.ndim() > 3 ? X.ndim() : 4, - dims.data(), - strides.data())); - } else { - vector dims = {N, C, H, W, D}; - vector strides = {H * W * C * D, 1, W * D * C, D * C, C}; - CUDNN_ENFORCE(cudnnSetTensorNdDescriptor( - data_desc_, - cudnnTypeWrapper::type, - X.ndim() > 3 ? X.ndim() : 4, - dims.data(), - strides.data())); - } - CUDNN_ENFORCE(cudnnDeriveBNTensorDescriptor( - bn_param_desc_, data_desc_, mode_)); - } - - auto* dX = Output(INPUT_GRAD); - dX->ResizeLike(X); - T* dX_data = dX->template mutable_data(); - auto* dScale = Output(SCALE_GRAD); - auto* dBias = Output(BIAS_GRAD); - dScale->ResizeLike(scale); - dBias->ResizeLike(scale); - auto* dScale_data = dScale->template mutable_data(); - auto* dBias_data = dBias->template mutable_data(); - - const auto& saved_mean = Input(SAVED_MEAN); - const auto& saved_var = Input(SAVED_INV_VAR); - const void* saved_mean_data = saved_mean.template data(); - const void* saved_var_data = saved_var.template data(); - if (N == 0) { - // set gradients to zeros - math::Set(C, 0, dScale_data, &context_); - math::Set(C, 0, dBias_data, &context_); - return true; - } - CUDNN_ENFORCE(cudnnBatchNormalizationBackward( - cudnn_wrapper_.inline_cudnn_handle(), - mode_, - cudnnTypeWrapper::kOne(), - cudnnTypeWrapper::kZero(), - cudnnTypeWrapper::kOne(), - cudnnTypeWrapper::kZero(), - data_desc_, - X.template data(), - data_desc_, - dY.template data(), - data_desc_, - dX_data, - bn_param_desc_, - scale.template data(), - dScale_data, - dBias_data, - epsilon_, - saved_mean_data, - saved_var_data)); - return true; -} - -bool CudnnSpatialBNGradientOp::RunOnDevice() { - if (Input(0).IsType()) { - return DoRunWithType(); - } else if (Input(0).IsType()) { - return DoRunWithType(); - } else { - LOG(FATAL) << "Unsupported input types"; - } - return true; -} - -// Since there is no default implementation for spatial batch normalization, -// we will register the cudnn version as the default as well. -REGISTER_CUDA_OPERATOR(SpatialBN, CudnnSpatialBNOp); -REGISTER_CUDA_OPERATOR(SpatialBNGradient, CudnnSpatialBNGradientOp); - -REGISTER_CUDNN_OPERATOR(SpatialBN, CudnnSpatialBNOp); -REGISTER_CUDNN_OPERATOR(SpatialBNGradient, CudnnSpatialBNGradientOp); -} // namespace caffe2 diff --git a/caffe2/operators/spatial_batch_norm_op_cudnn.cu b/caffe2/operators/spatial_batch_norm_op_cudnn.cu new file mode 100644 index 0000000000000..39adc54a17c26 --- /dev/null +++ b/caffe2/operators/spatial_batch_norm_op_cudnn.cu @@ -0,0 +1,345 @@ +#include "caffe2/operators/spatial_batch_norm_op.h" + +#include +#include +#include +#include + +#include "caffe2/core/context_gpu.h" +#include "caffe2/core/cudnn_wrappers.h" +#include "caffe2/operators/spatial_batch_norm_op_gpu_impl.cuh" +#include "caffe2/utils/math.h" + +#if CUDNN_VERSION_MIN(5, 0, 0) + +namespace caffe2 { + +namespace { + +void SetTensorDescriptor( + const cudnnDataType_t data_type, + const cudnnBatchNormMode_t mode, + const StorageOrder order, + const std::vector& input_dims, + cudnnTensorDescriptor_t data_desc, + cudnnTensorDescriptor_t param_desc) { + const int ndim = input_dims.size(); + const int N = input_dims[0]; + const int C = order == StorageOrder::NCHW ? input_dims[1] : input_dims.back(); + if (ndim == 3) { + const int H = 1; + const int W = order == StorageOrder::NCHW ? input_dims[2] : input_dims[1]; + CUDNN_ENFORCE(cudnnSetTensor4dDescriptor( + data_desc, GetCudnnTensorFormat(order), data_type, N, C, H, W)); + } else if (ndim == 4) { + const int H = order == StorageOrder::NCHW ? input_dims[2] : input_dims[1]; + const int W = order == StorageOrder::NCHW ? input_dims[3] : input_dims[2]; + CUDNN_ENFORCE(cudnnSetTensor4dDescriptor( + data_desc, GetCudnnTensorFormat(order), data_type, N, C, H, W)); + } else { + const int H = order == StorageOrder::NCHW ? input_dims[2] : input_dims[1]; + const int W = order == StorageOrder::NCHW ? input_dims[3] : input_dims[2]; + const auto l_iter = order == StorageOrder::NCHW ? input_dims.cbegin() + 4 + : input_dims.cbegin() + 3; + const auto r_iter = + order == StorageOrder::NCHW ? input_dims.cend() : input_dims.cend() - 1; + const int D = std::accumulate(l_iter, r_iter, 1, std::multiplies()); + const std::array dims = {N, C, H, W, D}; + const std::array strides = order == StorageOrder::NCHW + ? std::array{C * H * W * D, H * W * D, W * D, D, 1} + : std::array{C * H * W * D, 1, W * D * C, D * C, C}; + CUDNN_ENFORCE(cudnnSetTensorNdDescriptor( + data_desc, data_type, 5, dims.data(), strides.data())); + } + CUDNN_ENFORCE(cudnnDeriveBNTensorDescriptor(param_desc, data_desc, mode)); +} + +} // namespace + +class CuDNNSpatialBNOp final : public SpatialBNOp { + public: + USE_OPERATOR_FUNCTIONS(CUDAContext); + + CuDNNSpatialBNOp(const OperatorDef& operator_def, Workspace* ws) + : SpatialBNOp(operator_def, ws), + cudnn_wrapper_(&context_), +#if CUDNN_VERSION_MIN(7, 0, 0) + mode_(CUDNN_BATCHNORM_SPATIAL_PERSISTENT) { +#else + mode_(CUDNN_BATCHNORM_SPATIAL) { +#endif + CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&data_desc_)); + CUDNN_ENFORCE(cudnnCreateTensorDescriptor(¶m_desc_)); + if (epsilon_ < CUDNN_BN_MIN_EPSILON) { + LOG(ERROR) << "Provided epsilon is smaller than CUDNN_BN_MIN_EPSILON. " + "Setting it to CUDNN_BN_MIN_EPSILON instead."; + epsilon_ = CUDNN_BN_MIN_EPSILON; + } + } + + ~CuDNNSpatialBNOp() { + CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(data_desc_)); + CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(param_desc_)); + } + + bool RunOnDevice() override { + // CuDNN doesn't support multi-batch SpatialBN and it's NHWC order SpatialBN + // is much slower, so in such cases fallback to SpatialBNOp. + if (num_batches_ > 1 || order_ == StorageOrder::NHWC) { + return SpatialBNOp::RunOnDevice(); + } + return DispatchHelper>::call(this, Input(0)); + } + + template + bool DoRunWithType() { + typedef typename cudnnTypeWrapper::BNParamType BNParamType; + + const auto& X = Input(INPUT); + const auto& scale = Input(SCALE); + const auto& bias = Input(BIAS); + auto* Y = Output(OUTPUT); + const int ndim = X.ndim(); + CAFFE_ENFORCE_GE(ndim, 3); + const int N = X.dim32(0); + const int C = + (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(ndim - 1)); + CAFFE_ENFORCE_EQ(scale.size(), C); + CAFFE_ENFORCE_EQ(bias.size(), C); + Y->ResizeLike(X); + const T* X_data = X.data(); + const BNParamType* scale_data = scale.data(); + const BNParamType* bias_data = bias.data(); + T* Y_data = Y->mutable_data(); + + if (N > 0) { + const std::vector input_dims(X.dims().cbegin(), X.dims().cend()); + if (input_dims != data_dims_) { + data_dims_ = input_dims; + SetTensorDescriptor( + cudnnTypeWrapper::type, + mode_, + order_, + input_dims, + data_desc_, + param_desc_); + } + } + if (is_test_) { + const auto& mean = Input(EST_MEAN); + const auto& var = Input(EST_VAR); + CAFFE_ENFORCE_EQ(mean.size(), C); + CAFFE_ENFORCE_EQ(var.size(), C); + if (N == 0) { + return true; + } + CUDNN_ENFORCE(cudnnBatchNormalizationForwardInference( + cudnn_wrapper_.inline_cudnn_handle(), + // Note: PERSISTENT not implemented for inference + CUDNN_BATCHNORM_SPATIAL, + cudnnTypeWrapper::kOne(), + cudnnTypeWrapper::kZero(), + data_desc_, + X_data, + data_desc_, + Y_data, + param_desc_, + scale_data, + bias_data, + mean.data(), + var.data(), + epsilon_)); + } else { + auto* saved_mean = Output(SAVED_MEAN); + auto* saved_inv_std = Output(SAVED_INV_STD); + saved_mean->Resize(C); + saved_inv_std->Resize(C); + BNParamType* saved_mean_data = saved_mean->mutable_data(); + BNParamType* saved_inv_std_data = + saved_inv_std->mutable_data(); + auto* running_mean = Output(RUNNING_MEAN); + auto* running_var = Output(RUNNING_VAR); + if (running_mean->size() != C) { + running_mean->Resize(C); + math::Set( + C, + BNParamType(0), + running_mean->mutable_data(), + &context_); + } + if (running_var->size() != C) { + running_var->Resize(C); + math::Set( + C, + BNParamType(0), + running_var->mutable_data(), + &context_); + } + BNParamType* running_mean_data = + running_mean->mutable_data(); + BNParamType* running_var_data = running_var->mutable_data(); + if (N == 0) { + math::Set( + C, BNParamType(0), saved_mean_data, &context_); + math::Set( + C, BNParamType(0), saved_inv_std_data, &context_); + return true; + } + const double alpha = static_cast(1.0f - momentum_); + CUDNN_ENFORCE(cudnnBatchNormalizationForwardTraining( + cudnn_wrapper_.inline_cudnn_handle(), + mode_, + cudnnTypeWrapper::kOne(), + cudnnTypeWrapper::kZero(), + data_desc_, + X_data, + data_desc_, + Y_data, + param_desc_, + scale_data, + bias_data, + alpha, + running_mean_data, + running_var_data, + epsilon_, + saved_mean_data, + saved_inv_std_data)); + } + return true; + } + + private: + CuDNNWrapper cudnn_wrapper_; + cudnnTensorDescriptor_t data_desc_; + cudnnTensorDescriptor_t param_desc_; + cudnnBatchNormMode_t mode_; + + std::vector data_dims_; +}; + +class CuDNNSpatialBNGradientOp final : public SpatialBNGradientOp { + public: + USE_OPERATOR_FUNCTIONS(CUDAContext); + + CuDNNSpatialBNGradientOp(const OperatorDef& operator_def, Workspace* ws) + : SpatialBNGradientOp(operator_def, ws), + cudnn_wrapper_(&context_), +#if CUDNN_VERSION_MIN(7, 0, 0) + mode_(CUDNN_BATCHNORM_SPATIAL_PERSISTENT) { +#else + mode_(CUDNN_BATCHNORM_SPATIAL) { +#endif + CUDNN_ENFORCE(cudnnCreateTensorDescriptor(&data_desc_)); + CUDNN_ENFORCE(cudnnCreateTensorDescriptor(¶m_desc_)); + if (epsilon_ < CUDNN_BN_MIN_EPSILON) { + LOG(ERROR) << "Provided epsilon is smaller than CUDNN_BN_MIN_EPSILON. " + "Setting it to CUDNN_BN_MIN_EPSILON instead."; + epsilon_ = CUDNN_BN_MIN_EPSILON; + } + } + + ~CuDNNSpatialBNGradientOp() { + CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(data_desc_)); + CUDNN_ENFORCE(cudnnDestroyTensorDescriptor(param_desc_)); + } + + bool RunOnDevice() override { + // CuDNN doesn't support multi-batch SpatialBN and it's NHWC order SpatialBN + // is much slower, so in such cases fallback to SpatialBNOp. + if (num_batches_ > 1 || order_ == StorageOrder::NHWC) { + return SpatialBNGradientOp::RunOnDevice(); + } + return DispatchHelper>::call(this, Input(0)); + } + + template + bool DoRunWithType() { + typedef typename cudnnTypeWrapper::BNParamType BNParamType; + + const auto& X = Input(INPUT); + const auto& scale = Input(SCALE); + const auto& dY = Input(OUTPUT_GRAD); + const auto& saved_mean = Input(SAVED_MEAN); + const auto& saved_rstd = Input(SAVED_INV_STD); + auto* dX = Output(INPUT_GRAD); + auto* dscale = Output(SCALE_GRAD); + auto* dbias = Output(BIAS_GRAD); + const int ndim = X.ndim(); + CAFFE_ENFORCE_GE(ndim, 3); + const int N = X.dim32(0); + const int C = + (order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(ndim - 1)); + CAFFE_ENFORCE_EQ(scale.size(), C); + CAFFE_ENFORCE_EQ(saved_mean.size(), C); + CAFFE_ENFORCE_EQ(saved_rstd.size(), C); + dX->ResizeLike(X); + dscale->ResizeLike(scale); + dbias->ResizeLike(scale); + const T* X_data = X.template data(); + const T* scale_data = scale.template data(); + const T* dY_data = dY.template data(); + const BNParamType* saved_mean_data = + saved_mean.template data(); + const BNParamType* saved_rstd_data = + saved_rstd.template data(); + T* dX_data = dX->template mutable_data(); + BNParamType* dscale_data = dscale->template mutable_data(); + BNParamType* dbias_data = dbias->template mutable_data(); + if (N == 0) { + math::Set( + C, BNParamType(0), dscale_data, &context_); + math::Set( + C, BNParamType(0), dbias_data, &context_); + return true; + } + + const std::vector input_dims(X.dims().cbegin(), X.dims().cend()); + if (input_dims != data_dims_) { + data_dims_ = input_dims; + SetTensorDescriptor( + cudnnTypeWrapper::type, + mode_, + order_, + input_dims, + data_desc_, + param_desc_); + } + CUDNN_ENFORCE(cudnnBatchNormalizationBackward( + cudnn_wrapper_.inline_cudnn_handle(), + mode_, + cudnnTypeWrapper::kOne(), + cudnnTypeWrapper::kZero(), + cudnnTypeWrapper::kOne(), + cudnnTypeWrapper::kZero(), + data_desc_, + X_data, + data_desc_, + dY_data, + data_desc_, + dX_data, + param_desc_, + scale_data, + dscale_data, + dbias_data, + epsilon_, + saved_mean_data, + saved_rstd_data)); + + return true; + } + + private: + CuDNNWrapper cudnn_wrapper_; + cudnnTensorDescriptor_t data_desc_; + cudnnTensorDescriptor_t param_desc_; + cudnnBatchNormMode_t mode_; + + std::vector data_dims_; +}; + +REGISTER_CUDNN_OPERATOR(SpatialBN, CuDNNSpatialBNOp); +REGISTER_CUDNN_OPERATOR(SpatialBNGradient, CuDNNSpatialBNGradientOp); + +} // namespace caffe2 + +#endif // CUDNN_VERSION_MIN(5, 0, 0) diff --git a/caffe2/operators/spatial_batch_norm_op_gpu_impl.cuh b/caffe2/operators/spatial_batch_norm_op_gpu_impl.cuh new file mode 100644 index 0000000000000..7d87e81fb9679 --- /dev/null +++ b/caffe2/operators/spatial_batch_norm_op_gpu_impl.cuh @@ -0,0 +1,439 @@ +#include "caffe2/operators/spatial_batch_norm_op.h" + +#include + +#include "caffe2/core/context_gpu.h" +#include "caffe2/utils/math.h" + +namespace caffe2 { + +namespace { + +template +using BlockReduce = cub::BlockReduce; + +template +__global__ void ComputeFusedParamCUDAKernel( + const int C, + const T epsilon, + const T* scale, + const T* bias, + const T* mean, + const T* var, + T* alpha, + T* beta); + +template <> +__global__ void ComputeFusedParamCUDAKernel( + const int C, + const float epsilon, + const float* scale, + const float* bias, + const float* mean, + const float* var, + float* alpha, + float* beta) { + CUDA_1D_KERNEL_LOOP(i, C) { +#if __CUDA_ARCH__ >= 350 + const float scale_x_rstd = + __ldg(scale + i) * rsqrtf(__ldg(var + i) + epsilon); + alpha[i] = scale_x_rstd; + beta[i] = __ldg(bias + i) - scale_x_rstd * __ldg(mean + i); +#else + const float scale_x_rstd = scale[i] * rsqrtf(var[i] + epsilon); + alpha[i] = scale_x_rstd; + beta[i] = bias[i] - scale_x_rstd * mean[i]; +#endif + } +} + +template +__global__ void ComputeBatchMomentsCUDAKernel( + const int C, + const T scale, + const T* batch_mean_sum, + const T* batch_var_sum, + T* mean, + T* var) { + CUDA_1D_KERNEL_LOOP(i, C) { +#if __CUDA_ARCH__ >= 350 + const T mu = __ldg(batch_mean_sum + i) * scale; + mean[i] = mu; + var[i] = __ldg(batch_var_sum + i) * scale - mu * mu; +#else + const T mu = batch_mean_sum[i] * scale; + mean[i] = mu; + var[i] = batch_var_sum[i] * scale - mu * mu; +#endif + } +} + +template +__global__ void ComputeRunningMomentsAndFusedParamCUDAKernel( + const int C, + const T momentum, + const T epsilon, + const T* scale, + const T* bias, + const T* mean, + const T* var, + T* running_mean, + T* running_var, + T* rstd, + T* alpha, + T* beta); + +template <> +__global__ void ComputeRunningMomentsAndFusedParamCUDAKernel( + const int C, + const float momentum, + const float epsilon, + const float* scale, + const float* bias, + const float* mean, + const float* var, + float* running_mean, + float* running_var, + float* rstd, + float* alpha, + float* beta) { + const float a = 1.0f - momentum; + const float b = momentum; + CUDA_1D_KERNEL_LOOP(i, C) { +#if __CUDA_ARCH__ >= 350 + running_mean[i] = a * __ldg(mean + i) + b * __ldg(running_mean + i); + running_var[i] = a * __ldg(var + i) + b * __ldg(running_var + i); + const float rstd_val = rsqrtf(__ldg(var + i) + epsilon); + const float scale_x_rstd = __ldg(scale + i) * rstd_val; + rstd[i] = rstd_val; + alpha[i] = scale_x_rstd; + beta[i] = bias[i] - scale_x_rstd * __ldg(mean + i); +#else + running_mean[i] = a * mean[i] + b * running_mean[i]; + running_var[i] = a * var[i] + b * running_var[i]; + const float rstd_val = rsqrtf(var[i] + epsilon); + const float scale_x_rstd = scale[i] * rstd_val; + rstd[i] = rstd_val; + alpha[i] = scale_x_rstd; + beta[i] = bias[i] - scale_x_rstd * mean[i]; +#endif + } +} + +template +__global__ void ComputeMultiBatchScaleBiasGradientsAndFusedParamsCUDAKernel( + const int C, + const T inv_num_batches, + const T inv_nhw, + const T* scale, + const T* mean, + const T* rstd, + const T* dscale_sum, + const T* dbias_sum, + T* dscale, + T* dbias, + T* alpha, + T* beta, + T* gamma) { + CUDA_1D_KERNEL_LOOP(i, C) { +#if __CUDA_ARCH__ >= 350 + const T dscale_val = __ldg(dscale_sum + i) * inv_num_batches; + const T dbias_val = __ldg(dbias_sum + i) * inv_num_batches; + const T scale_x_rstd = __ldg(scale + i) * __ldg(rstd + i); + const T dscale_x_rstd = dscale_val * __ldg(rstd + i); + dscale[i] = dscale_val; + dbias[i] = dbias_val; + alpha[i] = scale_x_rstd; + beta[i] = -scale_x_rstd * dscale_x_rstd * inv_nhw; + gamma[i] = + scale_x_rstd * (__ldg(mean + i) * dscale_x_rstd - dbias_val) * inv_nhw; +#else + const T dscale_val = dscale_sum[i] * inv_num_batches; + const T dbias_val = dbias_sum[i] * inv_num_batches; + const T scale_x_rstd = scale[i] * rstd[i]; + const T dscale_x_rstd = dscale_val * rstd[i]; + dscale[i] = dscale_val; + dbias[i] = dbias_val; + alpha[i] = scale_x_rstd; + beta[i] = -scale_x_rstd * dscale_x_rstd * inv_nhw; + gamma[i] = scale_x_rstd * (mean[i] * dscale_x_rstd - dbias_val) * inv_nhw; +#endif + } +} + +template +__global__ void ComputeScaleBiasGradientsAndFusedParamsCUDAKernel( + const int N, + const int C, + const int HxW, + const T* dY, + const T* X, + const T* scale, + const T* mean, + const T* rstd, + T* dscale, + T* dbias, + T* alpha, + T* beta, + T* gamma) { + const int outer_size = C; + const int inner_size = N * HxW; + const T inv_nhw = T(1) / static_cast(N * HxW); + __shared__ typename BlockReduce::TempStorage ds_storage; + __shared__ typename BlockReduce::TempStorage db_storage; + for (int i = blockIdx.x; i < outer_size; i += gridDim.x) { + T ds_val = 0; + T db_val = 0; + for (int j = threadIdx.x; j < inner_size; j += blockDim.x) { + const int index = kOrder == StorageOrder::NCHW + ? (j / HxW * C + i) * HxW + j % HxW + : j * C + i; + ds_val += dY[index] * (X[index] - mean[i]) * rstd[i]; + db_val += dY[index]; + } + ds_val = BlockReduce(ds_storage).Sum(ds_val); + db_val = BlockReduce(db_storage).Sum(db_val); + if (threadIdx.x == 0) { +#if __CUDA_ARCH__ >= 350 + const T scale_x_rstd = __ldg(scale + i) * __ldg(rstd + i); + const T dscale_x_rstd = ds_val * __ldg(rstd + i); + dscale[i] = ds_val; + dbias[i] = db_val; + alpha[i] = scale_x_rstd; + beta[i] = -scale_x_rstd * dscale_x_rstd * inv_nhw; + gamma[i] = + scale_x_rstd * (__ldg(mean + i) * dscale_x_rstd - db_val) * inv_nhw; +#else + const T scale_x_rstd = scale[i] * rstd[i]; + const T dscale_x_rstd = ds_val * rstd[i]; + dscale[i] = ds_val; + dbias[i] = db_val; + alpha[i] = scale_x_rstd; + beta[i] = -scale_x_rstd * dscale_x_rstd * inv_nhw; + gamma[i] = scale_x_rstd * (mean[i] * dscale_x_rstd - db_val) * inv_nhw; +#endif + } + __syncthreads(); + } +} + +template +__global__ void ComputeXGradientCUDAKernel( + const int size, + const int C, + const int HxW, + const T* dY, + const T* X, + const T* alpha, + const T* beta, + const T* gamma, + T* dX) { + CUDA_1D_KERNEL_LOOP(i, size) { + const int c = kOrder == StorageOrder::NCHW ? i / HxW % C : i % C; +#if __CUDA_ARCH__ >= 350 + dX[i] = __ldg(alpha + c) * __ldg(dY + i) + __ldg(beta + c) * __ldg(X + i) + + __ldg(gamma + c); +#else + dX[i] = alpha[c] * dY[i] + beta[c] * X[i] + gamma[c]; +#endif + } +} + +} // namespace + +template <> +template +void SpatialBNOp::ComputeFusedParam( + const int C, + const T* scale, + const T* bias, + const T* mean, + const T* var, + T* alpha, + T* beta) { + ComputeFusedParamCUDAKernel + <<>>( + C, static_cast(epsilon_), scale, bias, mean, var, alpha, beta); +} + +template <> +template +void SpatialBNOp::ComputeBatchMoments( + const int N, + const int C, + const int HxW, + const T* batch_mean_sum, + const T* batch_var_sum, + T* mean, + T* var) { + const T scale = T(1) / static_cast(num_batches_ * N * HxW); + ComputeBatchMomentsCUDAKernel + <<>>( + C, scale, batch_mean_sum, batch_var_sum, mean, var); +} + +template <> +template +void SpatialBNOp::ComputeRunningMomentsAndFusedParam( + const int C, + const T* scale, + const T* bias, + const T* mean, + const T* var, + T* running_mean, + T* running_var, + T* rstd, + T* alpha, + T* beta) { + ComputeRunningMomentsAndFusedParamCUDAKernel + <<>>( + C, + static_cast(momentum_), + static_cast(epsilon_), + scale, + bias, + mean, + var, + running_mean, + running_var, + rstd, + alpha, + beta); +} + +template <> +template +void SpatialBNGradientOp:: + ComputeMultiBatchScaleBiasGradientsAndFusedParams( + const int N, + const int C, + const int HxW, + const T* scale, + const T* mean, + const T* rstd, + const T* dscale_sum, + const T* dbias_sum, + T* dscale, + T* dbias, + T* alpha, + T* beta, + T* gamma) { + const T inv_num_batches = T(1) / static_cast(num_batches_); + const T inv_nhw = T(1) / static_cast(N * HxW); + ComputeMultiBatchScaleBiasGradientsAndFusedParamsCUDAKernel + <<>>( + C, + inv_num_batches, + inv_nhw, + scale, + mean, + rstd, + dscale_sum, + dbias_sum, + dscale, + dbias, + alpha, + beta, + gamma); +} + +template <> +template +void SpatialBNGradientOp::ComputeScaleBiasGradientsAndFusedParams( + const int N, + const int C, + const int HxW, + const T* dY, + const T* X, + const T* scale, + const T* mean, + const T* rstd, + T* dscale, + T* dbias, + T* alpha, + T* beta, + T* gamma) { + if (order_ == StorageOrder::NCHW) { + ComputeScaleBiasGradientsAndFusedParamsCUDAKernel + <<>>( + N, + C, + HxW, + dY, + X, + scale, + mean, + rstd, + dscale, + dbias, + alpha, + beta, + gamma); + } else { + ComputeScaleBiasGradientsAndFusedParamsCUDAKernel + <<>>( + N, + C, + HxW, + dY, + X, + scale, + mean, + rstd, + dscale, + dbias, + alpha, + beta, + gamma); + } +} + +template <> +template +void SpatialBNGradientOp::ComputeXGradient( + const int N, + const int C, + const int HxW, + const T* dY, + const T* X, + const T* alpha, + const T* beta, + const T* gamma, + T* dX) { + const int size = N * C * HxW; + if (order_ == StorageOrder::NCHW) { + ComputeXGradientCUDAKernel + <<>>( + size, C, HxW, dY, X, alpha, beta, gamma, dX); + } else { + ComputeXGradientCUDAKernel + <<>>( + size, C, HxW, dY, X, alpha, beta, gamma, dX); + } +} + +} // namespace caffe2 diff --git a/caffe2/operators/utility_ops.cc b/caffe2/operators/utility_ops.cc index eb771974fbf39..6c287c37f3d6c 100644 --- a/caffe2/operators/utility_ops.cc +++ b/caffe2/operators/utility_ops.cc @@ -14,6 +14,31 @@ bool WeightedSumGradientOp::RunOnDevice() { return DoRunWithType(); } +std::vector WeightedSumShapeInference( + const OperatorDef& /* unused */, + const vector& in) { + vector out(1); + out[0] = in[0]; + return out; +} + +OpSchema::Cost CostInferenceForWeightedSum( + const OperatorDef& /* unused */, + const vector& in) { + CAFFE_ENFORCE_EQ( + in.size() % 2, 0, "WeightedSum requires an even number of inputs"); + struct OpSchema::Cost c; + + const auto& X0 = in[0]; + const auto& nElem = nElemFromDim(X0); + const auto& nInputs = in.size(); + c.flops = (nInputs - 1) * nElem; + c.bytes_read = (nInputs / 2) * (nElem + 1) * sizeof(X0.data_type()); + c.bytes_written = nElem * sizeof(X0.data_type()); + c.params_bytes = (nInputs / 2) * sizeof(X0.data_type()); + return c; +} + REGISTER_CPU_OPERATOR(WallClockTime, WallClockTimeOp); REGISTER_CPU_OPERATOR(Print, PrintOp); REGISTER_CPU_OPERATOR(FlattenToVec, FlattenToVecOp); @@ -261,6 +286,8 @@ OPERATOR_SCHEMA(SumInt) OPERATOR_SCHEMA(WeightedSum) .NumInputs([](int n) { return (n > 0 && n % 2 == 0); }) .NumOutputs(1) + .TensorInferenceFunction(WeightedSumShapeInference) + .CostInferenceFunction(CostInferenceForWeightedSum) .AllowInplace({{0, 0}}) .IdenticalTypeAndShapeOfInput(0) .SetDoc(R"DOC( diff --git a/caffe2/operators/utility_ops_gpu_test.cc b/caffe2/operators/utility_ops_gpu_test.cc index fb8baba549768..eb70a09aefb45 100644 --- a/caffe2/operators/utility_ops_gpu_test.cc +++ b/caffe2/operators/utility_ops_gpu_test.cc @@ -16,7 +16,7 @@ static void AddConstInput( const string& name, Workspace* ws) { DeviceOption option; - option.set_device_type(CUDA); + option.set_device_type(PROTO_CUDA); CUDAContext context(option); Blob* blob = ws->CreateBlob(name); auto* tensor = blob->GetMutableTensor(CUDA); @@ -37,7 +37,7 @@ TEST(UtilityOpGPUTest, testReshapeWithScalar) { def.add_output("XNew"); def.add_output("OldShape"); def.add_arg()->CopyFrom(MakeArgument("shape", vector{1})); - def.mutable_device_option()->set_device_type(CUDA); + def.mutable_device_option()->set_device_type(PROTO_CUDA); AddConstInput(vector(), 3.14, "X", &ws); // execute the op unique_ptr op(CreateOperator(def, &ws)); diff --git a/caffe2/opt/backend_cutting_test.cc b/caffe2/opt/backend_cutting_test.cc index 14cd435a77299..6335933d9c881 100644 --- a/caffe2/opt/backend_cutting_test.cc +++ b/caffe2/opt/backend_cutting_test.cc @@ -83,9 +83,9 @@ TEST(BackendCuttingTest, line) { EXPECT_EQ(3, net_opt.op_size()); } -// X0 -> CopyIn -> MyConv -\ +// X0 -> CopyIn -> MyConv -| // > Concat -> CopyOut -> Y -// N2 -> MyConv -> MyRelu -/ +// N2 -> MyConv -> MyRelu -| TEST(BackendCuttingTest, convergedPaths) { caffe2::NetDef net; net.add_external_input("X0"); @@ -118,8 +118,8 @@ TEST(BackendCuttingTest, convergedPaths) { }; // -> Random -> Relu -> MyConv4 -// / \ -// N0 -> MyConv -> MyRelu -> MyConv2 ---------- > Concat -> CopyOut -> Y +// | | +// N0 -> MyConv -> MyRelu -> MyConv2 ----------> Concat -> CopyOut -> Y TEST(BackendCuttingTest, skipPath) { caffe2::NetDef net; net.add_external_input("N0"); diff --git a/caffe2/opt/converter.h b/caffe2/opt/converter.h index d662967991463..31281ab90572e 100644 --- a/caffe2/opt/converter.h +++ b/caffe2/opt/converter.h @@ -54,7 +54,7 @@ class Caffe2Annotation : public nom::repr::Annotation { std::string Device = ""; caffe2::OperatorDef OpDef; bool OpDefExists = false; - int DeviceType = caffe2::DeviceType::CPU; + int DeviceType = caffe2::DeviceTypeProto::PROTO_CPU; }; CAFFE2_API nom::repr::NNModule convertToNNModule(caffe2::NetDef &net, std::unordered_map* blobMapOut = nullptr); diff --git a/caffe2/opt/dead_code_elim.cc b/caffe2/opt/dead_code_elim.cc new file mode 100644 index 0000000000000..5e7359329bfd7 --- /dev/null +++ b/caffe2/opt/dead_code_elim.cc @@ -0,0 +1,43 @@ +#include "caffe2/core/logging.h" +#include "caffe2/opt/converter.h" +#include "caffe2/opt/passes.h" + +namespace caffe2 { +namespace opt { + +using namespace nom; +using namespace nom::repr; + +void deadCodeElim(NNModule* nn) { + // Iteratively remove unconsumed non-external outputs. + bool changed = false; + do { + changed = false; + for (const auto& node : nn->dataFlow.getMutableNodes()) { + NOM_REQUIRE_OR_CONT(nn::is(node)); + + bool isUsed = false; + for (const auto& output : nn::getOutputs(node)) { + if (nn::hasConsumer(output) || nn->outputs.count(output)) { + isUsed = true; + break; + } + } + + NOM_REQUIRE_OR_CONT(!isUsed); + + // No outputs are used, delete them and the node itself. + for (const auto& output : nn::getOutputs(node)) { + nn->dataFlow.deleteNode(output); + } + nn->dataFlow.deleteNode(node); + changed = true; + break; + } + } while (changed); +} + +REGISTER_OPT_PASS_FROM_FUNC(DeadCodeElim, deadCodeElim); + +} // namespace opt +} // namespace caffe2 diff --git a/caffe2/opt/dead_code_elim_test.cc b/caffe2/opt/dead_code_elim_test.cc new file mode 100644 index 0000000000000..86b969cad5db5 --- /dev/null +++ b/caffe2/opt/dead_code_elim_test.cc @@ -0,0 +1,57 @@ +#include "caffe2/core/common.h" +#include "caffe2/opt/converter.h" +#include "caffe2/opt/passes.h" + +#include + +TEST(DeadCodeElim, BasicElim) { + caffe2::NetDef net; + { + caffe2::OperatorDef* def = net.add_op(); + def->set_type("Fake"); + def->add_input("X"); + def->add_output("Y"); + } + + auto nn = caffe2::convertToNNModule(net); + auto pass = caffe2::OptimizationPassRegistry()->Create("DeadCodeElim", &nn); + pass->run(); + auto optimized_net = caffe2::convertToCaffe2Proto(nn, net); + EXPECT_EQ(optimized_net.op().size(), 0); +} + +TEST(DeadCodeElim, BasicNoElim) { + caffe2::NetDef net; + { + caffe2::OperatorDef* def = net.add_op(); + def->set_type("Fake"); + def->add_input("X"); + def->add_output("Y"); + } + net.add_external_output("Y"); + + auto nn = caffe2::convertToNNModule(net); + auto pass = caffe2::OptimizationPassRegistry()->Create("DeadCodeElim", &nn); + pass->run(); + auto optimized_net = caffe2::convertToCaffe2Proto(nn, net); + EXPECT_EQ(optimized_net.op().size(), 1); +} + +TEST(DeadCodeElim, PartiallyUsedNoElim) { + caffe2::NetDef net; + { + caffe2::OperatorDef* def = net.add_op(); + def->set_type("Fake"); + def->add_input("X"); + def->add_output("Y"); + def->add_output("Z"); + } + net.add_external_output("Y"); + // Z is unused, but we should keep Fake because Y is + + auto nn = caffe2::convertToNNModule(net); + auto pass = caffe2::OptimizationPassRegistry()->Create("DeadCodeElim", &nn); + pass->run(); + auto optimized_net = caffe2::convertToCaffe2Proto(nn, net); + EXPECT_EQ(optimized_net.op().size(), 1); +} diff --git a/caffe2/opt/device_test.cc b/caffe2/opt/device_test.cc index d7c99131084a3..725516d8883b8 100644 --- a/caffe2/opt/device_test.cc +++ b/caffe2/opt/device_test.cc @@ -27,13 +27,13 @@ TEST(DeviceTest, InsertCopies) { ADD_ARG(def, "pad", i, 0); ADD_ARG(def, "order", s, "NCHW"); def->add_output("X"); - def->mutable_device_option()->set_device_type(caffe2::CPU); + def->mutable_device_option()->set_device_type(caffe2::PROTO_CPU); } else { caffe2::OperatorDef* def = net.add_op(); def->set_type("Relu"); def->add_input("X"); def->add_output("X"); - def->mutable_device_option()->set_device_type(caffe2::CPU); + def->mutable_device_option()->set_device_type(caffe2::PROTO_CPU); } } auto nn = caffe2::convertToNNModule(net); @@ -42,7 +42,7 @@ TEST(DeviceTest, InsertCopies) { if (nn::is(node)) { auto annot = nn::get(node)->getMutableAnnotation(); auto c2_annot = dyn_cast(annot); - c2_annot->setDeviceType(caffe2::OPENCL); + c2_annot->setDeviceType(caffe2::PROTO_OPENCL); } } @@ -57,7 +57,7 @@ TEST(DeviceTest, InsertCopies) { NOM_REQUIRE_OR_RET_FALSE(annot); auto c2_annot = dyn_cast(annot); NOM_REQUIRE_OR_RET_FALSE(c2_annot); - return c2_annot->getDeviceType() == caffe2::OPENCL; + return c2_annot->getDeviceType() == caffe2::PROTO_OPENCL; }, [](NNGraph& g) { return g.createNode(nom::util::make_unique()); diff --git a/caffe2/opt/fusion.h b/caffe2/opt/fusion.h index 8f9652d0d99e1..33dc2e4c54b1a 100644 --- a/caffe2/opt/fusion.h +++ b/caffe2/opt/fusion.h @@ -37,7 +37,7 @@ CAFFE2_API void fuseConvBN(repr::NNModule* nn, caffe2::Workspace* ws); // \param postprocess Functor to postprocess the conv node, // attaching additional attributes if necessary template -void fuseActivation( +CAFFE2_EXPORT void fuseActivation( repr::NNModule* nn, std::function should_fuse, std::function postprocess) { diff --git a/caffe2/opt/onnx_convert.h b/caffe2/opt/onnx_convert.h index 809f78c8d6bf1..707d41321f77e 100644 --- a/caffe2/opt/onnx_convert.h +++ b/caffe2/opt/onnx_convert.h @@ -1,4 +1,6 @@ -class OnnxAnnotation : public nom::repr::Annotation { +#include "caffe2/core/common.h" + +class CAFFE2_API OnnxAnnotation : public nom::repr::Annotation { public: OnnxAnnotation() : Annotation(AnnotationKind::Onnx) {} OnnxAnnotation(std::string device) diff --git a/caffe2/opt/optimize_ideep.cc b/caffe2/opt/optimize_ideep.cc index 6edd8de5d3570..33c350d8ef5cf 100644 --- a/caffe2/opt/optimize_ideep.cc +++ b/caffe2/opt/optimize_ideep.cc @@ -56,7 +56,7 @@ caffe2::OperatorDef* getMutableOpDef(repr::NeuralNetOperator& nnOp) { bool isOnIdeepDevice(const repr::NeuralNetOperator& nnOp) { // We only want to fuse for IDEEP convs const auto& op = getOpDef(nnOp); - return op.device_option().device_type() == DeviceType::IDEEP; + return op.device_option().device_type() == DeviceTypeProto::PROTO_IDEEP; } bool shouldFuseConv(const repr::Conv& conv) { diff --git a/caffe2/proto/caffe2.proto b/caffe2/proto/caffe2.proto index b561b43373765..8812b081a21b0 100644 --- a/caffe2/proto/caffe2.proto +++ b/caffe2/proto/caffe2.proto @@ -111,17 +111,17 @@ message Argument { // Note: if you add a device type, make sure you add the corresponding device // line in the DeviceTypeName() function in caffe2/utils/proto_utils.cc // and update ATen/core/DeviceType.h -enum DeviceType { - CPU = 0; // In default, we will use CPU. - CUDA = 1; // CUDA. - MKLDNN = 2; // Reserved for explicit MKLDNN - OPENGL = 3; // OpenGL - OPENCL = 4; // OpenCL - IDEEP = 5; // IDEEP. - HIP = 6; // AMD HIP +enum DeviceTypeProto { + PROTO_CPU = 0; // In default, we will use CPU. + PROTO_CUDA = 1; // CUDA. + PROTO_MKLDNN = 2; // Reserved for explicit MKLDNN + PROTO_OPENGL = 3; // OpenGL + PROTO_OPENCL = 4; // OpenCL + PROTO_IDEEP = 5; // IDEEP. + PROTO_HIP = 6; // AMD HIP // Change the following number if you add more devices in the code. - COMPILE_TIME_MAX_DEVICE_TYPES = 7; - ONLY_FOR_TEST = 20901701; // This device type is only for test. + PROTO_COMPILE_TIME_MAX_DEVICE_TYPES = 7; + PROTO_ONLY_FOR_TEST = 20901701; // This device type is only for test. } // Device-specific options. We do not distinguish DeviceOption protos for diff --git a/caffe2/proto/caffe2_pb.h b/caffe2/proto/caffe2_pb.h index 80f3825f619ba..291cb6bfe82b2 100644 --- a/caffe2/proto/caffe2_pb.h +++ b/caffe2/proto/caffe2_pb.h @@ -1,2 +1,86 @@ #pragma once +#include +#include #include + +namespace caffe2 { + +using DeviceType = at::DeviceType; +constexpr DeviceType CPU = DeviceType::CPU; +constexpr DeviceType CUDA = DeviceType::CUDA; +constexpr DeviceType OPENGL = DeviceType::OPENGL; +constexpr DeviceType OPENCL = DeviceType::OPENCL; +constexpr DeviceType MKLDNN = DeviceType::MKLDNN; +constexpr DeviceType IDEEP = DeviceType::IDEEP; +constexpr DeviceType HIP = DeviceType::HIP; +constexpr DeviceType COMPILE_TIME_MAX_DEVICE_TYPES = + DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES; +constexpr DeviceType ONLY_FOR_TEST = DeviceType::ONLY_FOR_TEST; + +inline CAFFE2_API DeviceType ProtoToType(const caffe2::DeviceTypeProto p) { + switch (p) { + case caffe2::PROTO_CPU: + return DeviceType::CPU; + case caffe2::PROTO_CUDA: + return DeviceType::CUDA; + case caffe2::PROTO_OPENGL: + return DeviceType::OPENGL; + case caffe2::PROTO_OPENCL: + return DeviceType::OPENCL; + case caffe2::PROTO_MKLDNN: + return DeviceType::MKLDNN; + case caffe2::PROTO_IDEEP: + return DeviceType::IDEEP; + case caffe2::PROTO_HIP: + return DeviceType::HIP; + case caffe2::PROTO_COMPILE_TIME_MAX_DEVICE_TYPES: + return DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES; + case caffe2::PROTO_ONLY_FOR_TEST: + return DeviceType::ONLY_FOR_TEST; + default: + AT_ERROR( + "Unknown device:", + static_cast(p), + ". If you have recently updated the caffe2.proto file to add a new " + "device type, did you forget to update the ProtoToType() and TypeToProto" + "function to reflect such recent changes?"); + // The below code won't run but is needed to suppress some compiler + // warnings. + return DeviceType::ONLY_FOR_TEST; + } +} + +inline CAFFE2_API DeviceTypeProto TypeToProto(const DeviceType& t) { + switch (t) { + case DeviceType::CPU: + return caffe2::PROTO_CPU; + case DeviceType::CUDA: + return caffe2::PROTO_CUDA; + case DeviceType::OPENGL: + return caffe2::PROTO_OPENGL; + case DeviceType::OPENCL: + return caffe2::PROTO_OPENCL; + case DeviceType::MKLDNN: + return caffe2::PROTO_MKLDNN; + case DeviceType::IDEEP: + return caffe2::PROTO_IDEEP; + case DeviceType::HIP: + return caffe2::PROTO_HIP; + case DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES: + return caffe2::PROTO_COMPILE_TIME_MAX_DEVICE_TYPES; + case DeviceType::ONLY_FOR_TEST: + return caffe2::PROTO_ONLY_FOR_TEST; + default: + AT_ERROR( + "Unknown device:", + static_cast(t), + ". If you have recently updated the caffe2.proto file to add a new " + "device type, did you forget to update the ProtoToType() and TypeToProto" + "function to reflect such recent changes?"); + // The below code won't run but is needed to suppress some compiler + // warnings. + return PROTO_ONLY_FOR_TEST; + } +} + +} // namespace caffe2 diff --git a/caffe2/python/__init__.py b/caffe2/python/__init__.py index e69de29bb2d1d..9401af575c9b7 100644 --- a/caffe2/python/__init__.py +++ b/caffe2/python/__init__.py @@ -0,0 +1,12 @@ +from __future__ import absolute_import, division, print_function, unicode_literals +from caffe2.proto import caffe2_pb2 +# TODO: refactor & remove the following alias +caffe2_pb2.CPU = caffe2_pb2.PROTO_CPU +caffe2_pb2.CUDA = caffe2_pb2.PROTO_CUDA +caffe2_pb2.MKLDNN = caffe2_pb2.PROTO_MKLDNN +caffe2_pb2.OPENGL = caffe2_pb2.PROTO_OPENGL +caffe2_pb2.OPENCL = caffe2_pb2.PROTO_OPENCL +caffe2_pb2.IDEEP = caffe2_pb2.PROTO_IDEEP +caffe2_pb2.HIP = caffe2_pb2.PROTO_HIP +caffe2_pb2.COMPILE_TIME_MAX_DEVICE_TYPES = caffe2_pb2.PROTO_COMPILE_TIME_MAX_DEVICE_TYPES +caffe2_pb2.ONLY_FOR_TEST = caffe2_pb2.PROTO_ONLY_FOR_TEST diff --git a/caffe2/python/data_parallel_model_test.py b/caffe2/python/data_parallel_model_test.py index b3e262272b1a5..91a08df74f165 100644 --- a/caffe2/python/data_parallel_model_test.py +++ b/caffe2/python/data_parallel_model_test.py @@ -447,6 +447,121 @@ def add_model_ops(model, loss_scale): self.assertTrue(transform.called) self.assertEqual(transform.call_count, 1) + @given(seed=st.integers(0, 65535), batch_size=st.integers(1, 20)) + def test_multi_device_spatial_bn_cpu(self, seed, batch_size): + self._spatial_bn_check("cpu", seed, batch_size) + + def _spatial_bn_check(self, device_type, seed, batch_size): + devices = [0, 1, 2] + epsilon = 1e-3 + tolerance = 1e-3 + + def _test_forward_pass(x, devices, device_type, scale, bias, epsilon): + x_concat = np.concatenate(x) + mean = np.mean(x_concat, axis=0) + var = np.var(x_concat, axis=0) + for device in devices: + x_i = x[device] + x_hat = (x_i - mean) / (np.sqrt(var + epsilon)) + expected_out = scale * x_hat + bias + spatial_out = workspace.FetchBlob( + "{}_{}/sp_out".format(device_type, device)) + rel_error = np.linalg.norm(spatial_out - expected_out) \ + / np.linalg.norm(expected_out) + self.assertTrue(rel_error < 0.005) + + def _test_backward_pass(x, devices, device_type, scale, tolerance): + dBias_arr = [] + dY_arr = [] + dGamma_arr = [] + num_devices = len(devices) + mean = np.array(workspace.FetchBlob( + "{}_0/sp_out_sm".format(device_type)), dtype=np.float32) + inv_var = np.array(workspace.FetchBlob( + "{}_0/sp_out_siv".format(device_type)), dtype=np.float32) + + # dBias + # Sum dBias values over all devices to find the average gradient + for device in devices: + dY_blob = workspace.FetchBlob( + "{}_{}/sp_out_grad".format(device_type, device)) + dY = np.array(dY_blob, dtype=np.float32) + dY_arr.append(dY) + dBias_arr.append(np.array(np.sum(dY, axis=0), dtype=np.float32)) + dBias = np.sum(dBias_arr, dtype=np.float32) + dBias_avg = dBias / num_devices + for device in devices: + dBiasActual = np.sum(workspace.FetchBlob("{}_{}/sp_out_b_grad" + .format(device_type, device)), dtype=np.float32) + self.assertTrue(np.isclose([dBiasActual], [dBias], atol=tolerance)) + + # dGamma + # Sum dGamma values over all devices to find the average gradient + for device in devices: + dGamma = np.sum((x[device] - mean) * inv_var * dY_arr[device], + axis=0, dtype=np.float32) + dGamma_arr.append(dGamma) + dGamma = np.sum(dGamma_arr, axis=0, dtype=np.float32) + dGamma_avg = dGamma / num_devices + for device in devices: + dGammaActual = workspace.FetchBlob( + "{}_{}/sp_out_s_grad".format(device_type, device)) + self.assertTrue(np.isclose([dGamma], [dGammaActual], atol=tolerance)) + + # dX + scale_inv_var = scale * inv_var / batch_size + for device in devices: + dX = scale_inv_var * (dY_arr[device] * batch_size - dBias_avg + - (x[device] - mean) * dGamma_avg * inv_var) + dX_actual = workspace.FetchBlob( + "{}_{}/tanh_grad".format(device_type, device)) + self.assertTrue(np.isclose([dX], [dX_actual], atol=tolerance).all()) + + def add_input_ops(model): + for device in devices: + data = np.random.rand(batch_size, 1, 1, 1).astype(np.float32) + workspace.FeedBlob("{}_{}/data".format(device_type, device), data) + + def add_model_ops(model, loss_scale): + model.Tanh("data", "tanh") + model.SpatialBN("tanh", "sp_out", 1, epsilon=epsilon, is_test=False) + model.Sqr("sp_out", "sqr") + loss = model.SumElements("sqr", "loss") + return [loss] + + def add_optimizer(model): + return optimizer.build_sgd(model, 0.1) + + np.random.seed(seed) + workspace.ResetWorkspace() + model = cnn.CNNModelHelper( + order="NCHW", + name="test" + ) + data_parallel_model.Parallelize( + model, + input_builder_fun=add_input_ops, + forward_pass_builder_fun=add_model_ops, + optimizer_builder_fun=add_optimizer, + devices=devices, + cpu_device=device_type == "cpu", + shared_model=False, + combine_spatial_bn=True, + ) + + workspace.RunNetOnce(model.param_init_net) + scale = workspace.FetchBlob("{}_0/sp_out_s".format(device_type)) + bias = workspace.FetchBlob("{}_0/sp_out_b".format(device_type)) + workspace.RunNetOnce(model.net) + + x = [] + for device in devices: + x_blob = workspace.FetchBlob("{}_{}/tanh".format(device_type, device)) + x_i = np.array(x_blob, dtype=np.float32) + x.append(x_i) + + _test_forward_pass(x, devices, device_type, scale, bias, epsilon) + _test_backward_pass(x, devices, device_type, scale, tolerance) class RecurrentNetworkParallelTest(TestCase): diff --git a/caffe2/python/data_workers.py b/caffe2/python/data_workers.py index 65866c3f71ff6..eb49da78c0afa 100644 --- a/caffe2/python/data_workers.py +++ b/caffe2/python/data_workers.py @@ -242,7 +242,7 @@ def put(self, chunk, data_input_coordinator): qsize = self._internal_queue.qsize() if qsize < 2 and (time.time() - self._last_warning) > LOG_INT_SECS: log.warning("Warning, data loading lagging behind: " + - "name={}".format(qsize, self._input_source_name)) + "queue size={}, name={}".format(qsize, self._input_source_name)) self._last_warning = time.time() self._counter += 1 self._internal_queue.put(chunk, block=True, timeout=0.5) diff --git a/caffe2/python/hypothesis_test.py b/caffe2/python/hypothesis_test.py index a84a747c282f5..64174f6e71c67 100644 --- a/caffe2/python/hypothesis_test.py +++ b/caffe2/python/hypothesis_test.py @@ -10,6 +10,7 @@ from hypothesis import assume, given, settings, HealthCheck import hypothesis.strategies as st import unittest +import os from caffe2.python import core, workspace, tt_core, dyndep import caffe2.python.hypothesis_test_util as hu @@ -193,6 +194,7 @@ def ref(x, y): _test_binary("Mul", ref, filter_=not_overflow, test_gradient=True)(self) _test_binary_broadcast("Mul", ref, filter_=not_overflow)(self) + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") def test_div(self): def ref(x, y): return (x / y, ) @@ -1823,6 +1825,7 @@ def ref(data): out, = self.assertReferenceChecks(gc, op, [a], ref) self.assertEqual(dst, out.dtype) + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(a=hu.tensor(), eps=st.floats(min_value=1e-4, max_value=1e-2), a_grad=hu.tensor(elements=st.floats(min_value=0.01, max_value=0.99)), diff --git a/caffe2/python/layers/concat.py b/caffe2/python/layers/concat.py index 126a421745139..6afc603f1a38d 100644 --- a/caffe2/python/layers/concat.py +++ b/caffe2/python/layers/concat.py @@ -11,10 +11,27 @@ ) from future.utils import viewitems import numpy as np +from collections import defaultdict import logging logger = logging.getLogger(__name__) + +def get_concatenated_feature_to_index(blobs_to_concat): + concat_feature_to_index = defaultdict(list) + start_pos = 0 + for scalar in blobs_to_concat: + num_dims = scalar.dtype.shape[0] + if hasattr(scalar, 'metadata') \ + and hasattr(scalar.metadata, 'feature_specs') \ + and hasattr(scalar.metadata.feature_specs, 'feature_to_index') \ + and isinstance(scalar.metadata.feature_specs.feature_to_index, dict): # noqa B950 + for k, v in scalar.metadata.feature_specs.feature_to_index.items(): + concat_feature_to_index[k].extend([start_pos + vi for vi in v]) + start_pos += num_dims + return dict(concat_feature_to_index) if concat_feature_to_index.keys() else None + + class Concat(ModelLayer): """ Construct Concat layer @@ -95,6 +112,19 @@ def __init__(self, model, input_record, axis=1, add_axis=0, (np.float32, output_dims), self.get_next_blob_reference('output')) + record_to_concat = input_record.fields.values() + concated_feature_to_index = get_concatenated_feature_to_index( + record_to_concat + ) + if concated_feature_to_index: + metadata = schema.Metadata( + feature_specs=schema.FeatureSpec( + feature_to_index=concated_feature_to_index + ) + ) + self.output_schema.set_metadata(metadata) + + def add_ops(self, net): net.Concat( self.input_record.field_blobs(), diff --git a/caffe2/python/model_helper.py b/caffe2/python/model_helper.py index 6dfc7dd3c4911..f8e3f32bb2c22 100644 --- a/caffe2/python/model_helper.py +++ b/caffe2/python/model_helper.py @@ -465,6 +465,9 @@ def GetCompleteNet(self): for op in new_net.Proto().op: op.debug_info = op.debug_info + "/param_init_net" new_net.AppendNet(self.net) + # keep the execution optimization + if self.net.Proto().HasField("type"): + new_net.Proto().type = self.net.Proto().type return new_net def ConstructInitTrainNetfromNet(self, net): diff --git a/caffe2/python/model_helper_test.py b/caffe2/python/model_helper_test.py index cff297e9f5d97..fcccddf401db6 100644 --- a/caffe2/python/model_helper_test.py +++ b/caffe2/python/model_helper_test.py @@ -8,6 +8,26 @@ class ModelHelperTest(unittest.TestCase): + def test_get_complete_net_type(self): + model = model_helper.ModelHelper("test_orig") + brew.conv( + model, + "input", + "conv", + dim_in=3, + dim_out=16, + weight_init=("MSRAFill", {}), + kernel=3, + stride=1, + pad=0, + ) + model.net.Proto().type = "async_scheduling" + net = model.GetCompleteNet() + model2 = model_helper.ModelHelper("test_new") + model2.ConstructInitTrainNetfromNet(net) + self.assertTrue(model2.net.Proto().type, "async_scheduling") + self.assertTrue(model2.param_init_net.Proto().type, "async_scheduling") + def test_get_complete_net(self): model = model_helper.ModelHelper("test_orig") conv = brew.conv( diff --git a/caffe2/python/modeling/gradient_clipping.py b/caffe2/python/modeling/gradient_clipping.py index 1c7fa5d291137..69fff42f9d200 100644 --- a/caffe2/python/modeling/gradient_clipping.py +++ b/caffe2/python/modeling/gradient_clipping.py @@ -26,7 +26,8 @@ class GradientClipping(NetModifier): def __init__(self, grad_clip_method, clip_norm_type='l2_norm', clip_threshold=0.1, use_parameter_norm=False, - compute_norm_ratio=False, clip_max=1, clip_min=-1): + compute_norm_ratio=False, clip_max=1, clip_min=-1, + blobs_to_include=None, blobs_to_exclude=None): """ Clips gradient to avoid gradient magnitude explosion or vanishing gradient. @@ -42,6 +43,9 @@ def __init__(self, grad_clip_method, clip_norm_type='l2_norm', clip_max will be clipped to clip_max clip_min: when clipping by_value, any value that is smaller than clip_min will be clipped to clip_min + blobs_to_include: names of blobs whose gradient is to be clipped. If it is set + to none, all param 's gradient in grad_map will be clipped. + blobs_to_exclude: names of blobs whose gradient is not to be clipped. """ assert grad_clip_method in self.GRAD_CLIP_METHODS, ( @@ -59,6 +63,8 @@ def __init__(self, grad_clip_method, clip_norm_type='l2_norm', self.compute_norm_ratio = compute_norm_ratio self.clip_max = float(clip_max) self.clip_min = float(clip_min) + self.blobs_to_include = blobs_to_include + self.blobs_to_exclude = blobs_to_exclude def modify_net(self, net, init_net=None, grad_map=None, blob_to_device=None, modify_output_record=False): @@ -67,8 +73,22 @@ def modify_net(self, net, init_net=None, grad_map=None, blob_to_device=None, CPU = core.DeviceOption(caffe2_pb2.CPU) - for param, grad in grad_map.items(): - + final_param_map = {} + if self.blobs_to_include is None: + final_param_map = grad_map + else: + for blob in self.blobs_to_include: + param = core.BlobReference(blob) + if not net.BlobIsDefined(param): + raise Exception('param {0} is not defined in net {1}'.format( + param, net.Name())) + final_param_map[param] = grad_map[param] + + if self.blobs_to_exclude is not None: + for blob in self.blobs_to_exclude: + final_param_map.pop(blob, None) + + for param, grad in final_param_map.items(): # currently sparse gradients won't be clipped # futher implementation is needed to enable it if isinstance(grad, core.GradientSlice): diff --git a/caffe2/python/modeling/gradient_clipping_test.py b/caffe2/python/modeling/gradient_clipping_test.py index e5fc2bab4e8f7..ca5c2ba8e22b6 100644 --- a/caffe2/python/modeling/gradient_clipping_test.py +++ b/caffe2/python/modeling/gradient_clipping_test.py @@ -199,3 +199,73 @@ def test_gradient_clipping_by_value(self): fc1_w_grad = workspace.FetchBlob('fc1_w_grad') self.assertLessEqual(np.amax(fc1_w_grad), clip_max) self.assertGreaterEqual(np.amin(fc1_w_grad), clip_min) + + def test_gradient_clipping_by_norm_including_blobs(self): + model = model_helper.ModelHelper(name="test") + data = model.net.AddExternalInput("data") + fc1 = brew.fc(model, data, "fc1", dim_in=4, dim_out=2) + + # no operator name set, will use default + fc2 = brew.fc(model, fc1, "fc2", dim_in=2, dim_out=1) + + sigm = model.net.Sigmoid(fc2, 'sigm') + sq = model.net.SquaredL2Distance([sigm, 'label'], 'sq') + loss = model.net.SumElements(sq, 'loss') + + grad_map = model.AddGradientOperators([loss]) + + grad_map_for_param = {key: grad_map[key] for key in ['fc1_w', 'fc2_w']} + + net_modifier = GradientClipping( + grad_clip_method='by_norm', + clip_norm_type='l2_norm', + clip_threshold=0.1, + blobs_to_include=['fc1_w'], + blobs_to_exclude=None + ) + + net_modifier(model.net, grad_map=grad_map_for_param) + + workspace.FeedBlob('data', np.random.rand(10, 4).astype(np.float32)) + workspace.FeedBlob('label', np.random.rand(10, 1).astype(np.float32)) + + workspace.RunNetOnce(model.param_init_net) + workspace.RunNetOnce(model.net) + + # 5 forward ops + 6 backward ops + 1 * (3 gradient clipping ops) + self.assertEqual(len(model.net.Proto().op), 14) + + def test_gradient_clipping_by_norm_excluding_blobs(self): + model = model_helper.ModelHelper(name="test") + data = model.net.AddExternalInput("data") + fc1 = brew.fc(model, data, "fc1", dim_in=4, dim_out=2) + + # no operator name set, will use default + fc2 = brew.fc(model, fc1, "fc2", dim_in=2, dim_out=1) + + sigm = model.net.Sigmoid(fc2, 'sigm') + sq = model.net.SquaredL2Distance([sigm, 'label'], 'sq') + loss = model.net.SumElements(sq, 'loss') + + grad_map = model.AddGradientOperators([loss]) + + grad_map_for_param = {key: grad_map[key] for key in ['fc1_w', 'fc2_w']} + + net_modifier = GradientClipping( + grad_clip_method='by_norm', + clip_norm_type='l2_norm', + clip_threshold=0.1, + blobs_to_include=None, + blobs_to_exclude=['fc1_w', 'fc2_w'] + ) + + net_modifier(model.net, grad_map=grad_map_for_param) + + workspace.FeedBlob('data', np.random.rand(10, 4).astype(np.float32)) + workspace.FeedBlob('label', np.random.rand(10, 1).astype(np.float32)) + + workspace.RunNetOnce(model.param_init_net) + workspace.RunNetOnce(model.net) + + # 5 forward ops + 6 backward ops + 0 * (3 gradient clipping ops) + self.assertEqual(len(model.net.Proto().op), 11) diff --git a/caffe2/python/nomnigraph.py b/caffe2/python/nomnigraph.py index 84142356fcb1f..708eae6b2a071 100644 --- a/caffe2/python/nomnigraph.py +++ b/caffe2/python/nomnigraph.py @@ -29,6 +29,14 @@ def __init__(self, net=None): def dataFlow(self): return self._NNModule.dataFlow() + def convertToCaffe2Proto(self, old_proto=None): + if not old_proto: + old_proto = caffe2_pb2.NetDef() + output = self._NNModule.convertToCaffe2Proto(old_proto) + new_proto = caffe2_pb2.NetDef() + new_proto.ParseFromString(output) + return new_proto + def match(self, pattern): for n in self.dataFlow.getMutableNodes(): m = C.matchSubgraph(n, pattern) diff --git a/caffe2/python/nomnigraph_test.py b/caffe2/python/nomnigraph_test.py index 0c9fd34d79a04..7739ac05f2979 100644 --- a/caffe2/python/nomnigraph_test.py +++ b/caffe2/python/nomnigraph_test.py @@ -151,3 +151,26 @@ def test_genericGraph(self): n2 = g.createNode("hello2") e = g.createEdge(n1, n2) ng.render(g) + + def test_convertToProto(self): + net = core.Net("name") + net.FC(["X", "W"], ["Y"]) + nn = ng.NNModule(net) + new_netdef = nn.convertToCaffe2Proto() + print(new_netdef) + print(net.Proto()) + assert len(new_netdef.op) == len(net.Proto().op) + for i in range(len(new_netdef.op)): + op = net.Proto().op[i] + new_op = new_netdef.op[i] + assert op.type == new_op.type + assert len(op.input) == len(new_op.input) + assert len(op.output) == len(new_op.output) + for a, b in zip(op.input, new_op.input): + assert a == b + for a, b in zip(op.output, new_op.output): + assert a == b + for a, b in zip(new_netdef.external_input, net.Proto().external_input): + assert a == b + for a, b in zip(new_netdef.external_output, net.Proto().external_output): + assert a == b diff --git a/caffe2/python/onnx/tests/c2_ref_test.py b/caffe2/python/onnx/tests/c2_ref_test.py index 8ff58a68ce108..3e297416770ba 100644 --- a/caffe2/python/onnx/tests/c2_ref_test.py +++ b/caffe2/python/onnx/tests/c2_ref_test.py @@ -648,6 +648,7 @@ def test_bvlc_googlenet(self): def test_bvlc_reference_caffenet(self): self._test_net('bvlc_reference_caffenet') + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") def test_bvlc_reference_rcnn_ilsvrc13(self): self._test_net('bvlc_reference_rcnn_ilsvrc13') diff --git a/caffe2/python/onnx/tests/onnx_backend_test.py b/caffe2/python/onnx/tests/onnx_backend_test.py index e04efb65f69d5..b3607c1426cfe 100644 --- a/caffe2/python/onnx/tests/onnx_backend_test.py +++ b/caffe2/python/onnx/tests/onnx_backend_test.py @@ -37,6 +37,7 @@ '|test_maxpool_with_argmax.*' # MaxPool outputs indices in different format. '|test_convtranspose.*' # ConvTranspose needs some more complicated translation '|test_mvn.*' # MeanVarianceNormalization is experimental and not supported. + '|test_dynamic_slice.*' # MeanVarianceNormalization is experimental and not supported. ')') # Quick patch to unbreak master CI, is working on the debugging. diff --git a/caffe2/python/operator_test/adagrad_test.py b/caffe2/python/operator_test/adagrad_test.py index c3c58d0e6a0b2..e4101e92cf01c 100644 --- a/caffe2/python/operator_test/adagrad_test.py +++ b/caffe2/python/operator_test/adagrad_test.py @@ -16,6 +16,8 @@ ref_adagrad, adagrad_sparse_test_helper ) +import unittest +import os class TestAdagrad(hu.HypothesisTestCase): @staticmethod @@ -158,6 +160,7 @@ def ref_sparse(param, momentum, indices, grad, lr): gc, op, [param_i, momentum_i, indices, grad, lr], ref_sparse ) + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") # Suppress filter_too_much health check. # Likely caused by `assume` call falling through too often. @settings(suppress_health_check=[HealthCheck.filter_too_much]) diff --git a/caffe2/python/operator_test/collect_and_distribute_fpn_rpn_proposals_op_test.py b/caffe2/python/operator_test/collect_and_distribute_fpn_rpn_proposals_op_test.py index b2c05a58d99ad..facb675a4944e 100644 --- a/caffe2/python/operator_test/collect_and_distribute_fpn_rpn_proposals_op_test.py +++ b/caffe2/python/operator_test/collect_and_distribute_fpn_rpn_proposals_op_test.py @@ -5,6 +5,7 @@ import numpy as np import unittest +import os from hypothesis import given, settings import hypothesis.strategies as st @@ -129,6 +130,7 @@ def collect_and_distribute_fpn_rpn_ref(*inputs): class TestCollectAndDistributeFpnRpnProposals(hu.HypothesisTestCase): + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(proposal_count=st.integers(min_value=1000, max_value=8000), rpn_min_level=st.integers(min_value=1, max_value=4), rpn_num_levels=st.integers(min_value=1, max_value=6), diff --git a/caffe2/python/operator_test/conv_test.py b/caffe2/python/operator_test/conv_test.py index d41add5848a18..8e65a9324535a 100644 --- a/caffe2/python/operator_test/conv_test.py +++ b/caffe2/python/operator_test/conv_test.py @@ -15,6 +15,8 @@ from caffe2.python.model_helper import ModelHelper import caffe2.python._import_c_extension as C +import unittest +import os def _cudnn_supports( dilation=False, @@ -430,6 +432,7 @@ def test_3d_convolution_cudnn_nchw(self, op_type, batch_size, stride, size, or "CUDNN_STATUS_NOT_SUPPORTED" not in es: raise e + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(op_type=st.sampled_from(["Conv", "Conv2D"]), stride=st.integers(1, 3), pad=st.integers(0, 3), diff --git a/caffe2/python/operator_test/deform_conv_test.py b/caffe2/python/operator_test/deform_conv_test.py index 03bea2c3f91fa..583e1db1679f8 100644 --- a/caffe2/python/operator_test/deform_conv_test.py +++ b/caffe2/python/operator_test/deform_conv_test.py @@ -9,6 +9,8 @@ from caffe2.proto import caffe2_pb2 from caffe2.python import core, workspace import caffe2.python.hypothesis_test_util as hu +import unittest +import os import unittest @@ -375,6 +377,7 @@ def reference_conv_op(*args): # CUDNN does NOT support different padding values and we skip it @unittest.skipIf(not workspace.has_gpu_support, "No gpu support") + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(stride_h=st.integers(1, 3), stride_w=st.integers(1, 3), pad_h=st.integers(0, 3), diff --git a/caffe2/python/operator_test/elementwise_ops_test.py b/caffe2/python/operator_test/elementwise_ops_test.py index c20aad4218f17..a3ae326303474 100644 --- a/caffe2/python/operator_test/elementwise_ops_test.py +++ b/caffe2/python/operator_test/elementwise_ops_test.py @@ -9,6 +9,8 @@ import hypothesis.strategies as st import numpy as np +import unittest +import os class TestElementwiseOps(hu.HypothesisTestCase): @@ -131,6 +133,7 @@ def sqr_op(X): self.assertGradientChecks( gc, op, [X], 0, [0], stepsize=1e-4, threshold=1e-2) + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given( X=hu.tensor( elements=st.floats(0.1, 10), @@ -333,6 +336,7 @@ def sigmoid_ref(X): self.assertDeviceChecks(dc, op, [X], [0]) self.assertGradientChecks(gc, op, [X], 0, [0]) + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(X=hu.tensor(dtype=np.float32), inplace=st.booleans(), alpha=st.floats(min_value=-100.0, max_value=100.0), diff --git a/caffe2/python/operator_test/group_conv_test.py b/caffe2/python/operator_test/group_conv_test.py index aac50d01713f0..fd4c5adf0075d 100644 --- a/caffe2/python/operator_test/group_conv_test.py +++ b/caffe2/python/operator_test/group_conv_test.py @@ -11,10 +11,12 @@ import caffe2.python.hypothesis_test_util as hu import unittest +import os class TestGroupConvolution(hu.HypothesisTestCase): + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(stride=st.integers(1, 3), pad=st.integers(0, 3), kernel=st.integers(1, 5), diff --git a/caffe2/python/operator_test/gru_test.py b/caffe2/python/operator_test/gru_test.py index d807b66e275f9..1292d843d0a82 100644 --- a/caffe2/python/operator_test/gru_test.py +++ b/caffe2/python/operator_test/gru_test.py @@ -15,6 +15,7 @@ import hypothesis.strategies as st import numpy as np import unittest +import os def gru_unit(*args, **kwargs): @@ -248,6 +249,7 @@ def generate_input_state(n, d): class GRUCellTest(hu.HypothesisTestCase): # Test just for GRUUnitOp + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given( seed=st.integers(0, 2**32 - 1), input_tensor=gru_unit_op_input(), diff --git a/caffe2/python/operator_test/instance_norm_test.py b/caffe2/python/operator_test/instance_norm_test.py index 4377b82be2f9f..a91154a4e45f6 100644 --- a/caffe2/python/operator_test/instance_norm_test.py +++ b/caffe2/python/operator_test/instance_norm_test.py @@ -9,6 +9,8 @@ from caffe2.python import core, model_helper, brew import caffe2.python.hypothesis_test_util as hu +import unittest +import os class TestInstanceNorm(hu.HypothesisTestCase): @@ -48,6 +50,7 @@ def _feed_inputs(self, input_blobs, device_option): for name, blob in zip(names, input_blobs): self.ws.create_blob(name).feed(blob, device_option=device_option) + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(gc=hu.gcs['gc'], dc=hu.gcs['dc'], N=st.integers(2, 3), diff --git a/caffe2/python/operator_test/layer_norm_op_test.py b/caffe2/python/operator_test/layer_norm_op_test.py index e8bdfc35510df..fa0958afd1f99 100644 --- a/caffe2/python/operator_test/layer_norm_op_test.py +++ b/caffe2/python/operator_test/layer_norm_op_test.py @@ -6,11 +6,14 @@ from caffe2.python import brew, core from hypothesis import given import caffe2.python.hypothesis_test_util as hu +import unittest +import os import numpy as np from caffe2.python.model_helper import ModelHelper class TestLayerNormOp(hu.HypothesisTestCase): + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(X=hu.tensors(n=1), **hu.gcs) def test_layer_norm_grad_op(self, X, gc, dc): X = X[0] @@ -82,6 +85,7 @@ def layer_norm_grad_ref(gout_full, norm, mean_full, stdev_full, X_full): outputs_to_check=[0], ) + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(X=hu.tensors(n=1), **hu.gcs) def test_layer_norm_op(self, X, gc, dc): X = X[0] diff --git a/caffe2/python/operator_test/rand_quantization_op_test.py b/caffe2/python/operator_test/rand_quantization_op_test.py index dda3a28fdf110..322b708daa46f 100644 --- a/caffe2/python/operator_test/rand_quantization_op_test.py +++ b/caffe2/python/operator_test/rand_quantization_op_test.py @@ -5,6 +5,8 @@ import numpy as np import struct +import unittest +import os from hypothesis import given, example import hypothesis.strategies as st @@ -15,6 +17,7 @@ np.set_printoptions(precision=6) class TestFloatToFusedRandRowwiseQuantized(hu.HypothesisTestCase): + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(X=hu.tensor(min_dim=2, max_dim=2, min_value=1, max_value=17), # only matrix is supported bitwidth_=st.sampled_from([1, 2, 4, 8]), diff --git a/caffe2/python/operator_test/recurrent_network_test.py b/caffe2/python/operator_test/recurrent_network_test.py index afb9e1ebf86fa..e37bd9092e8f8 100644 --- a/caffe2/python/operator_test/recurrent_network_test.py +++ b/caffe2/python/operator_test/recurrent_network_test.py @@ -10,8 +10,11 @@ import hypothesis.strategies as st import numpy as np +import unittest +import os class RecurrentNetworkTest(hu.HypothesisTestCase): + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(T=st.integers(1, 4), n=st.integers(1, 5), d=st.integers(1, 5)) diff --git a/caffe2/python/operator_test/sequence_ops_test.py b/caffe2/python/operator_test/sequence_ops_test.py index 6c35ce5689d34..28eb9278da1d3 100644 --- a/caffe2/python/operator_test/sequence_ops_test.py +++ b/caffe2/python/operator_test/sequence_ops_test.py @@ -8,6 +8,7 @@ import hypothesis.strategies as st import numpy as np import unittest +import os from functools import partial @@ -185,6 +186,7 @@ def test_remove_padding(self, start_pad_width, end_pad_width, args, gc, dc): inputs=[data, lengths], reference=partial(_remove_padding_ref, start_pad_width, end_pad_width)) + @unittest.skipIf("IN_CIRCLECI" in os.environ, "FIXME: flaky test in CircleCI") @given(start_pad_width=st.integers(min_value=0, max_value=2), end_pad_width=st.integers(min_value=0, max_value=2), args=_gen_test_add_padding(with_pad_data=True), diff --git a/caffe2/python/operator_test/spatial_bn_op_test.py b/caffe2/python/operator_test/spatial_bn_op_test.py index 7124e3bcf2307..bf87f06ef7fe7 100644 --- a/caffe2/python/operator_test/spatial_bn_op_test.py +++ b/caffe2/python/operator_test/spatial_bn_op_test.py @@ -16,7 +16,8 @@ def _run_in_hip(gc, dc): - return gc.device_type == caffe2_pb2.HIP or caffe2_pb2.HIP in {d.device_type for d in dc} + return (gc.device_type == caffe2_pb2.HIP) or ( + caffe2_pb2.HIP in {d.device_type for d in dc}) class TestSpatialBN(hu.HypothesisTestCase): @@ -27,12 +28,13 @@ class TestSpatialBN(hu.HypothesisTestCase): seed=st.integers(0, 65535), order=st.sampled_from(["NCHW", "NHWC"]), epsilon=st.floats(min_value=1e-5, max_value=1e-2), - inplace=st.sampled_from([True, False]), + inplace=st.booleans(), + engine=st.sampled_from(["", "CUDNN"]), # Currently HIP SpatialBN only supports 2D **hu.gcs_no_hip) def test_spatialbn_test_mode_3d( self, size, input_channels, batch_size, seed, order, epsilon, - inplace, gc, dc): + inplace, engine, gc, dc): op = core.CreateOperator( "SpatialBN", ["X", "scale", "bias", "mean", "var"], @@ -40,12 +42,13 @@ def test_spatialbn_test_mode_3d( order=order, is_test=True, epsilon=epsilon, - engine="CUDNN", + engine=engine, ) def reference_spatialbn_test(X, scale, bias, mean, var): if order == "NCHW": - scale = scale[np.newaxis, :, np.newaxis, np.newaxis, np.newaxis] + scale = scale[np.newaxis, :, + np.newaxis, np.newaxis, np.newaxis] bias = bias[np.newaxis, :, np.newaxis, np.newaxis, np.newaxis] mean = mean[np.newaxis, :, np.newaxis, np.newaxis, np.newaxis] var = var[np.newaxis, :, np.newaxis, np.newaxis, np.newaxis] @@ -58,7 +61,7 @@ def reference_spatialbn_test(X, scale, bias, mean, var): mean = np.random.randn(input_channels).astype(np.float32) var = np.random.rand(input_channels).astype(np.float32) + 0.5 X = np.random.rand(batch_size, input_channels, size, size, size)\ - .astype(np.float32) - 0.5 + .astype(np.float32) - 0.5 if order == "NHWC": X = X.transpose(0, 2, 3, 4, 1) @@ -66,19 +69,21 @@ def reference_spatialbn_test(X, scale, bias, mean, var): reference_spatialbn_test) self.assertDeviceChecks(dc, op, [X, scale, bias, mean, var], [0]) - @unittest.skipIf(not workspace.has_gpu_support and not workspace.has_hip_support, "No gpu support") + @unittest.skipIf((not workspace.has_gpu_support) and ( + not workspace.has_hip_support), "No gpu support") @given(size=st.integers(7, 10), input_channels=st.integers(1, 10), batch_size=st.integers(0, 3), seed=st.integers(0, 65535), order=st.sampled_from(["NCHW", "NHWC"]), epsilon=st.floats(min_value=1e-5, max_value=1e-2), - inplace=st.sampled_from([True, False]), + inplace=st.booleans(), + engine=st.sampled_from(["", "CUDNN"]), # Currently HIP SpatialBN only supports 2D **hu.gcs_no_hip) def test_spatialbn_test_mode_1d( self, size, input_channels, batch_size, seed, order, epsilon, - inplace, gc, dc): + inplace, engine, gc, dc): op = core.CreateOperator( "SpatialBN", ["X", "scale", "bias", "mean", "var"], @@ -86,7 +91,7 @@ def test_spatialbn_test_mode_1d( order=order, is_test=True, epsilon=epsilon, - engine="CUDNN", + engine=engine, ) def reference_spatialbn_test(X, scale, bias, mean, var): @@ -109,6 +114,7 @@ def reference_spatialbn_test(X, scale, bias, mean, var): X = X.swapaxes(1, 2) self.assertReferenceChecks(gc, op, [X, scale, bias, mean, var], reference_spatialbn_test) + self.assertDeviceChecks(dc, op, [X, scale, bias, mean, var], [0]) @given(size=st.integers(7, 10), input_channels=st.integers(1, 10), @@ -117,7 +123,7 @@ def reference_spatialbn_test(X, scale, bias, mean, var): order=st.sampled_from(["NCHW", "NHWC"]), epsilon=st.floats(min_value=1e-5, max_value=1e-2), engine=st.sampled_from(["", "CUDNN"]), - inplace=st.sampled_from([True, False]), + inplace=st.booleans(), **hu.gcs) def test_spatialbn_test_mode( self, size, input_channels, batch_size, seed, order, epsilon, @@ -165,12 +171,13 @@ def reference_spatialbn_test(X, scale, bias, mean, var): seed=st.integers(0, 65535), order=st.sampled_from(["NCHW", "NHWC"]), epsilon=st.floats(1e-5, 1e-2), + momentum=st.floats(0.5, 0.9), engine=st.sampled_from(["", "CUDNN"]), inplace=st.sampled_from([True, False]), **hu.gcs) def test_spatialbn_train_mode( self, size, input_channels, batch_size, seed, order, epsilon, - inplace, engine, gc, dc): + momentum, inplace, engine, gc, dc): # Currently HIP SpatialBN only supports NCHW if _run_in_hip(gc, dc) and (order != 'NCHW'): return @@ -183,6 +190,7 @@ def test_spatialbn_train_mode( order=order, is_test=False, epsilon=epsilon, + momentum=momentum, engine=engine, ) np.random.seed(1701) @@ -205,11 +213,12 @@ def test_spatialbn_train_mode( seed=st.integers(0, 65535), order=st.sampled_from(["NCHW", "NHWC"]), epsilon=st.floats(min_value=1e-5, max_value=1e-2), + momentum=st.floats(0.5, 0.9), engine=st.sampled_from(["", "CUDNN"]), **hu.gcs) def test_spatialbn_train_mode_gradient_check( self, size, input_channels, batch_size, seed, order, epsilon, - engine, gc, dc): + momentum, engine, gc, dc): # Currently HIP SpatialBN only supports NCHW if _run_in_hip(gc, dc) and (order != 'NCHW'): return @@ -221,6 +230,7 @@ def test_spatialbn_train_mode_gradient_check( order=order, is_test=False, epsilon=epsilon, + momentum=momentum, engine=engine ) np.random.seed(seed) @@ -243,11 +253,13 @@ def test_spatialbn_train_mode_gradient_check( seed=st.integers(0, 65535), order=st.sampled_from(["NCHW", "NHWC"]), epsilon=st.floats(min_value=1e-5, max_value=1e-2), + momentum=st.floats(min_value=0.5, max_value=0.9), + engine=st.sampled_from(["", "CUDNN"]), # Currently HIP SpatialBN only supports 2D **hu.gcs_no_hip) def test_spatialbn_train_mode_gradient_check_1d( self, size, input_channels, batch_size, seed, order, epsilon, - gc, dc): + momentum, engine, gc, dc): op = core.CreateOperator( "SpatialBN", ["X", "scale", "bias", "mean", "var"], @@ -255,7 +267,8 @@ def test_spatialbn_train_mode_gradient_check_1d( order=order, is_test=False, epsilon=epsilon, - engine="CUDNN", + momentum=momentum, + engine=engine, ) np.random.seed(seed) scale = np.random.rand(input_channels).astype(np.float32) + 0.5 @@ -271,6 +284,146 @@ def test_spatialbn_train_mode_gradient_check_1d( self.assertGradientChecks(gc, op, [X, scale, bias, mean, var], input_to_check, [0], stepsize=0.01) + @given(N=st.integers(0, 5), + C=st.integers(1, 10), + H=st.integers(1, 5), + W=st.integers(1, 5), + epsilon=st.floats(1e-5, 1e-2), + momentum=st.floats(0.5, 0.9), + order=st.sampled_from(["NCHW", "NHWC"]), + num_batches=st.integers(2, 5), + in_place=st.booleans(), + engine=st.sampled_from(["", "CUDNN"]), + **hu.gcs) + def test_spatial_bn_multi_batch( + self, N, C, H, W, epsilon, momentum, order, num_batches, in_place, + engine, gc, dc): + if in_place: + outputs = ["Y", "mean", "var", "batch_mean", "batch_var"] + else: + outputs = ["Y", "mean", "var", "saved_mean", "saved_var"] + op = core.CreateOperator( + "SpatialBN", + ["X", "scale", "bias", "mean", "var", "batch_mean", "batch_var"], + outputs, + order=order, + is_test=False, + epsilon=epsilon, + momentum=momentum, + num_batches=num_batches, + engine=engine, + ) + if order == "NCHW": + X = np.random.randn(N, C, H, W).astype(np.float32) + else: + X = np.random.randn(N, H, W, C).astype(np.float32) + scale = np.random.randn(C).astype(np.float32) + bias = np.random.randn(C).astype(np.float32) + mean = np.random.randn(C).astype(np.float32) + var = np.random.rand(C).astype(np.float32) + batch_mean = np.random.rand(C).astype(np.float32) - 0.5 + batch_var = np.random.rand(C).astype(np.float32) + 1.0 + inputs = [X, scale, bias, mean, var, batch_mean, batch_var] + + def spatial_bn_multi_batch_ref( + X, scale, bias, mean, var, batch_mean, batch_var): + if N == 0: + batch_mean = np.zeros(C).astype(np.float32) + batch_var = np.zeros(C).astype(np.float32) + else: + size = num_batches * N * H * W + batch_mean /= size + batch_var = batch_var / size - np.square(batch_mean) + mean = momentum * mean + (1.0 - momentum) * batch_mean + var = momentum * var + (1.0 - momentum) * batch_var + batch_var = 1.0 / np.sqrt(batch_var + epsilon) + if order == "NCHW": + scale = np.reshape(scale, (C, 1, 1)) + bias = np.reshape(bias, (C, 1, 1)) + batch_mean = np.reshape(batch_mean, (C, 1, 1)) + batch_var = np.reshape(batch_var, (C, 1, 1)) + Y = (X - batch_mean) * batch_var * scale + bias + if order == "NCHW": + batch_mean = np.reshape(batch_mean, (C)) + batch_var = np.reshape(batch_var, (C)) + return (Y, mean, var, batch_mean, batch_var) + + self.assertReferenceChecks( + device_option=gc, + op=op, + inputs=inputs, + reference=spatial_bn_multi_batch_ref, + ) + self.assertDeviceChecks(dc, op, inputs, [0, 1, 2, 3, 4]) + + @given(N=st.integers(0, 5), + C=st.integers(1, 10), + H=st.integers(1, 5), + W=st.integers(1, 5), + epsilon=st.floats(1e-5, 1e-2), + order=st.sampled_from(["NCHW", "NHWC"]), + num_batches=st.integers(2, 5), + in_place=st.booleans(), + engine=st.sampled_from(["", "CUDNN"]), + **hu.gcs) + def test_spatial_bn_multi_batch_grad( + self, N, C, H, W, epsilon, order, num_batches, in_place, engine, + gc, dc): + if in_place: + outputs = ["dX", "dscale_sum", "dbias_sum"] + else: + outputs = ["dX", "dscale", "dbias"] + op = core.CreateOperator( + "SpatialBNGradient", + ["X", "scale", "dY", "mean", "rstd", "dscale_sum", "dbias_sum"], + outputs, + order=order, + epsilon=epsilon, + num_batches=num_batches, + engine=engine, + ) + if order == "NCHW": + dY = np.random.randn(N, C, H, W).astype(np.float32) + X = np.random.randn(N, C, H, W).astype(np.float32) + else: + dY = np.random.randn(N, H, W, C).astype(np.float32) + X = np.random.randn(N, H, W, C).astype(np.float32) + scale = np.random.randn(C).astype(np.float32) + mean = np.random.randn(C).astype(np.float32) + rstd = np.random.rand(C).astype(np.float32) + dscale_sum = np.random.randn(C).astype(np.float32) + dbias_sum = np.random.randn(C).astype(np.float32) + inputs = [X, scale, dY, mean, rstd, dscale_sum, dbias_sum] + + def spatial_bn_multi_batch_grad_ref( + X, scale, dY, mean, rstd, dscale_sum, dbias_sum): + if N == 0: + dscale = np.zeros(C).astype(np.float32) + dbias = np.zeros(C).astype(np.float32) + alpha = np.zeros(C).astype(np.float32) + beta = np.zeros(C).astype(np.float32) + gamma = np.zeros(C).astype(np.float32) + else: + dscale = dscale_sum / num_batches + dbias = dbias_sum / num_batches + alpha = scale * rstd + beta = -alpha * dscale * rstd / (N * H * W) + gamma = alpha * (mean * dscale * rstd - dbias) / (N * H * W) + if order == "NCHW": + alpha = np.reshape(alpha, (C, 1, 1)) + beta = np.reshape(beta, (C, 1, 1)) + gamma = np.reshape(gamma, (C, 1, 1)) + dX = alpha * dY + beta * X + gamma + return (dX, dscale, dbias) + + self.assertReferenceChecks( + device_option=gc, + op=op, + inputs=inputs, + reference=spatial_bn_multi_batch_grad_ref, + ) + self.assertDeviceChecks(dc, op, inputs, [0, 1, 2]) + @given(size=st.integers(7, 10), input_channels=st.integers(1, 10), batch_size=st.integers(0, 3), diff --git a/caffe2/python/pybind_state.cc b/caffe2/python/pybind_state.cc index 25db6a2a026ea..9a29f73baef81 100644 --- a/caffe2/python/pybind_state.cc +++ b/caffe2/python/pybind_state.cc @@ -59,7 +59,7 @@ CAFFE_DEFINE_TYPED_REGISTRY( std::unique_ptr); CAFFE_DEFINE_TYPED_REGISTRY( BlobFeederRegistry, - int, + caffe2::DeviceType, BlobFeederBase, std::unique_ptr); @@ -368,7 +368,7 @@ void addObjectMethods(py::module& m) { [](DLPackWrapper* t) -> py::object { CAFFE_ENFORCE_EQ( t->device_option.device_type(), - CPU, + PROTO_CPU, "Expected CPU device option for CPU tensor"); return t->data(); }, @@ -378,7 +378,7 @@ void addObjectMethods(py::module& m) { [](DLPackWrapper* t, py::object obj) { CAFFE_ENFORCE_EQ( t->device_option.device_type(), - CPU, + PROTO_CPU, "Expected CPU device option for CPU tensor"); t->feed(obj); }, @@ -985,7 +985,7 @@ void addGlobalMethods(py::module& m) { m.def( "set_op_engine_pref", [](const std::string& op_type, - const CaffeMap& op_pref) -> void { + const CaffeMap& op_pref) -> void { caffe2::SetOpEnginePref(op_type, op_pref); }); diff --git a/caffe2/python/pybind_state.h b/caffe2/python/pybind_state.h index 967331a19c5af..1271d67f48e7f 100644 --- a/caffe2/python/pybind_state.h +++ b/caffe2/python/pybind_state.h @@ -73,13 +73,14 @@ inline unique_ptr CreateFetcher(TypeIdentifier id) { CAFFE_DECLARE_TYPED_REGISTRY( BlobFeederRegistry, - int, + DeviceType, BlobFeederBase, std::unique_ptr); #define REGISTER_BLOB_FEEDER(device_type, ...) \ CAFFE_REGISTER_TYPED_CLASS(BlobFeederRegistry, device_type, __VA_ARGS__) inline unique_ptr CreateFeeder(int device_type) { - return BlobFeederRegistry()->Create(device_type); + return BlobFeederRegistry()->Create( + caffe2::ProtoToType(static_cast(device_type))); } static_assert( @@ -311,7 +312,7 @@ class PythonOpBase : public Operator { py::gil_scoped_acquire g; DeviceOption cpu_option; - cpu_option.set_device_type(CPU); + cpu_option.set_device_type(PROTO_CPU); std::vector inputs; inputs.reserve(InputSize()); diff --git a/caffe2/python/pybind_state_dlpack.cc b/caffe2/python/pybind_state_dlpack.cc index abfe6bb343d26..3978e24658cab 100644 --- a/caffe2/python/pybind_state_dlpack.cc +++ b/caffe2/python/pybind_state_dlpack.cc @@ -7,8 +7,8 @@ namespace py = pybind11; const DLDeviceType* CaffeToDLDeviceType(int device_type) { static std::map dl_device_type_map{ - {CPU, kCPU}, - {CUDA, kGPU}, + {PROTO_CPU, kCPU}, + {PROTO_CUDA, kGPU}, }; const auto it = dl_device_type_map.find(device_type); return it == dl_device_type_map.end() ? nullptr : &it->second; diff --git a/caffe2/python/pybind_state_gpu.cc b/caffe2/python/pybind_state_gpu.cc index 42a28de781f9e..9c15bc2145d5a 100644 --- a/caffe2/python/pybind_state_gpu.cc +++ b/caffe2/python/pybind_state_gpu.cc @@ -22,10 +22,10 @@ namespace caffe2 { namespace python { -REGISTER_CUDA_OPERATOR(Python, GPUFallbackOp>); +REGISTER_CUDA_OPERATOR(Python, GPUFallbackOp); REGISTER_CUDA_OPERATOR( PythonGradient, - GPUFallbackOp>); + GPUFallbackOp); REGISTER_CUDA_OPERATOR(PythonDLPack, PythonOp); REGISTER_CUDA_OPERATOR( @@ -121,7 +121,7 @@ void addCUDAObjectMethods(py::module& m) { [](DLPackWrapper* t) -> py::object { CAFFE_ENFORCE_EQ( t->device_option.device_type(), - CUDA, + PROTO_CUDA, "Expected CUDA device option for CUDA tensor"); return t->data(); @@ -132,7 +132,7 @@ void addCUDAObjectMethods(py::module& m) { [](DLPackWrapper* t, py::object obj) { CAFFE_ENFORCE_EQ( t->device_option.device_type(), - CUDA, + PROTO_CUDA, "Expected CUDA device option for CUDA tensor"); t->feed(obj); }, diff --git a/caffe2/python/pybind_state_hip.cc b/caffe2/python/pybind_state_hip.cc index d5e8720c645fb..a5d443cb9a7a8 100644 --- a/caffe2/python/pybind_state_hip.cc +++ b/caffe2/python/pybind_state_hip.cc @@ -13,10 +13,10 @@ namespace caffe2 { namespace python { -REGISTER_HIP_OPERATOR(Python, GPUFallbackOp>); +REGISTER_HIP_OPERATOR(Python, GPUFallbackOp); REGISTER_HIP_OPERATOR( PythonGradient, - GPUFallbackOp>); + GPUFallbackOp); REGISTER_HIP_OPERATOR(PythonDLPack, PythonOp); REGISTER_HIP_OPERATOR(PythonDLPackGradient, PythonGradientOp); @@ -53,7 +53,7 @@ void addHIPObjectMethods(py::module& m) { [](DLPackWrapper* t) -> py::object { CAFFE_ENFORCE_EQ( t->device_option.device_type(), - HIP, + PROTO_HIP, "Expected HIP device option for HIP tensor"); return t->data(); @@ -64,7 +64,7 @@ void addHIPObjectMethods(py::module& m) { [](DLPackWrapper* t, py::object obj) { CAFFE_ENFORCE_EQ( t->device_option.device_type(), - HIP, + PROTO_HIP, "Expected HIP device option for HIP tensor"); t->feed(obj); }, diff --git a/caffe2/python/pybind_state_ideep.cc b/caffe2/python/pybind_state_ideep.cc index 056558c9a7333..ebad6cf8d9683 100644 --- a/caffe2/python/pybind_state_ideep.cc +++ b/caffe2/python/pybind_state_ideep.cc @@ -161,7 +161,7 @@ class IDeepFeeder : public BlobFeederBase { FeedTensor(option, original_array, blob->GetMutable()); } else { DeviceOption cpu_option(option); - cpu_option.set_device_type(DeviceType::CPU); + cpu_option.set_device_type(DeviceTypeProto::PROTO_CPU); TensorFeeder cpu_tensor_feeder; cpu_tensor_feeder.FeedTensor(cpu_option, original_array, blob->GetMutableTensor(CPU)); diff --git a/caffe2/python/pybind_state_nomni.cc b/caffe2/python/pybind_state_nomni.cc index e7e4dfcbf2e9b..bca9b69706987 100644 --- a/caffe2/python/pybind_state_nomni.cc +++ b/caffe2/python/pybind_state_nomni.cc @@ -104,7 +104,19 @@ void addNomnigraphMethods(pybind11::module& m) { .def( "dataFlow", [](NNModule* nn) -> NNGraph* { return &nn->dataFlow; }, - py::return_value_policy::reference_internal); + py::return_value_policy::reference_internal) + .def("convertToCaffe2Proto", [](NNModule& nn, py::object def) { + auto attr = def.attr("SerializeToString"); + CAFFE_ENFORCE( + attr, "convertToCaffe2Proto takes either no args", "a NetDef"); + auto str = attr(); + caffe2::NetDef proto; + proto.ParseFromString(py::bytes(str)); + auto new_proto = caffe2::convertToCaffe2Proto(nn, proto); + std::string out; + new_proto.SerializeToString(&out); + return py::bytes(out); + }); // NNGraph methods py::class_ nngraph(m, "NNGraph"); diff --git a/caffe2/sgd/lars_op.h b/caffe2/sgd/lars_op.h index 9a2c89a3b8782..7b65598cdffe1 100644 --- a/caffe2/sgd/lars_op.h +++ b/caffe2/sgd/lars_op.h @@ -9,7 +9,7 @@ namespace caffe2 { template -class CAFFE2_API LarsOp final : public Operator { +class LarsOp final : public Operator { public: USE_OPERATOR_CONTEXT_FUNCTIONS; LarsOp(const OperatorDef& operator_def, Workspace* ws) diff --git a/caffe2/transforms/conv_to_nnpack_transform.h b/caffe2/transforms/conv_to_nnpack_transform.h index 5d61be232c2dc..8563732f225e2 100644 --- a/caffe2/transforms/conv_to_nnpack_transform.h +++ b/caffe2/transforms/conv_to_nnpack_transform.h @@ -12,7 +12,7 @@ class CAFFE2_API ConvToNNPackTransform : public SingleOpTransform { // Specify what the op needs to be to match the pattern. bool MatchOperator(const OperatorDef& op) override { return ( - op.type() == "Conv" && op.device_option().device_type() == CPU && + op.type() == "Conv" && op.device_option().device_type() == PROTO_CPU && op.engine() != "NNPACK"); } diff --git a/caffe2/transforms/conv_to_nnpack_transform_test.cc b/caffe2/transforms/conv_to_nnpack_transform_test.cc index 4ab80fccd9314..92b5f1ade2407 100644 --- a/caffe2/transforms/conv_to_nnpack_transform_test.cc +++ b/caffe2/transforms/conv_to_nnpack_transform_test.cc @@ -15,7 +15,7 @@ TEST(ConvToNNPackTest, TestSimple) { op = AddOp(&netdef, "Conv", {"in"}, {"out"}); op = AddOp(&netdef, "Relu", {"out"}, {"out"}); op = AddOp(&netdef, "Conv", {"out"}, {"out"}); // if not CPU, won't transform - op->mutable_device_option()->set_device_type(CUDA); + op->mutable_device_option()->set_device_type(PROTO_CUDA); op = AddOp(&netdef, "Relu", {"out"}, {"out"}); op = AddOp(&netdef, "Conv", {"out"}, {"out"}); op->set_engine("NNPACK"); // does not need to be transformed @@ -28,7 +28,7 @@ TEST(ConvToNNPackTest, TestSimple) { int nnpack_count = 0; for (auto& op : transformed_netdef.op()) { - if (op.type() == "Conv" && op.device_option().device_type() == CPU) { + if (op.type() == "Conv" && op.device_option().device_type() == PROTO_CPU) { EXPECT_EQ(op.engine(), "NNPACK"); nnpack_count++; } diff --git a/caffe2/transforms/pattern_net_transform_test.cc b/caffe2/transforms/pattern_net_transform_test.cc index 36925d9d43fc4..8ac21af540739 100644 --- a/caffe2/transforms/pattern_net_transform_test.cc +++ b/caffe2/transforms/pattern_net_transform_test.cc @@ -250,19 +250,19 @@ TEST(PatternNetTransformTest, TestDeviceOptionMatching) { NetDef pdef; auto op = AddOp(&pdef, "DummyOp1", {"in"}, {"out"}); - op->mutable_device_option()->set_device_type(CPU); + op->mutable_device_option()->set_device_type(PROTO_CPU); NetDef rdef; op = AddOp(&rdef, "DummyOp1", {"in"}, {"out"}); - op->mutable_device_option()->set_device_type(CUDA); + op->mutable_device_option()->set_device_type(PROTO_CUDA); NetDef netdef; op = AddOp(&netdef, "DummyOp1", {"in"}, {"mid"}); - op->mutable_device_option()->set_device_type(CPU); + op->mutable_device_option()->set_device_type(PROTO_CPU); op = AddOp(&netdef, "DummyOp1", {"mid"}, {"mid"}); // should not match - op->mutable_device_option()->set_device_type(CUDA); + op->mutable_device_option()->set_device_type(PROTO_CUDA); op = AddOp(&netdef, "DummyOp1", {"mid"}, {"out"}); - op->mutable_device_option()->set_device_type(CPU); + op->mutable_device_option()->set_device_type(PROTO_CPU); PatternNetTransform t(pdef, rdef); transform::Graph g(netdef); @@ -272,7 +272,7 @@ TEST(PatternNetTransformTest, TestDeviceOptionMatching) { NetDef transformed_net = t.ApplyTo(netdef); for (const auto& opdef : transformed_net.op()) { EXPECT_TRUE(opdef.has_device_option()); - EXPECT_EQ(opdef.device_option().device_type(), CUDA); + EXPECT_EQ(opdef.device_option().device_type(), PROTO_CUDA); } } diff --git a/caffe2/utils/CMakeLists.txt b/caffe2/utils/CMakeLists.txt index 6bda0254c667d..ff8c323dd096d 100644 --- a/caffe2/utils/CMakeLists.txt +++ b/caffe2/utils/CMakeLists.txt @@ -1,11 +1,3 @@ -# TODO(orionr): Enable all of this for Windows DLL when we -# can figure out how to get it to build -if (MSVC AND BUILD_SHARED_LIBS) -list(APPEND Caffe2_CPU_SRCS utils/proto_wrap.cc) -set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} PARENT_SCOPE) -return() -endif() - list(APPEND Caffe2_CPU_SRCS utils/proto_wrap.cc utils/proto_utils.cc diff --git a/caffe2/utils/fatal_signal_asan_no_sig_test.cc b/caffe2/utils/fatal_signal_asan_no_sig_test.cc index b06231b01bbc3..a56539a1f1894 100644 --- a/caffe2/utils/fatal_signal_asan_no_sig_test.cc +++ b/caffe2/utils/fatal_signal_asan_no_sig_test.cc @@ -14,6 +14,7 @@ namespace { void* dummy_thread(void*) { while (1) { } + return nullptr; } bool forkAndPipe( diff --git a/caffe2/utils/filler.h b/caffe2/utils/filler.h index 4d04778e401e8..7016f09a3bab0 100644 --- a/caffe2/utils/filler.h +++ b/caffe2/utils/filler.h @@ -18,8 +18,12 @@ class TensorFiller { void Fill(Tensor* tensor, Context* context) const { CAFFE_ENFORCE(context, "context is null"); CAFFE_ENFORCE(tensor, "tensor is null"); - auto min = static_cast(min_); - auto max = static_cast(max_); + auto min = (min_ < std::numeric_limits::min()) + ? std::numeric_limits::min() + : static_cast(min_); + auto max = (max_ > std::numeric_limits::max()) + ? std::numeric_limits::max() + : static_cast(max_); CAFFE_ENFORCE_LE(min, max); Tensor temp_tensor(shape_, Context::GetDeviceType()); diff --git a/caffe2/utils/hip/math_blas_hip_test.cc b/caffe2/utils/hip/math_blas_hip_test.cc index 6aed0273f6771..911c2b09868fc 100644 --- a/caffe2/utils/hip/math_blas_hip_test.cc +++ b/caffe2/utils/hip/math_blas_hip_test.cc @@ -15,7 +15,7 @@ TEST(MathROCBLASTest, GemmNoTransNoTrans) { return; Workspace ws; DeviceOption option; - option.set_device_type(HIP); + option.set_device_type(PROTO_HIP); HIPContext context(option); Blob* blobX = ws.CreateBlob("X"); @@ -115,7 +115,7 @@ TEST(MathROCBLASTest, GemmNoTransTrans) { return; Workspace ws; DeviceOption option; - option.set_device_type(HIP); + option.set_device_type(PROTO_HIP); HIPContext context(option); Blob* blobX = ws.CreateBlob("X"); @@ -214,7 +214,7 @@ TEST(MathROCBLASTest, GemvNoTrans) { return; Workspace ws; DeviceOption option; - option.set_device_type(HIP); + option.set_device_type(PROTO_HIP); HIPContext context(option); Blob* blobA = ws.CreateBlob("A"); @@ -304,7 +304,7 @@ TEST(MathROCBLASTest, GemvTrans) { return; Workspace ws; DeviceOption option; - option.set_device_type(HIP); + option.set_device_type(PROTO_HIP); HIPContext context(option); Blob* blobA = ws.CreateBlob("A"); diff --git a/caffe2/utils/hip/math_hip.cc b/caffe2/utils/hip/math_hip.cc index 4fa375ed88d50..d1c1a9546e5ec 100644 --- a/caffe2/utils/hip/math_hip.cc +++ b/caffe2/utils/hip/math_hip.cc @@ -3361,6 +3361,47 @@ CAFFE2_SPECIALIZED_HIP_MOMENTS(float) namespace { +template +__global__ void +InvStdHIPKernel(const int N, const T epsilon, const T* var, T* inv_std); + +#define DELEGATE_INV_STD_KERNEL_FUNCTION(T, Func) \ + template <> \ + __global__ void InvStdHIPKernel( \ + const int N, const T epsilon, const T* var, T* inv_std) { \ + HIP_1D_KERNEL_LOOP(i, N) { \ + inv_std[i] = Func(var[i] + epsilon); \ + } \ + } +DELEGATE_INV_STD_KERNEL_FUNCTION(float, rsqrtf) +#undef DELEGATE_INV_STD_KERNEL_FUNCTION + +} // namespace + +#define CAFFE2_SPECIALIZED_HIP_INV_STD(T) \ + template <> \ + void InvStd( \ + const int N, \ + const T epsilon, \ + const T* var, \ + T* inv_std, \ + HIPContext* context) { \ + hipLaunchKernelGGL( \ + InvStdHIPKernel, \ + CAFFE_GET_BLOCKS(N), \ + CAFFE_HIP_NUM_THREADS, \ + 0, \ + context->hip_stream(), \ + N, \ + epsilon, \ + var, \ + inv_std); \ + } +CAFFE2_SPECIALIZED_HIP_INV_STD(float) +#undef CAFFE2_SPECIALIZED_HIP_INV_STD + +namespace { + template __global__ void TransposeHIPKernel( const int size, @@ -3434,5 +3475,55 @@ CAFFE2_SPECIALIZED_HIP_TRANSPOSE(double) CAFFE2_SPECIALIZED_HIP_TRANSPOSE(int) CAFFE2_SPECIALIZED_HIP_TRANSPOSE(TIndex) #undef CAFFE2_SPECIALIZED_HIP_TRANSPOSE + +namespace { + +template +__global__ void AffineChannelHIPKernel( + const int size, + const int C, + const int HxW, + const T* X, + const T* scale, + const T* bias, + T* Y) { + HIP_1D_KERNEL_LOOP(i, size) { + const int c = kOrder == StorageOrder::NCHW ? i / HxW % C : i % C; + Y[i] = scale[c] * X[i] + bias[c]; + } +} + +} // namespace + +#define CAFFE2_SPECIALIZED_HIP_AFFINE_CHANNEL(T, kOrder) \ + template <> \ + void AffineChannel( \ + const int N, \ + const int C, \ + const int HxW, \ + const T* X, \ + const T* scale, \ + const T* bias, \ + T* Y, \ + HIPContext* context) { \ + const int size = N * C * HxW; \ + hipLaunchKernelGGL( \ + AffineChannelHIPKernel, \ + CAFFE_GET_BLOCKS(size), \ + CAFFE_HIP_NUM_THREADS, \ + 0, \ + context->hip_stream(), \ + size, \ + C, \ + HxW, \ + X, \ + scale, \ + bias, \ + Y); \ + } +CAFFE2_SPECIALIZED_HIP_AFFINE_CHANNEL(float, StorageOrder::NCHW) +CAFFE2_SPECIALIZED_HIP_AFFINE_CHANNEL(float, StorageOrder::NHWC) +#undef CAFFE2_SPECIALIZED_HIP_AFFINE_CHANNEL + } // namespace math } // namespace caffe2 diff --git a/caffe2/utils/math.h b/caffe2/utils/math.h index fc17f824aef99..0b4e6dd0f6ffe 100644 --- a/caffe2/utils/math.h +++ b/caffe2/utils/math.h @@ -262,6 +262,15 @@ CAFFE2_API void Moments( T* variance, Context* context); +// Computes inv_std from variance. +template +CAFFE2_API void InvStd( + const int N, + const T epsilon, + const T* var, + T* inv_std, + Context* context); + // Adds batch sub-tensors elementwise to output. Stripe is the stripe length // and N is the number of elements to add (size of Y). template @@ -642,6 +651,17 @@ CAFFE2_API void CopyMatrix( template CAFFE2_API void CopyVector(const int N, const T* A, T* B, Context* context); +template +CAFFE2_API void AffineChannel( + const int N, + const int C, + const int HxW, + const T* X, + const T* scale, + const T* bias, + T* Y, + Context* context); + // Calculates ceil(a / b). User must be careful to ensure that there // is no overflow or underflow in the calculation. template diff --git a/caffe2/utils/math_cpu.cc b/caffe2/utils/math_cpu.cc index a5f44dd204a60..c54226af68ed7 100644 --- a/caffe2/utils/math_cpu.cc +++ b/caffe2/utils/math_cpu.cc @@ -76,7 +76,7 @@ namespace math { // (transpose) if the argument TransA or TransB is set to CblasNoTrans or // CblasTrans, respectively, for each of A and B. template <> -void Gemm( +CAFFE2_EXPORT void Gemm( const CBLAS_TRANSPOSE trans_A, const CBLAS_TRANSPOSE trans_B, const int M, @@ -134,7 +134,7 @@ void Gemm( } template <> -void GemmEx( +CAFFE2_EXPORT void GemmEx( const CBLAS_TRANSPOSE trans_A, const CBLAS_TRANSPOSE trans_B, const int M, @@ -206,7 +206,7 @@ void GemmEx( } template <> -void Gemv( +CAFFE2_EXPORT void Gemv( const CBLAS_TRANSPOSE trans_A, const int M, const int N, @@ -245,7 +245,7 @@ void Gemv( #define CAFFE2_SPECIALIZED_DOT(T) \ template <> \ - void Dot( \ + CAFFE2_EXPORT void Dot( \ const int N, const T* a, const T* b, T* y, CPUContext* context) { \ *y = ConstEigenVectorMap(a, N).dot(ConstEigenVectorMap(b, N)); \ } @@ -254,12 +254,12 @@ CAFFE2_SPECIALIZED_DOT(float) #define CAFFE2_SPECIALIZED_AXPY(T) \ template <> \ - void Axpy( \ + CAFFE2_EXPORT void Axpy( \ const int N, const T alpha, const T* x, T* Y, CPUContext* context) { \ EigenVectorMap(Y, N) += ConstEigenVectorMap(x, N) * alpha; \ } \ template <> \ - void Axpy( \ + CAFFE2_EXPORT void Axpy( \ const int N, const T* alpha, const T* x, T* Y, CPUContext* context) { \ EigenVectorMap(Y, N) += ConstEigenVectorMap(x, N) * (*alpha); \ } @@ -268,7 +268,7 @@ CAFFE2_SPECIALIZED_AXPY(float) #define CAFFE2_SPECIALIZED_AXPBY(T) \ template <> \ - void Axpby( \ + CAFFE2_EXPORT void Axpby( \ const int N, \ const T alpha, \ const T* x, \ @@ -279,7 +279,7 @@ CAFFE2_SPECIALIZED_AXPY(float) y_arr = y_arr * beta + ConstEigenVectorArrayMap(x, N) * alpha; \ } \ template <> \ - void Axpby( \ + CAFFE2_EXPORT void Axpby( \ const int N, \ const T* alpha, \ const T* x, \ @@ -295,7 +295,7 @@ CAFFE2_SPECIALIZED_AXPBY(float) #else // CAFFE2_USE_EIGEN_FOR_BLAS template <> -void Gemm( +CAFFE2_EXPORT void Gemm( const CBLAS_TRANSPOSE trans_A, const CBLAS_TRANSPOSE trans_B, const int M, @@ -328,7 +328,7 @@ void Gemm( } template <> -void GemmEx( +CAFFE2_EXPORT void GemmEx( const CBLAS_TRANSPOSE trans_A, const CBLAS_TRANSPOSE trans_B, const int M, @@ -361,7 +361,7 @@ void GemmEx( } template <> -void Gemv( +CAFFE2_EXPORT void Gemv( const CBLAS_TRANSPOSE trans_A, const int M, const int N, @@ -377,7 +377,7 @@ void Gemv( #define CAFFE2_SPECIALIZED_SCALE(TAlpha, TData, prefix) \ template <> \ - void Scale( \ + CAFFE2_EXPORT void Scale( \ const int n, \ const TAlpha alpha, \ const TData* x, \ @@ -391,7 +391,7 @@ void Gemv( } \ } \ template <> \ - void Scale( \ + CAFFE2_EXPORT void Scale( \ const int n, \ const TAlpha* alpha, \ const TData* x, \ @@ -411,7 +411,7 @@ CAFFE2_SPECIALIZED_SCALE(float, double, d) #define CAFFE2_SPECIALIZED_DOT(T, prefix) \ template <> \ - void Dot( \ + CAFFE2_EXPORT void Dot( \ const int N, const T* a, const T* b, T* y, CPUContext*) { \ *y = cblas_##prefix##dot(N, a, 1, b, 1); \ } @@ -420,12 +420,12 @@ CAFFE2_SPECIALIZED_DOT(float, s) #define CAFFE2_SPECIALIZED_AXPY(T, prefix) \ template <> \ - void Axpy( \ + CAFFE2_EXPORT void Axpy( \ const int N, const T alpha, const T* x, T* y, CPUContext*) { \ cblas_##prefix##axpy(N, alpha, x, 1, y, 1); \ } \ template <> \ - void Axpy( \ + CAFFE2_EXPORT void Axpy( \ const int N, const T* alpha, const T* x, T* y, CPUContext*) { \ cblas_##prefix##axpy(N, *alpha, x, 1, y, 1); \ } @@ -437,7 +437,7 @@ CAFFE2_SPECIALIZED_AXPY(float, s) #ifdef CAFFE2_USE_MKL #define CAFFE2_SPECIALIZED_AXPBY(T, prefix) \ template <> \ - void Axpby( \ + CAFFE2_EXPORT void Axpby( \ const int N, \ const T alpha, \ const T* x, \ @@ -447,7 +447,7 @@ CAFFE2_SPECIALIZED_AXPY(float, s) cblas_##prefix##axpby(N, alpha, x, 1, beta, y, 1); \ } \ template <> \ - void Axpby( \ + CAFFE2_EXPORT void Axpby( \ const int N, \ const T* alpha, \ const T* x, \ @@ -459,7 +459,7 @@ CAFFE2_SPECIALIZED_AXPY(float, s) #else // CAFFE2_USE_MKL #define CAFFE2_SPECIALIZED_AXPBY(T, prefix) \ template <> \ - void Axpby( \ + CAFFE2_EXPORT void Axpby( \ const int N, \ const T alpha, \ const T* x, \ @@ -470,7 +470,7 @@ CAFFE2_SPECIALIZED_AXPY(float, s) cblas_##prefix##axpy(N, alpha, x, 1, y, 1); \ } \ template <> \ - void Axpby( \ + CAFFE2_EXPORT void Axpby( \ const int N, \ const T* alpha, \ const T* x, \ @@ -488,7 +488,7 @@ CAFFE2_SPECIALIZED_AXPBY(float, s) #define CAFFE2_SPECIALIZED_SCALE(TAlpha, TData) \ template <> \ - void Scale( \ + CAFFE2_EXPORT void Scale( \ const int n, \ const TAlpha alpha, \ const TData* x, \ @@ -498,7 +498,7 @@ CAFFE2_SPECIALIZED_AXPBY(float, s) ConstEigenVectorMap(x, n) * static_cast(alpha); \ } \ template <> \ - void Scale( \ + CAFFE2_EXPORT void Scale( \ const int n, \ const TAlpha* alpha, \ const TData* x, \ @@ -517,7 +517,7 @@ CAFFE2_SPECIALIZED_SCALE(std::int64_t, std::int64_t) #undef CAFFE2_SPECIALIZED_SCALE template <> -void GemmBatched( +CAFFE2_EXPORT void GemmBatched( const CBLAS_TRANSPOSE trans_A, const CBLAS_TRANSPOSE trans_B, const int batch_size, @@ -563,7 +563,7 @@ void GemmBatched( } template <> -void GemmStridedBatched( +CAFFE2_EXPORT void GemmStridedBatched( const CBLAS_TRANSPOSE trans_A, const CBLAS_TRANSPOSE trans_B, const int batch_size, @@ -634,7 +634,7 @@ void GemmStridedBatched( #define DELEGATE_SIMPLE_UNARY_FUNCTION(T, Funcname, OriginalFunc, ...) \ template <> \ - void Funcname(const int N, const T* x, T* y, CPUContext*) { \ + CAFFE2_EXPORT void Funcname(const int N, const T* x, T* y, CPUContext*) { \ OriginalFunc(N, x, y, ##__VA_ARGS__); \ } DELEGATE_SIMPLE_UNARY_FUNCTION( @@ -683,7 +683,7 @@ DELEGATE_SIMPLE_UNARY_FUNCTION(double, Inv, vdInv) #define DELEGATE_SINCOS_FUNCTION(T, OriginalFunc) \ template <> \ - void SinCos( \ + CAFFE2_EXPORT void SinCos( \ const int N, const T* a, T* ys, T* yc, CPUContext*) { \ OriginalFunc(N, a, ys, yc); \ } @@ -693,7 +693,7 @@ DELEGATE_SINCOS_FUNCTION(double, vdSinCos) #define DELEGATE_POWX_FUNCTION(T, OriginalFunc) \ template <> \ - void Powx(const int N, const T* a, T b, T* y, CPUContext*) { \ + CAFFE2_EXPORT void Powx(const int N, const T* a, T b, T* y, CPUContext*) { \ OriginalFunc(N, a, b, y); \ } DELEGATE_POWX_FUNCTION(float, vsPowx) @@ -702,7 +702,7 @@ DELEGATE_POWX_FUNCTION(double, vdPowx) #define DELEGATE_SIMPLE_BINARY_FUNCTION(T, Func, FuncImpl) \ template <> \ - void Func( \ + CAFFE2_EXPORT void Func( \ const int N, const T* A, const T* B, T* C, CPUContext*) { \ FuncImpl(N, A, B, C); \ } @@ -720,7 +720,7 @@ DELEGATE_SIMPLE_BINARY_FUNCTION(double, Div, vdDiv) #define DELEGATE_SIMPLE_UNARY_FUNCTION(T, Funcname, expr) \ template <> \ - void Funcname(const int N, const T* x, T* y, CPUContext*) { \ + CAFFE2_EXPORT void Funcname(const int N, const T* x, T* y, CPUContext*) { \ EigenVectorMap(y, N) = ConstEigenVectorArrayMap(x, N).expr(); \ } DELEGATE_SIMPLE_UNARY_FUNCTION(float, Exp, exp) @@ -750,7 +750,7 @@ DELEGATE_SIMPLE_UNARY_FUNCTION(double, Rsqrt, rsqrt) #define DELEGATE_SINCOS_FUNCTION(T) \ template <> \ - void SinCos( \ + CAFFE2_EXPORT void SinCos( \ const int N, const T* x, T* ys, T* yc, CPUContext*) { \ EigenVectorMap(ys, N) = ConstEigenVectorArrayMap(x, N).sin(); \ EigenVectorMap(yc, N) = ConstEigenVectorArrayMap(x, N).cos(); \ @@ -761,7 +761,7 @@ DELEGATE_SINCOS_FUNCTION(double) #define DELEGATE_TANH_FUNCTION(T) \ template <> \ - void Tanh(const int N, const T* X, T* Y, CPUContext*) { \ + CAFFE2_EXPORT void Tanh(const int N, const T* X, T* Y, CPUContext*) { \ EigenVectorMap(Y, N) = T(1) - \ ((ConstEigenVectorArrayMap(X, N) * T(2)).exp() + T(1)).inverse() * \ T(2); \ @@ -772,7 +772,7 @@ DELEGATE_TANH_FUNCTION(double) #define DELEGATE_CBRT_FUNCTION(T) \ template <> \ - void Cbrt(const int N, const T* X, T* Y, CPUContext*) { \ + CAFFE2_EXPORT void Cbrt(const int N, const T* X, T* Y, CPUContext*) { \ std::transform(X, X + N, Y, [](const T x) { return cbrt(x); }); \ } DELEGATE_CBRT_FUNCTION(float) @@ -781,7 +781,7 @@ DELEGATE_CBRT_FUNCTION(double) #define DELEGATE_POWX_FUNCTION(T) \ template <> \ - void Powx( \ + CAFFE2_EXPORT void Powx( \ const int N, const T* a, const T b, T* y, CPUContext*) { \ EigenVectorMap(y, N) = ConstEigenVectorArrayMap(a, N).pow(b); \ } @@ -790,7 +790,7 @@ DELEGATE_POWX_FUNCTION(float) #define DELEGATE_SINH_FUNCTION(T) \ template <> \ - void Sinh(const int N, const T* X, T* Y, CPUContext*) { \ + CAFFE2_EXPORT void Sinh(const int N, const T* X, T* Y, CPUContext*) { \ ConstEigenVectorArrayMap X_arr(X, N); \ EigenVectorMap(Y, N) = (X_arr.exp() - (-X_arr).exp()) / 2; \ } @@ -800,7 +800,7 @@ DELEGATE_SINH_FUNCTION(double) #define DELEGATE_COSH_FUNCTION(T) \ template <> \ - void Cosh(const int N, const T* X, T* Y, CPUContext*) { \ + CAFFE2_EXPORT void Cosh(const int N, const T* X, T* Y, CPUContext*) { \ ConstEigenVectorArrayMap X_arr(X, N); \ EigenVectorMap(Y, N) = (X_arr.exp() + (-X_arr).exp()) / 2; \ } @@ -810,7 +810,7 @@ DELEGATE_COSH_FUNCTION(double) #define DELEGATE_INV_FUNCTION(T) \ template <> \ - void Inv(const int N, const T* x, T* y, CPUContext*) { \ + CAFFE2_EXPORT void Inv(const int N, const T* x, T* y, CPUContext*) { \ EigenVectorMap(y, N) = ConstEigenVectorArrayMap(x, N).inverse(); \ } DELEGATE_INV_FUNCTION(float) @@ -821,7 +821,7 @@ DELEGATE_INV_FUNCTION(double) #define DELEGATE_NEG_FUNCTION(T) \ template <> \ - void Neg(const int N, const T* x, T* y, CPUContext*) { \ + CAFFE2_EXPORT void Neg(const int N, const T* x, T* y, CPUContext*) { \ EigenVectorMap(y, N) = -ConstEigenVectorMap(x, N); \ } DELEGATE_NEG_FUNCTION(float) @@ -832,7 +832,7 @@ DELEGATE_NEG_FUNCTION(std::int64_t) #define DELEGATE_SIGN_FUNCTION(T) \ template <> \ - void Sign(const int N, const T* x, T* y, CPUContext*) { \ + CAFFE2_EXPORT void Sign(const int N, const T* x, T* y, CPUContext*) { \ EigenVectorMap(y, N) = ConstEigenVectorArrayMap(x, N).sign(); \ } DELEGATE_SIGN_FUNCTION(float) @@ -843,7 +843,7 @@ DELEGATE_SIGN_FUNCTION(std::int64_t) #define DELEGATE_ABS_FUNCTION(T) \ template <> \ - void Abs(const int N, const T* x, T* y, CPUContext*) { \ + CAFFE2_EXPORT void Abs(const int N, const T* x, T* y, CPUContext*) { \ EigenVectorMap(y, N) = ConstEigenVectorArrayMap(x, N).abs(); \ } #ifndef CAFFE2_USE_MKL @@ -856,7 +856,7 @@ DELEGATE_ABS_FUNCTION(std::int64_t) #define DELEGATE_CUBE_FUNCTION(T) \ template <> \ - void Cube(const int N, const T* X, T* Y, CPUContext*) { \ + CAFFE2_EXPORT void Cube(const int N, const T* X, T* Y, CPUContext*) { \ EigenVectorMap(Y, N) = ConstEigenVectorArrayMap(X, N).cube(); \ } DELEGATE_CUBE_FUNCTION(float) @@ -867,7 +867,7 @@ DELEGATE_CUBE_FUNCTION(std::int64_t) #define EIGEN_SIMPLE_BINARY_FUNCTION(T, Func, expr) \ template <> \ - void Func( \ + CAFFE2_EXPORT void Func( \ const int N, const T* A, const T* B, T* C, CPUContext*) { \ EigenVectorMap(C, N) = ConstEigenVectorArrayMap(A, N) \ expr ConstEigenVectorArrayMap(B, N); \ @@ -905,7 +905,7 @@ DEFINE_SIMPLE_BINARY_FUNCTION(Div, /) #define CAFFE2_SPECIALIZED_SET(T) \ template <> \ - void Set(const size_t N, const T alpha, T* Y, CPUContext*) { \ + CAFFE2_EXPORT void Set(const size_t N, const T alpha, T* Y, CPUContext*) { \ if (N == 0) { \ return; \ } \ @@ -932,7 +932,7 @@ CAFFE2_SPECIALIZED_SET(uint16_t); #define CAFFE2_SPECIALIZED_REDUCEMIN(T) \ template <> \ - void ReduceMin( \ + CAFFE2_EXPORT void ReduceMin( \ const int N, \ const T* x, \ T* y, \ @@ -945,7 +945,7 @@ CAFFE2_SPECIALIZED_REDUCEMIN(float) #define CAFFE2_SPECIALIZED_REDUCEMAX(T) \ template <> \ - void ReduceMax( \ + CAFFE2_EXPORT void ReduceMax( \ const int N, \ const T* x, \ T* y, \ @@ -991,7 +991,7 @@ struct SquaredL2NormFunctor { #define DELEGATE_ROWWISE_REDUCE_FUNCTION(Func, EigenOp) \ template \ - void Rowwise##Func( \ + CAFFE2_EXPORT void Rowwise##Func( \ const int rows, const int cols, const T alpha, const T* X, T* Y) { \ EigenVectorMap(Y, rows) = \ ConstEigenMatrixMap(X, cols, rows).colwise().EigenOp() * alpha; \ @@ -1006,7 +1006,7 @@ DELEGATE_ROWWISE_REDUCE_FUNCTION(ReduceL2, norm) #define DELEGATE_COLWISE_REDUCE_FUNCTION(Func, EigenOp) \ template \ - void Colwise##Func( \ + CAFFE2_EXPORT void Colwise##Func( \ const int rows, const int cols, const T alpha, const T* X, T* Y) { \ EigenVectorMap(Y, cols) = \ ConstEigenMatrixMap(X, cols, rows).rowwise().EigenOp() * alpha; \ @@ -1020,7 +1020,7 @@ DELEGATE_COLWISE_REDUCE_FUNCTION(ReduceL2, norm) #undef DELEGATE_COLWISE_REDUCE_FUNCTION template -void BothEndsReduceMin( +CAFFE2_EXPORT void BothEndsReduceMin( const int pre, const int mid, const int nxt, @@ -1044,7 +1044,7 @@ void BothEndsReduceMin( } template -void BothEndsReduceMax( +CAFFE2_EXPORT void BothEndsReduceMax( const int pre, const int mid, const int nxt, @@ -1066,7 +1066,7 @@ void BothEndsReduceMax( } template -void BothEndsReduceSum( +CAFFE2_EXPORT void BothEndsReduceSum( const int pre, const int mid, const int nxt, @@ -1087,7 +1087,7 @@ void BothEndsReduceSum( } template -void BothEndsReduceMean( +CAFFE2_EXPORT void BothEndsReduceMean( const int pre, const int mid, const int nxt, @@ -1108,7 +1108,7 @@ void BothEndsReduceMean( } template -void BothEndsReduceL1( +CAFFE2_EXPORT void BothEndsReduceL1( const int pre, const int mid, const int nxt, @@ -1135,7 +1135,7 @@ void BothEndsReduceL1( } template -void BothEndsReduceL2( +CAFFE2_EXPORT void BothEndsReduceL2( const int pre, const int mid, const int nxt, @@ -1155,7 +1155,7 @@ void BothEndsReduceL2( } template -void ReduceTensor( +CAFFE2_EXPORT void ReduceTensor( const int ndim, const int* X_dims, const int* Y_dims, @@ -1183,7 +1183,7 @@ void ReduceTensor( #define DELEGATE_REDUCE_FUNCTION(T, Func, reducer, init, is_norm) \ template <> \ - void Func( \ + CAFFE2_EXPORT void Func( \ const int num_dims, \ const int* dims, \ const int num_axes, \ @@ -1325,7 +1325,7 @@ DELEGATE_REDUCE_FUNCTION( #define CAFFE2_SPECIALIZED_REDUCE_MEAN(T) \ template <> \ - void ReduceMean( \ + CAFFE2_EXPORT void ReduceMean( \ const int num_dims, \ const int* dims, \ const int num_axes, \ @@ -1392,7 +1392,7 @@ CAFFE2_SPECIALIZED_REDUCE_MEAN(double) #define CAFFE2_SPECIALIZED_REDUCE_L2(T) \ template <> \ - void ReduceL2( \ + CAFFE2_EXPORT void ReduceL2( \ const int num_dims, \ const int* dims, \ const int num_axes, \ @@ -1462,7 +1462,7 @@ CAFFE2_SPECIALIZED_REDUCE_L2(double) namespace { template -void BroadcastImpl( +CAFFE2_EXPORT void BroadcastImpl( const int X_ndim, const int* X_dims, const int Y_ndim, @@ -1495,7 +1495,7 @@ void BroadcastImpl( #define CAFFE2_SPECIALIZED_BROADCAST(T) \ template <> \ - void Broadcast( \ + CAFFE2_EXPORT void Broadcast( \ const int X_ndim, \ const int* X_dims, \ const int Y_ndim, \ @@ -1515,64 +1515,68 @@ CAFFE2_SPECIALIZED_BROADCAST(double) namespace { template -void RowwiseMoments( +CAFFE2_EXPORT void RowwiseMoments( const int rows, const int cols, const T* X, T* mean, T* variance) { - ConstEigenArrayMap X_mat(X, cols, rows); + ConstEigenArrayMap X_arr(X, cols, rows); EigenVectorArrayMap mean_arr(mean, rows); - EigenVectorArrayMap variance_arr(variance, rows); - mean_arr = X_mat.colwise().mean(); - variance_arr = X_mat.array().square().colwise().mean(); - variance_arr -= mean_arr.square(); + EigenVectorArrayMap var_arr(variance, rows); + mean_arr = X_arr.colwise().mean(); + var_arr = X_arr.square().colwise().mean() - mean_arr.square().transpose(); } template -void ColwiseMoments( +CAFFE2_EXPORT void ColwiseMoments( const int rows, const int cols, const T* X, T* mean, T* variance) { - ConstEigenArrayMap X_mat(X, cols, rows); + std::memset(mean, 0, sizeof(T) * cols); + std::memset(variance, 0, sizeof(T) * cols); + ConstEigenArrayMap X_arr(X, cols, rows); EigenVectorArrayMap mean_arr(mean, cols); - EigenVectorArrayMap variance_arr(variance, cols); - mean_arr = X_mat.rowwise().mean(); - variance_arr = X_mat.array().square().rowwise().mean(); - variance_arr -= mean_arr.square(); + EigenVectorArrayMap var_arr(variance, cols); + // Eigen rowwise reduction is about 10 times slower than this for-loop. + for (int i = 0; i < rows; ++i) { + mean_arr += X_arr.col(i); + var_arr += X_arr.col(i).square(); + } + const T scale = T(1) / static_cast(rows); + mean_arr *= scale; + var_arr = var_arr * scale - mean_arr.square(); } template -void BothEndsMoments( +CAFFE2_EXPORT void BothEndsMoments( const int pre, const int mid, const int nxt, const T* X, T* mean, T* variance) { + std::memset(mean, 0, sizeof(T) * mid); + std::memset(variance, 0, sizeof(T) * mid); EigenVectorArrayMap mean_arr(mean, mid); - EigenVectorArrayMap variance_arr(variance, mid); - mean_arr = ConstEigenArrayMap(X, nxt, mid).colwise().mean(); - variance_arr = ConstEigenArrayMap(X, nxt, mid).square().colwise().mean(); - const int stride = mid * nxt; - const T* X_ptr = X + stride; - for (int i = 1; i < pre; ++i) { - mean_arr += ConstEigenArrayMap(X_ptr, nxt, mid).colwise().mean(); - variance_arr += - ConstEigenArrayMap(X_ptr, nxt, mid).square().colwise().mean(); - X_ptr += stride; - } - if (pre > 1) { - mean_arr /= static_cast(pre); - variance_arr /= static_cast(pre); + EigenVectorArrayMap var_arr(variance, mid); + ConstEigenArrayMap X_arr(X, nxt, pre * mid); + for (int i = 0; i < pre; ++i) { + for (int j = 0; j < mid; ++j) { + const int c = i * mid + j; + mean_arr(j) += X_arr.col(c).sum(); + var_arr(j) += X_arr.col(c).square().sum(); + } } - variance_arr -= mean_arr.square(); + const T scale = T(1) / static_cast(pre * nxt); + mean_arr *= scale; + var_arr = var_arr * scale - mean_arr.square(); } template -void MomentsImpl( +CAFFE2_EXPORT void MomentsImpl( const int num_dims, const int* dims, const int num_axes, @@ -1592,15 +1596,13 @@ void MomentsImpl( const int Y_size = std::accumulate(Y_dims, Y_dims + num_dims, 1, std::multiplies()); if (X_size == 0) { - if (Y_size > 0) { - memset(mean, 0, sizeof(T) * Y_size); - memset(variance, 0, sizeof(T) * Y_size); - } + std::memset(mean, 0, sizeof(T) * Y_size); + std::memset(variance, 0, sizeof(T) * Y_size); return; } if (std::equal(X_dims, X_dims + num_dims, Y_dims)) { - memcpy(mean, X, sizeof(T) * Y_size); - memset(variance, 0, sizeof(T) * Y_size); + std::memcpy(mean, X, sizeof(T) * Y_size); + std::memset(variance, 0, sizeof(T) * Y_size); return; } int rows; @@ -1630,17 +1632,18 @@ void MomentsImpl( utils::IncreaseIndexInDims(num_dims, dims, index.data()); } const T scale = static_cast(Y_size) / static_cast(X_size); - Scale(Y_size, scale, mean, mean, context); - EigenVectorArrayMap variance_arr(variance, Y_size); - variance_arr = - variance_arr * scale - ConstEigenVectorArrayMap(mean, Y_size).square(); + EigenVectorArrayMap mean_arr(mean, Y_size); + EigenVectorArrayMap var_arr(variance, Y_size); + mean_arr *= scale; + var_arr = + var_arr * scale - ConstEigenVectorArrayMap(mean, Y_size).square(); } } // namespace #define CAFFE2_SPECIALIZED_MOMENTS(T) \ template <> \ - void Moments( \ + CAFFE2_EXPORT void Moments( \ const int num_dims, \ const int* dims, \ const int num_axes, \ @@ -1655,9 +1658,23 @@ void MomentsImpl( CAFFE2_SPECIALIZED_MOMENTS(float) #undef CAFFE2_SPECIALIZED_MOMENTS +#define CAFFE2_SPECIALIZED_INV_STD(T) \ + template <> \ + void InvStd( \ + const int N, \ + const T epsilon, \ + const T* var, \ + T* inv_std, \ + CPUContext* context) { \ + EigenVectorArrayMap(inv_std, N) = \ + (ConstEigenVectorArrayMap(var, N) + epsilon).rsqrt(); \ + } +CAFFE2_SPECIALIZED_INV_STD(float) +#undef CAFFE2_SPECIALIZED_INV_STD + #define CAFFE2_SPECIALIZED_ROWWISEMAX(T) \ template <> \ - void RowwiseMax( \ + CAFFE2_EXPORT void RowwiseMax( \ const int N, const int D, const T* x, T* y, CPUContext*) { \ EigenVectorMap(y, N) = \ ConstEigenMatrixMap(x, D, N).colwise().maxCoeff(); \ @@ -1667,7 +1684,7 @@ CAFFE2_SPECIALIZED_ROWWISEMAX(float) #define CAFFE2_SPECIALIZED_COLWISEMAX(T) \ template <> \ - void ColwiseMax( \ + CAFFE2_EXPORT void ColwiseMax( \ const int N, const int D, const T* x, T* y, CPUContext*) { \ EigenVectorMap(y, D) = \ ConstEigenMatrixMap(x, D, N).rowwise().maxCoeff(); \ @@ -1677,7 +1694,7 @@ CAFFE2_SPECIALIZED_COLWISEMAX(float) #define CAFFE2_SPECIALIZED_ELEMWISEMAX(T) \ template <> \ - void ElemwiseMax( \ + CAFFE2_EXPORT void ElemwiseMax( \ const int N, const T* x, const T* y, T* z, CPUContext* /*context*/) { \ std::transform(x, x + N, y, z, [](const T& x_i, const T& y_i) { \ return std::max(x_i, y_i); \ @@ -1688,7 +1705,7 @@ CAFFE2_SPECIALIZED_ELEMWISEMAX(float) #define CAFFE2_SPECIALIZED_MAXIMUM(T) \ template <> \ - void Maximum( \ + CAFFE2_EXPORT void Maximum( \ const int N, const float alpha, const T* x, T* y, CPUContext* context) { \ std::transform( \ x, x + N, y, [&alpha](const T& x_i) { return std::max(x_i, alpha); }); \ @@ -1701,7 +1718,7 @@ CAFFE2_SPECIALIZED_MAXIMUM(float) #define DELEGATE_EIGEN_2D_BROADCAST_1ST_BINARY_FUNCTION(T, Func, expr) \ template <> \ - void Rowwise##Func( \ + CAFFE2_EXPORT void Rowwise##Func( \ const int rows, \ const int cols, \ const T* A, \ @@ -1718,7 +1735,7 @@ CAFFE2_SPECIALIZED_MAXIMUM(float) } \ } \ template <> \ - void Colwise##Func( \ + CAFFE2_EXPORT void Colwise##Func( \ const int rows, \ const int cols, \ const T* A, \ @@ -1738,7 +1755,7 @@ CAFFE2_SPECIALIZED_MAXIMUM(float) #define DELEGATE_EIGEN_2D_BROADCAST_2ND_BINARY_FUNCTION(T, Func, expr) \ template <> \ - void Rowwise##Func( \ + CAFFE2_EXPORT void Rowwise##Func( \ const int rows, \ const int cols, \ const T* A, \ @@ -1755,7 +1772,7 @@ CAFFE2_SPECIALIZED_MAXIMUM(float) } \ } \ template <> \ - void Colwise##Func( \ + CAFFE2_EXPORT void Colwise##Func( \ const int rows, \ const int cols, \ const T* A, \ @@ -1791,7 +1808,7 @@ DEFINE_EIGEN_2D_BROADCAST_BINARY_FUNCTION(Mul, *) #define DEFINE_EIGEN_2D_BROADCAST_SUB_FUNCTION(T) \ template <> \ - void RowwiseSub( \ + CAFFE2_EXPORT void RowwiseSub( \ const int rows, \ const int cols, \ const T* A, \ @@ -1803,7 +1820,7 @@ DEFINE_EIGEN_2D_BROADCAST_BINARY_FUNCTION(Mul, *) ConstEigenVectorArrayMap(A, cols); \ } \ template <> \ - void ColwiseSub( \ + CAFFE2_EXPORT void ColwiseSub( \ const int rows, \ const int cols, \ const T* A, \ @@ -1825,7 +1842,7 @@ DEFINE_EIGEN_2D_BROADCAST_SUB_FUNCTION(std::int64_t) #define DEFINE_EIGEN_2D_BROADCAST_DIV_FUNCTION(T) \ template <> \ - void RowwiseDiv( \ + CAFFE2_EXPORT void RowwiseDiv( \ const int rows, \ const int cols, \ const T* A, \ @@ -1837,7 +1854,7 @@ DEFINE_EIGEN_2D_BROADCAST_SUB_FUNCTION(std::int64_t) ConstEigenVectorArrayMap(A, cols); \ } \ template <> \ - void ColwiseDiv( \ + CAFFE2_EXPORT void ColwiseDiv( \ const int rows, \ const int cols, \ const T* A, \ @@ -1861,7 +1878,7 @@ DELEGATE_EIGEN_2D_BROADCAST_2ND_BINARY_FUNCTION(std::int64_t, Div, /) #undef DELEGATE_EIGEN_2D_BROADCAST_2ND_BINARY_FUNCTION template <> -void Not( +CAFFE2_EXPORT void Not( const int N, const bool* x, bool* y, @@ -1876,7 +1893,7 @@ void Not( #define CAFFE2_SPECIALIZED_CPU_ADD_STRIPED_BATCH(T) \ template <> \ - void AddStripedBatch( \ + CAFFE2_EXPORT void AddStripedBatch( \ const int N, \ const T* first, \ T* y, \ @@ -1894,7 +1911,7 @@ CAFFE2_SPECIALIZED_CPU_ADD_STRIPED_BATCH(float); namespace { template -void RowwiseBinaryOp( +CAFFE2_EXPORT void RowwiseBinaryOp( const int rows, const int cols, const BinaryOperator& op, @@ -1912,7 +1929,7 @@ void RowwiseBinaryOp( } template -void ColwiseBinaryOp( +CAFFE2_EXPORT void ColwiseBinaryOp( const int rows, const int cols, const BinaryOperator& op, @@ -1930,7 +1947,7 @@ void ColwiseBinaryOp( } template -void BroadcastBinaryOpImpl( +CAFFE2_EXPORT void BroadcastBinaryOpImpl( const int ndim, const int* A_dims, const int* B_dims, @@ -1954,7 +1971,7 @@ void BroadcastBinaryOpImpl( #define DELEGATE_1D_BINARY_FUNCTION(TIn, TOut, Func, Op) \ template <> \ - void Func( \ + CAFFE2_EXPORT void Func( \ const int N, const TIn* A, const TIn* B, TOut* C, CPUContext*) { \ std::transform(A, A + N, B, C, Op()); \ } @@ -1994,7 +2011,7 @@ DEFINE_1D_BITWISE_BINARY_FUNCTION(BitwiseXor, std::bit_xor) #define DELEGATE_2D_BROADCAST_BINARY_FUNCTION(TIn, TOut, Func, Op) \ template <> \ - void Rowwise##Func( \ + CAFFE2_EXPORT void Rowwise##Func( \ const int rows, \ const int cols, \ const TIn* A, \ @@ -2004,7 +2021,7 @@ DEFINE_1D_BITWISE_BINARY_FUNCTION(BitwiseXor, std::bit_xor) RowwiseBinaryOp, true>(rows, cols, Op(), A, B, C); \ } \ template <> \ - void Rowwise##Func( \ + CAFFE2_EXPORT void Rowwise##Func( \ const int rows, \ const int cols, \ const TIn* A, \ @@ -2015,7 +2032,7 @@ DEFINE_1D_BITWISE_BINARY_FUNCTION(BitwiseXor, std::bit_xor) rows, cols, Op(), A, B, C); \ } \ template <> \ - void Colwise##Func( \ + CAFFE2_EXPORT void Colwise##Func( \ const int rows, \ const int cols, \ const TIn* A, \ @@ -2025,7 +2042,7 @@ DEFINE_1D_BITWISE_BINARY_FUNCTION(BitwiseXor, std::bit_xor) ColwiseBinaryOp, true>(rows, cols, Op(), A, B, C); \ } \ template <> \ - void Colwise##Func( \ + CAFFE2_EXPORT void Colwise##Func( \ const int rows, \ const int cols, \ const TIn* A, \ @@ -2071,7 +2088,7 @@ DEFINE_2D_BROADCAST_BITWISE_BINARY_FUNCTION(BitwiseXor, std::bit_xor) #define DEFINE_2D_BROADCAST_1ST_DIV_FUNCTION(T) \ template <> \ - void RowwiseDiv( \ + CAFFE2_EXPORT void RowwiseDiv( \ const int rows, \ const int cols, \ const T* A, \ @@ -2082,7 +2099,7 @@ DEFINE_2D_BROADCAST_BITWISE_BINARY_FUNCTION(BitwiseXor, std::bit_xor) rows, cols, std::divides(), A, B, C); \ } \ template <> \ - void ColwiseDiv( \ + CAFFE2_EXPORT void ColwiseDiv( \ const int rows, \ const int cols, \ const T* A, \ @@ -2098,7 +2115,7 @@ DEFINE_2D_BROADCAST_1ST_DIV_FUNCTION(std::int64_t) #define DELEGATE_BROADCAST_BINARY_FUNCTION(TIn, TOut, Func, Op) \ template <> \ - void Func( \ + CAFFE2_EXPORT void Func( \ const int A_ndim, \ const int* A_dims, \ const int B_ndim, \ @@ -2241,7 +2258,7 @@ DEFINE_BROADCAST_BITWISE_BINARY_FUNCTION(BitwiseXor, std::bit_xor) #define CAFFE2_RAND_UNIFORM_REAL(T) \ template <> \ - void RandUniform( \ + CAFFE2_EXPORT void RandUniform( \ const size_t n, const T a, const T b, T* r, CPUContext* context) { \ std::uniform_real_distribution distribution(a, b); \ for (size_t i = 0; i < n; ++i) { \ @@ -2254,7 +2271,7 @@ CAFFE2_RAND_UNIFORM_REAL(double); #define CAFFE2_RAND_UNIFORM_CHAR(T) \ template <> \ - void RandUniform( \ + CAFFE2_EXPORT void RandUniform( \ const size_t n, const T a, const T b, T* r, CPUContext* context) { \ std::uniform_int_distribution distribution((short)a, (short)b); \ for (size_t i = 0; i < n; ++i) { \ @@ -2267,7 +2284,7 @@ CAFFE2_RAND_UNIFORM_CHAR(uint8_t); #define CAFFE2_RAND_UNIFORM_INT(T) \ template <> \ - void RandUniform( \ + CAFFE2_EXPORT void RandUniform( \ const size_t n, const T a, const T b, T* r, CPUContext* context) { \ std::uniform_int_distribution distribution(a, b); \ for (size_t i = 0; i < n; ++i) { \ @@ -2293,7 +2310,7 @@ CAFFE2_RAND_UNIFORM_INT(uint64_t); // each value. #define CAFFE2_RAND_FIXED_SUM(T) \ template <> \ - void RandFixedSum( \ + CAFFE2_EXPORT void RandFixedSum( \ const size_t n, \ const T a, \ const T b, \ @@ -2387,7 +2404,7 @@ Ind_t generate_stack_distance( } template -void generate_trace_lru( +CAFFE2_EXPORT void generate_trace_lru( std::vector& uni_ref, std::vector& cum_val, std::vector& cum_dis, @@ -2464,7 +2481,7 @@ void generate_trace_lru( // case we need to know the table id, to sample from the right distribution #define CAFFE2_RAND_SYNTHETIC_DATA(T) \ template <> \ - void RandSyntheticData( \ + CAFFE2_EXPORT void RandSyntheticData( \ const size_t n, const T a, const T b, T* r, CPUContext* context) { \ /* unique memory references */ \ std::vector mem_ref = {1, 2, 3, 4, 5, 6}; \ @@ -2503,7 +2520,7 @@ CAFFE2_RAND_SYNTHETIC_DATA(uint64_t); #define CAFFE2_SPECIALIZED_RAND_UNIFORM_UNIQUE(T) \ template <> \ - void RandUniformUnique( \ + CAFFE2_EXPORT void RandUniformUnique( \ const size_t n, \ const T a, \ const T b, \ @@ -2516,7 +2533,7 @@ CAFFE2_RAND_SYNTHETIC_DATA(uint64_t); std::unordered_set avoid_set(n); \ if (m) { \ avoid_set.insert(avoid, avoid + m); \ - CAFFE_ENFORCE_EQ(m, avoid_set.size(), "Avoid should be unique"); \ + CAFFE_ENFORCE_EQ(m, avoid_set.size(), "ACAFFE2_EXPORT void should be unique"); \ } \ std::uniform_int_distribution distribution(a, b); \ T v = 0; \ @@ -2534,7 +2551,7 @@ CAFFE2_SPECIALIZED_RAND_UNIFORM_UNIQUE(int64_t); #undef CAFFE2_SPECIALIZED_RAND_UNIFORM_UNIQUE template <> -void RandGaussian( +CAFFE2_EXPORT void RandGaussian( const size_t n, const float mean, const float std, @@ -2548,7 +2565,7 @@ void RandGaussian( #define CAFFE2_SPECIALIZED_SUM(T) \ template <> \ - void Sum( \ + CAFFE2_EXPORT void Sum( \ const int N, \ const T* x, \ T* y, \ @@ -2564,7 +2581,7 @@ CAFFE2_SPECIALIZED_SUM(int64_t); #undef CAFFE2_SPECIALIZED_SUM template <> -void SumSqr( +CAFFE2_EXPORT void SumSqr( const int N, const float* x, float* y, @@ -2574,7 +2591,7 @@ void SumSqr( } template <> -void Select( +CAFFE2_EXPORT void Select( const int N, const int D, const float* x, @@ -2588,7 +2605,7 @@ void Select( } template <> -void CopyMatrix( +CAFFE2_EXPORT void CopyMatrix( const size_t itemsize, const int M, const int N, @@ -2631,7 +2648,7 @@ void CopyMatrix( #define DELEGATE_COPY_MATRIX_FUNCTION(T, Func) \ template <> \ - void CopyMatrix( \ + CAFFE2_EXPORT void CopyMatrix( \ const int M, \ const int N, \ const T* A, \ @@ -2642,7 +2659,7 @@ void CopyMatrix( Func('R', 'N', M, N, T(1), A, lda, B, ldb); \ } \ template <> \ - void CopyMatrix( \ + CAFFE2_EXPORT void CopyMatrix( \ const int M, \ const int N, \ const T* A, \ @@ -2673,7 +2690,7 @@ DELEGATE_COPY_MATRIX_FUNCTION(double, mkl_domatcopy) #define CAFFE2_SPECIALIZED_COPY_MATRIX(T) \ template <> \ - void CopyMatrix( \ + CAFFE2_EXPORT void CopyMatrix( \ const int M, \ const int N, \ const T* A, \ @@ -2703,7 +2720,7 @@ DELEGATE_COPY_MATRIX_FUNCTION(double, mkl_domatcopy) } \ } \ template <> \ - void CopyMatrix( \ + CAFFE2_EXPORT void CopyMatrix( \ const int M, \ const int N, \ const T* A, \ @@ -2742,7 +2759,7 @@ CAFFE2_SPECIALIZED_COPY_MATRIX(std::uint16_t) namespace { template -void Im2ColZeroPaddingAndNoDilationNCHW( +CAFFE2_EXPORT void Im2ColZeroPaddingAndNoDilationNCHW( const int C, const int H, const int W, @@ -2789,7 +2806,7 @@ void Im2ColZeroPaddingAndNoDilationNCHW( } template -void Col2ImZeroPaddingAndNoDilationNCHW( +CAFFE2_EXPORT void Col2ImZeroPaddingAndNoDilationNCHW( const int C, const int H, const int W, @@ -2825,7 +2842,7 @@ void Col2ImZeroPaddingAndNoDilationNCHW( } template -void Im2ColZeroPaddingAndNoDilationNHWC( +CAFFE2_EXPORT void Im2ColZeroPaddingAndNoDilationNHWC( const int C, const int H, const int W, @@ -2850,7 +2867,7 @@ void Im2ColZeroPaddingAndNoDilationNHWC( } template -void Col2ImZeroPaddingAndNoDilationNHWC( +CAFFE2_EXPORT void Col2ImZeroPaddingAndNoDilationNHWC( const int C, const int H, const int W, @@ -2877,7 +2894,7 @@ void Col2ImZeroPaddingAndNoDilationNHWC( } template -void Im2ColNdNCHWImpl( +CAFFE2_EXPORT void Im2ColNdNCHWImpl( const int N, const int img_size, const int col_size, @@ -2933,7 +2950,7 @@ void Im2ColNdNCHWImpl( } // namespace template <> -void Im2ColNd( +CAFFE2_EXPORT void Im2ColNd( const int N, const int img_size, const int col_size, @@ -2961,7 +2978,7 @@ void Im2ColNd( } template <> -void Col2ImNd( +CAFFE2_EXPORT void Col2ImNd( const int N, const int img_size, const int col_size, @@ -2989,7 +3006,7 @@ void Col2ImNd( } template <> -void Im2Col( +CAFFE2_EXPORT void Im2Col( const int C, const int H, const int W, @@ -3055,7 +3072,7 @@ void Im2Col( } template <> -void Im2Col( +CAFFE2_EXPORT void Im2Col( const int C, const int H, const int W, @@ -3155,7 +3172,7 @@ void Im2Col( } template <> -void Col2Im( +CAFFE2_EXPORT void Col2Im( const int C, const int H, const int W, @@ -3222,7 +3239,7 @@ void Col2Im( } template <> -void Col2Im( +CAFFE2_EXPORT void Col2Im( const int C, const int H, const int W, @@ -3318,7 +3335,7 @@ void Col2Im( } template <> -void BiasCHW( +CAFFE2_EXPORT void BiasCHW( const float* bias, const float* /*bias_multiplier*/, const int bias_channels, @@ -3403,7 +3420,7 @@ void BiasCHW( #define CAFFE2_SPECIALIZED_COPYVECTOR(T) \ template <> \ - void CopyVector( \ + CAFFE2_EXPORT void CopyVector( \ const int N, const T* src, T* dst, CPUContext* /*context*/) { \ if (src != dst && N > 0) { \ memcpy(dst, src, sizeof(T) * N); \ @@ -3459,13 +3476,13 @@ bool TransposeWithHPTT( #endif // CAFFE2_USE_HPTT template -void Tranpose2D(const int rows, const int cols, const T* X, T* Y); +void Transpose2D(const int rows, const int cols, const T* X, T* Y); #ifdef CAFFE2_USE_MKL #define DELEGATE_TRANSPOSE_2D_FUNCTION(T, Func) \ template <> \ - void Tranpose2D(const int rows, const int cols, const T* X, T* Y) { \ + void Transpose2D(const int rows, const int cols, const T* X, T* Y) { \ Func('R', 'T', rows, cols, T(1), X, cols, Y, rows); \ } DELEGATE_TRANSPOSE_2D_FUNCTION(float, mkl_somatcopy); @@ -3476,7 +3493,7 @@ DELEGATE_TRANSPOSE_2D_FUNCTION(double, mkl_domatcopy); #define CAFFE2_SPECIALIZED_TRANSPOSE_2D(T) \ template <> \ - void Tranpose2D(const int rows, const int cols, const T* X, T* Y) { \ + void Transpose2D(const int rows, const int cols, const T* X, T* Y) { \ EigenMatrixMap(Y, rows, cols) = \ ConstEigenMatrixMap(X, cols, rows).transpose(); \ } @@ -3484,7 +3501,7 @@ DELEGATE_TRANSPOSE_2D_FUNCTION(double, mkl_domatcopy); #ifndef CAFFE2_USE_MKL template <> -void Tranpose2D( +void Transpose2D( const int rows, const int cols, const float* X, @@ -3581,7 +3598,7 @@ void TransposeCPUImpl( return; } if (ndim == 2) { - Tranpose2D(dims[0], dims[1], X, Y); + Transpose2D(dims[0], dims[1], X, Y); } else { TransposeND(ndim, dims, axes, X, Y); } @@ -3601,7 +3618,7 @@ void TransposeCPUImpl( return; } if (ndim == 2) { - Tranpose2D(dims[0], dims[1], X, Y); + Transpose2D(dims[0], dims[1], X, Y); } else { #ifdef CAFFE2_USE_HPTT if (TransposeWithHPTT(ndim, dims, axes, X, Y)) { @@ -3616,7 +3633,7 @@ void TransposeCPUImpl( #define CAFFE2_SPECIALIZED_TRANSPOSE(T) \ template <> \ - void Transpose( \ + CAFFE2_EXPORT void Transpose( \ const int ndim, \ const int* dims, \ const int* axes, \ @@ -3636,5 +3653,50 @@ CAFFE2_SPECIALIZED_TRANSPOSE(std::uint8_t) CAFFE2_SPECIALIZED_TRANSPOSE(std::uint16_t) #undef CAFFE2_SPECIALIZED_TRANSPOSE +#define CAFFE2_SPECIALIZED_AFFINE_CHANNEL(T) \ + template <> \ + void AffineChannel( \ + const int N, \ + const int C, \ + const int HxW, \ + const T* X, \ + const T* scale, \ + const T* bias, \ + T* Y, \ + CPUContext* /* context */) { \ + ConstEigenVectorArrayMap scale_arr(scale, C); \ + ConstEigenVectorArrayMap bias_arr(bias, C); \ + const int stride = C * HxW; \ + const T* X_ptr = X; \ + T* Y_ptr = Y; \ + for (int i = 0; i < N; ++i) { \ + EigenArrayMap(Y_ptr, HxW, C) = \ + (ConstEigenArrayMap(X_ptr, HxW, C).rowwise() * \ + scale_arr.transpose()) \ + .rowwise() + \ + bias_arr.transpose(); \ + X_ptr += stride; \ + Y_ptr += stride; \ + } \ + } \ + template <> \ + void AffineChannel( \ + const int N, \ + const int C, \ + const int HxW, \ + const T* X, \ + const T* scale, \ + const T* bias, \ + T* Y, \ + CPUContext* /* context */) { \ + EigenArrayMap(Y, C, N * HxW) = \ + (ConstEigenArrayMap(X, C, N * HxW).colwise() * \ + ConstEigenVectorArrayMap(scale, C)) \ + .colwise() + \ + ConstEigenVectorArrayMap(bias, C); \ + } +CAFFE2_SPECIALIZED_AFFINE_CHANNEL(float) +#undef CAFFE2_SPECIALIZED_AFFINE_CHANNEL + } // namespace math } // namespace caffe2 diff --git a/caffe2/utils/math_gpu.cu b/caffe2/utils/math_gpu.cu index 44b230c134c7a..920bffc0ae3bf 100644 --- a/caffe2/utils/math_gpu.cu +++ b/caffe2/utils/math_gpu.cu @@ -3573,6 +3573,7 @@ __global__ void RowwiseMomentsCUDAKernel( T* variance) { __shared__ typename BlockReduce::TempStorage m_storage; __shared__ typename BlockReduce::TempStorage v_storage; + const T scale = T(1) / static_cast(cols); for (int i = blockIdx.x; i < rows; i += gridDim.x) { T m_val = 0; T v_val = 0; @@ -3586,11 +3587,12 @@ __global__ void RowwiseMomentsCUDAKernel( v_val += X[X_index] * X[X_index]; #endif } - m_val = BlockReduce(m_storage).Reduce(m_val, cub::Sum()); - v_val = BlockReduce(v_storage).Reduce(v_val, cub::Sum()); + m_val = BlockReduce(m_storage).Sum(m_val); + v_val = BlockReduce(v_storage).Sum(v_val); if (threadIdx.x == 0) { - mean[i] = m_val / static_cast(cols); - variance[i] = v_val / static_cast(cols) - mean[i] * mean[i]; + const T mu = m_val * scale; + mean[i] = mu; + variance[i] = v_val * scale - mu * mu; } __syncthreads(); } @@ -3605,6 +3607,7 @@ __global__ void ColwiseMomentsCUDAKernel( T* variance) { __shared__ typename BlockReduce::TempStorage m_storage; __shared__ typename BlockReduce::TempStorage v_storage; + const T scale = T(1) / static_cast(rows); for (int i = blockIdx.x; i < cols; i += gridDim.x) { T m_val = 0; T v_val = 0; @@ -3618,11 +3621,12 @@ __global__ void ColwiseMomentsCUDAKernel( v_val += X[X_index] * X[X_index]; #endif } - m_val = BlockReduce(m_storage).Reduce(m_val, cub::Sum()); - v_val = BlockReduce(v_storage).Reduce(v_val, cub::Sum()); + m_val = BlockReduce(m_storage).Sum(m_val); + v_val = BlockReduce(v_storage).Sum(v_val); if (threadIdx.x == 0) { - mean[i] = m_val / static_cast(rows); - variance[i] = v_val / static_cast(rows) - mean[i] * mean[i]; + const T mu = m_val * scale; + mean[i] = mu; + variance[i] = v_val * scale - mu * mu; } __syncthreads(); } @@ -3639,6 +3643,7 @@ __global__ void MomentsCUDAKernel( T* variance) { __shared__ typename BlockReduce::TempStorage m_storage; __shared__ typename BlockReduce::TempStorage v_storage; + const T scale = T(1) / static_cast(inner_size); for (int i = blockIdx.x; i < outer_size; i += gridDim.x) { T m_val = 0; T v_val = 0; @@ -3659,11 +3664,12 @@ __global__ void MomentsCUDAKernel( v_val += X[X_index] * X[X_index]; #endif } - m_val = BlockReduce(m_storage).Reduce(m_val, cub::Sum()); - v_val = BlockReduce(v_storage).Reduce(v_val, cub::Sum()); + m_val = BlockReduce(m_storage).Sum(m_val); + v_val = BlockReduce(v_storage).Sum(v_val); if (threadIdx.x == 0) { - mean[i] = m_val / static_cast(inner_size); - variance[i] = v_val / static_cast(inner_size) - mean[i] * mean[i]; + const T mu = m_val * scale; + mean[i] = mu; + variance[i] = v_val * scale - mu * mu; } __syncthreads(); } @@ -3704,7 +3710,7 @@ void MomentsCUDA( T* variance, CUDAContext* context) { CAFFE_ENFORCE_LE(num_axes, num_dims); - std::vector Y_dims_vector(num_dims); + std::vector Y_dims_vector(dims, dims + num_dims); for (int i = 0; i < num_axes; ++i) { Y_dims_vector[axes[i]] = 1; } @@ -3794,6 +3800,42 @@ CAFFE2_SPECIALIZED_CUDA_MOMENTS(float) namespace { +template +__global__ void +InvStdCUDAKernel(const int N, const T epsilon, const T* var, T* inv_std); + +#define DELEGATE_INV_STD_KERNEL_FUNCTION(T, Func) \ + template <> \ + __global__ void InvStdCUDAKernel( \ + const int N, const T epsilon, const T* var, T* inv_std) { \ + CUDA_1D_KERNEL_LOOP(i, N) { \ + inv_std[i] = Func(var[i] + epsilon); \ + } \ + } +DELEGATE_INV_STD_KERNEL_FUNCTION(float, rsqrtf) +#undef DELEGATE_INV_STD_KERNEL_FUNCTION + +} // namespace + +#define CAFFE2_SPECIALIZED_CUDA_INV_STD(T) \ + template <> \ + void InvStd( \ + const int N, \ + const T epsilon, \ + const T* var, \ + T* inv_std, \ + CUDAContext* context) { \ + InvStdCUDAKernel \ + <<cuda_stream()>>>(N, epsilon, var, inv_std); \ + } +CAFFE2_SPECIALIZED_CUDA_INV_STD(float) +#undef CAFFE2_SPECIALIZED_CUDA_INV_STD + +namespace { + template __global__ void TransposeCUDAKernel( const int size, @@ -3866,5 +3908,50 @@ CAFFE2_SPECIALIZED_CUDA_TRANSPOSE(int) CAFFE2_SPECIALIZED_CUDA_TRANSPOSE(TIndex) #undef CAFFE2_SPECIALIZED_CUDA_TRANSPOSE +namespace { + +template +__global__ void AffineChannelCUDAKernel( + const int size, + const int C, + const int HxW, + const T* X, + const T* scale, + const T* bias, + T* Y) { + CUDA_1D_KERNEL_LOOP(i, size) { + const int c = kOrder == StorageOrder::NCHW ? i / HxW % C : i % C; +#if __CUDA_ARCH__ >= 350 + Y[i] = __ldg(scale + c) * __ldg(X + i) + __ldg(bias + c); +#else + Y[i] = scale[c] * X[i] + bias[c]; +#endif + } +} + +} // namespace + +#define CAFFE2_SPECIALIZED_CUDA_AFFINE_CHANNEL(T, kOrder) \ + template <> \ + void AffineChannel( \ + const int N, \ + const int C, \ + const int HxW, \ + const T* X, \ + const T* scale, \ + const T* bias, \ + T* Y, \ + CUDAContext* context) { \ + const int size = N * C * HxW; \ + AffineChannelCUDAKernel \ + <<cuda_stream()>>>(size, C, HxW, X, scale, bias, Y); \ + } +CAFFE2_SPECIALIZED_CUDA_AFFINE_CHANNEL(float, StorageOrder::NCHW) +CAFFE2_SPECIALIZED_CUDA_AFFINE_CHANNEL(float, StorageOrder::NHWC) +#undef CAFFE2_SPECIALIZED_CUDA_AFFINE_CHANNEL + } // namespace math } // namespace caffe2 diff --git a/caffe2/utils/math_gpu_test.cc b/caffe2/utils/math_gpu_test.cc index 5619c317525fa..b9f09706f9862 100644 --- a/caffe2/utils/math_gpu_test.cc +++ b/caffe2/utils/math_gpu_test.cc @@ -33,7 +33,7 @@ void executeGpuBinaryOpTest( return; Workspace ws; DeviceOption option; - option.set_device_type(CUDA); + option.set_device_type(PROTO_CUDA); CUDAContext context(option); Blob* blobx0 = ws.CreateBlob("X0"); @@ -85,7 +85,7 @@ TEST(MathUtilGPUTest, testAddStripedBatch) { return; Workspace ws; DeviceOption option; - option.set_device_type(CUDA); + option.set_device_type(PROTO_CUDA); CUDAContext context(option); Blob* blobx = ws.CreateBlob("X"); Blob* bloby = ws.CreateBlob("Y"); @@ -253,7 +253,7 @@ class GemmBatchedGPUTest if (!HasCudaGPU()) { return; } - option_.set_device_type(CUDA); + option_.set_device_type(PROTO_CUDA); cuda_context_ = make_unique(option_); Blob* X_blob = ws_.CreateBlob("X"); Blob* W_blob = ws_.CreateBlob("W"); @@ -377,7 +377,7 @@ class ReduceTensorGPUTest : public testing::Test { if (!HasCudaGPU()) { return; } - option_.set_device_type(CUDA); + option_.set_device_type(PROTO_CUDA); cuda_context_ = make_unique(option_); Blob* blob_x = ws_.CreateBlob("X"); Blob* blob_y = ws_.CreateBlob("Y"); @@ -660,7 +660,7 @@ class BroadcastGPUTest : public testing::Test { if (!HasCudaGPU()) { return; } - option_.set_device_type(CUDA); + option_.set_device_type(PROTO_CUDA); cuda_context_ = make_unique(option_); Blob* blob_x = ws_.CreateBlob("X"); Blob* blob_y = ws_.CreateBlob("Y"); @@ -736,7 +736,7 @@ class MomentsGPUTest : public testing::Test { if (!HasCudaGPU()) { return; } - option_.set_device_type(CUDA); + option_.set_device_type(PROTO_CUDA); cuda_context_ = make_unique(option_); Blob* blob_x = ws_.CreateBlob("X"); Blob* blob_mean = ws_.CreateBlob("mean"); @@ -864,7 +864,7 @@ class TransposeGPUTest : public testing::Test { if (!HasCudaGPU()) { return; } - option_.set_device_type(CUDA); + option_.set_device_type(PROTO_CUDA); cuda_context_ = make_unique(option_); Blob* blob_x = ws_.CreateBlob("X"); Blob* blob_y = ws_.CreateBlob("Y"); diff --git a/caffe2/utils/proto_utils.cc b/caffe2/utils/proto_utils.cc index b1c25328cba12..1daacff3eda2f 100644 --- a/caffe2/utils/proto_utils.cc +++ b/caffe2/utils/proto_utils.cc @@ -27,13 +27,13 @@ CAFFE2_EXPORT std::string DeviceTypeName(const int32_t& d) { CAFFE2_EXPORT int DeviceId(const DeviceOption& option) { switch (option.device_type()) { - case CPU: + case PROTO_CPU: return option.numa_node_id(); - case CUDA: + case PROTO_CUDA: return option.cuda_gpu_id(); - case MKLDNN: + case PROTO_MKLDNN: return option.numa_node_id(); - case HIP: + case PROTO_HIP: return option.hip_gpu_id(); default: CAFFE_THROW("Unknown device id for device type: ", option.device_type()); diff --git a/caffe2/utils/proto_utils_test.cc b/caffe2/utils/proto_utils_test.cc index 645e08019c65d..c9f37f4c98c29 100644 --- a/caffe2/utils/proto_utils_test.cc +++ b/caffe2/utils/proto_utils_test.cc @@ -15,8 +15,8 @@ TEST(ProtoUtilsTest, IsSameDevice) { EXPECT_FALSE(IsSameDevice(a, b)); a.set_cuda_gpu_id(2); EXPECT_TRUE(IsSameDevice(a, b)); - a.set_device_type(DeviceType::CUDA); - b.set_device_type(DeviceType::CPU); + a.set_device_type(DeviceTypeProto::PROTO_CUDA); + b.set_device_type(DeviceTypeProto::PROTO_CPU); EXPECT_FALSE(IsSameDevice(a, b)); } diff --git a/cmake/Caffe2Config.cmake.in b/cmake/Caffe2Config.cmake.in index 7e7ede760e264..7995dbc6ac467 100644 --- a/cmake/Caffe2Config.cmake.in +++ b/cmake/Caffe2Config.cmake.in @@ -106,6 +106,8 @@ if (@USE_CUDA@) endif() endif() +include("${CMAKE_CURRENT_LIST_DIR}/public/mkl.cmake") + # import targets include ("${CMAKE_CURRENT_LIST_DIR}/Caffe2Targets.cmake") diff --git a/cmake/Codegen.cmake b/cmake/Codegen.cmake index bb42109b770f6..ff838c58889e4 100644 --- a/cmake/Codegen.cmake +++ b/cmake/Codegen.cmake @@ -48,30 +48,41 @@ install(FILES ${CMAKE_BINARY_DIR}/caffe2/core/macros.h # ---[ ATen specific if (NOT BUILD_ATEN_MOBILE) + SET(OPT_FLAG "-O3 ") + IF(MSVC) + SET(OPT_FLAG "/Ox /fp:strict ") + ENDIF() + SET(VCOMP_LIB "vcomp") + + IF("${CMAKE_BUILD_TYPE}" MATCHES "Debug") + SET(OPT_FLAG " ") + SET(VCOMP_LIB "vcompd") + ENDIF() + # SET_SOURCE_FILES_PROPERTIES must be in the same CMakeLists.txt file as the target that includes the file # so we need to set these commands here rather than in src/TH IF(C_SSE4_1_FOUND AND C_SSE4_2_FOUND) IF(MSVC) - SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/generic/simd/convolve5x5_sse.cpp PROPERTIES COMPILE_FLAGS "${MSVC_OPT_FLAG}/fp:fast") + SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/generic/simd/convolve5x5_sse.cpp PROPERTIES COMPILE_FLAGS "${OPT_FLAG}/fp:fast") ELSE(MSVC) - SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/generic/simd/convolve5x5_sse.cpp PROPERTIES COMPILE_FLAGS "-O3 -ffast-math") + SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/generic/simd/convolve5x5_sse.cpp PROPERTIES COMPILE_FLAGS "${OPT_FLAG} -ffast-math") ENDIF(MSVC) ENDIF(C_SSE4_1_FOUND AND C_SSE4_2_FOUND) IF(C_AVX_FOUND) IF(MSVC) - SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/generic/simd/convolve5x5_avx.cpp PROPERTIES COMPILE_FLAGS "${MSVC_OPT_FLAG}/fp:fast ${CXX_AVX_FLAGS}") - SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/vector/AVX.cpp PROPERTIES COMPILE_FLAGS "${MSVC_OPT_FLAG}/arch:AVX ${CXX_AVX_FLAGS}") + SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/generic/simd/convolve5x5_avx.cpp PROPERTIES COMPILE_FLAGS "${OPT_FLAG}/fp:fast ${CXX_AVX_FLAGS}") + SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/vector/AVX.cpp PROPERTIES COMPILE_FLAGS "${OPT_FLAG}/arch:AVX ${CXX_AVX_FLAGS}") ELSE(MSVC) - SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/generic/simd/convolve5x5_avx.cpp PROPERTIES COMPILE_FLAGS "-O3 -ffast-math ${CXX_AVX_FLAGS}") - SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/vector/AVX.cpp PROPERTIES COMPILE_FLAGS "-O3 ${CXX_AVX_FLAGS}") + SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/generic/simd/convolve5x5_avx.cpp PROPERTIES COMPILE_FLAGS "${OPT_FLAG} -ffast-math ${CXX_AVX_FLAGS}") + SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/vector/AVX.cpp PROPERTIES COMPILE_FLAGS "${OPT_FLAG} ${CXX_AVX_FLAGS}") ENDIF(MSVC) ENDIF(C_AVX_FOUND) IF(C_AVX2_FOUND) IF(MSVC) - SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/vector/AVX2.cpp PROPERTIES COMPILE_FLAGS "${MSVC_OPT_FLAG}/arch:AVX2 ${CXX_AVX2_FLAGS}") + SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/vector/AVX2.cpp PROPERTIES COMPILE_FLAGS "${OPT_FLAG}/arch:AVX2 ${CXX_AVX2_FLAGS}") ELSE(MSVC) - SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/vector/AVX2.cpp PROPERTIES COMPILE_FLAGS "-O3 ${CXX_AVX2_FLAGS}") + SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_LIST_DIR}/../aten/src/TH/vector/AVX2.cpp PROPERTIES COMPILE_FLAGS "${OPT_FLAG} ${CXX_AVX2_FLAGS}") ENDIF(MSVC) ENDIF(C_AVX2_FOUND) @@ -81,28 +92,16 @@ if (NOT BUILD_ATEN_MOBILE) FILE(GLOB cpu_kernel_cpp_in "${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/cpu/*.cpp") - IF(MSVC AND NOT "${CMAKE_BUILD_TYPE}" MATCHES "Debug") - SET(MSVC_OPT_FLAG "/Ox /fp:strict ") - SET(VCOMP_LIB "vcomp") - ELSE() - SET(MSVC_OPT_FLAG " ") - SET(VCOMP_LIB "vcompd") - ENDIF() - LIST(APPEND CPU_CAPABILITY_NAMES "DEFAULT") - IF(MSVC) - LIST(APPEND CPU_CAPABILITY_FLAGS "${MSVC_OPT_FLAG}") - ELSE(MSVC) - LIST(APPEND CPU_CAPABILITY_FLAGS "-O3") - ENDIF(MSVC) + LIST(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG}") IF(CXX_AVX_FOUND) SET(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DHAVE_AVX_CPU_DEFINITION") LIST(APPEND CPU_CAPABILITY_NAMES "AVX") IF(MSVC) - LIST(APPEND CPU_CAPABILITY_FLAGS "${MSVC_OPT_FLAG}/arch:AVX") + LIST(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG}/arch:AVX") ELSE(MSVC) - LIST(APPEND CPU_CAPABILITY_FLAGS "-O3 -mavx") + LIST(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG} -mavx") ENDIF(MSVC) ENDIF(CXX_AVX_FOUND) @@ -110,9 +109,9 @@ if (NOT BUILD_ATEN_MOBILE) SET(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DHAVE_AVX2_CPU_DEFINITION") LIST(APPEND CPU_CAPABILITY_NAMES "AVX2") IF(MSVC) - LIST(APPEND CPU_CAPABILITY_FLAGS "${MSVC_OPT_FLAG}/arch:AVX2") + LIST(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG}/arch:AVX2") ELSE(MSVC) - LIST(APPEND CPU_CAPABILITY_FLAGS "-O3 -mavx2 -mfma") + LIST(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG} -mavx2 -mfma") ENDIF(MSVC) ENDIF(CXX_AVX2_FOUND) diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index 6d046d386c4d0..238dc3dff253f 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -83,9 +83,10 @@ elseif(BLAS STREQUAL "MKL") else() find_package(MKL QUIET) endif() + include(${CMAKE_CURRENT_LIST_DIR}/public/mkl.cmake) if(MKL_FOUND) include_directories(SYSTEM ${MKL_INCLUDE_DIR}) - list(APPEND Caffe2_PUBLIC_DEPENDENCY_LIBS ${MKL_LIBRARIES}) + list(APPEND Caffe2_PUBLIC_DEPENDENCY_LIBS caffe2::mkl) else() message(WARNING "MKL could not be found. Defaulting to Eigen") set(BLAS "Eigen" CACHE STRING "Selected BLAS library") @@ -756,9 +757,6 @@ if (USE_NNAPI AND NOT ANDROID) caffe2_update_option(USE_NNAPI OFF) endif() -# TODO(orionr): Enable all of this for Windows DLL when we -# can figure out how to get it to build -if (NOT (MSVC AND BUILD_SHARED_LIBS)) if (NOT BUILD_ATEN_MOBILE) if (CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO) list(APPEND Caffe2_DEPENDENCY_LIBS aten_op_header_gen) @@ -768,7 +766,6 @@ if (NOT BUILD_ATEN_MOBILE) include_directories(${PROJECT_BINARY_DIR}/caffe2/contrib/aten) endif() endif() -endif() if (USE_ZSTD) list(APPEND Caffe2_DEPENDENCY_LIBS libzstd_static) diff --git a/cmake/Modules/FindMKL.cmake b/cmake/Modules/FindMKL.cmake index faf2340b4205a..b296e5f2e47ae 100644 --- a/cmake/Modules/FindMKL.cmake +++ b/cmake/Modules/FindMKL.cmake @@ -310,7 +310,7 @@ if (USE_MKL AND USE_MKLML) set(CAFFE2_USE_MKL 1) endif() -if (USE_MKL AND USE_IDEEP) +if (USE_MKL AND USE_IDEEP) set(IDEEP_ROOT "${PROJECT_SOURCE_DIR}/third_party/ideep") set(MKLDNN_ROOT "${IDEEP_ROOT}/mkl-dnn") find_path(IDEEP_INCLUDE_DIR ideep.hpp PATHS ${IDEEP_ROOT} PATH_SUFFIXES include) @@ -319,7 +319,7 @@ if (USE_MKL AND USE_IDEEP) execute_process(COMMAND git submodule update --init mkl-dnn WORKING_DIRECTORY ${IDEEP_ROOT}) find_path(MKLDNN_INCLUDE_DIR mkldnn.hpp mkldnn.h PATHS ${MKLDNN_ROOT} PATH_SUFFIXES include) endif() - + if (MKLDNN_INCLUDE_DIR) list(APPEND IDEEP_INCLUDE_DIR ${MKLDNN_INCLUDE_DIR}) list(APPEND __ideep_looked_for MKLDNN_INCLUDE_DIR) diff --git a/cmake/ProtoBufPatch.cmake b/cmake/ProtoBufPatch.cmake index 91db791e9d075..f6f09b267ebdb 100644 --- a/cmake/ProtoBufPatch.cmake +++ b/cmake/ProtoBufPatch.cmake @@ -11,6 +11,13 @@ string( content "${content}") +string( + REPLACE + "PROTOBUF_CONSTEXPR" + "" + content + "${content}") + foreach(ns ${NAMESPACES}) # Insert "const ::std::string& GetEmptyStringAlreadyInited();" within # the namespace and make sure we only do it once in the file. Unfortunately diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake index 0e39a651a838b..846f4d6154b8b 100644 --- a/cmake/Summary.cmake +++ b/cmake/Summary.cmake @@ -81,10 +81,6 @@ function (caffe2_print_configuration_summary) message(STATUS " USE_FFMPEG : ${USE_FFMPEG}") message(STATUS " USE_GFLAGS : ${USE_GFLAGS}") message(STATUS " USE_GLOG : ${USE_GLOG}") - message(STATUS " USE_GLOO : ${USE_GLOO}") - if(${USE_GLOO}) - message(STATUS " USE_GLOO_IBVERBS : ${USE_GLOO_IBVERBS}") - endif() message(STATUS " USE_LEVELDB : ${USE_LEVELDB}") if(${USE_LEVELDB}) message(STATUS " LevelDB version : ${LEVELDB_VERSION}") @@ -102,7 +98,6 @@ function (caffe2_print_configuration_summary) message(STATUS " USE_IDEEP : ${USE_IDEEP}") endif() message(STATUS " USE_MOBILE_OPENGL : ${USE_MOBILE_OPENGL}") - message(STATUS " USE_MPI : ${USE_MPI}") message(STATUS " USE_NCCL : ${USE_NCCL}") if(${USE_NCCL}) message(STATUS " USE_SYSTEM_NCCL : ${USE_SYSTEM_NCCL}") @@ -125,7 +120,9 @@ function (caffe2_print_configuration_summary) message(STATUS " USE_ZMQ : ${USE_ZMQ}") message(STATUS " USE_DISTRIBUTED : ${USE_DISTRIBUTED}") if(${USE_DISTRIBUTED}) - message(STATUS " USE_DISTRIBUTED_MW : ${USE_DISTRIBUTED_MW}") + message(STATUS " USE_MPI : ${USE_MPI}") + message(STATUS " USE_GLOO : ${USE_GLOO}") + message(STATUS " USE_GLOO_IBVERBS : ${USE_GLOO_IBVERBS}") endif() message(STATUS " Public Dependencies : ${Caffe2_PUBLIC_DEPENDENCY_LIBS}") diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake index 72e6dd67a7b44..ee943791c2ad9 100644 --- a/cmake/public/LoadHIP.cmake +++ b/cmake/public/LoadHIP.cmake @@ -38,6 +38,13 @@ ELSE() SET(ROCBLAS_PATH $ENV{ROCBLAS_PATH}) ENDIF() +# ROCSPARSE_PATH +IF(NOT DEFINED ENV{ROCSPARSE_PATH}) + SET(ROCSPARSE_PATH ${ROCM_PATH}/rocsparse) +ELSE() + SET(ROCSPARSE_PATH $ENV{ROCSPARSE_PATH}) +ENDIF() + # ROCFFT_PATH IF(NOT DEFINED ENV{ROCFFT_PATH}) SET(ROCBLAS_PATH ${ROCM_PATH}/rocfft) @@ -47,7 +54,7 @@ ENDIF() # HIPSPARSE_PATH IF(NOT DEFINED ENV{HIPSPARSE_PATH}) - SET(HIPSPARSE_PATH ${ROCM_PATH}/hcsparse) + SET(HIPSPARSE_PATH ${ROCM_PATH}/hipsparse) ELSE() SET(HIPSPARSE_PATH $ENV{HIPSPARSE_PATH}) ENDIF() @@ -115,12 +122,14 @@ IF(HIP_FOUND) set(rocblas_DIR ${ROCBLAS_PATH}/lib/cmake/rocblas) set(rocfft_DIR ${ROCFFT_PATH}/lib/cmake/rocfft) set(hipsparse_DIR ${HIPSPARSE_PATH}/lib/cmake/hipsparse) + set(rocsparse_DIR ${ROCSPARSE_PATH}/lib/cmake/rocsparse) find_package(rocrand REQUIRED) find_package(hiprand REQUIRED) find_package(rocblas REQUIRED) find_package(rocfft REQUIRED) find_package(miopen REQUIRED) + #find_package(rocsparse REQUIRED) #find_package(hipsparse REQUIRED) # TODO: hip_hcc has an interface include flag "-hc" which is only @@ -132,6 +141,7 @@ IF(HIP_FOUND) # however currently it's just the lib name FIND_LIBRARY(PYTORCH_MIOPEN_LIBRARIES ${miopen_LIBRARIES} HINTS ${MIOPEN_PATH}/lib) FIND_LIBRARY(hiprand_LIBRARIES hiprand HINTS ${HIPRAND_PATH}/lib) + FIND_LIBRARY(rocsparse_LIBRARIES rocsparse HINTS ${ROCSPARSE_PATH}/lib) FIND_LIBRARY(hipsparse_LIBRARIES hipsparse HINTS ${HIPSPARSE_PATH}/lib) diff --git a/cmake/public/cuda.cmake b/cmake/public/cuda.cmake index 314f20b655ebf..5a4c4424ec082 100644 --- a/cmake/public/cuda.cmake +++ b/cmake/public/cuda.cmake @@ -178,7 +178,7 @@ add_library(caffe2::cudart INTERFACE IMPORTED) if(CAFFE2_STATIC_LINK_CUDA) set_property( TARGET caffe2::cudart PROPERTY INTERFACE_LINK_LIBRARIES - "${CUDA_TOOLKIT_ROOT_DIR}/lib64/libcudart_static.a") + "${CUDA_TOOLKIT_ROOT_DIR}/lib64/libcudart_static.a" rt) else() set_property( TARGET caffe2::cudart PROPERTY INTERFACE_LINK_LIBRARIES diff --git a/cmake/public/mkl.cmake b/cmake/public/mkl.cmake new file mode 100644 index 0000000000000..2fd0406de5308 --- /dev/null +++ b/cmake/public/mkl.cmake @@ -0,0 +1,9 @@ +find_package(MKL QUIET) + +add_library(caffe2::mkl INTERFACE IMPORTED) +set_property( + TARGET caffe2::mkl PROPERTY INTERFACE_INCLUDE_DIRECTORIES + ${MKL_INCLUDE_DIR}) +set_property( + TARGET caffe2::mkl PROPERTY INTERFACE_LINK_LIBRARIES + ${MKL_LIBRARIES}) diff --git a/conda/caffe2/full/build.sh b/conda/caffe2/full/build.sh index 1a5f93aad47e7..f6b1c6137e57b 100755 --- a/conda/caffe2/full/build.sh +++ b/conda/caffe2/full/build.sh @@ -61,6 +61,9 @@ CMAKE_ARGS+=("-DCMAKE_PREFIX_PATH=$PREFIX") mkdir -p build cd build cmake "${CMAKE_ARGS[@]}" $CONDA_CMAKE_ARGS $PYTHON_ARGS .. -make VERBOSE=1 "-j$(nproc)" +if [ -z "$MAX_JOBS" ]; then + MAX_JOBS=$(nproc) +fi +make VERBOSE=1 "-j${MAX_JOBS}" make install/fast diff --git a/conda/caffe2/normal/build.sh b/conda/caffe2/normal/build.sh index 9a9a7067d42fd..ee74c97ddaae8 100755 --- a/conda/caffe2/normal/build.sh +++ b/conda/caffe2/normal/build.sh @@ -27,10 +27,15 @@ cmake_args+=("-DCMAKE_PREFIX_PATH=$PREFIX") mkdir -p build cd build cmake "${cmake_args[@]}" $CAFFE2_CMAKE_ARGS .. -if [ "$(uname)" == 'Darwin' ]; then - make "-j$(sysctl -n hw.ncpu)" -else - make "-j$(nproc)" +if [ -z "$MAX_JOBS" ]; then + if [ "$(uname)" == 'Darwin' ]; then + MAX_JOBS=$(sysctl -n hw.ncpu) + else + MAX_JOBS=$(nproc) + fi fi +# Building with too many cores could cause OOM +MAX_JOBS=$(( ${MAX_JOBS} > 8 ? 8 : ${MAX_JOBS} )) +make "-j${MAX_JOBS}" make install/fast diff --git a/docker/caffe2/jenkins/common/install_rocm.sh b/docker/caffe2/jenkins/common/install_rocm.sh index 120db9d7452c1..c69d857118b2d 100644 --- a/docker/caffe2/jenkins/common/install_rocm.sh +++ b/docker/caffe2/jenkins/common/install_rocm.sh @@ -50,13 +50,6 @@ install_hip_thrust() { git clone --recursive https://github.com/ROCmSoftwarePlatform/cub-hip.git /data/Thrust/thrust/system/cuda/detail/cub-hip } -# This will be removed after merging an upcoming PR. -install_hcsparse() { - mkdir -p /opt/rocm/debians - curl https://s3.amazonaws.com/ossci-linux/hcsparse-master-907a505-Linux.deb -o /opt/rocm/debians/hcsparse.deb - dpkg -i /opt/rocm/debians/hcsparse.deb -} - # Install an updated version of rocRand that's PyTorch compatible. install_rocrand() { mkdir -p /opt/rocm/debians @@ -73,6 +66,21 @@ install_hipsparse() { dpkg -i /opt/rocm/debians/hipsparse.deb } +# Install custom hcc containing two compiler fixes relevant to PyTorch +install_customhcc() { + mkdir -p /opt/rocm/debians + curl https://s3.amazonaws.com/ossci-linux/hcc-1.2.18272-Linux.deb -o /opt/rocm/debians/hcc-1.2.18272-Linux.deb + curl https://s3.amazonaws.com/ossci-linux/hip_base-1.5.18276.deb -o /opt/rocm/debians/hip_base-1.5.18276.deb + curl https://s3.amazonaws.com/ossci-linux/hip_doc-1.5.18276.deb -o /opt/rocm/debians/hip_doc-1.5.18276.deb + curl https://s3.amazonaws.com/ossci-linux/hip_samples-1.5.18276.deb -o /opt/rocm/debians/hip_samples-1.5.18276.deb + curl https://s3.amazonaws.com/ossci-linux/hip_hcc-1.5.18276.deb -o /opt/rocm/debians/hip_hcc-1.5.18276.deb + dpkg -i /opt/rocm/debians/hcc-1.2.18272-Linux.deb + dpkg -i /opt/rocm/debians/hip_base-1.5.18276.deb + dpkg -i /opt/rocm/debians/hip_doc-1.5.18276.deb + dpkg -i /opt/rocm/debians/hip_samples-1.5.18276.deb + dpkg -i /opt/rocm/debians/hip_hcc-1.5.18276.deb +} + # Install Python packages depending on the base OS if [ -f /etc/lsb-release ]; then install_ubuntu @@ -85,5 +93,5 @@ fi install_hip_thrust install_rocrand -install_hcsparse install_hipsparse +install_customhcc diff --git a/docs/cpp/Doxyfile b/docs/cpp/Doxyfile index 715a013b61711..7625ce44a5a56 100644 --- a/docs/cpp/Doxyfile +++ b/docs/cpp/Doxyfile @@ -70,14 +70,6 @@ OUTPUT_DIRECTORY = build CREATE_SUBDIRS = NO -# If the ALLOW_UNICODE_NAMES tag is set to YES, doxygen will allow non-ASCII -# characters to appear in the names of generated files. If set to NO, non-ASCII -# characters will be escaped, for example _xE3_x81_x84 will be used for Unicode -# U+3044. -# The default value is: NO. - -ALLOW_UNICODE_NAMES = NO - # The OUTPUT_LANGUAGE tag is used to specify the language in which all # documentation generated by doxygen is written. Doxygen will use this # information to generate all constant output in the proper language. @@ -305,15 +297,6 @@ EXTENSION_MAPPING = MARKDOWN_SUPPORT = YES -# When the TOC_INCLUDE_HEADINGS tag is set to a non-zero value, all headings up -# to that level are automatically included in the table of contents, even if -# they do not have an id attribute. -# Note: This feature currently applies only to Markdown headings. -# Minimum value: 0, maximum value: 99, default value: 0. -# This tag requires that the tag MARKDOWN_SUPPORT is set to YES. - -TOC_INCLUDE_HEADINGS = 0 - # When enabled doxygen tries to link words that correspond to documented # classes, or namespaces to their corresponding documentation. Such a link can # be prevented in individual cases by putting a % sign in front of the word or @@ -364,13 +347,6 @@ IDL_PROPERTY_SUPPORT = YES DISTRIBUTE_GROUP_DOC = NO -# If one adds a struct or class to a group and this option is enabled, then also -# any nested class or struct is added to the same group. By default this option -# is disabled and one has to add nested compounds explicitly via \ingroup. -# The default value is: NO. - -GROUP_NESTED_COMPOUNDS = NO - # Set the SUBGROUPING tag to YES to allow class member groups of the same type # (for instance a group of public functions) to be put as a subgroup of that # type (e.g. under the Public Functions section). Set it to NO to prevent @@ -535,26 +511,12 @@ CASE_SENSE_NAMES = NO HIDE_SCOPE_NAMES = NO -# If the HIDE_COMPOUND_REFERENCE tag is set to NO (default) then doxygen will -# append additional text to a page's title, such as Class Reference. If set to -# YES the compound reference will be hidden. -# The default value is: NO. - -HIDE_COMPOUND_REFERENCE= NO - # If the SHOW_INCLUDE_FILES tag is set to YES then doxygen will put a list of # the files that are included by a file in the documentation of that file. # The default value is: YES. SHOW_INCLUDE_FILES = YES -# If the SHOW_GROUPED_MEMB_INC tag is set to YES then Doxygen will add for each -# grouped member an include statement to the documentation, telling the reader -# which file to include in order to use the member. -# The default value is: NO. - -SHOW_GROUPED_MEMB_INC = NO - # If the FORCE_LOCAL_INCLUDES tag is set to YES then doxygen will list include # files with double quotes in the documentation rather than with sharp brackets. # The default value is: NO. @@ -707,6 +669,8 @@ FILE_VERSION_FILTER = LAYOUT_FILE = +GENERATE_LATEX = NO + # The CITE_BIB_FILES tag can be used to specify one or more bib files containing # the reference definitions. This must be a list of .bib files. The .bib # extension is automatically appended if omitted. This requires the bibtex tool @@ -726,7 +690,7 @@ CITE_BIB_FILES = # messages are off. # The default value is: NO. -QUIET = NO +QUIET = YES # The WARNINGS tag can be used to turn on/off the warning messages that are # generated to standard error (stderr) by doxygen. If WARNINGS is set to YES @@ -760,12 +724,6 @@ WARN_IF_DOC_ERROR = YES WARN_NO_PARAMDOC = NO -# If the WARN_AS_ERROR tag is set to YES then doxygen will immediately stop when -# a warning is encountered. -# The default value is: NO. - -WARN_AS_ERROR = NO - # The WARN_FORMAT tag determines the format of the warning messages that doxygen # can produce. The string should contain the $file, $line, and $text tags, which # will be replaced by the file and line number from which the warning originated @@ -810,9 +768,7 @@ INPUT = ../../torch/csrc/api/include \ ../../aten/src/ATen/core/ScalarType.h \ ../../aten/src/ATen/cuda/CUDAGuard.h \ ../../aten/src/ATen/cuda/CUDAStream.h \ - ../../aten/src/ATen/cuda/CUDAHalf.h \ ../../aten/src/ATen/cuda/CUDAContext.h \ - ../../aten/src/ATen/cuda/PinnedMemoryAllocator.h \ ../../aten/src/ATen/cudnn/Descriptors.h \ ../../aten/src/ATen/cudnn/Handles.h \ ../../aten/src/ATen/cudnn/Types.h \ @@ -1210,17 +1166,6 @@ HTML_COLORSTYLE_GAMMA = 80 HTML_TIMESTAMP = NO -# If the HTML_DYNAMIC_MENUS tag is set to YES then the generated HTML -# documentation will contain a main index with vertical navigation menus that -# are dynamically created via Javascript. If disabled, the navigation index will -# consists of multiple levels of tabs that are statically embedded in every HTML -# page. Disable this option to support browsers that do not have Javascript, -# like the Qt help browser. -# The default value is: YES. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_DYNAMIC_MENUS = YES - # If the HTML_DYNAMIC_SECTIONS tag is set to YES then the generated HTML # documentation will contain sections that can be hidden and shown after the # page has loaded. @@ -1634,285 +1579,6 @@ EXTERNAL_SEARCH_ID = EXTRA_SEARCH_MAPPINGS = -#--------------------------------------------------------------------------- -# Configuration options related to the LaTeX output -#--------------------------------------------------------------------------- - -# If the GENERATE_LATEX tag is set to YES, doxygen will generate LaTeX output. -# The default value is: YES. - -GENERATE_LATEX = NO - -# The LATEX_OUTPUT tag is used to specify where the LaTeX docs will be put. If a -# relative path is entered the value of OUTPUT_DIRECTORY will be put in front of -# it. -# The default directory is: latex. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_OUTPUT = latex - -# The LATEX_CMD_NAME tag can be used to specify the LaTeX command name to be -# invoked. -# -# Note that when enabling USE_PDFLATEX this option is only used for generating -# bitmaps for formulas in the HTML output, but not in the Makefile that is -# written to the output directory. -# The default file is: latex. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_CMD_NAME = latex - -# The MAKEINDEX_CMD_NAME tag can be used to specify the command name to generate -# index for LaTeX. -# The default file is: makeindex. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -MAKEINDEX_CMD_NAME = makeindex - -# If the COMPACT_LATEX tag is set to YES, doxygen generates more compact LaTeX -# documents. This may be useful for small projects and may help to save some -# trees in general. -# The default value is: NO. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -COMPACT_LATEX = NO - -# The PAPER_TYPE tag can be used to set the paper type that is used by the -# printer. -# Possible values are: a4 (210 x 297 mm), letter (8.5 x 11 inches), legal (8.5 x -# 14 inches) and executive (7.25 x 10.5 inches). -# The default value is: a4. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -PAPER_TYPE = a4 - -# The EXTRA_PACKAGES tag can be used to specify one or more LaTeX package names -# that should be included in the LaTeX output. The package can be specified just -# by its name or with the correct syntax as to be used with the LaTeX -# \usepackage command. To get the times font for instance you can specify : -# EXTRA_PACKAGES=times or EXTRA_PACKAGES={times} -# To use the option intlimits with the amsmath package you can specify: -# EXTRA_PACKAGES=[intlimits]{amsmath} -# If left blank no extra packages will be included. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -EXTRA_PACKAGES = - -# The LATEX_HEADER tag can be used to specify a personal LaTeX header for the -# generated LaTeX document. The header should contain everything until the first -# chapter. If it is left blank doxygen will generate a standard header. See -# section "Doxygen usage" for information on how to let doxygen write the -# default header to a separate file. -# -# Note: Only use a user-defined header if you know what you are doing! The -# following commands have a special meaning inside the header: $title, -# $datetime, $date, $doxygenversion, $projectname, $projectnumber, -# $projectbrief, $projectlogo. Doxygen will replace $title with the empty -# string, for the replacement values of the other commands the user is referred -# to HTML_HEADER. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_HEADER = - -# The LATEX_FOOTER tag can be used to specify a personal LaTeX footer for the -# generated LaTeX document. The footer should contain everything after the last -# chapter. If it is left blank doxygen will generate a standard footer. See -# LATEX_HEADER for more information on how to generate a default footer and what -# special commands can be used inside the footer. -# -# Note: Only use a user-defined footer if you know what you are doing! -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_FOOTER = - -# The LATEX_EXTRA_STYLESHEET tag can be used to specify additional user-defined -# LaTeX style sheets that are included after the standard style sheets created -# by doxygen. Using this option one can overrule certain style aspects. Doxygen -# will copy the style sheet files to the output directory. -# Note: The order of the extra style sheet files is of importance (e.g. the last -# style sheet in the list overrules the setting of the previous ones in the -# list). -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_EXTRA_STYLESHEET = - -# The LATEX_EXTRA_FILES tag can be used to specify one or more extra images or -# other source files which should be copied to the LATEX_OUTPUT output -# directory. Note that the files will be copied as-is; there are no commands or -# markers available. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_EXTRA_FILES = - -# If the PDF_HYPERLINKS tag is set to YES, the LaTeX that is generated is -# prepared for conversion to PDF (using ps2pdf or pdflatex). The PDF file will -# contain links (just like the HTML output) instead of page references. This -# makes the output suitable for online browsing using a PDF viewer. -# The default value is: YES. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -PDF_HYPERLINKS = YES - -# If the USE_PDFLATEX tag is set to YES, doxygen will use pdflatex to generate -# the PDF file directly from the LaTeX files. Set this option to YES, to get a -# higher quality PDF documentation. -# The default value is: YES. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -USE_PDFLATEX = YES - -# If the LATEX_BATCHMODE tag is set to YES, doxygen will add the \batchmode -# command to the generated LaTeX files. This will instruct LaTeX to keep running -# if errors occur, instead of asking the user for help. This option is also used -# when generating formulas in HTML. -# The default value is: NO. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_BATCHMODE = NO - -# If the LATEX_HIDE_INDICES tag is set to YES then doxygen will not include the -# index chapters (such as File Index, Compound Index, etc.) in the output. -# The default value is: NO. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_HIDE_INDICES = NO - -# If the LATEX_SOURCE_CODE tag is set to YES then doxygen will include source -# code with syntax highlighting in the LaTeX output. -# -# Note that which sources are shown also depends on other settings such as -# SOURCE_BROWSER. -# The default value is: NO. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_SOURCE_CODE = NO - -# The LATEX_BIB_STYLE tag can be used to specify the style to use for the -# bibliography, e.g. plainnat, or ieeetr. See -# https://en.wikipedia.org/wiki/BibTeX and \cite for more info. -# The default value is: plain. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_BIB_STYLE = plain - -# If the LATEX_TIMESTAMP tag is set to YES then the footer of each generated -# page will contain the date and time when the page was generated. Setting this -# to NO can help when comparing the output of multiple runs. -# The default value is: NO. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_TIMESTAMP = NO - -#--------------------------------------------------------------------------- -# Configuration options related to the RTF output -#--------------------------------------------------------------------------- - -# If the GENERATE_RTF tag is set to YES, doxygen will generate RTF output. The -# RTF output is optimized for Word 97 and may not look too pretty with other RTF -# readers/editors. -# The default value is: NO. - -GENERATE_RTF = NO - -# The RTF_OUTPUT tag is used to specify where the RTF docs will be put. If a -# relative path is entered the value of OUTPUT_DIRECTORY will be put in front of -# it. -# The default directory is: rtf. -# This tag requires that the tag GENERATE_RTF is set to YES. - -RTF_OUTPUT = rtf - -# If the COMPACT_RTF tag is set to YES, doxygen generates more compact RTF -# documents. This may be useful for small projects and may help to save some -# trees in general. -# The default value is: NO. -# This tag requires that the tag GENERATE_RTF is set to YES. - -COMPACT_RTF = NO - -# If the RTF_HYPERLINKS tag is set to YES, the RTF that is generated will -# contain hyperlink fields. The RTF file will contain links (just like the HTML -# output) instead of page references. This makes the output suitable for online -# browsing using Word or some other Word compatible readers that support those -# fields. -# -# Note: WordPad (write) and others do not support links. -# The default value is: NO. -# This tag requires that the tag GENERATE_RTF is set to YES. - -RTF_HYPERLINKS = NO - -# Load stylesheet definitions from file. Syntax is similar to doxygen's config -# file, i.e. a series of assignments. You only have to provide replacements, -# missing definitions are set to their default value. -# -# See also section "Doxygen usage" for information on how to generate the -# default style sheet that doxygen normally uses. -# This tag requires that the tag GENERATE_RTF is set to YES. - -RTF_STYLESHEET_FILE = - -# Set optional variables used in the generation of an RTF document. Syntax is -# similar to doxygen's config file. A template extensions file can be generated -# using doxygen -e rtf extensionFile. -# This tag requires that the tag GENERATE_RTF is set to YES. - -RTF_EXTENSIONS_FILE = - -# If the RTF_SOURCE_CODE tag is set to YES then doxygen will include source code -# with syntax highlighting in the RTF output. -# -# Note that which sources are shown also depends on other settings such as -# SOURCE_BROWSER. -# The default value is: NO. -# This tag requires that the tag GENERATE_RTF is set to YES. - -RTF_SOURCE_CODE = NO - -#--------------------------------------------------------------------------- -# Configuration options related to the man page output -#--------------------------------------------------------------------------- - -# If the GENERATE_MAN tag is set to YES, doxygen will generate man pages for -# classes and files. -# The default value is: NO. - -GENERATE_MAN = NO - -# The MAN_OUTPUT tag is used to specify where the man pages will be put. If a -# relative path is entered the value of OUTPUT_DIRECTORY will be put in front of -# it. A directory man3 will be created inside the directory specified by -# MAN_OUTPUT. -# The default directory is: man. -# This tag requires that the tag GENERATE_MAN is set to YES. - -MAN_OUTPUT = man - -# The MAN_EXTENSION tag determines the extension that is added to the generated -# man pages. In case the manual section does not start with a number, the number -# 3 is prepended. The dot (.) at the beginning of the MAN_EXTENSION tag is -# optional. -# The default value is: .3. -# This tag requires that the tag GENERATE_MAN is set to YES. - -MAN_EXTENSION = .3 - -# The MAN_SUBDIR tag determines the name of the directory created within -# MAN_OUTPUT in which the man pages are placed. If defaults to man followed by -# MAN_EXTENSION with the initial . removed. -# This tag requires that the tag GENERATE_MAN is set to YES. - -MAN_SUBDIR = - -# If the MAN_LINKS tag is set to YES and doxygen generates man output, then it -# will generate one additional man file for each entity documented in the real -# man page(s). These additional files only source the real man page, but without -# them the man command would be unable to find the correct page. -# The default value is: NO. -# This tag requires that the tag GENERATE_MAN is set to YES. - -MAN_LINKS = NO - #--------------------------------------------------------------------------- # Configuration options related to the XML output #--------------------------------------------------------------------------- @@ -1940,83 +1606,6 @@ XML_OUTPUT = xml XML_PROGRAMLISTING = YES -#--------------------------------------------------------------------------- -# Configuration options related to the DOCBOOK output -#--------------------------------------------------------------------------- - -# If the GENERATE_DOCBOOK tag is set to YES, doxygen will generate Docbook files -# that can be used to generate PDF. -# The default value is: NO. - -GENERATE_DOCBOOK = NO - -# The DOCBOOK_OUTPUT tag is used to specify where the Docbook pages will be put. -# If a relative path is entered the value of OUTPUT_DIRECTORY will be put in -# front of it. -# The default directory is: docbook. -# This tag requires that the tag GENERATE_DOCBOOK is set to YES. - -DOCBOOK_OUTPUT = docbook - -# If the DOCBOOK_PROGRAMLISTING tag is set to YES, doxygen will include the -# program listings (including syntax highlighting and cross-referencing -# information) to the DOCBOOK output. Note that enabling this will significantly -# increase the size of the DOCBOOK output. -# The default value is: NO. -# This tag requires that the tag GENERATE_DOCBOOK is set to YES. - -DOCBOOK_PROGRAMLISTING = NO - -#--------------------------------------------------------------------------- -# Configuration options for the AutoGen Definitions output -#--------------------------------------------------------------------------- - -# If the GENERATE_AUTOGEN_DEF tag is set to YES, doxygen will generate an -# AutoGen Definitions (see http://autogen.sourceforge.net/) file that captures -# the structure of the code including all documentation. Note that this feature -# is still experimental and incomplete at the moment. -# The default value is: NO. - -GENERATE_AUTOGEN_DEF = NO - -#--------------------------------------------------------------------------- -# Configuration options related to the Perl module output -#--------------------------------------------------------------------------- - -# If the GENERATE_PERLMOD tag is set to YES, doxygen will generate a Perl module -# file that captures the structure of the code including all documentation. -# -# Note that this feature is still experimental and incomplete at the moment. -# The default value is: NO. - -GENERATE_PERLMOD = NO - -# If the PERLMOD_LATEX tag is set to YES, doxygen will generate the necessary -# Makefile rules, Perl scripts and LaTeX code to be able to generate PDF and DVI -# output from the Perl module output. -# The default value is: NO. -# This tag requires that the tag GENERATE_PERLMOD is set to YES. - -PERLMOD_LATEX = NO - -# If the PERLMOD_PRETTY tag is set to YES, the Perl module output will be nicely -# formatted so it can be parsed by a human reader. This is useful if you want to -# understand what is going on. On the other hand, if this tag is set to NO, the -# size of the Perl module output will be much smaller and Perl will parse it -# just the same. -# The default value is: YES. -# This tag requires that the tag GENERATE_PERLMOD is set to YES. - -PERLMOD_PRETTY = YES - -# The names of the make variables in the generated doxyrules.make file are -# prefixed with the string contained in PERLMOD_MAKEVAR_PREFIX. This is useful -# so different doxyrules.make files included by the same Makefile don't -# overwrite each other's variables. -# This tag requires that the tag GENERATE_PERLMOD is set to YES. - -PERLMOD_MAKEVAR_PREFIX = - #--------------------------------------------------------------------------- # Configuration options related to the preprocessor #--------------------------------------------------------------------------- @@ -2056,7 +1645,10 @@ SEARCH_INCLUDES = YES # preprocessor. # This tag requires that the tag SEARCH_INCLUDES is set to YES. -INCLUDE_PATH = +INCLUDE_PATH = ../../ \ + ../../torch/csrc/api/include/ \ + ../../aten/src/ \ + ../../build/aten/src/ # You can use the INCLUDE_FILE_PATTERNS tag to specify one or more wildcard # patterns (like *.h and *.hpp) to filter out the header-files in the @@ -2170,13 +1762,6 @@ CLASS_DIAGRAMS = YES MSCGEN_PATH = -# You can include diagrams made with dia in doxygen documentation. Doxygen will -# then run dia to produce the diagram and insert it in the documentation. The -# DIA_PATH tag allows you to specify the directory where the dia binary resides. -# If left empty dia is assumed to be found in the default search path. - -DIA_PATH = - # If set to YES the inheritance and collaboration graphs will hide inheritance # and usage relations if the target is undocumented or is not a class. # The default value is: YES. @@ -2383,30 +1968,6 @@ DOTFILE_DIRS = MSCFILE_DIRS = -# The DIAFILE_DIRS tag can be used to specify one or more directories that -# contain dia files that are included in the documentation (see the \diafile -# command). - -DIAFILE_DIRS = - -# When using plantuml, the PLANTUML_JAR_PATH tag should be used to specify the -# path where java can find the plantuml.jar file. If left blank, it is assumed -# PlantUML is not used or called during a preprocessing step. Doxygen will -# generate a warning when it encounters a \startuml command in this case and -# will not generate output for the diagram. - -PLANTUML_JAR_PATH = - -# When using plantuml, the PLANTUML_CFG_FILE tag can be used to specify a -# configuration file for plantuml. - -PLANTUML_CFG_FILE = - -# When using plantuml, the specified paths are searched for files specified by -# the !include statement in a plantuml block. - -PLANTUML_INCLUDE_PATH = - # The DOT_GRAPH_MAX_NODES tag can be used to set the maximum number of nodes # that will be shown in the graph. If the number of nodes in a graph becomes # larger than this value, doxygen will truncate the graph, which is visualized diff --git a/docs/cpp/check-doxygen.sh b/docs/cpp/check-doxygen.sh new file mode 100755 index 0000000000000..58e65b403e39b --- /dev/null +++ b/docs/cpp/check-doxygen.sh @@ -0,0 +1,34 @@ +#!/bin/bash + +set -ex + +ignore_warning() { + # Invert match to filter out $1. + grep -v "$1" doxygen-log.txt > temp.txt + mv temp.txt doxygen-log.txt +} + +# Run doxygen and log all output. +doxygen 2> original-doxygen-log.txt +cp original-doxygen-log.txt doxygen-log.txt + +echo "Original output" +cat original-doxygen-log.txt + +# Filter out some warnings. +ignore_warning "warning: no uniquely matching class member found for" +ignore_warning "warning: source ../../build/aten/src/ is not a readable file" +ignore_warning "warning: source ../../build/aten/src/ATen/Tensor.h is not a readable file" +ignore_warning "warning: source ../../build/aten/src/ATen/Functions.h is not a readable file" + +# Count the number of remaining warnings. +warnings=$(grep 'warning:' doxygen-log.txt | wc -l) + +if [[ $warnings != 0 ]]; then + echo "Filtered output" + cat doxygen-log.txt + rm -f doxygen-log.txt original-doxygen-log.txt + exit 1 +fi + +rm -f doxygen-log.txt original-doxygen-log.txt diff --git a/docs/cpp/conf.py b/docs/cpp/conf.py index aeea35a4f331e..c494decdf9dbb 100644 --- a/docs/cpp/conf.py +++ b/docs/cpp/conf.py @@ -23,7 +23,6 @@ import sys import textwrap -import torch import sphinx_rtd_theme # -- General configuration ------------------------------------------------ @@ -109,7 +108,7 @@ # # The short X.Y version. # TODO: change to [:2] at v1.0 -version = 'master (' + torch.__version__ + ' )' +version = 'master' # The full version, including alpha/beta/rc tags. # TODO: verify this works as expected release = 'master' diff --git a/docs/cpp/requirements.txt b/docs/cpp/requirements.txt index ef1bcc7e57448..153044aca6baa 100644 --- a/docs/cpp/requirements.txt +++ b/docs/cpp/requirements.txt @@ -1,4 +1,4 @@ sphinx>=1.7.5 --e git://github.com/snide/sphinx_rtd_theme.git#egg=sphinx_rtd_theme +sphinx_rtd_theme breathe exhale diff --git a/docs/source/torch.rst b/docs/source/torch.rst index 9a022bed9a6a2..18d21f0e2a1a9 100644 --- a/docs/source/torch.rst +++ b/docs/source/torch.rst @@ -266,6 +266,7 @@ Other Operations .. autofunction:: histc .. autofunction:: meshgrid .. autofunction:: renorm +.. autofunction:: tensordot .. autofunction:: trace .. autofunction:: tril .. autofunction:: triu diff --git a/scripts/build_android.sh b/scripts/build_android.sh index c7c843eac161a..6c353f68da5e4 100755 --- a/scripts/build_android.sh +++ b/scripts/build_android.sh @@ -98,8 +98,11 @@ cmake "$CAFFE2_ROOT" \ "${CMAKE_ARGS[@]}" # Cross-platform parallel build -if [ "$(uname)" == "Darwin" ]; then - cmake --build . -- "-j$(sysctl -n hw.ncpu)" -else - cmake --build . -- "-j$(nproc)" +if [ -z "$MAX_JOBS" ]; then + if [ "$(uname)" == 'Darwin' ]; then + MAX_JOBS=$(sysctl -n hw.ncpu) + else + MAX_JOBS=$(nproc) + fi fi +cmake --build . -- "-j${MAX_JOBS}" diff --git a/scripts/build_host_protoc.sh b/scripts/build_host_protoc.sh index 77944a3f0f979..cd37db3b31713 100755 --- a/scripts/build_host_protoc.sh +++ b/scripts/build_host_protoc.sh @@ -49,8 +49,11 @@ fi cmake "$CAFFE2_ROOT/third_party/protobuf/cmake" ${CMAKE_ARGS[@]} -if [ "$(uname)" == 'Darwin' ]; then - cmake --build . -- "-j$(sysctl -n hw.ncpu)" install -else - cmake --build . -- "-j$(nproc)" install +if [ -z "$MAX_JOBS" ]; then + if [ "$(uname)" == 'Darwin' ]; then + MAX_JOBS=$(sysctl -n hw.ncpu) + else + MAX_JOBS=$(nproc) + fi fi +cmake --build . -- "-j${MAX_JOBS}" install diff --git a/scripts/build_local.sh b/scripts/build_local.sh index f2c40a60024b3..8f95ffda131bc 100755 --- a/scripts/build_local.sh +++ b/scripts/build_local.sh @@ -65,7 +65,9 @@ else # Determine the number of CPUs to build with. # If the `CAFFE_MAKE_NCPUS` variable is not specified, use them all. - if [ -n "${CAFFE_MAKE_NCPUS}" ]; then + if [ -n "${MAX_JOBS}" ]; then + CAFFE_MAKE_NCPUS="$MAX_JOBS" + elif [ -n "${CAFFE_MAKE_NCPUS}" ]; then CAFFE_MAKE_NCPUS="$CAFFE_MAKE_NCPUS" elif [ "$(uname)" == 'Darwin' ]; then CAFFE_MAKE_NCPUS="$(sysctl -n hw.ncpu)" diff --git a/setup.py b/setup.py index 2481f2cfe0c2f..cf1b0c9696ccd 100644 --- a/setup.py +++ b/setup.py @@ -37,7 +37,7 @@ # disables NNPACK build # # NO_DISTRIBUTED -# disables THD (distributed) build +# disables distributed (c10d, gloo, mpi, etc.) build # # NO_SYSTEM_NCCL # disables use of system-wide nccl (we will use our submoduled @@ -188,7 +188,10 @@ def hotpatch_var(var): lib_path = os.path.join(cwd, "torch", "lib") third_party_path = os.path.join(cwd, "third_party") tmp_install_path = lib_path + "/tmp_install" +caffe2_build_dir = os.path.join(cwd, "build") +# lib/pythonx.x/site-packages rel_site_packages = distutils.sysconfig.get_python_lib(prefix='') +# full absolute path to the dir above full_site_packages = distutils.sysconfig.get_python_lib() @@ -465,6 +468,15 @@ def run(self): setuptools.command.develop.develop.run(self) self.create_compile_commands() + # Copy Caffe2's Python proto files (generated during the build with the + # protobuf python compiler) from the build folder to the root folder + # cp root/build/caffe2/proto/proto.py root/caffe2/proto/proto.py + for src in glob.glob( + os.path.join(caffe2_build_dir, 'caffe2', 'proto', '*.py')): + dst = os.path.join( + cwd, os.path.relpath(src, caffe2_build_dir)) + self.copy_file(src, dst) + def create_compile_commands(self): def load(filename): with open(filename) as f: @@ -593,39 +605,40 @@ def build_extensions(self): # platform dependent build folder created by the "build" command of # setuptools. Only the contents of this folder are installed in the # "install" command by default. - if FULL_CAFFE2: - # We only make this copy for Caffe2's pybind extensions - caffe2_pybind_exts = [ - 'caffe2.python.caffe2_pybind11_state', - 'caffe2.python.caffe2_pybind11_state_gpu', - 'caffe2.python.caffe2_pybind11_state_hip', - ] - i = 0 - while i < len(self.extensions): - ext = self.extensions[i] - if ext.name not in caffe2_pybind_exts: - i += 1 - continue - fullname = self.get_ext_fullname(ext.name) - filename = self.get_ext_filename(fullname) - - src = os.path.join(tmp_install_path, rel_site_packages, filename) - if not os.path.exists(src): - print("{} does not exist".format(src)) - del self.extensions[i] - else: - dst = os.path.join(os.path.realpath(self.build_lib), filename) - dst_dir = os.path.dirname(dst) - if not os.path.exists(dst_dir): - os.makedirs(dst_dir) - self.copy_file(src, dst) - i += 1 + # We only make this copy for Caffe2's pybind extensions + caffe2_pybind_exts = [ + 'caffe2.python.caffe2_pybind11_state', + 'caffe2.python.caffe2_pybind11_state_gpu', + 'caffe2.python.caffe2_pybind11_state_hip', + ] + i = 0 + while i < len(self.extensions): + ext = self.extensions[i] + if ext.name not in caffe2_pybind_exts: + i += 1 + continue + fullname = self.get_ext_fullname(ext.name) + filename = self.get_ext_filename(fullname) + print("\nCopying extension {}".format(ext.name)) + + src = os.path.join(tmp_install_path, rel_site_packages, filename) + if not os.path.exists(src): + print("{} does not exist".format(src)) + del self.extensions[i] + else: + dst = os.path.join(os.path.realpath(self.build_lib), filename) + print("Copying {} from {} to {}".format(ext.name, src, dst)) + dst_dir = os.path.dirname(dst) + if not os.path.exists(dst_dir): + os.makedirs(dst_dir) + self.copy_file(src, dst) + i += 1 distutils.command.build_ext.build_ext.build_extensions(self) def get_outputs(self): outputs = distutils.command.build_ext.build_ext.get_outputs(self) - if FULL_CAFFE2: - outputs.append(os.path.join(self.build_lib, "caffe2")) + outputs.append(os.path.join(self.build_lib, "caffe2")) + print("setup.py::get_outputs returning {}".format(outputs)) return outputs @@ -831,6 +844,7 @@ def run(self): "torch/csrc/jit/ivalue.cpp", "torch/csrc/jit/passes/onnx.cpp", "torch/csrc/jit/passes/onnx/fixup_onnx_loop.cpp", + "torch/csrc/jit/passes/onnx/prepare_division_for_onnx.cpp", "torch/csrc/jit/passes/onnx/peephole.cpp", "torch/csrc/jit/passes/to_batch.cpp", "torch/csrc/jit/python_arg_flatten.cpp", @@ -929,8 +943,8 @@ def run(self): rocm_include_path = '/opt/rocm/include' hcc_include_path = '/opt/rocm/hcc/include' rocblas_include_path = '/opt/rocm/rocblas/include' + hipsparse_include_path = '/opt/rocm/hipsparse/include' rocfft_include_path = '/opt/rocm/rocfft/include' - hipsparse_include_path = '/opt/rocm/hcsparse/include' hiprand_include_path = '/opt/rocm/hiprand/include' rocrand_include_path = '/opt/rocm/rocrand/include' hip_lib_path = '/opt/rocm/hip/lib' @@ -1007,10 +1021,7 @@ def make_relative_rpath(path): ################################################################################ extensions = [] -if FULL_CAFFE2: - packages = find_packages(exclude=('tools', 'tools.*')) -else: - packages = find_packages(exclude=('tools', 'tools.*', 'caffe2', 'caffe2.*')) +packages = find_packages(exclude=('tools', 'tools.*')) C = Extension("torch._C", libraries=main_libraries, sources=main_sources, @@ -1054,19 +1065,18 @@ def make_relative_rpath(path): ) extensions.append(THNVRTC) -if FULL_CAFFE2: - # If building Caffe2 python as well, these extensions are built by cmake - # copied manually in build_extensions() inside the build_ext implementaiton - extensions.append( - setuptools.Extension( - name=str('caffe2.python.caffe2_pybind11_state'), - sources=[]), - ) - extensions.append( - setuptools.Extension( - name=str('caffe2.python.caffe2_pybind11_state_gpu'), - sources=[]), - ) +# These extensions are built by cmake and copied manually in build_extensions() +# inside the build_ext implementaiton +extensions.append( + setuptools.Extension( + name=str('caffe2.python.caffe2_pybind11_state'), + sources=[]), +) +extensions.append( + setuptools.Extension( + name=str('caffe2.python.caffe2_pybind11_state_gpu'), + sources=[]), +) cmdclass = { 'create_version_file': create_version_file, @@ -1082,14 +1092,12 @@ def make_relative_rpath(path): } cmdclass.update(build_dep_cmds) -entry_points = {} -if FULL_CAFFE2: - entry_points = { - 'console_scripts': [ - 'convert-caffe2-to-onnx = caffe2.python.onnx.bin.conversion:caffe2_to_onnx', - 'convert-onnx-to-caffe2 = caffe2.python.onnx.bin.conversion:onnx_to_caffe2', - ] - } +entry_points = { + 'console_scripts': [ + 'convert-caffe2-to-onnx = caffe2.python.onnx.bin.conversion:caffe2_to_onnx', + 'convert-onnx-to-caffe2 = caffe2.python.onnx.bin.conversion:onnx_to_caffe2', + ] +} if __name__ == '__main__': setup( @@ -1132,6 +1140,9 @@ def make_relative_rpath(path): 'lib/include/torch/csrc/utils/*.h', 'lib/include/torch/csrc/cuda/*.h', 'lib/include/torch/torch.h', + ], + 'caffe2': [ + rel_site_packages + '/caffe2/**/*.py' ] }, ) diff --git a/test/cpp/api/jit.cpp b/test/cpp/api/jit.cpp new file mode 100644 index 0000000000000..8879d6f0007fd --- /dev/null +++ b/test/cpp/api/jit.cpp @@ -0,0 +1,31 @@ +#include + +#include +#include + +#include + +TEST_CASE("torch script") { + SECTION("multiple functions") { + auto module = torch::jit::compile(R"JIT( + def test_mul(a, b): + return a * b + def test_relu(a, b): + return torch.relu(a + b) + def test_while(a, i): + while i < 10: + a += a + i += 1 + return a + )JIT"); + auto a = torch::ones(1); + auto b = torch::ones(1); + + REQUIRE(1 == module->run_method("test_mul", a, b).toTensor().toCLong()); + + REQUIRE(2 == module->run_method("test_relu", a, b).toTensor().toCLong()); + + REQUIRE( + 0x200 == module->run_method("test_while", a, b).toTensor().toCLong()); + } +} diff --git a/test/cpp/api/optim.cpp b/test/cpp/api/optim.cpp index c7047a6c8d193..ab278180b12b3 100644 --- a/test/cpp/api/optim.cpp +++ b/test/cpp/api/optim.cpp @@ -35,8 +35,7 @@ bool test_optimizer_xor(Options options) { const int64_t kBatchSize = 4; const int64_t kMaximumNumberOfEpochs = 3000; - auto optimizer = OptimizerClass(std::vector(), options); - optimizer.add_parameters(model->parameters()); + OptimizerClass optimizer(model->parameters(), options); float running_loss = 1; int epoch = 0; @@ -152,6 +151,9 @@ TEST_CASE("Optim/BasicInterface") { REQUIRE(optimizer.size() == 0); optimizer.add_parameters(parameters); REQUIRE(optimizer.size() == parameters.size()); + for (size_t p = 0; p < parameters.size(); ++p) { + REQUIRE(optimizer.parameters()[p].allclose(parameters[p])); + } } { Linear linear(3, 4); diff --git a/test/expect/TestCudaSparse.test_print.expect b/test/expect/TestCudaSparse.test_print.expect new file mode 100644 index 0000000000000..cae534bd16e8d --- /dev/null +++ b/test/expect/TestCudaSparse.test_print.expect @@ -0,0 +1,294 @@ +# shape: torch.Size([]) +# nnz: 2 +# sparseDim: 0 +# indices shape: torch.Size([0, 2]) +# values shape: torch.Size([2]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0, 1]), + device='cuda:0', size=(), nnz=2, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 2), dtype=torch.int64) +# _values +tensor([0, 1], device='cuda:0', dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 1.]), + device='cuda:0', size=(), nnz=2, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 1.]), + device='cuda:0', size=(), nnz=2, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 4)), + values=tensor([0., 1., 0., 1.]), + device='cuda:0', size=(), nnz=4, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 2), dtype=torch.int64, + grad_fn=) +# _values +tensor([0., 1.], device='cuda:0', dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([0]) +# nnz: 10 +# sparseDim: 0 +# indices shape: torch.Size([0, 10]) +# values shape: torch.Size([10, 0]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + device='cuda:0', size=(0,), nnz=10, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 10), dtype=torch.int64) +# _values +tensor([], device='cuda:0', size=(10, 0), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + device='cuda:0', size=(0,), nnz=10, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + device='cuda:0', size=(0,), nnz=10, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 20)), + values=tensor([], size=(20, 0)), + device='cuda:0', size=(0,), nnz=20, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 10), dtype=torch.int64, + grad_fn=) +# _values +tensor([], device='cuda:0', size=(10, 0), dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([2]) +# nnz: 3 +# sparseDim: 0 +# indices shape: torch.Size([0, 3]) +# values shape: torch.Size([3, 2]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0, 0], + [0, 1], + [1, 1]]), + device='cuda:0', size=(2,), nnz=3, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int64) +# _values +tensor([[0, 0], + [0, 1], + [1, 1]], device='cuda:0', dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + device='cuda:0', size=(2,), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + device='cuda:0', size=(2,), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 6)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667], + [0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + device='cuda:0', size=(2,), nnz=6, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int64, + grad_fn=) +# _values +tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]], device='cuda:0', dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([100, 3]) +# nnz: 3 +# sparseDim: 1 +# indices shape: torch.Size([1, 3]) +# values shape: torch.Size([3, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([[0, 1, 2]]), + values=tensor([[0, 0, 0], + [0, 0, 1], + [1, 1, 1]]), + device='cuda:0', size=(100, 3), nnz=3, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([[0, 1, 2]], device='cuda:0') +# _values +tensor([[0, 0, 0], + [0, 0, 1], + [1, 1, 1]], device='cuda:0', dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([[0, 1, 2]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + device='cuda:0', size=(100, 3), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([[0, 1, 2]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + device='cuda:0', size=(100, 3), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([[0, 1, 2, 0, 1, 2]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778], + [0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + device='cuda:0', size=(100, 3), nnz=6, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([[0, 1, 2]], device='cuda:0', grad_fn=) +# _values +tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]], device='cuda:0', dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([100, 20, 3]) +# nnz: 0 +# sparseDim: 2 +# indices shape: torch.Size([2, 0]) +# values shape: torch.Size([0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + device='cuda:0', size=(100, 20, 3), nnz=0, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(2, 0), dtype=torch.int64) +# _values +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + device='cuda:0', size=(100, 20, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + device='cuda:0', size=(100, 20, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + device='cuda:0', size=(100, 20, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(2, 0), dtype=torch.int64, + grad_fn=) +# _values +tensor([], device='cuda:0', size=(0, 3), dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([10, 0, 3]) +# nnz: 3 +# sparseDim: 0 +# indices shape: torch.Size([0, 3]) +# values shape: torch.Size([3, 10, 0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=3, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int64) +# _values +tensor([], device='cuda:0', size=(3, 10, 0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 6)), + values=tensor([], size=(6, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=6, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int64, + grad_fn=) +# _values +tensor([], device='cuda:0', size=(3, 10, 0, 3), dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([10, 0, 3]) +# nnz: 0 +# sparseDim: 0 +# indices shape: torch.Size([0, 0]) +# values shape: torch.Size([0, 10, 0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=0, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 0), dtype=torch.int64) +# _values +tensor([], device='cuda:0', size=(0, 10, 0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 0), dtype=torch.int64, + grad_fn=) +# _values +tensor([], device='cuda:0', size=(0, 10, 0, 3), dtype=torch.float32, + grad_fn=) diff --git a/test/expect/TestCudaUncoalescedSparse.test_print.expect b/test/expect/TestCudaUncoalescedSparse.test_print.expect new file mode 100644 index 0000000000000..e6207c6da3d27 --- /dev/null +++ b/test/expect/TestCudaUncoalescedSparse.test_print.expect @@ -0,0 +1,294 @@ +# shape: torch.Size([]) +# nnz: 2 +# sparseDim: 0 +# indices shape: torch.Size([0, 2]) +# values shape: torch.Size([2]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0, 1]), + device='cuda:0', size=(), nnz=2, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 2), dtype=torch.int64) +# _values +tensor([0, 1], device='cuda:0', dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 1.]), + device='cuda:0', size=(), nnz=2, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 1.]), + device='cuda:0', size=(), nnz=2, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 4)), + values=tensor([0., 1., 0., 1.]), + device='cuda:0', size=(), nnz=4, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 2), dtype=torch.int64, + grad_fn=) +# _values +tensor([0., 1.], device='cuda:0', dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([0]) +# nnz: 10 +# sparseDim: 0 +# indices shape: torch.Size([0, 10]) +# values shape: torch.Size([10, 0]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + device='cuda:0', size=(0,), nnz=10, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 10), dtype=torch.int64) +# _values +tensor([], device='cuda:0', size=(10, 0), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + device='cuda:0', size=(0,), nnz=10, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + device='cuda:0', size=(0,), nnz=10, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 20)), + values=tensor([], size=(20, 0)), + device='cuda:0', size=(0,), nnz=20, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 10), dtype=torch.int64, + grad_fn=) +# _values +tensor([], device='cuda:0', size=(10, 0), dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([2]) +# nnz: 3 +# sparseDim: 0 +# indices shape: torch.Size([0, 3]) +# values shape: torch.Size([3, 2]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0, 0], + [0, 1], + [1, 1]]), + device='cuda:0', size=(2,), nnz=3, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int64) +# _values +tensor([[0, 0], + [0, 1], + [1, 1]], device='cuda:0', dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + device='cuda:0', size=(2,), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + device='cuda:0', size=(2,), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 6)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667], + [0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + device='cuda:0', size=(2,), nnz=6, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int64, + grad_fn=) +# _values +tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]], device='cuda:0', dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([100, 3]) +# nnz: 3 +# sparseDim: 1 +# indices shape: torch.Size([1, 3]) +# values shape: torch.Size([3, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([[0, 1, 0]]), + values=tensor([[0, 0, 0], + [0, 0, 1], + [1, 1, 1]]), + device='cuda:0', size=(100, 3), nnz=3, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([[0, 1, 0]], device='cuda:0') +# _values +tensor([[0, 0, 0], + [0, 0, 1], + [1, 1, 1]], device='cuda:0', dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([[0, 1, 0]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + device='cuda:0', size=(100, 3), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([[0, 1, 0]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + device='cuda:0', size=(100, 3), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([[0, 1, 0, 0, 1, 0]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778], + [0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + device='cuda:0', size=(100, 3), nnz=6, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([[0, 1, 0]], device='cuda:0', grad_fn=) +# _values +tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]], device='cuda:0', dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([100, 20, 3]) +# nnz: 0 +# sparseDim: 2 +# indices shape: torch.Size([2, 0]) +# values shape: torch.Size([0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + device='cuda:0', size=(100, 20, 3), nnz=0, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(2, 0), dtype=torch.int64) +# _values +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + device='cuda:0', size=(100, 20, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + device='cuda:0', size=(100, 20, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + device='cuda:0', size=(100, 20, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(2, 0), dtype=torch.int64, + grad_fn=) +# _values +tensor([], device='cuda:0', size=(0, 3), dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([10, 0, 3]) +# nnz: 3 +# sparseDim: 0 +# indices shape: torch.Size([0, 3]) +# values shape: torch.Size([3, 10, 0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=3, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int64) +# _values +tensor([], device='cuda:0', size=(3, 10, 0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=3, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 6)), + values=tensor([], size=(6, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=6, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 3), dtype=torch.int64, + grad_fn=) +# _values +tensor([], device='cuda:0', size=(3, 10, 0, 3), dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([10, 0, 3]) +# nnz: 0 +# sparseDim: 0 +# indices shape: torch.Size([0, 0]) +# values shape: torch.Size([0, 10, 0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=0, dtype=torch.int32, + layout=torch.sparse_coo) +# _indices +tensor([], device='cuda:0', size=(0, 0), dtype=torch.int64) +# _values +tensor([], device='cuda:0', size=(0, 10, 0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo, requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + device='cuda:0', size=(10, 0, 3), nnz=0, dtype=torch.float32, + layout=torch.sparse_coo, grad_fn=) +# _indices +tensor([], device='cuda:0', size=(0, 0), dtype=torch.int64, + grad_fn=) +# _values +tensor([], device='cuda:0', size=(0, 10, 0, 3), dtype=torch.float32, + grad_fn=) diff --git a/test/expect/TestJit.test_concat_fusion.expect b/test/expect/TestJit.test_concat_fusion_cuda.expect similarity index 100% rename from test/expect/TestJit.test_concat_fusion.expect rename to test/expect/TestJit.test_concat_fusion_cuda.expect diff --git a/test/expect/TestJit.test_cpp.expect b/test/expect/TestJit.test_cpp_cuda.expect similarity index 100% rename from test/expect/TestJit.test_cpp.expect rename to test/expect/TestJit.test_cpp_cuda.expect diff --git a/test/expect/TestJit.test_fuse_last_device.expect b/test/expect/TestJit.test_fuse_last_device_cuda.expect similarity index 100% rename from test/expect/TestJit.test_fuse_last_device.expect rename to test/expect/TestJit.test_fuse_last_device_cuda.expect diff --git a/test/expect/TestJit.test_fusion_distribute.expect b/test/expect/TestJit.test_fusion_distribute_cuda.expect similarity index 100% rename from test/expect/TestJit.test_fusion_distribute.expect rename to test/expect/TestJit.test_fusion_distribute_cuda.expect diff --git a/test/expect/TestJit.test_lstm_fusion_concat.expect b/test/expect/TestJit.test_lstm_fusion_concat_cuda.expect similarity index 100% rename from test/expect/TestJit.test_lstm_fusion_concat.expect rename to test/expect/TestJit.test_lstm_fusion_concat_cuda.expect diff --git a/test/expect/TestScript.test_loop_unroll_unused_counter.expect b/test/expect/TestScript.test_loop_unroll_unused_counter.expect index 8627e64448a84..292b251c75da3 100644 --- a/test/expect/TestScript.test_loop_unroll_unused_counter.expect +++ b/test/expect/TestScript.test_loop_unroll_unused_counter.expect @@ -4,7 +4,7 @@ graph(%x : Dynamic) { %3 : int = prim::TensorToNum(%x) %4 : int = prim::Constant[value=1]() %5 : int = prim::Constant[value=8]() - %6 : int = aten::div(%3, %5) + %6 : int = aten::floordiv(%3, %5) %7 : int = prim::Constant[value=8]() %8 : int = aten::mul(%6, %7) %9 : int = aten::sub(%3, %8) diff --git a/test/expect/TestScript.test_loop_unrolling.expect b/test/expect/TestScript.test_loop_unrolling.expect index c2982c3f4b3c8..29546b5c1e971 100644 --- a/test/expect/TestScript.test_loop_unrolling.expect +++ b/test/expect/TestScript.test_loop_unrolling.expect @@ -4,7 +4,7 @@ graph(%x : Dynamic) { %3 : int = prim::Constant[value=1]() %4 : int = prim::Constant[value=0]() %5 : int = prim::Constant[value=8]() - %6 : int = aten::div(%2, %5) + %6 : int = aten::floordiv(%2, %5) %7 : int = prim::Constant[value=8]() %8 : int = aten::mul(%6, %7) %9 : int = aten::sub(%2, %8) diff --git a/test/expect/TestScript.test_loop_unrolling_nested.expect b/test/expect/TestScript.test_loop_unrolling_nested.expect index b6e57f0bf9b5d..a107eb81f7291 100644 --- a/test/expect/TestScript.test_loop_unrolling_nested.expect +++ b/test/expect/TestScript.test_loop_unrolling_nested.expect @@ -8,7 +8,7 @@ graph(%x : Dynamic) { %8 : int = prim::Constant[value=1]() %9 : int = prim::Constant[value=0]() %10 : int = prim::Constant[value=8]() - %11 : int = aten::div(%7, %10) + %11 : int = aten::floordiv(%7, %10) %12 : int = prim::Constant[value=8]() %13 : int = aten::mul(%11, %12) %14 : int = aten::sub(%7, %13) diff --git a/test/expect/TestScript.test_onnx_export_script_truediv.expect b/test/expect/TestScript.test_onnx_export_script_truediv.expect new file mode 100644 index 0000000000000..35d45d5b62fdd --- /dev/null +++ b/test/expect/TestScript.test_onnx_export_script_truediv.expect @@ -0,0 +1,23 @@ +ModelProto { + producer_name: "pytorch" + domain: "" + doc_string: "" + graph: + GraphProto { + name: "torch-jit-export" + inputs: [{name: "x", type:Tensor dims: 1 2 3}] + outputs: [{name: "8", type:Tensor dims: 1 2 3}] + initializers: [] + nodes: [ + Node {type: "Constant", inputs: [], outputs: [1], attributes: [{ name: 'value', type: tensor, value:TensorProto shape: []}]}, + Node {type: "Constant", inputs: [], outputs: [2], attributes: [{ name: 'value', type: tensor, value:TensorProto shape: []}]}, + Node {type: "Shape", inputs: [x], outputs: [3], attributes: []}, + Node {type: "Gather", inputs: [3,2], outputs: [4], attributes: [{ name: 'axis', type: int, value: 0}]}, + Node {type: "Cast", inputs: [4], outputs: [5], attributes: [{ name: 'to', type: int, value: 1}]}, + Node {type: "Cast", inputs: [1], outputs: [6], attributes: [{ name: 'to', type: int, value: 1}]}, + Node {type: "Div", inputs: [5,6], outputs: [7], attributes: []}, + Node {type: "Add", inputs: [x,7], outputs: [8], attributes: []} + ] + } + opset_import: [OperatorSetIdProto { domain: }], +} diff --git a/test/expect/TestScript.test_onnx_raw_export_script_truediv.expect b/test/expect/TestScript.test_onnx_raw_export_script_truediv.expect new file mode 100644 index 0000000000000..2eb94a744dd04 --- /dev/null +++ b/test/expect/TestScript.test_onnx_raw_export_script_truediv.expect @@ -0,0 +1,25 @@ +ModelProto { + producer_name: "pytorch" + domain: "" + doc_string: "" + graph: + GraphProto { + name: "torch-jit-export" + inputs: [{name: "x", type:Tensor dims: 1 2 3}] + outputs: [{name: "10", type:Tensor dims: 1 2 3}] + initializers: [] + nodes: [ + Node {type: "Constant", inputs: [], outputs: [1], attributes: [{ name: 'value', type: tensor, value:TensorProto shape: []}]}, + Node {type: "Constant", inputs: [], outputs: [2], attributes: [{ name: 'value', type: tensor, value:TensorProto shape: []}]}, + Node {type: "size", inputs: [x,2], outputs: [3], attributes: []}, + Node {type: "Constant", inputs: [], outputs: [4], attributes: [{ name: 'value', type: tensor, value:TensorProto shape: []}]}, + Node {type: "_cast_Float", inputs: [3,4], outputs: [5], attributes: []}, + Node {type: "Constant", inputs: [], outputs: [6], attributes: [{ name: 'value', type: tensor, value:TensorProto shape: []}]}, + Node {type: "_cast_Float", inputs: [1,6], outputs: [7], attributes: []}, + Node {type: "div", inputs: [5,7], outputs: [z], attributes: []}, + Node {type: "Constant", inputs: [], outputs: [9], attributes: [{ name: 'value', type: tensor, value:TensorProto shape: []}]}, + Node {type: "add", inputs: [x,z,9], outputs: [10], attributes: []} + ] + } + opset_import: [OperatorSetIdProto { domain: }], +} diff --git a/test/expect/TestScript.test_scalar_fusion.expect b/test/expect/TestScript.test_scalar_fusion.expect new file mode 100644 index 0000000000000..9d45a9f765d63 --- /dev/null +++ b/test/expect/TestScript.test_scalar_fusion.expect @@ -0,0 +1,12 @@ +graph(%x : Float() + %y : Float()) { + %2 : Float() = prim::FusionGroup_0[device=-1](%x, %y) + return (%2); +} +with prim::FusionGroup_0 = graph(%0 : Float() + %1 : Float()) { + %2 : Float() = aten::type_as(%1, %0) + %3 : int = prim::Constant[value=1]() + %4 : Float() = aten::add(%0, %2, %3) + return (%4); +} diff --git a/test/expect/TestSparse.test_print.expect b/test/expect/TestSparse.test_print.expect new file mode 100644 index 0000000000000..c0223c1483741 --- /dev/null +++ b/test/expect/TestSparse.test_print.expect @@ -0,0 +1,262 @@ +# shape: torch.Size([]) +# nnz: 2 +# sparseDim: 0 +# indices shape: torch.Size([0, 2]) +# values shape: torch.Size([2]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0, 1]), + size=(), nnz=2, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 2), dtype=torch.int64) +# _values +tensor([0, 1], dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 1.]), + size=(), nnz=2, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 1.]), + size=(), nnz=2, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 2.]), + size=(), nnz=2, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 2), dtype=torch.int64, grad_fn=) +# _values +tensor([0., 1.], dtype=torch.float32, grad_fn=) + +# shape: torch.Size([0]) +# nnz: 10 +# sparseDim: 0 +# indices shape: torch.Size([0, 10]) +# values shape: torch.Size([10, 0]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + size=(0,), nnz=10, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 10), dtype=torch.int64) +# _values +tensor([], size=(10, 0), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + size=(0,), nnz=10, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + size=(0,), nnz=10, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + size=(0,), nnz=10, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 10), dtype=torch.int64, grad_fn=) +# _values +tensor([], size=(10, 0), dtype=torch.float32, grad_fn=) + +# shape: torch.Size([2]) +# nnz: 3 +# sparseDim: 0 +# indices shape: torch.Size([0, 3]) +# values shape: torch.Size([3, 2]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0, 0], + [0, 1], + [1, 1]]), + size=(2,), nnz=3, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 3), dtype=torch.int64) +# _values +tensor([[0, 0], + [0, 1], + [1, 1]], dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + size=(2,), nnz=3, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + size=(2,), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.6667], + [1.3333, 2.0000], + [2.6667, 3.3333]]), + size=(2,), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 3), dtype=torch.int64, grad_fn=) +# _values +tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]], dtype=torch.float32, grad_fn=) + +# shape: torch.Size([100, 3]) +# nnz: 3 +# sparseDim: 1 +# indices shape: torch.Size([1, 3]) +# values shape: torch.Size([3, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([[0, 1, 2]]), + values=tensor([[0, 0, 0], + [0, 0, 1], + [1, 1, 1]]), + size=(100, 3), nnz=3, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([[0, 1, 2]]) +# _values +tensor([[0, 0, 0], + [0, 0, 1], + [1, 1, 1]], dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([[0, 1, 2]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + size=(100, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([[0, 1, 2]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + size=(100, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([[0, 1, 2]]), + values=tensor([[0.0000, 0.4444, 0.8889], + [1.3333, 1.7778, 2.2222], + [2.6667, 3.1111, 3.5556]]), + size=(100, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([[0, 1, 2]], grad_fn=) +# _values +tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]], dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([100, 20, 3]) +# nnz: 0 +# sparseDim: 2 +# indices shape: torch.Size([2, 0]) +# values shape: torch.Size([0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + size=(100, 20, 3), nnz=0, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(2, 0), dtype=torch.int64) +# _values +tensor([], size=(0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + size=(100, 20, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + size=(100, 20, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + size=(100, 20, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(2, 0), dtype=torch.int64, grad_fn=) +# _values +tensor([], size=(0, 3), dtype=torch.float32, grad_fn=) + +# shape: torch.Size([10, 0, 3]) +# nnz: 3 +# sparseDim: 0 +# indices shape: torch.Size([0, 3]) +# values shape: torch.Size([3, 10, 0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + size=(10, 0, 3), nnz=3, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 3), dtype=torch.int64) +# _values +tensor([], size=(3, 10, 0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + size=(10, 0, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + size=(10, 0, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + size=(10, 0, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 3), dtype=torch.int64, grad_fn=) +# _values +tensor([], size=(3, 10, 0, 3), dtype=torch.float32, grad_fn=) + +# shape: torch.Size([10, 0, 3]) +# nnz: 0 +# sparseDim: 0 +# indices shape: torch.Size([0, 0]) +# values shape: torch.Size([0, 10, 0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + size=(10, 0, 3), nnz=0, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 0), dtype=torch.int64) +# _values +tensor([], size=(0, 10, 0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + size=(10, 0, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + size=(10, 0, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + size=(10, 0, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 0), dtype=torch.int64, grad_fn=) +# _values +tensor([], size=(0, 10, 0, 3), dtype=torch.float32, grad_fn=) diff --git a/test/expect/TestUncoalescedSparse.test_print.expect b/test/expect/TestUncoalescedSparse.test_print.expect new file mode 100644 index 0000000000000..244442de0cc73 --- /dev/null +++ b/test/expect/TestUncoalescedSparse.test_print.expect @@ -0,0 +1,262 @@ +# shape: torch.Size([]) +# nnz: 2 +# sparseDim: 0 +# indices shape: torch.Size([0, 2]) +# values shape: torch.Size([2]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0, 1]), + size=(), nnz=2, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 2), dtype=torch.int64) +# _values +tensor([0, 1], dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 1.]), + size=(), nnz=2, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 1.]), + size=(), nnz=2, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 2)), + values=tensor([0., 2.]), + size=(), nnz=2, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 2), dtype=torch.int64, grad_fn=) +# _values +tensor([0., 1.], dtype=torch.float32, grad_fn=) + +# shape: torch.Size([0]) +# nnz: 10 +# sparseDim: 0 +# indices shape: torch.Size([0, 10]) +# values shape: torch.Size([10, 0]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + size=(0,), nnz=10, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 10), dtype=torch.int64) +# _values +tensor([], size=(10, 0), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + size=(0,), nnz=10, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + size=(0,), nnz=10, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 10)), + values=tensor([], size=(10, 0)), + size=(0,), nnz=10, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 10), dtype=torch.int64, grad_fn=) +# _values +tensor([], size=(10, 0), dtype=torch.float32, grad_fn=) + +# shape: torch.Size([2]) +# nnz: 3 +# sparseDim: 0 +# indices shape: torch.Size([0, 3]) +# values shape: torch.Size([3, 2]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0, 0], + [0, 1], + [1, 1]]), + size=(2,), nnz=3, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 3), dtype=torch.int64) +# _values +tensor([[0, 0], + [0, 1], + [1, 1]], dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + size=(2,), nnz=3, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]]), + size=(2,), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 3)), + values=tensor([[0.0000, 0.6667], + [1.3333, 2.0000], + [2.6667, 3.3333]]), + size=(2,), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 3), dtype=torch.int64, grad_fn=) +# _values +tensor([[0.0000, 0.3333], + [0.6667, 1.0000], + [1.3333, 1.6667]], dtype=torch.float32, grad_fn=) + +# shape: torch.Size([100, 3]) +# nnz: 3 +# sparseDim: 1 +# indices shape: torch.Size([1, 3]) +# values shape: torch.Size([3, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([[0, 1, 0]]), + values=tensor([[0, 0, 0], + [0, 0, 1], + [1, 1, 1]]), + size=(100, 3), nnz=3, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([[0, 1, 0]]) +# _values +tensor([[0, 0, 0], + [0, 0, 1], + [1, 1, 1]], dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([[0, 1, 0]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + size=(100, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([[0, 1, 0]]), + values=tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]]), + size=(100, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([[0, 1, 0]]), + values=tensor([[0.0000, 0.4444, 0.8889], + [1.3333, 1.7778, 2.2222], + [2.6667, 3.1111, 3.5556]]), + size=(100, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([[0, 1, 0]], grad_fn=) +# _values +tensor([[0.0000, 0.2222, 0.4444], + [0.6667, 0.8889, 1.1111], + [1.3333, 1.5556, 1.7778]], dtype=torch.float32, + grad_fn=) + +# shape: torch.Size([100, 20, 3]) +# nnz: 0 +# sparseDim: 2 +# indices shape: torch.Size([2, 0]) +# values shape: torch.Size([0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + size=(100, 20, 3), nnz=0, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(2, 0), dtype=torch.int64) +# _values +tensor([], size=(0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + size=(100, 20, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + size=(100, 20, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(2, 0)), + values=tensor([], size=(0, 3)), + size=(100, 20, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(2, 0), dtype=torch.int64, grad_fn=) +# _values +tensor([], size=(0, 3), dtype=torch.float32, grad_fn=) + +# shape: torch.Size([10, 0, 3]) +# nnz: 3 +# sparseDim: 0 +# indices shape: torch.Size([0, 3]) +# values shape: torch.Size([3, 10, 0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + size=(10, 0, 3), nnz=3, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 3), dtype=torch.int64) +# _values +tensor([], size=(3, 10, 0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + size=(10, 0, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + size=(10, 0, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 3)), + values=tensor([], size=(3, 10, 0, 3)), + size=(10, 0, 3), nnz=3, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 3), dtype=torch.int64, grad_fn=) +# _values +tensor([], size=(3, 10, 0, 3), dtype=torch.float32, grad_fn=) + +# shape: torch.Size([10, 0, 3]) +# nnz: 0 +# sparseDim: 0 +# indices shape: torch.Size([0, 0]) +# values shape: torch.Size([0, 10, 0, 3]) +########## torch.int32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + size=(10, 0, 3), nnz=0, dtype=torch.int32, layout=torch.sparse_coo) +# _indices +tensor([], size=(0, 0), dtype=torch.int64) +# _values +tensor([], size=(0, 10, 0, 3), dtype=torch.int32) +########## torch.float32 ########## +# sparse tensor +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + size=(10, 0, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo) +# after requires_grad_ +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + size=(10, 0, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo, + requires_grad=True) +# after addition +tensor(indices=tensor([], size=(0, 0)), + values=tensor([], size=(0, 10, 0, 3)), + size=(10, 0, 3), nnz=0, dtype=torch.float32, layout=torch.sparse_coo, + grad_fn=) +# _indices +tensor([], size=(0, 0), dtype=torch.int64, grad_fn=) +# _values +tensor([], size=(0, 10, 0, 3), dtype=torch.float32, grad_fn=) diff --git a/test/onnx/expect/TestOperators.test_batchnorm_noaffine.expect b/test/onnx/expect/TestOperators.test_batchnorm_noaffine.expect new file mode 100644 index 0000000000000..d807b30800d39 --- /dev/null +++ b/test/onnx/expect/TestOperators.test_batchnorm_noaffine.expect @@ -0,0 +1,151 @@ +ir_version: 3 +producer_name: "pytorch" +producer_version: "0.4" +graph { + node { + output: "4" + op_type: "Constant" + attribute { + name: "value" + t { + dims: 128 + data_type: FLOAT + raw_data: "\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?" + } + type: TENSOR + } + } + node { + output: "5" + op_type: "Constant" + attribute { + name: "value" + t { + dims: 128 + data_type: FLOAT + raw_data: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000" + } + type: TENSOR + } + } + node { + input: "0" + input: "4" + input: "5" + input: "1" + input: "2" + output: "6" + op_type: "BatchNormalization" + attribute { + name: "epsilon" + f: 1e-05 + type: FLOAT + } + attribute { + name: "momentum" + f: 1 + type: FLOAT + } + } + name: "torch-jit-export" + initializer { + dims: 128 + data_type: FLOAT + name: "1" + raw_data: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000" + } + initializer { + dims: 128 + data_type: FLOAT + name: "2" + raw_data: "\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?" + } + initializer { + data_type: INT64 + name: "3" + raw_data: "\000\000\000\000\000\000\000\000" + } + input { + name: "0" + type { + tensor_type { + elem_type: FLOAT + shape { + dim { + dim_value: 128 + } + dim { + dim_value: 128 + } + dim { + dim_value: 1 + } + dim { + dim_value: 1 + } + } + } + } + } + input { + name: "1" + type { + tensor_type { + elem_type: FLOAT + shape { + dim { + dim_value: 128 + } + } + } + } + } + input { + name: "2" + type { + tensor_type { + elem_type: FLOAT + shape { + dim { + dim_value: 128 + } + } + } + } + } + input { + name: "3" + type { + tensor_type { + elem_type: INT64 + shape { + } + } + } + } + output { + name: "6" + type { + tensor_type { + elem_type: FLOAT + shape { + dim { + dim_value: 128 + } + dim { + dim_value: 128 + } + dim { + dim_value: 1 + } + dim { + dim_value: 1 + } + } + } + } + } +} +opset_import { + version: 7 +} diff --git a/test/onnx/test_operators.py b/test/onnx/test_operators.py index 736b681ab850c..b996e49f7de85 100644 --- a/test/onnx/test_operators.py +++ b/test/onnx/test_operators.py @@ -436,6 +436,10 @@ def test_unsqueeze(self): x = Variable(torch.randn(3, 4), requires_grad=True) self.assertONNX(lambda x: x.unsqueeze(len(x.shape)), x) + def test_batchnorm_noaffine(self): + x = Variable(torch.randn(128, 128, 1, 1), requires_grad=True) + self.assertONNX(nn.BatchNorm2d(128, affine=False), x) + def test_symbolic_override(self): """Lifted from fast-neural-style: custom implementation of instance norm to be mapped to ONNX operator""" diff --git a/test/onnx/test_pytorch_onnx_caffe2.py b/test/onnx/test_pytorch_onnx_caffe2.py index 94a2fa933b402..f8038d2465dbc 100644 --- a/test/onnx/test_pytorch_onnx_caffe2.py +++ b/test/onnx/test_pytorch_onnx_caffe2.py @@ -488,6 +488,11 @@ def test_batchnorm1d_special(self): model = nn.BatchNorm1d(224) self.run_model_test(model, train=True, input=c, batch_size=BATCH_SIZE) + def test_batchnorm2d_noaffine(self): + c = Variable(torch.randn(128, 128, 1, 1)) + model = nn.BatchNorm2d(128, affine=False) + self.run_model_test(model, train=False, input=c, batch_size=BATCH_SIZE) + def test_constant(self): c = Variable(torch.randn(BATCH_SIZE, 3, 224, 224)) diff --git a/test/test_c10d.py b/test/test_c10d.py index 13f7b779d0473..395f8135f9922 100644 --- a/test/test_c10d.py +++ b/test/test_c10d.py @@ -4,6 +4,7 @@ import sys import tempfile import unittest +from datetime import timedelta from functools import wraps from collections import namedtuple @@ -98,7 +99,9 @@ def tearDown(self): self.file.close() def _create_store(self): - return c10d.FileStore(self.file.name) + store = c10d.FileStore(self.file.name) + store.set_timeout(timedelta(seconds=300)) + return store class PrefixFileStoreTest(TestCase, StoreTestBase): @@ -106,6 +109,7 @@ def setUp(self): self.file = tempfile.NamedTemporaryFile() self.filestore = c10d.FileStore(self.file.name) self.prefix = "test_prefix" + self.filestore.set_timeout(timedelta(seconds=300)) def tearDown(self): self.file.close() @@ -118,7 +122,9 @@ class TCPStoreTest(TestCase, StoreTestBase): def _create_store(self): addr = 'localhost' port = common.find_free_port() - return c10d.TCPStore(addr, port, True) + store = c10d.TCPStore(addr, port, True) + store.set_timeout(timedelta(seconds=300)) + return store class PrefixTCPStoreTest(TestCase, StoreTestBase): @@ -127,6 +133,7 @@ def setUp(self): port = common.find_free_port() self.tcpstore = c10d.TCPStore(addr, port, True) self.prefix = "test_prefix" + self.tcpstore.set_timeout(timedelta(seconds=300)) def _create_store(self): return c10d.PrefixStore(self.prefix, self.tcpstore) diff --git a/test/test_cpp_extensions.py b/test/test_cpp_extensions.py index 3db7a42ffd236..4d35c0ced3c78 100755 --- a/test/test_cpp_extensions.py +++ b/test/test_cpp_extensions.py @@ -266,6 +266,113 @@ def test_lenient_flag_handling_in_jit_extensions(self): z = module.tanh_add(x, y).cpu() self.assertEqual(z, x.tanh() + y.tanh()) + def test_complex_registration(self): + cpp_source = ''' + #include + #include + #include + #include + + #include "ATen/TensorImpl.h" + #include "ATen/CPUGenerator.h" + #include "ATen/TensorImpl.h" + #include "ATen/Allocator.h" + #include "ATen/DeviceGuard.h" + #include "ATen/NativeFunctions.h" + #include "ATen/UndefinedTensor.h" + #include "ATen/Utils.h" + #include "ATen/WrapDimUtils.h" + #include "ATen/core/Half.h" + #include "ATen/core/optional.h" + + #include + #include + #include + #include + + #include "ATen/Config.h" + + namespace at { + + struct CPUComplexFloatType : public at::CPUTypeDefault { + + CPUComplexFloatType() + : CPUTypeDefault(CPUTensorId(), /*is_variable=*/false, /*is_undefined=*/false) {} + + ScalarType scalarType() const override; + Backend backend() const override; + const char * toString() const override; + size_t elementSizeInBytes() const override; + TypeID ID() const override; + Tensor & s_copy_(Tensor & self, const Tensor & src, bool non_blocking) const override; + Tensor & _s_copy_from(const Tensor & self, Tensor & dst, bool non_blocking) const override; + + Tensor tensor(IntList size) const override { + // TODO: Upstream this + int64_t numel = 1; + for (auto s : size) { + numel *= s; + } + Storage s{c10::make_intrusive( + scalarTypeToDataType(ScalarType::ComplexFloat), + numel, + getCPUAllocator(), + /* resizable */ true)}; + Tensor t{c10::make_intrusive( + std::move(s), + at::CPUTensorId(), + /* is_variable */ false)}; + return t; + } + }; + + struct ComplexHooks : public at::ComplexHooksInterface { + ComplexHooks(ComplexHooksArgs) {} + void registerComplexTypes(Context* context) const override { + context->registerType(Backend::CPU, ScalarType::ComplexFloat, new CPUComplexFloatType()); + } + }; + + ScalarType CPUComplexFloatType::scalarType() const { + return ScalarType::ComplexFloat; + } + + Backend CPUComplexFloatType::backend() const { + return Backend::CPU; + } + + const char * CPUComplexFloatType::toString() const { + return "CPUComplexFloatType"; + } + TypeID CPUComplexFloatType::ID() const { + return TypeID::CPUComplexFloat; + } + + size_t CPUComplexFloatType::elementSizeInBytes() const { + return sizeof(float); + } + + Tensor & CPUComplexFloatType::s_copy_(Tensor & dst, const Tensor & src, bool non_blocking) const { + AT_ERROR("not yet supported"); + } + + Tensor & CPUComplexFloatType::_s_copy_from(const Tensor & src, Tensor & dst, bool non_blocking) const { + AT_ERROR("not yet supported"); + } + + REGISTER_COMPLEX_HOOKS(ComplexHooks); + + } // namespace at + ''' + + module = torch.utils.cpp_extension.load_inline( + name='complex_registration_extension', + cpp_sources=cpp_source, + functions=[], + verbose=True) + + torch.empty(2, 2, dtype=torch.complex64) + if __name__ == '__main__': common.run_tests() diff --git a/test/test_distributed.py b/test/test_distributed.py index 3b53d69bfab08..dd21a597e6680 100644 --- a/test/test_distributed.py +++ b/test/test_distributed.py @@ -170,6 +170,12 @@ def _init_group_test(self): return (group, group_id, rank) + def _init_full_group_test(self): + group = [i for i in range(0, dist.get_world_size())] + group_id = dist.new_group() + rank = dist.get_rank() + return (group, group_id, rank) + def _init_global_test(self): group = [i for i in range(0, dist.get_world_size())] group_id = dist.group.WORLD @@ -223,24 +229,65 @@ def test_get_rank(self): self._barrier() + # GET default group + def test_get_default_group(self): + default_grp = dist.get_default_group() + self.assertNotEqual(default_grp, None) + + # Test destroy + def test_destroy_group(self): + if dist.get_world_size() > 2: + group = [1, 2] + else: + group = [0, 1] + group_id = dist.new_group(group) + dist.destroy_process_group(group_id) + + # Test get rank and size of group + def test_get_rank_size_group(self): + if dist.get_world_size() > 2: + group = [1, 2] + else: + group = [0, 1] + group_id = dist.new_group(group) + if dist.get_rank() in group: + self.assertEqual(dist.get_world_size(group_id), 2) + self.assertTrue(dist.get_rank(group_id) in list(range(2))) + else: + self.assertEqual(dist.get_world_size(group_id), -1) + self.assertEqual(dist.get_rank(group_id), -1) + + # Test destroy full groups + def test_destroy_full_group(self): + _, group_id, _ = self._init_full_group_test() + dist.destroy_process_group(group_id) + + # Test get rank and size of full group + def test_get_rank_size_full_group(self): + _, group_id, _ = self._init_full_group_test() + self.assertEqual(dist.get_world_size(group_id), dist.get_world_size()) + self.assertEqual(dist.get_rank(group_id), dist.get_rank()) + # SEND RECV @unittest.skipIf(BACKEND == "gloo", "Gloo does not support send/recv") @unittest.skipIf(BACKEND == "nccl", "Nccl does not support send/recv") def test_send_recv(self): rank = dist.get_rank() tensor = _build_tensor(rank + 1) - for dest in range(0, dist.get_world_size()): - if dest == rank: - continue - dist.send(tensor, dest) for src in range(0, dist.get_world_size()): if src == rank: - continue - tensor = _build_tensor(src + 1, value=-1) - expected_tensor = _build_tensor(src + 1) - dist.recv(tensor, src) - self.assertEqual(tensor, expected_tensor) + # Send mode + for dst in range(0, dist.get_world_size()): + if dst == rank: + continue + dist.send(tensor, dst) + else: + # Recv mode + expected_tensor = _build_tensor(src + 1) + output_tensor = _build_tensor(src + 1, value=-1) + dist.recv(output_tensor, src) + self.assertEqual(output_tensor, expected_tensor) self._barrier() @@ -253,20 +300,26 @@ def test_send_recv(self): ) def test_send_recv_any_source(self): rank = dist.get_rank() - tensor = _build_tensor(10, rank) - for dest in range(0, dist.get_world_size()): - if dest == rank: - continue - dist.send(tensor, dest) - + tensor = _build_tensor(10, value=rank) recv_ranks = set() - for src in range(0, dist.get_world_size()): - if src == rank: - continue - tensor = _build_tensor(10, value=-1) - sender = dist.recv(tensor) - self.assertTrue(tensor.eq(sender).all()) - recv_ranks.add(sender) + + for dst in range(0, dist.get_world_size()): + if dst == rank: + # Recv mode + for dst in range(0, dist.get_world_size()): + if dst == rank: + continue + output_tensor = _build_tensor(10, value=-1) + sender = dist.recv(output_tensor) + + # Assert the scalar value "sender" that should be + # equal to the rank of the sender is equal to all + # values in the received tensor. + self.assertTrue(output_tensor.eq(sender).all()) + recv_ranks.add(sender) + else: + # Send mode + dist.send(tensor, dst) self.assertEqual(len(recv_ranks), dist.get_world_size() - 1) self._barrier() @@ -370,6 +423,11 @@ def test_broadcast_group(self): group, group_id, rank = self._init_group_test() self._test_broadcast_helper(group, group_id, rank) + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_broadcast_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_broadcast_helper(group, group_id, rank) + # REDUCE def _test_reduce_helper( self, @@ -500,6 +558,46 @@ def test_reduce_group_max(self): group, group_id, rank = self._init_group_test() self._test_reduce_helper(group, group_id, rank, dist.reduce_op.MAX, -1, 10, 10) + @unittest.skipIf(BACKEND == "gloo", "Gloo does not support reduce") + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_reduce_full_group_sum(self): + group, group_id, rank = self._init_full_group_test() + self._test_reduce_helper( + group, + group_id, + rank, + dist.reduce_op.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + ) + + @unittest.skipIf(BACKEND == "gloo", "Gloo does not support reduce") + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_reduce_full_group_product(self): + group, group_id, rank = self._init_full_group_test() + self._test_reduce_helper( + group, + group_id, + rank, + dist.reduce_op.PRODUCT, + 2, + 10, + reduce((lambda x, y: x * y), [10] * (len(group) - 1), 2), + ) + + @unittest.skipIf(BACKEND == "gloo", "Gloo does not support reduce") + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_reduce_full_group_min(self): + group, group_id, rank = self._init_full_group_test() + self._test_reduce_helper(group, group_id, rank, dist.reduce_op.MIN, 1010, 1, 1) + + @unittest.skipIf(BACKEND == "gloo", "Gloo does not support reduce") + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_reduce_full_group_max(self): + group, group_id, rank = self._init_full_group_test() + self._test_reduce_helper(group, group_id, rank, dist.reduce_op.MAX, -1, 10, 10) + # ALL REDUCE def _test_all_reduce_helper( self, @@ -634,6 +732,46 @@ def test_all_reduce_group_max(self): group, group_id, rank, dist.reduce_op.MAX, -1, 10, 10 ) + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_all_reduce_full_group_sum(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.reduce_op.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + ) + + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_all_reduce_full_group_product(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.reduce_op.PRODUCT, + 2, + 10, + reduce((lambda x, y: x * y), [10] * (len(group) - 1), 2), + ) + + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_all_reduce_full_group_min(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_helper( + group, group_id, rank, dist.reduce_op.MIN, 1010, 1, 1 + ) + + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_all_reduce_full_group_max(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_helper( + group, group_id, rank, dist.reduce_op.MAX, -1, 10, 10 + ) + # SCATTER def _test_scatter_helper(self, group, group_id, rank): for dest in group: @@ -660,6 +798,12 @@ def test_scatter_group(self): group, group_id, rank = self._init_group_test() self._test_scatter_helper(group, group_id, rank) + @unittest.skipIf(BACKEND == "gloo", "Gloo does not support scatter") + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support scatter") + def test_scatter_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_scatter_helper(group, group_id, rank) + # GATHER def _test_gather_helper(self, group, group_id, rank): for dest in group: @@ -688,6 +832,12 @@ def test_gather_group(self): group, group_id, rank = self._init_group_test() self._test_gather_helper(group, group_id, rank) + @unittest.skipIf(BACKEND == "gloo", "Gloo does not support gather") + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_gather_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_gather_helper(group, group_id, rank) + # ALL GATHER def _test_all_gather_helper( self, group, group_id, rank, cuda=False, rank_to_GPU=None @@ -727,6 +877,12 @@ def test_all_gather_group(self): group, group_id, rank = self._init_group_test() self._test_all_gather_helper(group, group_id, rank) + @unittest.skipIf(BACKEND == "gloo", "Gloo does not support gather") + @unittest.skipIf(BACKEND == "nccl", "Nccl does not support CPU tensors") + def test_all_gather_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_gather_helper(group, group_id, rank) + # BARRIER def _test_barrier_helper(self, group, group_id, rank): WAIT_TIME = 0.3 # seconds @@ -756,6 +912,11 @@ def test_barrier_group(self): group, group_id, rank = self._init_group_test() self._test_barrier_helper(group, group_id, rank) + @unittest.skipIf(BACKEND != "mpi", "Only MPI supports barrier") + def test_barrier_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_barrier_helper(group, group_id, rank) + def _test_broadcast_multigpu_helper(self, group, group_id, rank, rank_to_GPU): for src in group: expected_tensor = _build_tensor(src + 1) @@ -1013,6 +1174,31 @@ def test_DistributedDataParallel(self): ) self._barrier() + @unittest.skipIf( + BACKEND == "nccl", "nccl does not support DistributedDataParallelCPU" + ) + def test_DistributedDataParallelCPU(self): + # Run a simple end to end DDP-CPU model, use result of single node + # model as baseline + group, group_id, rank = self._init_global_test() + + # cpu training setup + model_base = self._create_Net() + + # DDP-CPU training setup + model_DDP = copy.deepcopy(model_base) + model_DDP = nn.parallel._DistributedDataParallelC10dCPU(model_DDP) + + # dummy data initialization + local_bs = 2 + global_bs, input_cpu, target, loss = self._prepare_dummy_data(local_bs) + + # check two model parameters over 2 iterations + self._test_DDP_2iter( + model_base, model_DDP, input_cpu, target, loss, local_bs, rank, global_bs + ) + self._barrier() + if BACKEND == "gloo" or BACKEND == "nccl": WORLD_SIZE = os.environ["WORLD_SIZE"] @@ -1086,6 +1272,7 @@ def _run(self, rank): # self.id() == e.g. '__main__.TestDistributed.test_get_rank' # We're retreiving a corresponding test and executing it. getattr(self, self.id().split(".")[2])() + dist.destroy_process_group() sys.exit(0) def _join_and_reduce(self, fn): diff --git a/test/test_distributions.py b/test/test_distributions.py index 0419206b82754..c72f2bb2eca7d 100644 --- a/test/test_distributions.py +++ b/test/test_distributions.py @@ -93,6 +93,7 @@ def is_all_nan(tensor): {'probs': torch.tensor([0.7, 0.2, 0.4], requires_grad=True)}, {'probs': torch.tensor([0.3], requires_grad=True)}, {'probs': 0.3}, + {'logits': torch.tensor([0.], requires_grad=True)}, ]), Example(Geometric, [ {'probs': torch.tensor([0.7, 0.2, 0.4], requires_grad=True)}, @@ -112,6 +113,7 @@ def is_all_nan(tensor): Example(Categorical, [ {'probs': torch.tensor([[0.1, 0.2, 0.3], [0.5, 0.3, 0.2]], requires_grad=True)}, {'probs': torch.tensor([[1.0, 0.0], [0.0, 1.0]], requires_grad=True)}, + {'logits': torch.tensor([[0.0, 0.0], [0.0, 0.0]], requires_grad=True)}, ]), Example(Binomial, [ {'probs': torch.tensor([[0.1, 0.2, 0.3], [0.5, 0.3, 0.2]], requires_grad=True), 'total_count': 10}, @@ -322,6 +324,7 @@ def is_all_nan(tensor): Example(OneHotCategorical, [ {'probs': torch.tensor([[0.1, 0.2, 0.3], [0.5, 0.3, 0.2]], requires_grad=True)}, {'probs': torch.tensor([[1.0, 0.0], [0.0, 1.0]], requires_grad=True)}, + {'logits': torch.tensor([[0.0, 0.0], [0.0, 0.0]], requires_grad=True)}, ]), Example(Pareto, [ { @@ -713,6 +716,12 @@ def _check_enumerate_support(self, dist, examples): actual = dist(param).enumerate_support() self.assertEqual(actual, expected) + def test_repr(self): + for Dist, params in EXAMPLES: + for param in params: + dist = Dist(**param) + self.assertTrue(repr(dist).startswith(dist.__class__.__name__)) + def test_sample_detached(self): for Dist, params in EXAMPLES: for i, param in enumerate(params): diff --git a/test/test_jit.py b/test/test_jit.py index 9613ced33e765..93c9027123bd4 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -1,3 +1,4 @@ +from __future__ import division import torch import torch.jit import torch.nn as nn @@ -29,6 +30,10 @@ from torch.jit.frontend import NotSupportedError from torch.jit import BatchTensor +# For testing truediv in python 2 +from test_module.future_div import div_int_future, div_float_future +from test_module.no_future_div import div_int_nofuture, div_float_nofuture + try: import torchvision HAS_TORCHVISION = True @@ -83,12 +88,12 @@ def LSTMCellC(*args, **kwargs): def LSTMCellS(x, hx, cx, w_ih, w_hh, b_ih, b_hh): gates = x.mm(w_ih.t()) + hx.mm(w_hh.t()) + b_ih + b_hh ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1) - ingate = F.sigmoid(ingate) - forgetgate = F.sigmoid(forgetgate) - cellgate = F.tanh(cellgate) - outgate = F.sigmoid(outgate) + ingate = torch.sigmoid(ingate) + forgetgate = torch.sigmoid(forgetgate) + cellgate = torch.tanh(cellgate) + outgate = torch.sigmoid(outgate) cy = (forgetgate * cx) + (ingate * cellgate) - hy = outgate * F.tanh(cy) + hy = outgate * torch.tanh(cy) return hy, cy @@ -261,8 +266,8 @@ def run_pass(self, name, trace): return graph def checkTrace(self, func, reference_tensors, input_tensors=None, - optimize=True, drop=None, allow_unused=False, - verbose=False, inputs_require_grads=True, check_tolerance=1e-5): + optimize=True, drop=None, allow_unused=False, verbose=False, + inputs_require_grads=True, check_tolerance=1e-5, export_import=True): # TODO: check gradients for parameters, not just inputs def allSum(vs): # drop allows us to remove some values from ever being used @@ -286,6 +291,9 @@ def allSum(vs): else: ge = torch.jit.trace(func, input_tensors, optimize=optimize, check_tolerance=check_tolerance) + if export_import: + ge = self.getExportImportCopy(ge) + if verbose: print(ge.graph) @@ -522,6 +530,30 @@ def forward(self, x): self.assertExportImport(trace, (t,) + tuple(model.parameters())) self.assertExpectedONNXGraph(trace) + @unittest.skipIf(not IS_WINDOWS, "Testing Fuse skipped on windows") + @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") + def test_windows_fuse(self): + def scaleshift(x, scale, shift): + return x * scale + shift + + graph = torch.jit.script(scaleshift).graph + + inputs = [ + torch.randn(4, 4, dtype=torch.float, device='cuda'), + torch.randn(4, dtype=torch.float, device='cuda'), + torch.randn(4, dtype=torch.float, device='cuda'), + ] + + ge = self.checkTrace(scaleshift, inputs) + fuse_graph = ge.graph_for(*inputs) + + def run_graph(graph, inputs): + m = torch.jit.ScriptModule() + m._create_method_from_graph("forward", graph) + return m(*inputs) + + self.assertEqual(run_graph(graph, inputs), run_graph(fuse_graph, inputs)) + @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm @@ -565,7 +597,7 @@ def test_lstm_fusion_cpu(self): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm - def test_lstm_fusion_concat(self): + def test_lstm_fusion_concat_cuda(self): inputs = get_lstm_inputs('cuda') ge = self.checkTrace(LSTMCellC, inputs) self.assertExpectedGraph(ge.graph_for(*inputs)) @@ -573,7 +605,7 @@ def test_lstm_fusion_concat(self): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm - def test_concat_fusion(self): + def test_concat_fusion_cuda(self): hx = torch.randn(3, 20, dtype=torch.float, device='cuda') cx = torch.randn(3, 20, dtype=torch.float, device='cuda') @@ -604,7 +636,7 @@ def fn(x, y, z): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm - def test_fusion_distribute(self): + def test_fusion_distribute_cuda(self): def f(x, y): z1, z2 = (x + y).chunk(2, dim=1) return z1 * z2 @@ -618,7 +650,7 @@ def f(x, y): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm - def test_fusion_rand(self): + def test_fusion_rand_cuda(self): class M(torch.jit.ScriptModule): __constants__ = ['d'] @@ -642,7 +674,7 @@ def create(self, x): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm - def test_fusion_arg_configurations(self): + def test_fusion_arg_configurations_cuda(self): # A smoke test to make sure we won't use the same kernel for contiguous # and non-contiguous arguments. # TODO: add optionally enabled debug counters to the fuser to verify @@ -667,7 +699,7 @@ def fn_test_comparison_gt_lt(x, y): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm - def test_comparison_gt_lt(self): + def test_comparison_gt_lt_cuda(self): x = torch.randn(4, 4, dtype=torch.float, device='cuda') y = torch.randn(4, 4, dtype=torch.float, device='cuda') @@ -676,7 +708,7 @@ def test_comparison_gt_lt(self): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm - def test_comparison_ge_le(self): + def test_comparison_ge_le_cuda(self): def f(x, y): mask = (x >= 0).type_as(x) z = x * mask + y @@ -696,7 +728,7 @@ def fn_test_relu(x, y): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm - def test_relu(self): + def test_relu_cuda(self): x = torch.randn(4, 4, dtype=torch.float, device='cuda') y = torch.randn(4, 4, dtype=torch.float, device='cuda') @@ -704,7 +736,7 @@ def test_relu(self): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") - def test_small_constant(self): + def test_small_constant_cuda(self): def fn_test_small_constant(x, y): return (1e-8 * x + 5e-9 * y) * 1e8 x = torch.randn(4, 4, dtype=torch.float, device='cuda') @@ -719,7 +751,7 @@ def fn_test_exp(x, y): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm - def test_exp(self): + def test_exp_cuda(self): x = torch.randn(4, 4, dtype=torch.float, device='cuda') y = torch.randn(4, 4, dtype=torch.float, device='cuda') @@ -858,7 +890,7 @@ def broadcast(a, b): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA_MULTI_GPU, "needs non-zero device") @skipIfRocm - def test_fuse_last_device(self): + def test_fuse_last_device_cuda(self): device = 'cuda:' + str(1) x = torch.tensor([0.4], dtype=torch.float, device=device) y = torch.tensor([0.7], dtype=torch.float, device=device) @@ -1134,7 +1166,7 @@ def doit(x, y): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "cpp tests require CUDA") @skipIfRocm - def test_cpp(self): + def test_cpp_cuda(self): # rather than rebuild assertExpected in cpp, # just glob all the cpp outputs into one file for now self.assertExpected(torch._C._jit_run_cpp_tests()) @@ -1291,7 +1323,7 @@ def foo(a): @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "calls .cuda()") @skipIfRocm - def test_traced_module(self): + def test_traced_module_cuda(self): class Model(nn.Module): def __init__(self, num_features, num_layers): super(Model, self).__init__() @@ -1398,7 +1430,10 @@ def addmm(a, b, c): a, b, c = [torch.tensor(e) for e in (1, [[2.]], [[3.]])] addmm(a, b, c) - self.assertExpectedGraph(addmm.graph_for(a, b, c)) + graph = addmm.graph_for(a, b, c) + # graph fusion skipped in windows, which runs cse + self.run_pass('cse', graph) + self.assertExpectedGraph(graph) def test_index_put(self): ten = torch.zeros(3, 3) @@ -2333,24 +2368,24 @@ def func(x): x = torch.rand(10, dtype=torch.float, requires_grad=True) self.checkScript(func, [x], optimize=True) + def _check_code(self, code_str, fn_name, inputs): + scope = {} + exec(code_str, globals(), scope) + cu = torch.jit.CompilationUnit(code_str) + self.assertEqual(cu.func(*inputs), scope[fn_name](*inputs)) + def test_index(self): def consec(size, start=0): numel = torch.tensor(size).prod().item() return torch.arange(numel).view(size) - def check_code(code_str, fn_name, inputs): - scope = {} - exec(code_str, globals(), scope) - cu = torch.jit.CompilationUnit(code_str) - self.assertEqual(cu.func(*inputs), scope['func'](*inputs)) - def check_indexing(indexing, tensor): template = dedent(""" def func(x): return x{} """) - check_code(template.format(indexing), "func", [tensor]) + self._check_code(template.format(indexing), "func", [tensor]) def check_dynamic_indexing(indexing, tensor, value1, value2): value1 = torch.tensor(value1) @@ -2363,7 +2398,7 @@ def func(x, value1, value2): return x{} """) - check_code(template.format(indexing), "func", [tensor, value1, value2]) + self._check_code(template.format(indexing), "func", [tensor, value1, value2]) # basic slices check_indexing('[0]', consec((3, 3))) @@ -2462,6 +2497,93 @@ def test_func(func, x, tensor): for i in range(len(script_funs)): self.assertEqual(test_func(script_funs[i], x, tensor), test_func(funs[i], x, tensor)) + def test_advancedindex(self): + def consec(size, start=0): + numel = torch.tensor(size).prod().item() + return torch.arange(numel).view(size) + + def check_indexing(indexing, tensor, **kwargs): + indices_dict = kwargs + + template = dedent(""" + def func(x{formals}): + return x{expr} + """) + + formals = [] + values = [] + for formal, value in indices_dict.items(): + formals.append(formal) + values.append(value) + + formals = ''.join(map(', {}'.format, formals)) + inputs = [tensor] + values + + self._check_code(template.format(formals=formals, expr=indexing), + "func", inputs) + + # Indexing with tensor (basic) + check_indexing('[i]', consec((3, 3)), i=torch.tensor([0])) + check_indexing('[i]', consec((3, 3)), i=torch.tensor(1)) + check_indexing('[i]', consec((3, 3)), i=torch.tensor([-2])) + check_indexing('[i]', consec((3, 3), 2), i=torch.tensor([0, 0])) + check_indexing('[i]', consec((3, 3, 2, 2)), i=torch.tensor([0, -2, 1])) + + # NB: indexing with tensors and indexing with sequences can be implemented + # in a very similar way (sequences are converted to tensors), so only one + # case needs to be tested extensively. + # XXX: When we can index with sequences, replace these cases with + # sequence indexing expressions; those are much easier to read. + + # Misc sequence advanced indexing + inp = consec((4, 8, 5)) + to_check = [ + # [[0, 2], [1, 3]] + ['[i, j]', dict(i=[0, 2], j=[1, 3])], + # [[0, 2], [1, 3], [1, 1]] + ['[i, j, k]', dict(i=[0, 2], j=[1, 3], k=[1, 1])], + # [[0, 2], 1, [1, 1]] + ['[i, j, k]', dict(i=[0, 2], j=1, k=[1, 1])], + # [:, :, [0, 3, 4]] + ['[:, :, i]', dict(i=[0, 3, 4])], + # [:, [2, 4, 5, 7], 2:4] + ['[:, i, 2:4]', dict(i=[0, 2, 3])], + # [[2, 3], :, :] + ['[i, :, :]', dict(i=[2, 3])], + # [:, [0, 2, 3], [1, 3, 4]] + ['[:, i, j]', dict(i=[0, 2, 3], j=[1, 3, 4])], + # [:, [0], [1, 2, 4]] + ['[:, i, j]', dict(i=[0], j=[1, 2, 4])], + # [:, [0, 1, 3], [4]] + ['[:, i, j]', dict(i=[0, 1, 3], j=[4])], + # [:, [[0, 1], [1, 0]], [[2, 3]]] + ['[:, i, j]', dict(i=[[0, 1], [1, 0]], j=[[2, 3]])], + # [:, [[0, 1], [2, 3]], [[0]]] + ['[:, i, j]', dict(i=[[0, 1], [2, 3]], j=[[0]])], + # [:, [[5, 6]], [[0, 3], [4, 4]]] + ['[:, i, j]', dict(i=[[5, 6]], j=[[0, 3], [4, 4]])], + # [[0, 2, 3], [1, 3, 4], :] + ['[i, j, :]', dict(i=[0, 2, 3], j=[1, 3, 4])], + # [0, [1, 2, 4], :] + ['[i, j, :]', dict(i=0, j=[1, 2, 4])], + # [[0, 1, 3], 4, :] + ['[i, j, :]', dict(i=[0, 1, 3], j=4)], + # [[[0, 1], [1, 0]], [[2, 1], [3, 5]], :] + ['[i, j, :]', dict(i=[[0, 1], [1, 0]], j=[[2, 1], [3, 5]])], + # [[[0, 1], [1, 0]], [[2, 3]], :] + ['[i, j, :]', dict(i=[[0, 1], [1, 0]], j=[[2, 3]])], + # [[[0, 1], [2, 3]], [[0]], :] + ['[i, j, :]', dict(i=[[0, 1], [2, 3]], j=[[0]])], + # [[[2, 1]], [[0, 3], [4, 4]], :] + ['[i, j, :]', dict(i=[[2, 1]], j=[[0, 3], [4, 4]])], + # [[[2]], [[0, 3], [4, 1]], 0:2] + ['[i, j, 0:2]', dict(i=[[2]], j=[[0, 3], [4, 1]])], + ] + + for expr, argdict in to_check: + tensordict = {k: torch.tensor(v) for (k, v) in argdict.items()} + check_indexing(expr, inp, **tensordict) + def test_keyword(self): @torch.jit.script def func(x): @@ -2915,6 +3037,16 @@ def test_fuser_multiple_blocks(this, that, theother, meme): self.assertEqual(cu.test_fuser_multiple_blocks(*inputs), outputs) + @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") + def test_scalar_fusion(self): + def fn(x, y): + return x + y.type_as(x) + + x = torch.tensor(0.1, dtype=torch.float, device='cpu') + y = torch.tensor(1, dtype=torch.float, device='cpu') + ge = self.checkScript(fn, (x, y)) + self.assertExpectedGraph(ge.graph_for(x, y)) + @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows") @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @skipIfRocm @@ -3347,7 +3479,7 @@ def foo(a): @torch.jit.script def bar(c, b): - return c / b + return c + b with self.assertRaisesRegex(RuntimeError, "failed in interpreter"): bar(Variable(torch.rand(10), requires_grad=True), Variable(torch.rand(9), requires_grad=True)) @@ -3376,10 +3508,6 @@ def func4(): return 3.14 {op} 3.14 ''') ops = ['+', '-', '*', '<', '<=', '>', '>=', '==', '!='] - # TODO: turn this on for py3 (and add PY3 division semantics) - ops_py2_only = ['/'] - if PY2: - ops.extend(ops_py2_only) for op in ops: code = template.format(op=op) @@ -3392,6 +3520,19 @@ def func4(): self.assertEqual(cu.func3(), scope['func3']()) self.assertEqual(cu.func4(), scope['func4']()) + def test_number_div(self): + self.checkScript(div_int_future, (), optimize=True) + self.checkScript(div_float_future, (), optimize=True) + + if PY2: + with self.assertRaisesRegex(RuntimeError, 'from __future__ import division'): + torch.jit.script(div_int_nofuture) + with self.assertRaisesRegex(RuntimeError, 'from __future__ import division'): + torch.jit.script(div_float_nofuture) + else: + self.checkScript(div_int_nofuture, (), optimize=True) + self.checkScript(div_float_nofuture, (), optimize=True) + def test_number_augassign(self): def func(): z = 1 @@ -4095,26 +4236,6 @@ def forward(self, x, seq_lens): f = io.BytesIO() torch.onnx._export(m, (x, seq_lens), f, verbose=False) - def test_pack_padded_wrong_types(self): - from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence - - class PackPaddedWrapper(torch.nn.Module): - def __init__(self): - super(PackPaddedWrapper, self).__init__() - self.seq_lens = [3, 3, 3, 3] - - __constants__ = ['seq_lens'] - - def forward(self, x): - return pack_padded_sequence(x, self.seq_lens) - - m = PackPaddedWrapper() - - x = torch.rand(3, 4, 5) - f = io.BytesIO() - with self.assertRaisesRegex(RuntimeError, 'PackPadded requires `lengths` to be a Tensor'): - torch.onnx._export(m, (x,), f) - def test_script_outputs(self): with self.assertRaisesRegex(RuntimeError, "cannot be used as a tuple"): @torch.jit.script @@ -4750,6 +4871,38 @@ def forward(self, x): mte, (torch.zeros(1, 2, 3),), None, verbose=False, example_outputs=outputs)) + def test_onnx_export_script_truediv(self): + class ModuleToExport(torch.jit.ScriptModule): + def __init__(self): + super(ModuleToExport, self).__init__() + + @torch.jit.script_method + def forward(self, x): + z = x.size(0) / 2 + return x + z + + mte = ModuleToExport() + outputs = mte(torch.zeros(1, 2, 3)) + self.assertExpected(torch.onnx.export_to_pretty_string( + mte, (torch.zeros(1, 2, 3),), None, verbose=False, + example_outputs=outputs)) + + def test_onnx_raw_export_script_truediv(self): + class ModuleToExport(torch.jit.ScriptModule): + def __init__(self): + super(ModuleToExport, self).__init__() + + @torch.jit.script_method + def forward(self, x): + z = x.size(0) / 2 + return x + z + + mte = ModuleToExport() + outputs = mte(torch.zeros(1, 2, 3)) + self.assertExpected(torch.onnx.export_to_pretty_string( + mte, (torch.zeros(1, 2, 3),), None, verbose=False, + example_outputs=outputs, export_raw_ir=True)) + def test_onnx_export_script_module_if(self): class ModuleToExport(torch.jit.ScriptModule): def __init__(self): @@ -6239,6 +6392,7 @@ def foo(x): y = torch.arange(0, x.shape[0]).double() return x + y.unsqueeze(1) + @suppress_warnings def test_trace_checker_dot_data(self): with self.assertRaisesRegex(torch.jit.TracingCheckError, r'Tensor-valued Constant nodes differed in value ' r'across invocations'): @@ -6249,13 +6403,15 @@ def foo(x): @suppress_warnings def test_trace_checker_control_flow(self): + def foo(x): + for _ in range(x.size(0)): + x = torch.neg(x) + return x + with self.assertRaisesRegex(torch.jit.TracingCheckError, r'Graphs differed across invocations!'): - @_trace(torch.rand(3, 4), check_inputs=[(torch.rand(4, 4),)]) - def foo(x): - for _ in range(x.size(0)): - x = torch.neg(x) - return x + torch.jit.trace(foo, torch.randn(3, 4), check_inputs=[torch.randn(4, 4)]) + @suppress_warnings def test_trace_checker_memoization(self): with self.assertRaisesRegex(torch.jit.TracingCheckError, r'Graphs differed across invocations!'): def foo(x): @@ -6277,13 +6433,19 @@ def foo(x): for i in range(3): x[i, :] = torch.zeros(4) return x - self.checkTracerWarning(foo, torch.rand(3, 4)) + + self.assertWarnsRegex(lambda: torch.jit.trace(foo, torch.rand(3, 4), check_inputs=[torch.rand(3, 4)]), + 'Output nr 1. of the traced function does not match the ' + 'corresponding output of the Python function') def test_trace_checker_inplace_on_view(self): def foo(x): x.view(-1).add_(-x.view(-1)) return x - self.checkTracerWarning(foo, torch.rand(3, 4), check_trace=False) + + self.assertWarnsRegex(lambda: torch.jit.trace(foo, torch.rand(3, 4), check_inputs=[torch.rand(5, 6)]), + 'Output nr 1. of the traced function does not match the ' + 'corresponding output of the Python function') def test_lhs_index_fails(self): def foo(x): @@ -6297,11 +6459,22 @@ def foo(y, x): return y self.checkTrace(foo, (torch.rand(3, 4), torch.rand(4)), inputs_require_grads=False) + def test_inplace_warn(self): + def foo(x): + x.view(-1).add_(-x.view(-1)) + return x + self.checkTracerWarning(foo, torch.rand(3, 4)) + + @suppress_warnings def test_trace_checker_dropout_train(self): - with self.assertRaisesRegex(torch.jit.TracingCheckError, r'Trace had nondeterministic nodes'): - @_trace(torch.rand(3, 4)) - def foo(x): - return torch.dropout(x, p=0.5, train=True) + def foo(x): + return torch.dropout(x, p=0.5, train=True) + + self.assertWarnsRegex(lambda: torch.jit.trace(foo, torch.rand(3, 4), check_inputs=[torch.rand(5, 6)]), + 'Output nr 1. of the traced function does not match the ' + 'corresponding output of the Python function') + self.assertWarnsRegex(lambda: torch.jit.trace(foo, torch.rand(3, 4), check_inputs=[torch.rand(5, 6)]), + 'Trace had nondeterministic nodes') def test_trace_checker_dropout_notrain(self): input = torch.rand(3, 4) @@ -6496,8 +6669,7 @@ def forward(self, x): x = self.fc2(x) return F.log_softmax(x, dim=1) - # FIXME: eval() is present because it works around the issue described - # in https://github.com/pytorch/pytorch/issues/8448 + # eval() is present because dropout makes this nondeterministic self.checkTrace(Net().eval(), (torch.rand(5, 1, 28, 28),)) def test_reinforcement_learning(self): @@ -6516,9 +6688,6 @@ def forward(self, x): @skipIfRocm def test_snli(self): - # TODO: - # 1) nn.LSTM is called as a Python function https://github.com/pytorch/pytorch/issues/8449 - # 2) Dropout is called as a Python function https://github.com/pytorch/pytorch/issues/8450 class Bottle(nn.Module): def forward(self, input): @@ -6686,7 +6855,8 @@ def forward(self, input): outputs = torch.cat((outputs, output), 1) return outputs - self.checkTrace(Sequence(), (torch.rand(3, 4),)) + # TODO: toggle export_import once above issues are fixed + self.checkTrace(Sequence(), (torch.rand(3, 4),), export_import=False) def test_vae(self): class VAE(nn.Module): @@ -6720,8 +6890,7 @@ def forward(self, x): z = self.reparameterize(mu, logvar) return self.decode(z), mu, logvar - # FIXME: this fails under training because of the call to `randn_like` - # https://github.com/pytorch/pytorch/issues/8443 + # eval() is present because randn_like makes this nondeterministic self.checkTrace(VAE().eval(), (torch.rand(128, 1, 28, 28),)) diff --git a/test/test_module/__init__.py b/test/test_module/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/test/test_module/future_div.py b/test/test_module/future_div.py new file mode 100644 index 0000000000000..4329365ff28b0 --- /dev/null +++ b/test/test_module/future_div.py @@ -0,0 +1,10 @@ +from __future__ import division +import torch + + +def div_int_future(): + return 1 / 2 + + +def div_float_future(): + return 3.14 / 0.125 diff --git a/test/test_module/no_future_div.py b/test/test_module/no_future_div.py new file mode 100644 index 0000000000000..1b1b9f0cd1316 --- /dev/null +++ b/test/test_module/no_future_div.py @@ -0,0 +1,9 @@ +import torch + + +def div_int_nofuture(): + return 1 / 2 + + +def div_float_nofuture(): + return 3.14 / 0.125 diff --git a/test/test_nn.py b/test/test_nn.py index 209b2cd0efa93..f0a8f86017d96 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -5353,6 +5353,7 @@ def test_affine_grid(self): @unittest.skipIf((not TEST_NUMPY) or (not TEST_SCIPY) or (scipy.__version__ < '1.0.0'), "Scipy v1.0 and/or numpy not found") + @skipIfRocm def test_affine_2d_rotate0(self): # scipy before 1.0.0 do not support homogeneous coordinate # scipy.ndimage.affine_transform, so we need to skip. @@ -5390,6 +5391,7 @@ def test_affine_2d_rotate0(self): @unittest.skipIf((not TEST_NUMPY) or (not TEST_SCIPY) or (scipy.__version__ < '1.0.0'), "Scipy v1.0 and/or numpy not found") + @skipIfRocm def test_affine_2d_rotate90(self): # scipy before 1.0.0 do not support homogeneous coordinate # scipy.ndimage.affine_transform, so we need to skip. @@ -5435,6 +5437,7 @@ def test_affine_2d_rotate90(self): @unittest.skipIf((not TEST_NUMPY) or (not TEST_SCIPY) or (scipy.__version__ < '1.0.0'), "Scipy v1.0 and/or numpy not found") + @skipIfRocm def test_affine_2d_rotate45(self): # scipy before 1.0.0 do not support homogeneous coordinate # scipy.ndimage.affine_transform, so we need to skip. @@ -5473,6 +5476,7 @@ def test_affine_2d_rotate45(self): @unittest.skipIf((not TEST_NUMPY) or (not TEST_SCIPY) or (scipy.__version__ < '1.0.0'), "Scipy v1.0 and/or numpy not found") + @skipIfRocm def test_affine_2d_rotateRandom(self): # scipy before 1.0.0 do not support homogeneous coordinate # scipy.ndimage.affine_transform, so we need to skip. @@ -5521,6 +5525,7 @@ def test_affine_2d_rotateRandom(self): @unittest.skipIf((not TEST_NUMPY) or (not TEST_SCIPY) or (scipy.__version__ < '1.0.0'), "Scipy v1.0 and/or numpy not found") + @skipIfRocm def test_affine_3d_rotateRandom(self): # scipy before 1.0.0 do not support homogeneous coordinate # scipy.ndimage.affine_transform, so we need to skip. diff --git a/test/test_sparse.py b/test/test_sparse.py index 01a382f3901ff..fab6ad978176b 100644 --- a/test/test_sparse.py +++ b/test/test_sparse.py @@ -107,6 +107,61 @@ def randn(self, *args, **kwargs): # TODO: Put this in torch.cuda.randn return self.ValueTensor(*args, **kwargs).normal_() + @skipIfRocm # ROCm stack doesn't like the x + x call + def test_print(self): + shape_sparseDim_nnz = [ + ((), 0, 2), + ((0,), 0, 10), + ((2,), 0, 3), + ((100, 3), 1, 3), + ((100, 20, 3), 2, 0), + ((10, 0, 3), 0, 3), + ((10, 0, 3), 0, 0), + ] + + printed = [] + for shape, sparseDim, nnz in shape_sparseDim_nnz: + indices_shape = torch.Size((sparseDim, nnz)) + values_shape = torch.Size((nnz,) + shape[sparseDim:]) + printed.append("# shape: {}".format(torch.Size(shape))) + printed.append("# nnz: {}".format(nnz)) + printed.append("# sparseDim: {}".format(sparseDim)) + printed.append("# indices shape: {}".format(indices_shape)) + printed.append("# values shape: {}".format(values_shape)) + + indices = torch.arange(indices_shape.numel(), dtype=self.IndexTensor.dtype, + device=self.device).view(indices_shape) + for d in range(sparseDim): + indices[d].clamp_(max=(shape[d] - 1)) # make it valid index + if self.is_uncoalesced and indices.numel() > 0: + indices[:, -1] = indices[:, 0] # make it uncoalesced + values_numel = values_shape.numel() + values = torch.arange(values_numel, dtype=self.ValueTensor.dtype, + device=self.device).view(values_shape).div_(values_numel / 2.) + sp_tensor = self.SparseTensor(indices, values, shape) + + dtypes = [torch.int32] + if values.dtype == torch.double: + dtypes.append(torch.float) + else: + dtypes.append(torch.double) + for dtype in dtypes: + printed.append("########## {} ##########".format(dtype)) + x = sp_tensor.detach().to(dtype) + printed.append("# sparse tensor") + printed.append(str(x)) + if x.dtype.is_floating_point: + printed.append("# after requires_grad_") + printed.append(str(x.requires_grad_())) + printed.append("# after addition") + printed.append(str(x + x)) + printed.append("# _indices") + printed.append(str(x._indices())) + printed.append("# _values") + printed.append(str(x._values())) + printed.append('') + self.assertExpected('\n'.join(printed)) + @skipIfRocm def test_basic(self): x, i, v = self._gen_sparse(3, 10, 100) @@ -997,6 +1052,10 @@ def test_factory_size_check(self): with self.assertRaisesRegex(RuntimeError, "sizes is inconsistent with indices"): torch.sparse_coo_tensor(indices, values, sizes) + indices.fill_(-1) + with self.assertRaisesRegex(RuntimeError, "found negative index"): + torch.sparse_coo_tensor(indices, values, sizes) + indices = self.IndexTensor([[1, 2], [0, 2]]) values = self.ValueTensor([[1, 1, 1], [1, 1, 1]]) sizes = torch.Size([3, 3, 2]) diff --git a/test/test_torch.py b/test/test_torch.py index 2f0aecbc4b164..02ffe7f1048f7 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -218,8 +218,7 @@ def test_namespace(ns, *skips): 'sparse_resize_and_clear_', ) test_namespace(torch.nn) - test_namespace(torch.nn.functional, 'assert_int_or_pair', 'bilinear', - 'dropout', 'dropout2d', 'dropout3d', 'feature_alpha_dropout') + test_namespace(torch.nn.functional, 'assert_int_or_pair', 'bilinear', 'feature_alpha_dropout') # TODO: add torch.* tests when we have proper namespacing on ATen functions # test_namespace(torch) @@ -3410,6 +3409,26 @@ def test_sort(self): # Test that we still have proper sorting with duplicate keys self.assertIsOrdered('descending', x, res2val, res2ind, 'random with duplicate keys') + @unittest.skipIf(not TEST_NUMPY, 'Numpy not found') + def test_tensordot(self): + devices = ['cpu'] if not torch.cuda.is_available() else ['cpu', 'cuda'] + for d in devices: + a = torch.arange(60., device=d).reshape(3, 4, 5) + b = torch.arange(24., device=d).reshape(4, 3, 2) + c = torch.tensordot(a, b, dims=([1, 0], [0, 1])).cpu() + cn = torch.from_numpy(np.tensordot(a.cpu().numpy(), b.cpu().numpy(), + axes=([1, 0], [0, 1]))) + self.assertEqual(c, cn) + a = torch.randn(2, 3, 4, 5, device=d) + b = torch.randn(4, 5, 6, 7, device=d) + c = torch.tensordot(a, b, dims=2).cpu() + cn = torch.from_numpy(np.tensordot(a.cpu().numpy(), b.cpu().numpy(), + axes=2)) + self.assertEqual(c, cn) + c = torch.tensordot(a, b).cpu() + cn = torch.from_numpy(np.tensordot(a.cpu().numpy(), b.cpu().numpy())) + self.assertEqual(c, cn) + def test_topk(self): def topKViaSort(t, k, dim, dir): sorted, indices = t.sort(dim, dir) @@ -7743,6 +7762,35 @@ def test_serialization_filelike_uses_readinto(self): b = torch.load(data) self.assertTrue(data.was_called('readinto')) + def test_serialization_storage_slice(self): + # Generated using: + # + # t = torch.zeros(2); + # s1 = t.storage()[:1] + # s2 = t.storage()[1:] + # torch.save((s1, s2), 'foo.ser') + # + # with PyTorch 0.3.1 + serialized = (b'\x80\x02\x8a\nl\xfc\x9cF\xf9 j\xa8P\x19.\x80\x02M\xe9\x03' + b'.\x80\x02}q\x00(X\n\x00\x00\x00type_sizesq\x01}q\x02(X\x03' + b'\x00\x00\x00intq\x03K\x04X\x05\x00\x00\x00shortq\x04K\x02X' + b'\x04\x00\x00\x00longq\x05K\x04uX\x10\x00\x00\x00protocol_versionq' + b'\x06M\xe9\x03X\r\x00\x00\x00little_endianq\x07\x88u.\x80\x02' + b'(X\x07\x00\x00\x00storageq\x00ctorch\nFloatStorage\nq\x01X\x0e' + b'\x00\x00\x0094279043900432q\x02X\x03\x00\x00\x00cpuq\x03K\x02' + b'X\x0e\x00\x00\x0094279029750368q\x04K\x00K\x01\x87q\x05tq\x06' + b'Q(h\x00h\x01X\x0e\x00\x00\x0094279043900432q\x07h\x03K\x02X' + b'\x0e\x00\x00\x0094279029750432q\x08K\x01K\x01\x87q\ttq\nQ' + b'\x86q\x0b.\x80\x02]q\x00X\x0e\x00\x00\x0094279043900432q' + b'\x01a.\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' + b'\x00\x00\x00\x00') + + buf = io.BytesIO(serialized) + (s1, s2) = torch.load(buf) + self.assertEqual(s1[0], 0) + self.assertEqual(s2[0], 0) + self.assertEqual(s1.data_ptr() + 4, s2.data_ptr()) + def test_load_error_msg(self): expected_err_msg = (".*You can only torch.load from a file that is seekable. " + "Please pre-load the data into a buffer like io.BytesIO and " + @@ -8276,14 +8324,18 @@ def test_numpy_array_interface(self): @unittest.skipIf(not TEST_NUMPY, "Numpy not found") def test_multiplication_numpy_scalar(self): - np_sc = np.float64(2.0) - t = torch.ones(2, requires_grad=True) - r1 = np_sc * t - self.assertIsInstance(r1, torch.Tensor) - self.assertTrue(r1.requires_grad) - r2 = t * np_sc - self.assertIsInstance(r2, torch.Tensor) - self.assertTrue(r2.requires_grad) + for np_dtype in [np.float32, np.float64, np.int32, np.int64, np.int16, np.uint8]: + for t_dtype in [torch.float, torch.double]: + np_sc = np_dtype(2.0) + t = torch.ones(2, requires_grad=True, dtype=t_dtype) + r1 = t * np_sc + self.assertIsInstance(r1, torch.Tensor) + self.assertTrue(r1.dtype == t_dtype) + self.assertTrue(r1.requires_grad) + r2 = np_sc * t + self.assertIsInstance(r2, torch.Tensor) + self.assertTrue(r2.dtype == t_dtype) + self.assertTrue(r2.requires_grad) def test_error_msg_type_translation(self): with self.assertRaisesRegex( diff --git a/third_party/onnx b/third_party/onnx index 1b09eb14c2c78..bff0b8835870c 160000 --- a/third_party/onnx +++ b/third_party/onnx @@ -1 +1 @@ -Subproject commit 1b09eb14c2c781fae078fa6b1c0390ba6fc0898c +Subproject commit bff0b8835870c7df7762ef43498d000d2d8ffb52 diff --git a/tools/amd_build/pyHIPIFY/cuda_to_hip_mappings.py b/tools/amd_build/pyHIPIFY/cuda_to_hip_mappings.py index 6bf931c589396..113403fd87bbf 100644 --- a/tools/amd_build/pyHIPIFY/cuda_to_hip_mappings.py +++ b/tools/amd_build/pyHIPIFY/cuda_to_hip_mappings.py @@ -2171,7 +2171,7 @@ CUDA_SPARSE_MAP = { "cusparseStatus_t": ("hipsparseStatus_t", CONV_MATH_FUNC, API_SPARSE), "cusparseHandle_t": ("hipsparseHandle_t", CONV_MATH_FUNC, API_SPARSE), - "cusparseOperation_t": ("hcsparseOperation_t", CONV_TYPE, API_SPARSE), + "cusparseOperation_t": ("hipsparseOperation_t", CONV_TYPE, API_SPARSE), "cusparseCreate": ("hipsparseCreate", CONV_MATH_FUNC, API_SPARSE), "cusparseDestroy": ("hipsparseDestroy", CONV_MATH_FUNC, API_SPARSE), "CUSPARSE_STATUS_SUCCESS": ("HIPSPARSE_STATUS_SUCCESS", CONV_NUMERIC_LITERAL, API_SPARSE), @@ -2202,6 +2202,7 @@ "common_gpu" : ("hip/common_hip", API_CAFFE2), "mixed_utils" : ("hip/mixed_utils_hip", API_CAFFE2), "operator_fallback_gpu" : ("hip/operator_fallback_hip", API_CAFFE2), + "spatial_batch_norm_op_gpu_impl" : ("hip/spatial_batch_norm_op_hip_impl", API_CAFFE2), "recurrent_network_executor_gpu" : ("hip/recurrent_network_executor_hip", API_CAFFE2), "max_pool_with_index_gpu": ("hip/max_pool_with_index_hip", API_CAFFE2), "THCCachingAllocator_gpu": ("hip/THCCachingAllocator_hip", API_CAFFE2), diff --git a/tools/amd_build/pyHIPIFY/hipify-python.py b/tools/amd_build/pyHIPIFY/hipify-python.py index 169a8cb6201df..e4201f48cef0f 100755 --- a/tools/amd_build/pyHIPIFY/hipify-python.py +++ b/tools/amd_build/pyHIPIFY/hipify-python.py @@ -718,7 +718,7 @@ def get_hip_file_path(filepath, hipify_caffe2): def is_caffe2_gpu_file(filepath): filename = os.path.basename(filepath) _, ext = os.path.splitext(filename) - return 'gpu' in filename or ext in ['.cu', '.cuh'] + return ('gpu' in filename or ext in ['.cu', '.cuh']) and ('cudnn' not in filename) def preprocessor(filepath, stats, hipify_caffe2): diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml index 0f5bfe8d9995e..67d9634b9babc 100644 --- a/tools/autograd/derivatives.yaml +++ b/tools/autograd/derivatives.yaml @@ -283,7 +283,7 @@ self: grad - name: gather(Tensor self, int64_t dim, Tensor index) - self: at::zeros(self.sizes(), grad.type()).scatter_add_(dim, index, grad) + self: at::zeros(self.sizes(), grad.options()).scatter_add_(dim, index, grad) - name: ge_(Tensor self, Scalar other) self: zeros_like(self) @@ -346,7 +346,7 @@ value: grad.index_select(dim, index).sum() - name: index_select(Tensor self, int64_t dim, Tensor index) - self: at::zeros(self.sizes(), grad.type()).index_add_(dim, index, grad) + self: at::zeros(self.sizes(), grad.options()).index_add_(dim, index, grad) - name: inverse(Tensor self) self: -at::mm(result.t(), at::mm(grad, result.t())) @@ -511,14 +511,14 @@ self: zeros_like(grad) - name: normal(Tensor mean, double std, Generator generator) - mean: at::zeros(mean.sizes(), grad.type()) + mean: at::zeros(mean.sizes(), grad.options()) - name: normal(double mean, Tensor std, Generator generator) - std: at::zeros(std.sizes(), grad.type()) + std: at::zeros(std.sizes(), grad.options()) - name: normal(Tensor mean, Tensor std, Generator generator) - mean: at::zeros(mean.sizes(), grad.type()) - std: at::zeros(std.sizes(), grad.type()) + mean: at::zeros(mean.sizes(), grad.options()) + std: at::zeros(std.sizes(), grad.options()) - name: orgqr(Tensor self, Tensor input2) self: not_implemented("orgqr") @@ -1259,3 +1259,7 @@ - name: _thnn_fused_gru_cell(Tensor input_gates, Tensor hidden_gates, Tensor hx, Tensor input_bias, Tensor hidden_bias) input_gates, hidden_gates, hx, input_bias, hidden_bias: _thnn_fused_gru_cell_backward(grad, result1, input_bias.defined()) + +# PackedSequence helpers +- name: _pack_padded_sequence(Tensor input, Tensor lengths, bool batch_first) + input: _pack_padded_sequence_backward(grad, input.sizes(), result1, batch_first) diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py index 1688828669ded..05affcbaa6b71 100644 --- a/tools/autograd/gen_variable_type.py +++ b/tools/autograd/gen_variable_type.py @@ -80,7 +80,7 @@ } METHOD_DECLARATION = CodeTemplate("""\ -virtual ${return_type} ${method_prefix_derived}${api_name}(${type_method_formals}) const override; +${return_type} ${method_prefix_derived}${api_name}(${type_method_formals}) const override; """) METHOD_DEFINITION = CodeTemplate("""\ @@ -302,16 +302,12 @@ def is_differentiable(arg): print('WARNING: derivative ignored for {}'.format(name), file=sys.stderr) def setup_derivative(): - def error_msg(): - name = declaration['api_name'] - return '"the derivative for {} is not implemented"'.format(name) - args_with_derivatives = find_args_with_derivatives() env = {} env['args_with_derivatives'] = reference_args(args_with_derivatives) - env['op'] = func['op'] if func is not None else 'Error' - env['op_ctor'] = '' if func is not None else error_msg() + env['op'] = func['op'] if func is not None else 'NotImplemented' + env['op_ctor'] = '' if func is not None else '"{}"'.format(declaration['api_name']) if is_out_fn: setup = ['throw_error_out_requires_grad("{}");'.format(base_name)] diff --git a/tools/autograd/templates/Functions.cpp b/tools/autograd/templates/Functions.cpp index e0022647c91ac..cc4695f3df504 100644 --- a/tools/autograd/templates/Functions.cpp +++ b/tools/autograd/templates/Functions.cpp @@ -121,7 +121,7 @@ Tensor pow_backward(Tensor grad, const Tensor & self, const Scalar & exponent_) } Tensor pow_backward_self(Tensor grad, const Tensor & self, const Tensor & exponent) { - return at::where(exponent == 0.0, at::zeros({}, grad.type()), grad * exponent * self.pow(exponent - 1)); + return at::where(exponent == 0.0, at::zeros({}, grad.options()), grad * exponent * self.pow(exponent - 1)); } Tensor pow_backward_exponent(Tensor grad, const Tensor & self, const Tensor & exponent) { @@ -164,7 +164,7 @@ Tensor sum_backward(const Tensor & grad, IntList sizes, IntList dims, bool keepd } Tensor reverse_dim(const Tensor& t, int64_t dim) { - Tensor index = at::arange(t.size(dim) - 1, -1, -1, t.type().toScalarType(at::ScalarType::Long)); + Tensor index = at::arange(t.size(dim) - 1, -1, -1, t.options().dtype(at::kLong)); return t.index_select(dim, index); } @@ -175,7 +175,7 @@ Tensor prod_safe_zeros_backward(const Tensor &grad, const Tensor& inp, int64_t d auto ones_size = inp.sizes().vec(); ones_size[dim] = 1; - Tensor ones = at::ones(ones_size, grad.type()); + Tensor ones = at::ones(ones_size, grad.options()); Tensor exclusive_normal_nocp = at::cat({ones, inp.narrow(dim, 0, inp.size(dim) - 1)}, dim); Tensor exclusive_normal = exclusive_normal_nocp.cumprod(dim); @@ -328,8 +328,8 @@ Tensor cumprod_backward(const Tensor &grad, const Tensor &input, int64_t dim) { auto ones_size = input.sizes().vec(); ones_size[dim] = 1; - Tensor ones = at::ones({1}, grad.type()).expand(ones_size); - Tensor grad_input = at::zeros(input.sizes(), grad.type()); + Tensor ones = at::ones({1}, grad.options()).expand(ones_size); + Tensor grad_input = at::zeros(input.sizes(), grad.options()); Tensor prods_from_k_plus_1; Tensor omitted_products; for (int k = 0; k < dim_size; ++k) { @@ -433,7 +433,7 @@ std::vector cat_tensors_backward(const Tensor & grad, const std::vector< auto& shape = sizes[i]; // If input was empty tensor, gradInput should be empty tensor. if (shape == std::vector({0})) { - grad_inputs[i] = at::zeros({0}, grad.type()); + grad_inputs[i] = at::zeros({0}, grad.options()); continue; } auto size = shape[dim]; @@ -546,7 +546,7 @@ Tensor index_select_backward(Tensor grad, int64_t dim, Tensor indices, IntList s grad = grad.unsqueeze(dim); indices = indices.unsqueeze(dim); } - return at::zeros(sizes, grad.type()).scatter_(dim, indices, grad); + return at::zeros(sizes, grad.options()).scatter_(dim, indices, grad); } Tensor slice_backward(Tensor grad, IntList input_sizes, int64_t dim, int64_t start, int64_t end, int64_t step) { @@ -566,25 +566,22 @@ Tensor trace_backward(const Tensor & grad, IntList sizes) { throw std::runtime_error("expected matrix input"); } - auto& long_type = grad.type().toScalarType(at::kLong); - - auto grad_input = at::zeros(sizes[0] * sizes[1], grad.type()); - auto indices = at::arange(0, grad_input.numel(), sizes[1] + 1, long_type); + auto grad_input = at::zeros(sizes[0] * sizes[1], grad.options()); + auto indices = at::arange(0, grad_input.numel(), sizes[1] + 1, grad.options().dtype(at::kLong)); grad_input.index_fill_(0, indices, grad); return grad_input.view(sizes); } Tensor unfold_backward(const Tensor & grad, IntList input_sizes, int64_t dim, int64_t size, int64_t step) { - auto& long_type = grad.type().toScalarType(at::kLong); int64_t numel = 1; for (auto size : input_sizes) { numel *= size; } - auto idx = at::arange(0, numel, long_type).view(input_sizes); + auto idx = at::arange(0, numel, grad.options().dtype(at::kLong)).view(input_sizes); auto idx_unfolded = idx.unfold(dim, size, step).contiguous().view(-1); - auto grad_input = at::zeros({numel}, grad.type()); + auto grad_input = at::zeros({numel}, grad.options()); grad_input.index_add_(0, idx_unfolded, grad.contiguous().view(-1)); return grad_input.view(input_sizes); } @@ -613,7 +610,7 @@ Tensor masked_scatter_backward(const Tensor & grad, const Tensor & mask, IntList if (diff_nelem > 0) { // because mask_selected returns a 1-d tensor with size of masked elements that are 1, // we need to fill out the rest with zeros then reshape back to tensor2's size. - auto zeros_fillin = at::zeros({diff_nelem}, grad.type()); + auto zeros_fillin = at::zeros({diff_nelem}, grad.options()); mask_selected = at::cat({mask_selected, zeros_fillin}, 0); } return mask_selected.view(sizes); @@ -710,7 +707,7 @@ Tensor glu_double_backward_grad_output(const Tensor & grad, const Tensor & input if (dim < 0) dim += input.dim(); auto sizes = input.sizes().vec(); sizes[dim] /= 2; - auto tmp = grad * glu_backward(at::ones(sizes, input.type()), input, dim); + auto tmp = grad * glu_backward(at::ones(sizes, input.options()), input, dim); return tmp.narrow(dim, 0, sizes[dim]) + tmp.narrow(dim, sizes[dim], sizes[dim]); } @@ -799,14 +796,14 @@ Tensor diag_backward(const Tensor & grad, IntList input_sizes, int64_t diagonal) } // Input was a matrix but was not square - auto grad_input = at::zeros(input_sizes, grad.type()); + auto grad_input = at::zeros(input_sizes, grad.options()); auto diag = grad_input.diagonal(diagonal); diag.copy_(grad); return grad_input; } Tensor diagonal_backward(const Tensor & grad, IntList input_sizes, int64_t offset, int64_t dim1, int64_t dim2) { - auto grad_input = at::zeros(input_sizes, grad.type()); + auto grad_input = at::zeros(input_sizes, grad.options()); auto diag = grad_input.diagonal(offset, dim1, dim2); diag.copy_(grad); return grad_input; @@ -1306,7 +1303,7 @@ Tensor as_strided_backward(Tensor grad, TensorGeometry input_geometry, IntList s auto size_i = sizes[i]; auto stride_i = strides[i]; if (size_i == 0) { - return at::zeros(input_geometry.sizes(), grad.type()); + return at::zeros(input_geometry.sizes(), grad.options()); } else if (size_i == 1) { grad = grad.squeeze(i); } else if (stride_i == 0) { @@ -1334,7 +1331,7 @@ Tensor as_strided_backward(Tensor grad, TensorGeometry input_geometry, IntList s auto size_i = inp_sizes[i]; auto stride_i = inp_strides[i]; if (size_i == 0) { - return at::zeros(input_geometry.sizes(), grad.type()); + return at::zeros(input_geometry.sizes(), grad.options()); } else if (size_i != 1) { inp_sizes_.insert(inp_sizes_.begin(), size_i); inp_strides_.insert(inp_strides_.begin(), stride_i); @@ -1362,13 +1359,12 @@ Tensor as_strided_backward(Tensor grad, TensorGeometry input_geometry, IntList s _min_storage_size(inp_sizes_, inp_strides_, inp_effective_offset), _min_storage_size(out_sizes_, out_strides_, out_effective_offset) ); - auto& ty = grad.type(); - auto storage = at::zeros({base_size}, ty); + auto storage = at::zeros({base_size}, grad.options()); // prepare indices tensor if we will do index_add_ later at::optional flatten_full_indices; if (inp_maybe_overlap || out_maybe_overlap) { - flatten_full_indices = at::arange(0, base_size, ty.toScalarType(at::kLong)); + flatten_full_indices = at::arange(0, base_size, grad.options().dtype(at::kLong)); } // Step (2): use output geometry to scatter gradients into storage @@ -1385,7 +1381,7 @@ Tensor as_strided_backward(Tensor grad, TensorGeometry input_geometry, IntList s if (inp_maybe_overlap) { auto count = at::zeros_like(storage); auto inp_indices = flatten_full_indices->as_strided(inp_sizes_, inp_strides_, inp_effective_offset).reshape(-1); - count.index_add_(0, inp_indices, at::ones({1}, ty).expand_as(inp_indices)); + count.index_add_(0, inp_indices, at::ones({1}, grad.options()).expand_as(inp_indices)); storage.div_(count); // this will give nan outside visible range } // Step (4): return as_strided view of the storage tensor with input geometry @@ -1411,9 +1407,9 @@ std::tuple prelu_double_backward( std::array output_mask) { // Zero-fill undefined grads (TODO: do this more efficiently) - auto ggI = mb_ggI.defined() ? mb_ggI : input.type().zeros_like(input); - auto ggW = mb_ggW.defined() ? mb_ggW : weight.type().zeros_like(weight); - auto gO = mb_gO.defined() ? mb_gO : input.type().zeros_like(input); + auto ggI = mb_ggI.defined() ? mb_ggI : at::zeros_like(input); + auto ggW = mb_ggW.defined() ? mb_ggW : at::zeros_like(weight); + auto gO = mb_gO.defined() ? mb_gO : at::zeros_like(input); auto positive_mask = (input > 0).type_as(ggI); auto nonpositive_mask = (input <= 0).type_as(ggW); @@ -1515,7 +1511,7 @@ Tensor svd_backward(const std::vector &grads, const T if (gsigma.defined()) { sigma_term = u.mm(gsigma.diag()).mm(vt); } else { - sigma_term = at::zeros({1}, self.type()).expand_as(self); + sigma_term = at::zeros({1}, self.options()).expand_as(self); } // in case that there are no gu and gv, we can avoid the series of kernel // calls below @@ -1524,8 +1520,8 @@ Tensor svd_backward(const std::vector &grads, const T } auto ut = u.t(); - auto im = eye(m, self.type()); - auto in = eye(n, self.type()); + auto im = eye(m, self.options()); + auto in = eye(n, self.options()); auto sigma_mat = sigma.diag(); auto sigma_mat_inv = sigma.pow(-1).diag(); auto sigma_expanded_sq = sigma.pow(2).expand_as(sigma_mat); @@ -1545,7 +1541,7 @@ Tensor svd_backward(const std::vector &grads, const T } u_term = u_term.mm(vt); } else { - u_term = at::zeros({1}, self.type()).expand_as(self); + u_term = at::zeros({1}, self.options()).expand_as(self); } if (gv.defined()) { @@ -1556,7 +1552,7 @@ Tensor svd_backward(const std::vector &grads, const T } v_term = u.mm(v_term); } else { - v_term = at::zeros({1}, self.type()).expand_as(self); + v_term = at::zeros({1}, self.options()).expand_as(self); } return u_term + sigma_term + v_term; @@ -1665,10 +1661,10 @@ std::tuple trtrs_backward( } } if (!grad_a.defined()) { - grad_a = at::zeros({1}, a.type()).expand_as(a); + grad_a = at::zeros({1}, a.options()).expand_as(a); } if (!grad_b.defined()) { - grad_b = at::zeros({1}, b.type()).expand_as(b); + grad_b = at::zeros({1}, b.options()).expand_as(b); } if (output_mask[1] && grad_m.defined()) { grad_a = grad_a.add(grad_m); @@ -1708,7 +1704,7 @@ Tensor fft_backward(const Tensor& self, const Tensor& grad, int64_t signal_ndim, } zero_grad_shape[signal_ndim] = zero_length; zero_grad_shape[signal_ndim + 1] = 2; - complex_full_grad = at::cat({ grad, at::zeros(zero_grad_shape, grad.type()) }, signal_ndim); + complex_full_grad = at::cat({ grad, at::zeros(zero_grad_shape, grad.options()) }, signal_ndim); } gI = _fft_with_size(complex_full_grad, signal_ndim, /* complex_input */ true, /* complex_output */ true, @@ -1841,7 +1837,7 @@ std::tuple batchnorm_double_backward( ggB_expanded = expand_as_dim1(ggB, input); } } else { - gamma_expanded = input.type().tensor({}).fill_(1); + gamma_expanded = at::ones({}, input.options()); } // define some terms we will reuse @@ -1900,7 +1896,7 @@ std::tuple batchnorm_double_backward( if (affine && ggI.defined()) { if (training) { // gG is just the first backwards with the gamma term removed (then shaped properly) - gG = ggI * first_back_grad_input(gO, sigma2_eps_neg_1_2.type().tensor({}).fill_(1)); + gG = ggI * first_back_grad_input(gO, at::ones({}, sigma2_eps_neg_1_2.options())); gG = sum_exclude_dim1(gG, false); } else { gG = sum_exclude_dim1(ggI * gO * sigma2_eps_neg_1_2, false); diff --git a/tools/autograd/templates/VariableType.cpp b/tools/autograd/templates/VariableType.cpp index 3eafab9ea1da8..c2d7cc39bdeae 100644 --- a/tools/autograd/templates/VariableType.cpp +++ b/tools/autograd/templates/VariableType.cpp @@ -55,10 +55,12 @@ ScalarType VariableType::scalarType() const { Backend VariableType::backend() const { return baseType->backend(); } -bool VariableType::is_cuda() const { return baseType->is_cuda(); } -bool VariableType::is_sparse() const { return baseType->is_sparse(); } -bool VariableType::is_distributed() const { return baseType->is_distributed(); } - +Allocator* VariableType::allocator() const { + return baseType->allocator(); +} +Device VariableType::getDeviceFromPtr(void * data) const { + return baseType->getDeviceFromPtr(data); +} Storage VariableType::storage(bool resizable) const { return baseType->storage(); } @@ -97,10 +99,6 @@ TypeID VariableType::ID() const { return static_cast(id_); } -const char * VariableType::typeString() { - return "VariableType"; -} - std::vector> type_to_variable_type; // XXX - this is not threadsafe with uses of Variables diff --git a/tools/autograd/templates/VariableType.h b/tools/autograd/templates/VariableType.h index 7b04b78e53ee7..b307a1459da1d 100644 --- a/tools/autograd/templates/VariableType.h +++ b/tools/autograd/templates/VariableType.h @@ -34,32 +34,30 @@ void register_variable_type_for(at::Type* baseType); struct TORCH_API VariableType final : public at::TypeDefault { VariableType(Context* context, at::Type* baseType); - virtual at::ScalarType scalarType() const override; - virtual at::Backend backend() const override; - virtual bool is_cuda() const override; - virtual bool is_sparse() const override; - virtual bool is_distributed() const override; - virtual Storage storage(bool resizable = false) const override; - virtual Storage storage(size_t size, bool resizable = false) const override; - virtual Storage storageFromBlob(void * data, int64_t size, const std::function & deleter) const override; - virtual Storage storageWithAllocator(int64_t size, at::Allocator* allocator) const override; - virtual std::unique_ptr generator() const override; - virtual const char * toString() const override; - virtual at::TypeID ID() const override; - virtual size_t elementSizeInBytes() const override; - virtual at::Type & toBackend(at::Backend b) const override; - virtual at::Type & toScalarType(at::ScalarType s) const override; - static const char * typeString(); - virtual Storage unsafeStorageFromTH(void * th_pointer, bool retain) const override; - virtual at::Tensor unsafeTensorFromTH(void * th_pointer, bool retain) const override; + at::ScalarType scalarType() const override; + at::Backend backend() const override; + at::Allocator* allocator() const override; + at::Device getDeviceFromPtr(void * data) const override; + Storage storage(bool resizable = false) const override; + Storage storage(size_t size, bool resizable = false) const override; + Storage storageFromBlob(void * data, int64_t size, const std::function & deleter) const override; + Storage storageWithAllocator(int64_t size, at::Allocator* allocator) const override; + std::unique_ptr generator() const override; + const char * toString() const override; + at::TypeID ID() const override; + size_t elementSizeInBytes() const override; + at::Type & toBackend(at::Backend b) const override; + at::Type & toScalarType(at::ScalarType s) const override; + Storage unsafeStorageFromTH(void * th_pointer, bool retain) const override; + at::Tensor unsafeTensorFromTH(void * th_pointer, bool retain) const override; static at::Type* getVariableTypeFromBaseType(const at::Type& baseType); static bool isVariableType(const at::Type& type); static std::vector allCUDATypes(); static std::vector allCPUTypes(); - virtual Tensor & s_copy_(Tensor & self, const Tensor & src, bool non_blocking) const override; - virtual Tensor & _s_copy_from(const Tensor & self, Tensor & dst, bool non_blocking) const override; + Tensor & s_copy_(Tensor & self, const Tensor & src, bool non_blocking) const override; + Tensor & _s_copy_from(const Tensor & self, Tensor & dst, bool non_blocking) const override; ${type_derived_method_declarations} private: diff --git a/tools/build_pytorch_libs.bat b/tools/build_pytorch_libs.bat index 1743ea3b253cd..b122df45051e3 100755 --- a/tools/build_pytorch_libs.bat +++ b/tools/build_pytorch_libs.bat @@ -181,6 +181,7 @@ goto:eof -DUSE_NNPACK=%USE_NNPACK% ^ -DUSE_GLOG=OFF ^ -DUSE_GFLAGS=OFF ^ + -DUSE_SYSTEM_EIGEN_INSTALL=OFF ^ -DCUDNN_INCLUDE_DIR="%CUDNN_INCLUDE_DIR%" ^ -DCUDNN_LIB_DIR="%CUDNN_LIB_DIR%" ^ -DCUDNN_LIBRARY="%CUDNN_LIBRARY%" ^ diff --git a/tools/build_pytorch_libs.sh b/tools/build_pytorch_libs.sh index 37f23dbcafadc..a8b0b9381d67c 100755 --- a/tools/build_pytorch_libs.sh +++ b/tools/build_pytorch_libs.sh @@ -17,6 +17,7 @@ USE_NNPACK=0 USE_MKLDNN=0 USE_GLOO_IBVERBS=0 FULL_CAFFE2=0 +CAFFE2_STATIC_LINK_CUDA=0 while [[ $# -gt 0 ]]; do case "$1" in --use-cuda) @@ -170,6 +171,7 @@ function build() { -DTH_LIB_PATH="$INSTALL_DIR/lib" \ -DTH_LIBRARIES="$INSTALL_DIR/lib/libTH$LD_POSTFIX" \ -DCAFFE2_LIBRARIES="$INSTALL_DIR/lib/libcaffe2$LD_POSTFIX" \ + -DCAFFE2_STATIC_LINK_CUDA=$CAFFE2_STATIC_LINK_CUDA \ -DTHNN_LIBRARIES="$INSTALL_DIR/lib/libTHNN$LD_POSTFIX" \ -DTHCUNN_LIBRARIES="$INSTALL_DIR/lib/libTHCUNN$LD_POSTFIX" \ -DTHS_LIBRARIES="$INSTALL_DIR/lib/libTHS$LD_POSTFIX" \ @@ -238,6 +240,9 @@ function build_nccl() { # detected them (to ensure that we have a consistent view between the # PyTorch and Caffe2 builds.) function build_caffe2() { + # pwd is pytorch_root/build + + # TODO change these to CMAKE_ARGS for consistency if [[ -z $EXTRA_CAFFE2_CMAKE_FLAGS ]]; then EXTRA_CAFFE2_CMAKE_FLAGS=() fi @@ -255,7 +260,7 @@ function build_caffe2() { -DBUILDING_WITH_TORCH_LIBS=ON \ -DCMAKE_BUILD_TYPE=$BUILD_TYPE \ -DBUILD_TORCH=$BUILD_TORCH \ - -DBUILD_PYTHON=$FULL_CAFFE2 \ + -DBUILD_PYTHON=ON \ -DBUILD_SHARED_LIBS=$BUILD_SHARED_LIBS \ -DBUILD_BINARY=$FULL_CAFFE2 \ -DBUILD_TEST=$FULL_CAFFE2 \ @@ -267,6 +272,7 @@ function build_caffe2() { -DUSE_NNPACK=$USE_NNPACK \ -DUSE_GLOG=OFF \ -DUSE_GFLAGS=OFF \ + -DUSE_SYSTEM_EIGEN_INSTALL=OFF \ -DCUDNN_INCLUDE_DIR=$CUDNN_INCLUDE_DIR \ -DCUDNN_LIB_DIR=$CUDNN_LIB_DIR \ -DCUDNN_LIBRARY=$CUDNN_LIBRARY \ @@ -286,21 +292,15 @@ function build_caffe2() { # This is needed by the aten tests built with caffe2 if [ -f "${INSTALL_DIR}/lib/libnccl.so" ] && [ ! -f "lib/libnccl.so.1" ]; then + # cp root/torch/lib/tmp_install/libnccl root/build/lib/libnccl cp "${INSTALL_DIR}/lib/libnccl.so.1" "lib/libnccl.so.1" fi ${CMAKE_INSTALL} -j"$MAX_JOBS" - # Install Python proto files - if [[ $FULL_CAFFE2 -ne 0 ]]; then - find . -name proto - for proto_file in ./caffe2/proto/*.py; do - cp $proto_file "../caffe2/proto/" - done - fi - # Fix rpaths of shared libraries if [[ $(uname) == 'Darwin' ]]; then + # root/torch/lib/tmp_install/lib pushd "$INSTALL_DIR/lib" for lib in *.dylib; do echo "Updating install_name for $lib" diff --git a/tools/clang_tidy.py b/tools/clang_tidy.py index 77b101dedf0f3..22c2accb6be6c 100644 --- a/tools/clang_tidy.py +++ b/tools/clang_tidy.py @@ -210,7 +210,7 @@ def parse_options(): def main(): options = parse_options() - paths = map(normalize_directory_path, options.paths) + paths = list(map(normalize_directory_path, options.paths)) if options.revision: files = get_changed_files(options.revision, paths, options.verbose) else: diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt index 08ba356e1c6f3..1fc87e10be15a 100644 --- a/torch/CMakeLists.txt +++ b/torch/CMakeLists.txt @@ -209,6 +209,7 @@ if (NOT NO_API AND NOT USE_ROCM) ${TORCH_SRC_DIR}/csrc/api/src/optim/lbfgs.cpp ${TORCH_SRC_DIR}/csrc/api/src/optim/rmsprop.cpp ${TORCH_SRC_DIR}/csrc/api/src/optim/sgd.cpp + ${TORCH_SRC_DIR}/csrc/api/src/jit.cpp ) endif() @@ -349,7 +350,7 @@ if(USE_ROCM) /opt/rocm/include /opt/rocm/hcc/include /opt/rocm/rocblas/include - /opt/rocm/hcsparse/include + /opt/rocm/hipsparse/include ) endif() @@ -433,6 +434,7 @@ if (BUILD_TORCH_TEST AND NOT NO_API AND NOT USE_ROCM) ${TORCH_API_TEST_DIR}/static.cpp ${TORCH_API_TEST_DIR}/tensor_cuda.cpp ${TORCH_API_TEST_DIR}/tensor.cpp + ${TORCH_API_TEST_DIR}/jit.cpp # Temporary until ATen tests are built with Caffe2 ${TORCH_API_TEST_DIR}/tensor_options.cpp ${TORCH_API_TEST_DIR}/tensor_options_cuda.cpp diff --git a/torch/_tensor_str.py b/torch/_tensor_str.py index 71f2afb45b49f..3d0c6b0c1ae1e 100644 --- a/torch/_tensor_str.py +++ b/torch/_tensor_str.py @@ -174,7 +174,7 @@ def _vector_str(self, indent, formatter, summarize): return '[' + (',' + '\n' + ' ' * (indent + 1)).join(lines) + ']' -def _tensor_str(self, indent, formatter, summarize): +def _tensor_str_with_formatter(self, indent, formatter, summarize): dim = self.dim() if dim == 0: @@ -183,24 +183,42 @@ def _tensor_str(self, indent, formatter, summarize): return _vector_str(self, indent, formatter, summarize) if summarize and self.size(0) > 2 * PRINT_OPTS.edgeitems: - slices = ([_tensor_str(self[i], indent + 1, formatter, summarize) + slices = ([_tensor_str_with_formatter(self[i], indent + 1, formatter, summarize) for i in range(0, PRINT_OPTS.edgeitems)] + ['...'] + - [_tensor_str(self[i], indent + 1, formatter, summarize) + [_tensor_str_with_formatter(self[i], indent + 1, formatter, summarize) for i in range(len(self) - PRINT_OPTS.edgeitems, len(self))]) else: - slices = [_tensor_str(self[i], indent + 1, formatter, summarize) for i in range(0, self.size(0))] + slices = [_tensor_str_with_formatter(self[i], indent + 1, formatter, summarize) + for i in range(0, self.size(0))] tensor_str = (',' + '\n' * (dim - 1) + ' ' * (indent + 1)).join(slices) return '[' + tensor_str + ']' -def _maybe_wrap_suffix(suffix, indent, tensor_str): - suffix_len = len(suffix) +def _tensor_str(self, indent): + if self.numel() == 0: + return '[]' + + summarize = self.numel() > PRINT_OPTS.threshold + formatter = _Formatter(get_summarized_data(self) if summarize else self) + return _tensor_str_with_formatter(self, indent, formatter, summarize) + + +def _add_suffixes(tensor_str, suffixes, indent, force_newline): + tensor_strs = [tensor_str] last_line_len = len(tensor_str) - tensor_str.rfind('\n') + 1 - if suffix_len > 2 and last_line_len + suffix_len > PRINT_OPTS.linewidth: - return ',\n' + ' ' * indent + suffix[2:] - return suffix + for suffix in suffixes: + suffix_len = len(suffix) + if force_newline or last_line_len + suffix_len + 2 > PRINT_OPTS.linewidth: + tensor_strs.append(',\n' + ' ' * indent + suffix) + last_line_len = indent + suffix_len + force_newline = False + else: + tensor_strs.append(', ' + suffix) + last_line_len += suffix_len + 2 + tensor_strs.append(')') + return ''.join(tensor_strs) def get_summarized_data(self): @@ -222,50 +240,60 @@ def get_summarized_data(self): def _str(self): - if self.is_sparse: - size_str = str(tuple(self.shape)).replace(' ', '') - return '{} of size {} with indices:\n{}\nand values:\n{}'.format( - self.type(), size_str, self._indices(), self._values()) - prefix = 'tensor(' indent = len(prefix) - summarize = self.numel() > PRINT_OPTS.threshold - suffix = '' + suffixes = [] if not torch._C._is_default_type_cuda(): if self.device.type == 'cuda': - suffix += ', device=\'' + str(self.device) + '\'' + suffixes.append('device=\'' + str(self.device) + '\'') else: if self.device.type == 'cpu' or torch.cuda.current_device() != self.device.index: - suffix += ', device=\'' + str(self.device) + '\'' + suffixes.append('device=\'' + str(self.device) + '\'') - if self.numel() == 0: - # Explicitly print the shape if it is not (0,), to match NumPy behavior - if self.dim() != 1: - suffix += ', size=' + str(tuple(self.shape)) - - # In an empty tensor, there are no elements to infer if the dtype should be int64, - # so it must be shown explicitly. - if self.dtype != torch.get_default_dtype(): - suffix += ', dtype=' + str(self.dtype) - tensor_str = '[]' + has_default_dtype = self.dtype == torch.get_default_dtype() or self.dtype == torch.int64 + + if self.is_sparse: + suffixes.append('size=' + str(tuple(self.shape))) + suffixes.append('nnz=' + str(self._nnz())) + if not has_default_dtype: + suffixes.append('dtype=' + str(self.dtype)) + indices_prefix = 'indices=tensor(' + indices = self._indices().detach() + indices_str = _tensor_str(indices, indent + len(indices_prefix)) + if indices.numel() == 0: + indices_str += ', size=' + str(tuple(indices.shape)) + values_prefix = 'values=tensor(' + values = self._values().detach() + values_str = _tensor_str(values, indent + len(values_prefix)) + if values.numel() == 0: + values_str += ', size=' + str(tuple(values.shape)) + tensor_str = indices_prefix + indices_str + '),\n' + ' ' * indent + values_prefix + values_str + ')' else: - if self.dtype != torch.get_default_dtype() and self.dtype != torch.int64: - suffix += ', dtype=' + str(self.dtype) + if self.numel() == 0 and not self.is_sparse: + # Explicitly print the shape if it is not (0,), to match NumPy behavior + if self.dim() != 1: + suffixes.append('size=' + str(tuple(self.shape))) + + # In an empty tensor, there are no elements to infer if the dtype + # should be int64, so it must be shown explicitly. + if self.dtype != torch.get_default_dtype(): + suffixes.append('dtype=' + str(self.dtype)) + tensor_str = '[]' + else: + if not has_default_dtype: + suffixes.append('dtype=' + str(self.dtype)) + tensor_str = _tensor_str(self, indent) - formatter = _Formatter(get_summarized_data(self) if summarize else self) - tensor_str = _tensor_str(self, indent, formatter, summarize) + if self.layout != torch.strided: + suffixes.append('layout=' + str(self.layout)) if self.grad_fn is not None: name = type(self.grad_fn).__name__ if name == 'CppFunction': name = self.grad_fn.name().rsplit('::', maxsplit=1)[-1] - suffix += ', grad_fn=<{}>'.format(name) + suffixes.append('grad_fn=<{}>'.format(name)) elif self.requires_grad: - suffix += ', requires_grad=True' - - suffix += ')' - - suffix = _maybe_wrap_suffix(suffix, indent, tensor_str) + suffixes.append('requires_grad=True') - return prefix + tensor_str + suffix + return _add_suffixes(prefix + tensor_str, suffixes, indent, force_newline=self.is_sparse) diff --git a/torch/csrc/Size.cpp b/torch/csrc/Size.cpp index 6bc458592f957..d11a718634473 100644 --- a/torch/csrc/Size.cpp +++ b/torch/csrc/Size.cpp @@ -135,6 +135,22 @@ static PyMappingMethods THPSize_as_mapping = { 0 }; +static PyObject *THPSize_numel(THPSize *self) +{ + HANDLE_TH_ERRORS + int64_t numel = 1; + for (Py_ssize_t i = 0; i < PyTuple_Size((PyObject*)self); ++i) { + numel *= PyLong_AsLong(PyTuple_GET_ITEM(self, i)); + } + return THPUtils_packInt64(numel); + END_HANDLE_TH_ERRORS +} + +static PyMethodDef THPSize_methods[] = { + {"numel", (PyCFunction)THPSize_numel, METH_NOARGS, nullptr}, + {nullptr} +}; + PyTypeObject THPSizeType = { PyVarObject_HEAD_INIT(nullptr, 0) @@ -157,14 +173,14 @@ PyTypeObject THPSizeType = { 0, /* tp_setattro */ 0, /* tp_as_buffer */ Py_TPFLAGS_DEFAULT, /* tp_flags */ - nullptr, /* tp_doc */ + nullptr, /* tp_doc */ 0, /* tp_traverse */ 0, /* tp_clear */ 0, /* tp_richcompare */ 0, /* tp_weaklistoffset */ 0, /* tp_iter */ 0, /* tp_iternext */ - 0, /* tp_methods */ + THPSize_methods, /* tp_methods */ 0, /* tp_members */ 0, /* tp_getset */ &PyTuple_Type, /* tp_base */ diff --git a/torch/csrc/Size.h b/torch/csrc/Size.h index 03a4c48a1c280..2ce259248ae91 100644 --- a/torch/csrc/Size.h +++ b/torch/csrc/Size.h @@ -2,7 +2,7 @@ #include "torch/csrc/python_headers.h" #include "torch/csrc/autograd/variable.h" -#include "stdint.h" +#include "cstdint" extern PyTypeObject THPSizeType; diff --git a/torch/csrc/Types.h b/torch/csrc/Types.h index 8342c7a94b0aa..f1bc466dc053d 100644 --- a/torch/csrc/Types.h +++ b/torch/csrc/Types.h @@ -5,7 +5,7 @@ #include #ifndef INT64_MAX -#include "stdint.h" +#include "cstdint" #endif template struct THPTypeInfo {}; diff --git a/torch/csrc/api/include/torch/detail/ordered_dict.h b/torch/csrc/api/include/torch/detail/ordered_dict.h index 1b576922da3cb..cb7cc59a3df66 100644 --- a/torch/csrc/api/include/torch/detail/ordered_dict.h +++ b/torch/csrc/api/include/torch/detail/ordered_dict.h @@ -72,7 +72,9 @@ class OrderedDict { // Move works by default, because you can move-construct vectors of const // values.. - OrderedDict(OrderedDict&& other) = default; + OrderedDict(OrderedDict&& other) noexcept( + noexcept(std::unordered_map()) && + noexcept(std::vector())) = default; OrderedDict& operator=(OrderedDict&& other) = default; ~OrderedDict() = default; diff --git a/torch/csrc/api/include/torch/detail/static.h b/torch/csrc/api/include/torch/detail/static.h index e33e2c27e126a..45ef0c029f2fb 100644 --- a/torch/csrc/api/include/torch/detail/static.h +++ b/torch/csrc/api/include/torch/detail/static.h @@ -50,7 +50,7 @@ inline constexpr bool check_not_lvalue_references() { return true; } -/// A type trait whose `::value` member is true if `M` derives from `Module`. +/// A type trait whose `value` member is true if `M` derives from `Module`. template using is_module = std::is_base_of::type>; diff --git a/torch/csrc/api/include/torch/jit.h b/torch/csrc/api/include/torch/jit.h new file mode 100644 index 0000000000000..dfcf26dcda2ae --- /dev/null +++ b/torch/csrc/api/include/torch/jit.h @@ -0,0 +1,34 @@ +#pragma once +#include +#include +#include + +#include +#include + +namespace torch { +namespace jit { + +/// Compiles Python JIT code into a graph that can be executed. +/// +/// For example: +/// @code +/// auto module = torch::jit::compile(R"JIT( +/// def relu_script(a, b): +/// return torch.relu(a + b) +/// def test_while(a, i): +/// while i < 10: +/// a += a +/// i += 1 +/// return a +/// )JIT"); +/// IValue output = module->run_method("relu_script", a, b); +/// @endcode +/// +/// @param source A string containing functions containing script code to +/// compile +/// @return A module containing the compiled functions +std::shared_ptr compile(const std::string& source); + +} // namespace jit +} // namespace torch diff --git a/torch/csrc/api/include/torch/nn/modules/linear.h b/torch/csrc/api/include/torch/nn/modules/linear.h index 5642a9acc58c7..1d1a25da36163 100644 --- a/torch/csrc/api/include/torch/nn/modules/linear.h +++ b/torch/csrc/api/include/torch/nn/modules/linear.h @@ -12,24 +12,41 @@ namespace torch { namespace nn { struct LinearOptions { LinearOptions(int64_t in, int64_t out); + /// The number of input features (columns of the input matrix). TORCH_ARG(int64_t, in); + /// The number of output features to produce (columns of the output matrix). TORCH_ARG(int64_t, out); + /// Whether to learn and add a bias after the linear transformation. TORCH_ARG(bool, with_bias) = true; }; +/// Applies a linear transformation with optional bias. class LinearImpl : public Cloneable { public: LinearImpl(int64_t in, int64_t out) : LinearImpl(LinearOptions(in, out)) {} explicit LinearImpl(LinearOptions options); void reset() override; - Tensor forward(Tensor); + /// Transforms the `input` tensor by multiplying with the `weight` and + /// optionally adding the `bias`, if `with_bias` is true in the options. + Tensor forward(Tensor input); + + /// The options used to configure this module. LinearOptions options; + + /// The learned weight. Tensor weight; + + /// The learned bias. If `with_bias` is false in the `options`, this tensor is + /// undefined. Tensor bias; }; +/// A `ModuleHolder` subclass for `LinearImpl`. +/// See the documentation for `LinearImpl` class to learn what methods it +/// provides, or the documentation for `ModuleHolder` to learn about PyTorch's +/// module storage semantics. TORCH_MODULE(Linear); } // namespace nn diff --git a/torch/csrc/api/include/torch/nn/modules/rnn.h b/torch/csrc/api/include/torch/nn/modules/rnn.h index 17ff80b66197c..326e9e267a9f3 100644 --- a/torch/csrc/api/include/torch/nn/modules/rnn.h +++ b/torch/csrc/api/include/torch/nn/modules/rnn.h @@ -17,29 +17,47 @@ namespace torch { namespace nn { +/// The output of a single invocation of an RNN module's `forward()` method. struct RNNOutput { + /// The result of applying the specific RNN algorithm + /// to the input tensor and input state. Tensor output; + /// The new, updated state that can be fed into the RNN + /// in the next forward step. Tensor state; }; namespace detail { + +/// Common options for LSTM and GRU modules. struct RNNOptionsBase { RNNOptionsBase(int64_t input_size, int64_t hidden_size); virtual ~RNNOptionsBase() = default; + /// The number of features of a single sample in the input sequence `x`. TORCH_ARG(int64_t, input_size); + /// The number of features in the hidden state `h`. TORCH_ARG(int64_t, hidden_size); + /// The number of recurrent layers (cells) to use. TORCH_ARG(int64_t, layers) = 1; + /// Whether a bias term should be added to all linear operations. TORCH_ARG(bool, with_bias) = true; + /// If non-zero, adds dropout with the given probability to the output of each + /// RNN layer, except the final layer. TORCH_ARG(double, dropout) = 0.0; + /// Whether to make the RNN bidirectional. TORCH_ARG(bool, bidirectional) = false; + /// If true, the input sequence should be provided as `(batch, sequence, + /// features)`. If false (default), the expected layout is `(sequence, batch, + /// features)`. TORCH_ARG(bool, batch_first) = false; }; +/// Base class for all RNN implementations (intended for code sharing). template class RNNImplBase : public torch::nn::Cloneable { public: - // These must line up with the CUDNN mode codes: - // https://docs.nvidia.com/deeplearning/sdk/cudnn-developer-guide/index.html#cudnnRNNMode_t + /// These must line up with the CUDNN mode codes: + /// https://docs.nvidia.com/deeplearning/sdk/cudnn-developer-guide/index.html#cudnnRNNMode_t enum class CuDNNMode { RNN_RELU = 0, RNN_TANH = 1, LSTM = 2, GRU = 3 }; explicit RNNImplBase( @@ -121,72 +139,110 @@ class RNNImplBase : public torch::nn::Cloneable { enum class RNNActivation { ReLU, Tanh }; +/// Options for RNN modules. struct RNNOptions { RNNOptions(int64_t input_size, int64_t hidden_size); + /// Sets the activation after linear operations to `tanh`. RNNOptions& tanh(); + /// Sets the activation after linear operations to `relu`. RNNOptions& relu(); + /// The number of features of a single sample in the input sequence `x`. TORCH_ARG(int64_t, input_size); + /// The number of features in the hidden state `h`. TORCH_ARG(int64_t, hidden_size); + /// The number of recurrent layers (cells) to use. TORCH_ARG(int64_t, layers) = 1; + /// Whether a bias term should be added to all linear operations. TORCH_ARG(bool, with_bias) = true; + /// If non-zero, adds dropout with the given probability to the output of each + /// RNN layer, except the final layer. TORCH_ARG(double, dropout) = 0.0; + /// Whether to make the RNN bidirectional. TORCH_ARG(bool, bidirectional) = false; + /// If true, the input sequence should be provided as `(batch, sequence, + /// features)`. If false (default), the expected layout is `(sequence, batch, + /// features)`. TORCH_ARG(bool, batch_first) = false; + /// The activation to use after linear operations. TORCH_ARG(RNNActivation, activation) = RNNActivation::ReLU; }; +/// A multi-layer Elman RNN module with Tanh or ReLU activation. +/// See https://pytorch.org/docs/master/nn.html#torch.nn.RNN for more +/// documenation. class RNNImpl : public detail::RNNImplBase { public: RNNImpl(int64_t input_size, int64_t hidden_size) : RNNImpl(RNNOptions(input_size, hidden_size)) {} explicit RNNImpl(RNNOptions options); + /// Applies the `RNN` module to an input sequence and input state. + /// The `input` should follow a `(sequence, batch, features)` layout unless + /// `batch_first` is true, in which case the layout should be `(batch, + /// sequence, features)`. RNNOutput forward(Tensor input, Tensor state = {}); RNNOptions options; }; -/// A multi-layer Elman RNN module with Tanh or ReLU activation. -/// See https://pytorch.org/docs/master/nn.html#torch.nn.RNN for more -/// documenation. +/// A `ModuleHolder` subclass for `RNNImpl`. +/// See the documentation for `RNNImpl` class to learn what methods it provides, +/// or the documentation for `ModuleHolder` to learn about PyTorch's module +/// storage semantics. TORCH_MODULE(RNN); // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ LSTM ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ using LSTMOptions = detail::RNNOptionsBase; +/// A multi-layer long-short-term-memory (LSTM) module. +/// See https://pytorch.org/docs/master/nn.html#torch.nn.LSTM for more +/// documenation. class LSTMImpl : public detail::RNNImplBase { public: LSTMImpl(int64_t input_size, int64_t hidden_size) : LSTMImpl(LSTMOptions(input_size, hidden_size)) {} explicit LSTMImpl(LSTMOptions options); + /// Applies the `LSTM` module to an input sequence and input state. + /// The `input` should follow a `(sequence, batch, features)` layout unless + /// `batch_first` is true, in which case the layout should be `(batch, + /// sequence, features)`. RNNOutput forward(Tensor input, Tensor state = {}); }; -/// A multi-layer long-short-term-memory (LSTM) module. -/// See https://pytorch.org/docs/master/nn.html#torch.nn.LSTM for more -/// documenation. +/// A `ModuleHolder` subclass for `LSTMImpl`. +/// See the documentation for `LSTMImpl` class to learn what methods it +/// provides, or the documentation for `ModuleHolder` to learn about PyTorch's +/// module storage semantics. TORCH_MODULE(LSTM); // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ GRU ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ using GRUOptions = detail::RNNOptionsBase; +/// A multi-layer gated recurrent unit (GRU) module. +/// See https://pytorch.org/docs/master/nn.html#torch.nn.GRU for more +/// documenation. class GRUImpl : public detail::RNNImplBase { public: GRUImpl(int64_t input_size, int64_t hidden_size) : GRUImpl(GRUOptions(input_size, hidden_size)) {} explicit GRUImpl(GRUOptions options); + /// Applies the `GRU` module to an input sequence and input state. + /// The `input` should follow a `(sequence, batch, features)` layout unless + /// `batch_first` is true, in which case the layout should be `(batch, + /// sequence, features)`. RNNOutput forward(Tensor input, Tensor state = {}); }; -/// A multi-layer gated recurrent unit (GRU) module. -/// See https://pytorch.org/docs/master/nn.html#torch.nn.GRU for more -/// documenation. +/// A `ModuleHolder` subclass for `GRUImpl`. +/// See the documentation for `GRUImpl` class to learn what methods it provides, +/// or the documentation for `ModuleHolder` to learn about PyTorch's module +/// storage semantics. TORCH_MODULE(GRU); } // namespace nn diff --git a/torch/csrc/api/include/torch/nn/pimpl.h b/torch/csrc/api/include/torch/nn/pimpl.h index 772e19b5e5307..af83f2e0ed731 100644 --- a/torch/csrc/api/include/torch/nn/pimpl.h +++ b/torch/csrc/api/include/torch/nn/pimpl.h @@ -12,7 +12,7 @@ namespace detail { /// This class exists only to do SFINAE on abstract types `T` that are really /// `ModuleHolder`, because there's no good way to say that `T` is a /// `ModuleHolder` over some unknown type `ModuleType`. With this, you can do -/// enable_if_t::value>::type. +/// `enable_if_t>`. struct ModuleHolderIndicator {}; template @@ -155,31 +155,31 @@ class ModuleHolder : torch::detail::ModuleHolderIndicator { } // namespace nn } // namespace torch -#define TORCH_ARG(T, name) \ - auto name(const T& new_##name)->decltype(*this) { \ - this->name##_ = new_##name; \ - return *this; \ - } \ - auto name(T&& new_##name)->decltype(*this) { \ - this->name##_ = std::move(new_##name); \ - return *this; \ - } \ - const T& name() const noexcept { \ - return this->name##_; \ - } \ - T name##_ +#define TORCH_ARG(T, name) \ + auto name(const T& new_##name)->decltype(*this) { /* NOLINT */ \ + this->name##_ = new_##name; \ + return *this; \ + } \ + auto name(T&& new_##name)->decltype(*this) { /* NOLINT */ \ + this->name##_ = std::move(new_##name); \ + return *this; \ + } \ + const T& name() const noexcept { /* NOLINT */ \ + return this->name##_; \ + } \ + T name##_ /* NOLINT */ /// Defines a class `Name` which inherits from `nn::ModuleHolder` to provide a /// wrapper over a `std::shared_ptr`. -#define TORCH_MODULE_IMPL(Name, Impl) \ - class Name : public torch::nn::ModuleHolder { \ - public: \ - using torch::nn::ModuleHolder::ModuleHolder; \ - Name(const Name&) = default; \ - Name(Name&&) = default; \ - Name(Name& other) : Name(static_cast(other)) {} \ - Name& operator=(const Name&) = default; \ - Name& operator=(Name&&) = default; \ +#define TORCH_MODULE_IMPL(Name, Impl) \ + class Name : public torch::nn::ModuleHolder { /* NOLINT */ \ + public: \ + using torch::nn::ModuleHolder::ModuleHolder; \ + Name(const Name&) = default; /* NOLINT */ \ + Name(Name&&) = default; /* NOLINT */ \ + Name(Name& other) : Name(static_cast(other)) {} /* NOLINT */ \ + Name& operator=(const Name&) = default; /* NOLINT */ \ + Name& operator=(Name&&) = default; /* NOLINT */ \ } /// Like `TORCH_MODULE_IMPL`, but defaults the `Impl` name to `Impl`. diff --git a/torch/csrc/api/include/torch/optim/optimizer.h b/torch/csrc/api/include/torch/optim/optimizer.h index cb86c842843d8..feb8ef617f95f 100644 --- a/torch/csrc/api/include/torch/optim/optimizer.h +++ b/torch/csrc/api/include/torch/optim/optimizer.h @@ -31,21 +31,20 @@ class OptimizerBase { virtual ~OptimizerBase() = default; /// Adds the given vector of parameters to the optimizer's parameter list. - /// Override this method if you want to modify the way parameters are added to - /// the `Optimizer`. - virtual void add_parameters(const std::vector& parameters); + void add_parameters(const std::vector& parameters); /// Adds the `ParameterCursor`'s parameters to the optimizer's parameter list. - /// NOTE: Calls the `vector` overload of `add_parameters` -- override - /// that method if you want to modify the behavior of `add_parameters`. - virtual void add_parameters(const ParameterCursor& cursor); + void add_parameters(const ParameterCursor& cursor); /// Zeros out the gradients of all parameters. virtual void zero_grad(); - /// Provides a reference to the parameters this optimizer holds. + /// Provides a const reference to the parameters this optimizer holds. const std::vector& parameters() const noexcept; + /// Provides a reference to the parameters this optimizer holds. + std::vector& parameters() noexcept; + /// Returns the number of parameters referenced by the optimizer. size_t size() const noexcept; diff --git a/torch/csrc/api/src/jit.cpp b/torch/csrc/api/src/jit.cpp new file mode 100644 index 0000000000000..7d7dedd33ec9c --- /dev/null +++ b/torch/csrc/api/src/jit.cpp @@ -0,0 +1,19 @@ +#include + +#include +#include + +#include +#include + +namespace torch { +namespace jit { + +std::shared_ptr compile(const std::string& source) { + auto module = std::make_shared(); + defineMethodsInModule(*module, source, script::nativeResolver, /*self=*/nullptr); + return module; +} + +} // namespace jit +} // namespace torch diff --git a/torch/csrc/api/src/nn/modules/functional.cpp b/torch/csrc/api/src/nn/modules/functional.cpp index 878a58f059770..591634db23640 100644 --- a/torch/csrc/api/src/nn/modules/functional.cpp +++ b/torch/csrc/api/src/nn/modules/functional.cpp @@ -7,7 +7,7 @@ namespace torch { namespace nn { -FunctionalImpl::FunctionalImpl(std::function function) +FunctionalImpl::FunctionalImpl(Function function) : function_(std::move(function)) {} void FunctionalImpl::reset() {} diff --git a/torch/csrc/api/src/optim/optimizer.cpp b/torch/csrc/api/src/optim/optimizer.cpp index 9eed489f1c648..b57bf6746167e 100644 --- a/torch/csrc/api/src/optim/optimizer.cpp +++ b/torch/csrc/api/src/optim/optimizer.cpp @@ -36,6 +36,14 @@ void OptimizerBase::zero_grad() { } } +const std::vector& OptimizerBase::parameters() const noexcept { + return parameters_; +} + +std::vector& OptimizerBase::parameters() noexcept { + return parameters_; +} + size_t OptimizerBase::size() const noexcept { return parameters_.size(); } diff --git a/torch/csrc/autograd/anomaly_mode.cpp b/torch/csrc/autograd/anomaly_mode.cpp index 7392286671e9f..edf3476b81396 100644 --- a/torch/csrc/autograd/anomaly_mode.cpp +++ b/torch/csrc/autograd/anomaly_mode.cpp @@ -2,6 +2,6 @@ namespace torch { namespace autograd { -bool AnomalyMode::_enabled = 0; +bool AnomalyMode::_enabled = false; }} diff --git a/torch/csrc/autograd/engine.cpp b/torch/csrc/autograd/engine.cpp index cb024a029620e..a8dc01fa5ec3a 100644 --- a/torch/csrc/autograd/engine.cpp +++ b/torch/csrc/autograd/engine.cpp @@ -63,7 +63,7 @@ struct FunctionTask { FunctionTask(GraphTask* base, std::shared_ptr fn, InputBuffer inputs) : base(base) - , fn(fn) + , fn(std::move(fn)) , inputs(std::move(inputs)) {} }; @@ -170,15 +170,10 @@ struct GraphTask { } GraphTask(bool keep_graph, bool grad_mode) - : exception() - , has_error(false) + : has_error(false) , outstanding_tasks(0) , keep_graph(keep_graph) , grad_mode(grad_mode) - , mutex() - , not_done() - , not_ready() - , dependencies() , owner(NO_DEVICE) {} }; @@ -194,12 +189,12 @@ auto ReadyQueue::push(FunctionTask item) -> void { auto ReadyQueue::pop() -> FunctionTask { std::unique_lock lock(mutex); not_empty.wait(lock, [this]{ return !heap.empty(); }); + // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) auto task = std::move(const_cast(heap.top())); heap.pop(); return task; } -Engine::Engine() : ready_queues() { -} +Engine::Engine() = default; // This Engine's ReadyQueues and their corresponding threads are leaked here Engine::~Engine() = default; @@ -376,6 +371,7 @@ static variable_list call_function(FunctionTask& task) { checkpoint_valid = prev_checkpoint_valid_state; if(has_post_hooks){ + // NOLINTNEXTLINE(bugprone-use-after-move) return call_post_hooks(fn, std::move(outputs), std::move(inputs)); } return outputs; @@ -478,7 +474,7 @@ auto Engine::compute_dependencies(Function* root, GraphTask& task) -> void { // Queue contains all nodes that will start propagating gradients. // We no longer have to expand functions that don't require grad. auto& dependencies = task.dependencies; - while (queue.size() > 0) { + while (!queue.empty()) { auto fn = queue.back(); queue.pop_back(); for (const auto& edge : fn->next_edges()) { if (auto next_ptr = edge.function.get()) { @@ -513,6 +509,7 @@ auto Engine::execute(const edge_list& roots, const edge_list& outputs) -> variable_list { std::call_once(start_threads_flag, &Engine::start_threads, this); + // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) validate_outputs(roots, const_cast(inputs), [](const std::string& msg) { return msg; }); @@ -559,6 +556,9 @@ auto Engine::execute(const edge_list& roots, // more callbacks (or they can be registered from other threads // while it's waiting. std::unique_lock cb_lock(post_callbacks_lock); + // WARNING: Don't use a range-for loop here because more callbacks may be + // added in between callback calls, so iterators may become invalidated. + // NOLINTNEXTLINE(modernize-loop-convert) for (size_t i = 0; i < final_callbacks.size(); ++i) { cb_lock.unlock(); final_callbacks[i](); diff --git a/torch/csrc/autograd/engine.h b/torch/csrc/autograd/engine.h index 94490303ccc24..0c7a5829da2ec 100644 --- a/torch/csrc/autograd/engine.h +++ b/torch/csrc/autograd/engine.h @@ -67,7 +67,7 @@ struct TORCH_API Engine { }; // allow python_engine to override the default engine when it loads -typedef Engine& (*EngineStub)(void); +using EngineStub = Engine& (*)(); TORCH_API void set_default_engine_stub(EngineStub stub); }} // namespace torch::autograd diff --git a/torch/csrc/autograd/function.cpp b/torch/csrc/autograd/function.cpp index 56ea7f7d29071..e077fadcb8214 100644 --- a/torch/csrc/autograd/function.cpp +++ b/torch/csrc/autograd/function.cpp @@ -114,14 +114,14 @@ void deleteFunction(Function* function) { delete function; - if (deleteFunctionQueue.size() == 0) { + if (deleteFunctionQueue.empty()) { return; } if (recursion_depth.value() != kDeleteFunctionMaxRecursionDepth) { AT_ERROR("Only one deleter per thread should be able to process " "the delete queue. Please open an issue."); } - while (deleteFunctionQueue.size() > 0) { + while (!deleteFunctionQueue.empty()) { auto queued_function = deleteFunctionQueue.front(); deleteFunctionQueue.pop_front(); delete queued_function; diff --git a/torch/csrc/autograd/functions/basic_ops.h b/torch/csrc/autograd/functions/basic_ops.h index ad7a9e52d3757..c8812c23f3a7c 100644 --- a/torch/csrc/autograd/functions/basic_ops.h +++ b/torch/csrc/autograd/functions/basic_ops.h @@ -24,6 +24,18 @@ struct TORCH_API Error : public Function { std::string msg; }; +// We print grad_fn names in tensor printing. For functions with backward +// NYI, grad_fn= will be printed if we use Error, which is confusing. So +// special case with a new NotImplemented function here. +struct TORCH_API NotImplemented : public Error { + NotImplemented(std::string forward_fn, edge_list&& next_edges) + : Error("derivative for " + forward_fn + " is not implemented", + std::move(next_edges)) {} + + NotImplemented(std::string forward_fn) + : Error("derivative for " + forward_fn + " is not implemented") {} +}; + // Identity in forward, Error in backward. Used to implement @once_differentiable struct TORCH_API DelayedError : public Function { DelayedError(std::string msg, int num_inputs) diff --git a/torch/csrc/autograd/functions/init.cpp b/torch/csrc/autograd/functions/init.cpp index 6988e650f888b..b25df7542d675 100644 --- a/torch/csrc/autograd/functions/init.cpp +++ b/torch/csrc/autograd/functions/init.cpp @@ -93,6 +93,9 @@ void THPAutograd_initFunctions() static PyTypeObject ErrorClass; addClass(module, ErrorClass, "Error"); + static PyTypeObject NotImplementedClass; + addClass(module, NotImplementedClass, "NotImplemented"); + static PyTypeObject DelayedErrorClass; addClass(module, DelayedErrorClass, "DelayedError"); diff --git a/torch/csrc/autograd/grad_mode.cpp b/torch/csrc/autograd/grad_mode.cpp index fc438dfad3d6a..8f5745668a77a 100644 --- a/torch/csrc/autograd/grad_mode.cpp +++ b/torch/csrc/autograd/grad_mode.cpp @@ -2,7 +2,7 @@ namespace torch { namespace autograd { -thread_local bool GradMode_enabled = 1; +thread_local bool GradMode_enabled = true; bool GradMode::is_enabled() { return GradMode_enabled; diff --git a/torch/csrc/autograd/python_function.cpp b/torch/csrc/autograd/python_function.cpp index 228fefe710288..2301df22e1a66 100644 --- a/torch/csrc/autograd/python_function.cpp +++ b/torch/csrc/autograd/python_function.cpp @@ -54,7 +54,7 @@ VariableInfo::VariableInfo(const Variable& var) Variable VariableInfo::zeros(at::DeviceGuard& device_guard) const { device_guard.set_index(device); - return at::zeros(size, *type); + return at::zeros(size, type->options()); } auto PyFunction::legacy_apply(const variable_list& inputs) -> variable_list { diff --git a/torch/csrc/autograd/variable.cpp b/torch/csrc/autograd/variable.cpp index b67b4b22cd003..6ec348dc01a0d 100644 --- a/torch/csrc/autograd/variable.cpp +++ b/torch/csrc/autograd/variable.cpp @@ -85,10 +85,6 @@ void Variable::Impl::set_storage_offset(int64_t storage_offset) { AT_ERROR("variable impl does not have set_storage_offset"); } -const char* Variable::Impl::typeString() { - return "VariableType"; -} - const at::Storage& Variable::Impl::storage() const { return data_.storage(); } diff --git a/torch/csrc/autograd/variable.h b/torch/csrc/autograd/variable.h index ba283d9136bb5..a1dce7e6c792a 100644 --- a/torch/csrc/autograd/variable.h +++ b/torch/csrc/autograd/variable.h @@ -276,8 +276,6 @@ struct TORCH_API Variable::Impl : public at::TensorImpl { const at::Storage& storage() const override; int64_t storage_offset() const override; - static const char* typeString(); - std::shared_ptr get_grad_accumulator(); virtual std::shared_ptr& get_grad_fn() { return grad_fn_; diff --git a/torch/csrc/byte_order.h b/torch/csrc/byte_order.h index 0b34730825fb4..9699556c69747 100644 --- a/torch/csrc/byte_order.h +++ b/torch/csrc/byte_order.h @@ -1,7 +1,7 @@ #ifndef THP_BYTE_ORDER_H #define THP_BYTE_ORDER_H -#include +#include #include #include diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp index d94301f1b41c3..10cd00976ee75 100644 --- a/torch/csrc/distributed/c10d/init.cpp +++ b/torch/csrc/distributed/c10d/init.cpp @@ -96,9 +96,23 @@ PyObject* c10d_init(PyObject* _unused) { "add", &::c10d::Store::add, py::call_guard()) + .def( + "set_timeout", + &::c10d::Store::setTimeout, + py::call_guard()) + .def( + "wait", + [](::c10d::Store& store, const std::vector& keys) { + store.wait(keys); + }, + py::call_guard()) .def( "wait", - &::c10d::Store::wait, + [](::c10d::Store& store, + const std::vector& keys, + const std::chrono::milliseconds& timeout) { + store.wait(keys, timeout); + }, py::call_guard()); shared_ptr_class_<::c10d::FileStore>(module, "FileStore", store) diff --git a/torch/csrc/generic/Storage.cpp b/torch/csrc/generic/Storage.cpp index 94c0b85bdb5ec..888b3d787410f 100644 --- a/torch/csrc/generic/Storage.cpp +++ b/torch/csrc/generic/Storage.cpp @@ -160,8 +160,32 @@ static PyObject * THPStorage_(get)(THPStorage *self, PyObject *index) return THPUtils_(newReal)(value); /* Slice index */ } else if (PySlice_Check(index)) { - THPUtils_setError("storages don't support slicing"); - return nullptr; + Py_ssize_t start, stop, slicelength, step; + int64_t len = THWStorage_(size)(LIBRARY_STATE self->cdata); + if (!THPUtils_parseSlice(index, len, &start, &stop, &step, &slicelength)) + return NULL; + if (step != 1) { + THPUtils_setError("Trying to slice with a step of %" PRId64 ", but only a step of " + "1 is supported", (int64_t)step); + return NULL; + } + + scalar_t *data = THWStorage_(data)(LIBRARY_STATE self->cdata); + + at::StorageImpl* old_storage = self->cdata; + c10::raw::intrusive_ptr::incref(old_storage); + at::Storage new_storage(c10::make_intrusive( + old_storage->dtype(), + slicelength, + at::DataPtr(static_cast(data + start), + old_storage, + [](void* s) { c10::raw::intrusive_ptr::decref(static_cast(s)); }, + old_storage->device()), + old_storage->allocator(), + /* resizable */ false)); + + PyObject *_ret = THPStorage_(New)(new_storage.unsafeReleaseStorageImpl()); + return _ret; } PyErr_Format(PyExc_TypeError, "can't index a " THPStorageStr " with %s", THPUtils_typename(index)); diff --git a/torch/csrc/jit/attributes.h b/torch/csrc/jit/attributes.h index 7199e0ae5787e..0a9f7e3fd2687 100644 --- a/torch/csrc/jit/attributes.h +++ b/torch/csrc/jit/attributes.h @@ -1,6 +1,6 @@ #pragma once #include -#include +#include #include #include #include @@ -36,7 +36,7 @@ struct ScalarAttributeValue : public AttributeValue { using ConstructorType = T; using ValueType = T; ScalarAttributeValue(Symbol name, ConstructorType value_) - : AttributeValue(name), value_(value_) {} + : AttributeValue(name), value_(std::move(value_)) {} ValueType & value() { return value_; } @@ -222,7 +222,7 @@ struct Attributes { typename T::ValueType & get(Symbol name) const { JIT_ASSERT(name.is_attr()); auto it = find(name, true); - T* child = dynamic_cast(it->get()); + auto* child = dynamic_cast(it->get()); if(child == nullptr) { throw AttributeError(name, true); } diff --git a/torch/csrc/jit/export.cpp b/torch/csrc/jit/export.cpp index 55747fa800fd4..98141d3382939 100644 --- a/torch/csrc/jit/export.cpp +++ b/torch/csrc/jit/export.cpp @@ -493,17 +493,18 @@ void ModuleEncoder::EncodeTypeInfo( } else if (kind == TypeKind::CompleteTensorType) { type_proto->set_denotation("CompleteTensorType"); CompleteTensorTypePtr node_type = type->cast(); - const std::vector& sizes = node_type->sizes(); // store the sizes and strides in the dims field of TensorShapeProto - for (size_t i = 0; i < sizes.size(); i++) { + size_t i = 0; + for (auto &size : node_type->sizes()) { shape_proto->add_dim(); - shape_proto->mutable_dim(i)->set_dim_value(sizes[i]); + shape_proto->mutable_dim(i)->set_dim_value(size); + i++; } - const std::vector& strides = node_type->strides(); - for (size_t i = 0; i < strides.size(); i++) { + for (auto &stride : node_type->strides()) { shape_proto->add_dim(); - shape_proto->mutable_dim(i)->set_dim_value(strides[i]); + shape_proto->mutable_dim(i)->set_dim_value(stride); + i++; } tensortype_proto->set_elem_type(ATenTypeToOnnxType(node_type->scalarType())); } else if (kind == TypeKind::TupleType) { @@ -612,6 +613,10 @@ void ModuleEncoder::EncodeMethod( script::Method &method, const std::string prefix) { node_proto->set_name(prefix + method.name()); + if (method.is_optimized()) { + // mark that this method was optimized + node_proto->set_domain("optimized"); + } // We store the schema string in the docstring. node_proto->set_doc_string(getExportableSchemaStringForMethod(method)); diff --git a/torch/csrc/jit/function_schema.h b/torch/csrc/jit/function_schema.h index ee5c5e5ca2555..c7b53abf46c2a 100644 --- a/torch/csrc/jit/function_schema.h +++ b/torch/csrc/jit/function_schema.h @@ -18,8 +18,8 @@ struct Argument { bool kwarg_only = false) : name(std::move(name)), type(type? type : DynamicType::get()), - N(N), - default_value(default_value), + N(std::move(N)), + default_value(std::move(default_value)), kwarg_only(kwarg_only) {} std::string name; TypePtr type; diff --git a/torch/csrc/jit/fusion_compiler.cpp b/torch/csrc/jit/fusion_compiler.cpp index 4f890599cc813..e21d00f00f40d 100644 --- a/torch/csrc/jit/fusion_compiler.cpp +++ b/torch/csrc/jit/fusion_compiler.cpp @@ -177,6 +177,10 @@ struct TensorInfo { IndexType sizes[N]; IndexType strides[N]; }; +template +struct TensorInfo { + T * data; +}; )"); // We rewrite the code for philox RNG from curand as nvrtc couldn't resolve the @@ -732,7 +736,9 @@ void compressContiguous( c_strides[compressed_dims] = strides[cur-1]; compressed_dims++; } - JIT_ASSERT(!cont.back() || strides.back() == 1); + if (ndim > 0) { + JIT_ASSERT(!cont.back() || strides.back() == 1); + } } } // anonymous namespace diff --git a/torch/csrc/jit/fusion_compiler.h b/torch/csrc/jit/fusion_compiler.h index 936b2aac5cf56..a24fe0817e051 100644 --- a/torch/csrc/jit/fusion_compiler.h +++ b/torch/csrc/jit/fusion_compiler.h @@ -29,7 +29,11 @@ struct TensorDesc { TensorDesc(const at::ScalarType& type, const std::vector& contiguity) : scalar_type(type), contiguity(contiguity) { - nDim_ = std::count(contiguity.begin(), contiguity.end(), false) + (lastIsContiguous() ? 1 : 0); + if (contiguity.size() == 0) { + nDim_ = 0; + } else { + nDim_ = std::count(contiguity.begin(), contiguity.end(), false) + (lastIsContiguous() ? 1 : 0); + } } TensorDesc(const at::ScalarType& type, const at::IntList& sizes, const at::IntList& strides) diff --git a/torch/csrc/jit/import.cpp b/torch/csrc/jit/import.cpp index 4ea87e895be73..ba6fe9fe0a27a 100644 --- a/torch/csrc/jit/import.cpp +++ b/torch/csrc/jit/import.cpp @@ -224,11 +224,19 @@ TypePtr ModuleDecoder::buildType(const onnx::TypeProto& type_proto) { if (kind == "DynamicType") { return DynamicType::get(); } else if (kind == "TensorType") { - // TODO: Don't use DynamicType here - return DynamicType::get(); + auto dims = shape_proto.dim_size(); + return TensorType::create(onnxTypeToATenType(tensortype_proto.elem_type()), -1, dims); } else if (kind == "CompleteTensorType") { - // TODO: Don't use DynamicType here - return DynamicType::get(); + // first half of the dims are sizes and the second half are strides + auto total = shape_proto.dim_size(); + std::vector sizes, strides; + for (int i = 0; i < total / 2; i++) { + sizes.push_back(shape_proto.dim(i).dim_value()); + } + for (int i = total / 2; i < total; i++) { + strides.push_back(shape_proto.dim(i).dim_value()); + } + return CompleteTensorType::create(onnxTypeToATenType(tensortype_proto.elem_type()), -1, sizes, strides); } else if (kind == "TupleType") { std::vector elems; for (auto &subkind : shape_proto.dim()) { @@ -371,6 +379,8 @@ ModuleDecoder::ModuleDecoder( } auto graph = buildGraph(node_proto.attribute(0).g()); + // has_domain field has a string iff the method was optimized + parent_module->set_optimized(node_proto.has_domain()); parent_module->create_method(name, graph, member_inputs); // We store the schema in the docstring so we can parse the schema and // assign it to the method. diff --git a/torch/csrc/jit/init.cpp b/torch/csrc/jit/init.cpp index e816961d86dd3..eb4956087c6c1 100644 --- a/torch/csrc/jit/init.cpp +++ b/torch/csrc/jit/init.cpp @@ -11,6 +11,7 @@ #include "torch/csrc/jit/passes/onnx.h" #include "torch/csrc/jit/passes/dead_code_elimination.h" #include "torch/csrc/jit/passes/erase_number_types.h" +#include "torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h" #include "torch/csrc/jit/passes/common_subexpression_elimination.h" #include "torch/csrc/jit/passes/peephole.h" #include "torch/csrc/jit/passes/canonicalize.h" @@ -90,6 +91,7 @@ void initJITBindings(PyObject *module) { }) .def("_jit_pass_remove_expands", RemoveExpands) .def("_jit_pass_erase_number_types", EraseNumberTypes) + .def("_jit_pass_prepare_division_for_onnx", PrepareDivisionForONNX) .def("_jit_pass_loop_unrolling", UnrollLoops) .def("_jit_pass_constant_propagation", [](std::shared_ptr& g) { return ConstantPropagation(g); diff --git a/torch/csrc/jit/interned_strings.cpp b/torch/csrc/jit/interned_strings.cpp index c1f018e611c52..bf487484b25e0 100644 --- a/torch/csrc/jit/interned_strings.cpp +++ b/torch/csrc/jit/interned_strings.cpp @@ -1,5 +1,5 @@ #include "torch/csrc/jit/interned_strings.h" -#include +#include #include #include #include diff --git a/torch/csrc/jit/interned_strings.h b/torch/csrc/jit/interned_strings.h index 899c3c8063664..32216971834e8 100644 --- a/torch/csrc/jit/interned_strings.h +++ b/torch/csrc/jit/interned_strings.h @@ -1,6 +1,6 @@ #pragma once #include -#include +#include #include #include #include @@ -57,6 +57,7 @@ namespace torch { namespace jit { _(prim, FusedConcat) \ _(prim, ConstantChunk) \ _(prim, NoneGenerator) \ + _(aten, floordiv) \ _(aten, __not__) \ FORALL_ATEN_BASE_SYMBOLS(_) \ _(onnx, Add) \ diff --git a/torch/csrc/jit/interned_strings_class.h b/torch/csrc/jit/interned_strings_class.h index 7b82c8034a98a..7dbf497d2739a 100644 --- a/torch/csrc/jit/interned_strings_class.h +++ b/torch/csrc/jit/interned_strings_class.h @@ -1,4 +1,4 @@ -#include +#include #include #include #include diff --git a/torch/csrc/jit/ir.h b/torch/csrc/jit/ir.h index 196ec9a1edb1c..9ac5059a2d801 100644 --- a/torch/csrc/jit/ir.h +++ b/torch/csrc/jit/ir.h @@ -150,6 +150,7 @@ struct Scope { } Scope* parent = this->parent_; while (!parent->isRoot()) { + // NOLINTNEXTLINE(performance-inefficient-string-concatenation) out = std::string(parent->name_.toUnqualString()) + separator + out; parent = parent->parent_; } @@ -181,7 +182,7 @@ struct Value { std::string unique_name_; TypePtr type_; public: - Value* setType(const TypePtr type); + Value* setType(TypePtr type); void inferTypeFrom(const at::Tensor& output) { setType(CompleteTensorType::create(output)); } @@ -368,7 +369,7 @@ struct Node : public Attributes { } bool hasUses() const { for(auto o : outputs()) { - if(o->uses().size() > 0) + if(!o->uses().empty()) return true; } return false; @@ -679,7 +680,7 @@ struct Node : public Attributes { } // XXX: this function is meant to be used with string literals only! - bool matches(const char *signature_literal, at::ArrayRef const_inputs={}) const; + TORCH_API bool matches(const char *signature_literal, at::ArrayRef const_inputs={}) const; const FunctionSchema& schema() const { if (!schema_) @@ -890,7 +891,7 @@ friend struct Block; Graph(std::shared_ptr scope_root) : next_unique_(0) , new_node_stage_(0) - , scope_root_(scope_root) + , scope_root_(std::move(scope_root)) , current_scope_(scope_root_.get()) , block_(new Block(this, nullptr)) , insert_before_(return_node()) {} @@ -1261,7 +1262,7 @@ inline Node::Node(Graph * graph_, NodeKind kind_) : inline void Node::eraseOutput(size_t i) { JIT_ASSERT(i < outputs_.size()); - JIT_ASSERT(outputs_[i]->uses().size() == 0); + JIT_ASSERT(outputs_[i]->uses().empty()); schema_ = nullptr; Value * n = outputs_[i]; outputs_.erase(outputs_.begin() + i); @@ -1286,9 +1287,9 @@ inline void Node::eraseBlock(size_t i) { } inline void Node::destroy() { - while(outputs().size() > 0) + while(!outputs().empty()) eraseOutput(outputs().size() - 1); - while(blocks().size() > 0) + while(!blocks().empty()) eraseBlock(blocks().size() - 1); removeAllInputs(); if(inBlockList()) @@ -1422,13 +1423,13 @@ inline Node* Graph::createPythonOp( } inline graph_node_list_iterator Node::iterator() { - return graph_node_list_iterator(this, 0); + return {this, 0}; } inline graph_node_list_iterator Node::reverseIterator() { return iterator().reverse(); } inline const_graph_node_list_iterator Node::iterator() const { - return const_graph_node_list_iterator(this, 0); + return {this, 0}; } inline const_graph_node_list_iterator Node::reverseIterator() const { return iterator().reverse(); diff --git a/torch/csrc/jit/ivalue.h b/torch/csrc/jit/ivalue.h index fff0abdea0a96..c8475b7ff8618 100644 --- a/torch/csrc/jit/ivalue.h +++ b/torch/csrc/jit/ivalue.h @@ -17,8 +17,8 @@ struct TORCH_API ConstantString : c10::intrusive_ptr_target { private: const std::string str_; public: - ConstantString(const std::string & str) - : str_(str) {} + ConstantString(std::string str) + : str_(std::move(str)) {} static c10::intrusive_ptr create(const std::string str_) { return c10::make_intrusive(str_); } @@ -88,7 +88,7 @@ struct TORCH_API IValue { c10::raw::intrusive_ptr::decref(as_intrusive_ptr); } } - IValue & operator=(IValue && rhs) & { + IValue & operator=(IValue && rhs) & noexcept { rhs.swap(*this); return *this; } @@ -96,7 +96,7 @@ struct TORCH_API IValue { IValue(rhs).swap(*this); return *this; } - void swap(IValue & rhs) { + void swap(IValue & rhs) noexcept { std::swap(payload, rhs.payload); std::swap(is_intrusive_ptr, rhs.is_intrusive_ptr); std::swap(tag, rhs.tag); diff --git a/torch/csrc/jit/passes/graph_fuser.cpp b/torch/csrc/jit/passes/graph_fuser.cpp index ea8ab51b938f8..81a1768f48a30 100644 --- a/torch/csrc/jit/passes/graph_fuser.cpp +++ b/torch/csrc/jit/passes/graph_fuser.cpp @@ -882,9 +882,14 @@ struct GraphFuser { } // anonymous namespace void FuseGraph(std::shared_ptr& graph) { + // NYI on Windows + #ifndef _WIN32 + GraphFuser(graph->block()).run(); // After FuseGraph some common subexpressions may come back EliminateCommonSubexpression(graph); + + #endif } }} diff --git a/torch/csrc/jit/passes/graph_fuser.h b/torch/csrc/jit/passes/graph_fuser.h index 1c1bbfe292ff5..f7135359f1e07 100644 --- a/torch/csrc/jit/passes/graph_fuser.h +++ b/torch/csrc/jit/passes/graph_fuser.h @@ -6,6 +6,7 @@ namespace torch { namespace jit { // NB: Be sure to run DCE before fusion, because dead instructions // can prevent fusion opportunities from being exploited. +// On Windows will noop, NYI TORCH_API void FuseGraph(std::shared_ptr& graph); }} diff --git a/torch/csrc/jit/passes/loop_unrolling.cpp b/torch/csrc/jit/passes/loop_unrolling.cpp index 4dbce8afea96d..015ac1f9b47d5 100644 --- a/torch/csrc/jit/passes/loop_unrolling.cpp +++ b/torch/csrc/jit/passes/loop_unrolling.cpp @@ -188,7 +188,7 @@ void unroll(Node *loop) { // Change the iteration counts of both loops Value* iter_count = loop->inputs().at(0); - Value* unrolled_iter_count = graph->insert(aten::div, {iter_count, kUnrollFactor}); + Value* unrolled_iter_count = graph->insert(aten::floordiv, {iter_count, kUnrollFactor}); loop->replaceInput(0, unrolled_iter_count); loop_epilogue->replaceInput(0, graph->insert(aten::sub, {iter_count, graph->insert(aten::mul,{unrolled_iter_count , kUnrollFactor})})); } diff --git a/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.cpp b/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.cpp new file mode 100644 index 0000000000000..8b204b66e7eb0 --- /dev/null +++ b/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.cpp @@ -0,0 +1,35 @@ +#include "torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h" +#include "torch/csrc/jit/constants.h" + +namespace torch { namespace jit { + +static void PrepareDivisionForONNXOnBlock(Block* block) { + for (auto it = block->nodes().begin(); it != block->nodes().end(); ++it) { + for (auto sub : it->blocks()) { + PrepareDivisionForONNXOnBlock(sub); + } + WithInsertPoint guard(*it); + auto* subgraph = it->owningGraph(); + + if (it->matches("aten::div(int a, int b) -> float")) { + // Cast to Float before dividing + std::vector floattensor_inputs = fmap(it->inputs(), [&](Value* input) { + auto* longtensor = subgraph->insertNode(subgraph->createNumToTensor(input))->output(); + auto* nonblocking = subgraph->insertConstant(0); + auto* cast = subgraph->create(aten::_cast_Float, {longtensor, nonblocking}); + return subgraph->insertNode(cast)->output(); + }); + + it->replaceInput(0, floattensor_inputs[0]); + it->replaceInput(1, floattensor_inputs[1]); + it->output()->setType(CompleteTensorType::fromNumberType(FloatType::get())); + } + } +} + +void PrepareDivisionForONNX(const std::shared_ptr& graph) { + PrepareDivisionForONNXOnBlock(graph->block()); +} + +}} + diff --git a/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h b/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h new file mode 100644 index 0000000000000..1041877e48862 --- /dev/null +++ b/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h @@ -0,0 +1,17 @@ +#pragma once + +#include "torch/csrc/jit/ir.h" + +namespace torch { namespace jit { + +// Prepare division ops for ONNX export. This is necessary for and only used +// by ONNX export. +// +// The pass corrects the following: +// +// - aten::div(int, int) -> float is the python truediv operator. This doesn't +// exist in ONNX so we cast the ints to FloatTensors +// +TORCH_API void PrepareDivisionForONNX(const std::shared_ptr& graph); + +}} diff --git a/torch/csrc/jit/pybind_utils.h b/torch/csrc/jit/pybind_utils.h index f2fab997a9037..f7285e2e3d30a 100644 --- a/torch/csrc/jit/pybind_utils.h +++ b/torch/csrc/jit/pybind_utils.h @@ -59,7 +59,8 @@ inline IValue toIValue(py::handle input) { } return Tuple::create(s); } else { - AT_ERROR("Only tensors and tuples of tensors are supported as inputs to traced functions"); + AT_ERROR("Only tensors and (possibly nested) tuples of tensors are supported " + "as inputs or outputs of traced functions"); } } diff --git a/torch/csrc/jit/python_tracer.cpp b/torch/csrc/jit/python_tracer.cpp index 7d069cb91dddf..dabe88991073a 100644 --- a/torch/csrc/jit/python_tracer.cpp +++ b/torch/csrc/jit/python_tracer.cpp @@ -52,7 +52,11 @@ std::shared_ptr createGraphByTracing( py_inputs[i] = py::cast(enter_info.second[i]); } auto out = func(*py_inputs); - if(!PyTuple_Check(out.ptr())) { + if (out.ptr() == Py_None) { + AT_ERROR("The traced function didn't return any values! Side-effects are not " + "captured in traces, so it would be a no-op."); + } + if (!PyTuple_Check(out.ptr())) { out = py::make_tuple(out); } tracer::exit(toStack(out)); diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp index c68c2aea516f1..7e70550a38094 100644 --- a/torch/csrc/jit/register_prim_ops.cpp +++ b/torch/csrc/jit/register_prim_ops.cpp @@ -616,9 +616,32 @@ RegisterOperators reg2({ DEFINE_BINARY_OP(aten::add, a + b) DEFINE_BINARY_OP(aten::sub, a - b) DEFINE_BINARY_OP(aten::mul, a * b) - DEFINE_BINARY_OP(aten::div, a / b) DEFINE_BINARY_OP(aten::pow, static_cast(pow(a, b))) + // TODO: Support python floordiv (//) + // Right now aten::floordiv is only used by loop unrolling + DEFINE_INT_OP(aten::floordiv, a / b) + + // NB: This is the python truediv operation + Operator("aten::div(int a, int b) -> float", + [](Node* node) { + return [=](Stack& stack) { + int64_t a, b; + pop(stack, a, b); + push(stack, static_cast(a) / static_cast(b)); + return 0; + }; + }), + Operator("aten::div(float a, float b) -> float", + [](Node* node) { + return [=](Stack& stack) { + double a, b; + pop(stack, a, b); + push(stack, a / b); + return 0; + }; + }), + DEFINE_COMPARISON_OP(aten::ne, a != b) DEFINE_COMPARISON_OP(aten::eq, a == b) DEFINE_COMPARISON_OP(aten::lt, a < b) diff --git a/torch/csrc/jit/script/compiler.cpp b/torch/csrc/jit/script/compiler.cpp index 937351f30cbf9..8aae072a02a1f 100644 --- a/torch/csrc/jit/script/compiler.cpp +++ b/torch/csrc/jit/script/compiler.cpp @@ -1525,8 +1525,7 @@ struct to_ir { if (slice_exprs[0].kind() == TK_SLICE_EXPR) { return emitBasicSlice(subscript); } else { - return emitBasicGather( - subscript.range(), TreeList{subscript.value(), slice_exprs[0]}); + return emitBasicGather(subscript); } } break; case TK_IF_EXPR: { @@ -1620,26 +1619,34 @@ struct to_ir { return emitBuiltinCall(loc, *graph, aten::slice, args, {step}, true); } - // Desugars multidim slicing into aten::slice / aten::select calls. - // - // XXX: Errors in user code are not elegantly reported. - // Let's say someone were to do the following: - // @torch.jit.script - // def fn(x): - // return x[0, 1] - // fn(torch.randn(5)) - // Because we desugar this into two aten::select ops, the error message - // complains about aten::select failing rather than there "not being - // enough dimensions to index". - Value * emitMultidimSlicing(const Subscript& subscript) { - const auto& loc = subscript.range(); - auto * sliceable = emitExpr(subscript.value()); - if (!sliceable->type()->isSubtypeOf(DynamicType::get())) { - throw ErrorReport(loc) - << "Unsupported operation: attempted to use multidimensional " - << "indexing on a non-tensor type."; - } + Value* emitIndex( + const SourceRange& loc, + Value* input, + at::ArrayRef indices) { + auto* index = graph->insertNode( + graph->createList(DynamicType::get(), indices))->output(); + return emitBuiltinCall(loc, *graph, aten::index, {input, index}, {}, true); + } + + // Emits multidimensional slicing with int and slice indices. + // Returns: + // - Value*: the input after it has been indexed by int and slice indices. + // - vector: A list of tensor Value* indices that have not been applied yet. + // Should be NULL at indices where sliceable (post-slicing) isn't indexed by a tensor. + std::pair> emitIntAndSliceIndexing( + const SourceRange& loc, + Value* sliceable, + const Subscript& subscript) { + std::vector tensor_indices; size_t dim = 0; + + auto handle_tensor = [&](Value* tensor) { + // NB: tensor_indices can have NULL holes because of how at::index works. + tensor_indices.resize(dim + 1); + tensor_indices[dim] = tensor; + dim++; + }; + for (const auto & subscript_expr : subscript.subscript_exprs()) { if (subscript_expr.kind() == TK_SLICE_EXPR) { sliceable = emitSlice(loc, sliceable, dim, SliceExpr(subscript_expr)); @@ -1651,13 +1658,64 @@ struct to_ir { sliceable = emitSelect(loc, sliceable, dim, index); continue; } else if (index->type()->isSubtypeOf(DynamicType::get())) { - throw std::runtime_error("NYI: advanced indexing"); + handle_tensor(index); + continue; } + throw ErrorReport(loc) + << "Unsupported operation: indexing tensor with unsupported index type " + << index->type()->str() << ". Only ints, slices, and tensors are supported."; + } + return std::make_pair(sliceable, tensor_indices); + } + + // The strategy is to slice and select the tensor for int and slices first + // in one pass and then apply at::index on the result of the slicing/selecting. + // Call the tensor after we've applied slice / select the `sliced`. + // tensor_indices should have the same size as sliced.dim(): + // - tensor_indices[i] = NULL if we should not index `sliced` at dim i + // - tensor_indices[i] = t if we should index `sliced` at dim i with tensor t. + Value* emitMultidimSlicing( + const SourceRange& loc, + Value* sliceable, + const Subscript& subscript) { + std::vector tensor_indices; + std::tie(sliceable, tensor_indices) = emitIntAndSliceIndexing(loc, sliceable, subscript); + + if (tensor_indices.empty()) { + // XXX: Might need to at::alias this when we support mutability + return sliceable; + } + + // at::index takes in a TensorList where some tensors can be undefined. + // Convert NULL tensor_indices to undefined tensors to pass to at::index. + for (auto& index : tensor_indices) { + if (index == nullptr) { + index = graph->insertNode(graph->createUndefined())->output(); + } + } + return emitIndex(loc, sliceable, tensor_indices); + } + + // Desugars multidim slicing into slice/select/index calls. + // + // XXX: Errors in user code are not elegantly reported. + // Let's say someone were to do the following: + // @torch.jit.script + // def fn(x): + // return x[0, 1] + // fn(torch.randn(5)) + // Because we desugar this into two aten::select ops, the error message + // complains about aten::select failing rather than there "not being + // enough dimensions to index". + Value* emitMultidimSlicing(const Subscript& subscript) { + const auto& loc = subscript.range(); + auto* sliceable = emitExpr(subscript.value()); + if (!sliceable->type()->isSubtypeOf(DynamicType::get())) { throw ErrorReport(loc) << "Unsupported operation: attempted to use multidimensional " - << "indexing with unsupported index type."; + << "indexing on a non-tensor type."; } - return sliceable; + return emitMultidimSlicing(loc, sliceable, subscript); } // Desugars slice syntactic sugar tensor[begin:end] -> tensor.slice(begin, @@ -1677,24 +1735,19 @@ struct to_ir { } // Desugars gather syntactic sugar foo[i] - Value* emitBasicGather( - const SourceRange& loc, - TreeList&& inputs) { - const auto applyInputs = - Compound::create(TK_LIST, loc, std::move(inputs)); - auto input_values = getNamedValues(applyInputs->trees(), - /*maybe_unpack*/false); - NamedValue gatherable = input_values[0]; - NamedValue idx = input_values[1]; - if (gatherable.value(*graph)->type()->kind() == TypeKind::ListType) { + Value* emitBasicGather(const Subscript& subscript) { + const auto& loc = subscript.range(); + JIT_ASSERT(subscript.subscript_exprs().size() == 1); + auto* gatherable = emitExpr(subscript.value()); + + if (gatherable->type()->kind() == TypeKind::ListType) { // if it's a list, emit a regular index selection op + auto* idx = emitExpr(subscript.subscript_exprs()[0]); return emitBuiltinCall( loc, *graph, aten::select, {gatherable, idx}, {}, true); - } else { - JIT_ASSERT(gatherable.value(*graph)->type()->isSubtypeOf(DynamicType::get())); - return emitSelect( - loc, gatherable.value(*graph), /*dim=*/0, idx.value(*graph)); + JIT_ASSERT(gatherable->type()->isSubtypeOf(DynamicType::get())); + return emitMultidimSlicing(loc, gatherable, subscript); } } }; @@ -1729,6 +1782,13 @@ std::shared_ptr SimpleValue::attr(SourceRange loc, Method & m, con Symbol::aten(field), NamedValue(loc, "self", value)); } +std::shared_ptr nativeResolver(const std::string& name, Method& m, const SourceRange& loc) { + if (name == "torch") { + return std::make_shared(name); + } + return nullptr; +} + std::vector inlineCallTo(Graph& g, Graph& callee, ArrayRef inputs) { std::unordered_map value_map; auto value_map_func = [&](Value* v) { return value_map.at(v); }; diff --git a/torch/csrc/jit/script/compiler.h b/torch/csrc/jit/script/compiler.h index ab362ec8ec5e6..deef6a5c2ca8f 100644 --- a/torch/csrc/jit/script/compiler.h +++ b/torch/csrc/jit/script/compiler.h @@ -122,8 +122,25 @@ struct TORCH_API BuiltinFunction : public SugaredValue { size_t n_binders) override; }; +struct TORCH_API BuiltinModule : public SugaredValue { + BuiltinModule(const std::string& name) + : name(name) {} + std::string name; + + std::string kind() const override { + return "builtin module"; + } + + std::shared_ptr attr(SourceRange loc, Method & m, const std::string& field) override { + return std::make_shared(Symbol::aten(field), at::nullopt); + } +}; + using Resolver = std::function(const std::string& name, Method& m, const SourceRange& loc)>; + +std::shared_ptr nativeResolver(const std::string& name, Method& m, const SourceRange& loc); + TORCH_API void defineMethodsInModule( Module & m, const std::vector& definitions, diff --git a/torch/csrc/jit/script/module.h b/torch/csrc/jit/script/module.h index 6af1d16a06d6e..70a3a8a7b9016 100644 --- a/torch/csrc/jit/script/module.h +++ b/torch/csrc/jit/script/module.h @@ -163,6 +163,10 @@ struct Method { return get_executor().getDebugState(); } + bool is_optimized() { + return optimize; + } + private: std::string name_; std::shared_ptr graph_; // for debugging and for inlining @@ -348,6 +352,24 @@ struct Module { return nullptr; } + /// Run a method from this module. + /// + /// For example: + /// @code + /// IValue output = module->run("relu_script", a, b); + /// @endcode + /// + /// To get a compile a module from a source string, see torch::jit::compile + /// + /// @param method_name The name of the method to run + /// @param args Arguments to be passed to the method + /// @return An IValue containing the return value (or values if it is a tuple) + /// from the method + template + IValue run_method(const std::string& method_name, Types&&... args) { + return get_method(method_name)({IValue(std::forward(args))...}); + } + void save(const std::string& filename); private: diff --git a/torch/csrc/jit/source_location.h b/torch/csrc/jit/source_location.h index ec55ce8f23ed1..6b86391a8f144 100644 --- a/torch/csrc/jit/source_location.h +++ b/torch/csrc/jit/source_location.h @@ -17,7 +17,7 @@ struct SourceLocation { void wrapAndRethrowException(const std::exception & e, const std::string & additional = "") { std::stringstream msg; msg << "\n" << e.what() << ":\n"; - if(additional.size() != 0) { + if(!additional.empty()) { msg << additional << ":\n"; } highlight(msg); diff --git a/torch/csrc/jit/source_range.h b/torch/csrc/jit/source_range.h index b84729f5dd1c6..0139c2527513a 100644 --- a/torch/csrc/jit/source_range.h +++ b/torch/csrc/jit/source_range.h @@ -10,10 +10,10 @@ namespace jit { // range. struct SourceRange : public SourceLocation { SourceRange( - const std::shared_ptr& file_, + std::shared_ptr file_, size_t start_, size_t end_) - : file_(file_), start_(start_), end_(end_) {} + : file_(std::move(file_)), start_(start_), end_(end_) {} const std::string text() const { return file().substr(start(), end() - start()); } @@ -22,7 +22,7 @@ struct SourceRange : public SourceLocation { } static const size_t CONTEXT = 10; - virtual void highlight(std::ostream& out) const override { + void highlight(std::ostream& out) const override { const std::string& str = file(); size_t begin_line = start(); // beginning of line to highlight size_t end_line = start(); // end of line to highlight @@ -57,7 +57,7 @@ struct SourceRange : public SourceLocation { out << std::string(len, '~') << (len < size() ? "... <--- HERE" : " <--- HERE"); out << str.substr(end_line, end_highlight - end_line); - if (str.size() > 0 && str.back() != '\n') + if (!str.empty() && str.back() != '\n') out << "\n"; } const std::string& file() const { diff --git a/torch/csrc/jit/test_jit.cpp b/torch/csrc/jit/test_jit.cpp index eeaf08b39d409..bec6b0459a081 100644 --- a/torch/csrc/jit/test_jit.cpp +++ b/torch/csrc/jit/test_jit.cpp @@ -702,7 +702,7 @@ void testCreateAutodiffSubgraphs(std::ostream & out) { } autograd::Variable var(at::Type & t, at::IntList sizes, bool requires_grad) { - return autograd::make_variable(at::rand(sizes, t), requires_grad); + return autograd::make_variable(at::rand(sizes, t.options()), requires_grad); } autograd::Variable undef() { return autograd::Variable(); diff --git a/torch/csrc/jit/type.h b/torch/csrc/jit/type.h index ed1a9205286e0..2c9c1d550ae2a 100644 --- a/torch/csrc/jit/type.h +++ b/torch/csrc/jit/type.h @@ -129,7 +129,7 @@ struct TORCH_API DynamicType : public Type { static constexpr bool is_singleton = true; template static DynamicTypePtr create( T&& ... all ) { - return DynamicTypePtr(new DynamicType( std::forward(all)... )); + return DynamicTypePtr(new DynamicType( std::forward(all)... )); // NOLINT(modernize-make-shared) } bool operator==(const Type& rhs) const override { @@ -156,7 +156,7 @@ struct TORCH_API TensorType : public Type { template static TensorTypePtr create( T&& ... all ) { - return TensorTypePtr(new TensorType( std::forward(all)... )); + return TensorTypePtr(new TensorType( std::forward(all)... )); // NOLINT(modernize-make-shared) } at::ScalarType scalarType() const { return scalar_type_; } @@ -215,15 +215,15 @@ struct TORCH_API CompleteTensorType : public TensorType { friend struct Type; template static CompleteTensorTypePtr create( T&& ... all ) { - return CompleteTensorTypePtr(new CompleteTensorType( std::forward(all)... )); + return CompleteTensorTypePtr(new CompleteTensorType( std::forward(all)... )); // NOLINT(modernize-make-shared) } // overloaded create variadic template argument as it could not distinguish initializer list static CompleteTensorTypePtr create(at::ScalarType scalar_type, int device, at::IntList sizes) { - return CompleteTensorTypePtr(new CompleteTensorType(scalar_type, device, sizes)); + return CompleteTensorTypePtr(new CompleteTensorType(scalar_type, device, sizes)); // NOLINT(modernize-make-shared) } static CompleteTensorTypePtr create(at::ScalarType scalar_type, int device, at::IntList sizes, at::IntList strides) { - return CompleteTensorTypePtr(new CompleteTensorType(scalar_type, device, sizes, strides)); + return CompleteTensorTypePtr(new CompleteTensorType(scalar_type, device, sizes, strides)); // NOLINT(modernize-make-shared) } static const TypeKind Kind = TypeKind::CompleteTensorType; @@ -295,7 +295,7 @@ struct TORCH_API CompleteTensorType : public TensorType { static std::vector contiguousStridesOf(at::IntList sizes) { std::vector strides(sizes.size()); - if(sizes.size() == 0) // zero-dim case + if(sizes.empty()) // zero-dim case return strides; strides.back() = 1; for(size_t i = strides.size() - 1; i > 0; i--) { @@ -318,7 +318,7 @@ struct TORCH_API ListType : public Type { friend struct Type; template static ListTypePtr create( T&& ... all ) { - return ListTypePtr(new ListType( std::forward(all)... )); + return ListTypePtr(new ListType( std::forward(all)... )); // NOLINT(modernize-make-shared) } bool operator==(const Type& rhs) const override { if(auto rhs_ = rhs.cast()) { @@ -340,7 +340,7 @@ struct TORCH_API ListType : public Type { static ListTypePtr ofFloats(); private: ListType(TypePtr elem) - : Type(TypeKind::ListType), elem(elem) {} + : Type(TypeKind::ListType), elem(std::move(elem)) {} static const TypeKind Kind = TypeKind::ListType; TypePtr elem; }; @@ -352,7 +352,7 @@ struct TORCH_API TupleType : public Type { static constexpr bool is_singleton = false; friend struct Type; static TupleTypePtr create(std::vector types) { - return TupleTypePtr(new TupleType( std::move(types) )); + return TupleTypePtr(new TupleType( std::move(types) )); // NOLINT(modernize-make-shared) } at::ArrayRef elements() const { return elements_; @@ -408,7 +408,7 @@ struct TORCH_API NumberType : public Type { static constexpr bool is_singleton = true; template static NumberTypePtr create( T&& ... all ) { - return NumberTypePtr(new NumberType( std::forward(all)... )); + return NumberTypePtr(new NumberType( std::forward(all)... )); // NOLINT(modernize-make-shared) } bool operator==(const Type& rhs) const override { return rhs.kind() == kind(); @@ -431,7 +431,7 @@ struct TORCH_API FloatType : public Type { static constexpr bool is_singleton = true; template static FloatTypePtr create( T&& ... all ) { - return FloatTypePtr(new FloatType( std::forward(all)... )); + return FloatTypePtr(new FloatType( std::forward(all)... )); // NOLINT(modernize-make-shared) } bool operator==(const Type& rhs) const override { return rhs.kind() == kind(); @@ -457,7 +457,7 @@ struct TORCH_API IntType : public Type { static constexpr bool is_singleton = true; template static IntTypePtr create( T&& ... all ) { - return IntTypePtr(new IntType( std::forward(all)... )); + return IntTypePtr(new IntType( std::forward(all)... )); // NOLINT(modernize-make-shared) } bool operator==(const Type& rhs) const override { return rhs.kind() == kind(); @@ -483,7 +483,7 @@ struct TORCH_API StringType : public Type { static constexpr bool is_singleton = true; template static StringTypePtr create( T&& ... all ) { - return StringTypePtr(new StringType( std::forward(all)... )); + return StringTypePtr(new StringType( std::forward(all)... )); // NOLINT(modernize-make-shared) } bool operator==(const Type& rhs) const override { return rhs.kind() == kind(); @@ -509,12 +509,12 @@ struct NoneType : public Type { static constexpr bool is_singleton = true; template static NoneTypePtr create( T&& ... all ) { - return NoneTypePtr(new NoneType( std::forward(all)... )); + return NoneTypePtr(new NoneType( std::forward(all)... )); // NOLINT(modernize-make-shared) } - virtual bool operator==(const Type& rhs) const override { + bool operator==(const Type& rhs) const override { return rhs.kind() == kind(); } - virtual std::string str() const override { + std::string str() const override { return "None"; } static const TypeKind Kind = TypeKind::NoneType; @@ -531,12 +531,12 @@ struct GeneratorType : public Type { static constexpr bool is_singleton = true; template static GeneratorTypePtr create( T&& ... all) { - return GeneratorTypePtr(new GeneratorType( std::forward(all)... )); + return GeneratorTypePtr(new GeneratorType( std::forward(all)... )); // NOLINT(modernize-make-shared) } - virtual bool operator==(const Type& rhs) const override { + bool operator==(const Type& rhs) const override { return rhs.kind() == kind(); } - virtual std::string str() const override { + std::string str() const override { return "Generator"; } static const TypeKind Kind = TypeKind::GeneratorType; diff --git a/torch/csrc/utils/object_ptr.h b/torch/csrc/utils/object_ptr.h index 14991b8779d9c..8a3b362d20267 100644 --- a/torch/csrc/utils/object_ptr.h +++ b/torch/csrc/utils/object_ptr.h @@ -6,16 +6,16 @@ template class THPPointer { public: THPPointer(): ptr(nullptr) {}; - explicit THPPointer(T *ptr): ptr(ptr) {}; - THPPointer(THPPointer &&p) { free(); ptr = p.ptr; p.ptr = nullptr; }; + explicit THPPointer(T *ptr) noexcept : ptr(ptr) {}; + THPPointer(THPPointer &&p) noexcept { free(); ptr = p.ptr; p.ptr = nullptr; }; ~THPPointer() { free(); }; T * get() { return ptr; } const T * get() const { return ptr; } T * release() { T *tmp = ptr; ptr = nullptr; return tmp; } operator T*() { return ptr; } - THPPointer& operator =(T *new_ptr) { free(); ptr = new_ptr; return *this; } - THPPointer& operator =(THPPointer &&p) { free(); ptr = p.ptr; p.ptr = nullptr; return *this; } + THPPointer& operator =(T *new_ptr) noexcept { free(); ptr = new_ptr; return *this; } + THPPointer& operator =(THPPointer &&p) noexcept { free(); ptr = p.ptr; p.ptr = nullptr; return *this; } T * operator ->() { return ptr; } explicit operator bool() const { return ptr != nullptr; } @@ -35,4 +35,4 @@ class THPPointer { * out the GIL itself. Easiest way to avoid this problem is to * not use THPPointer in this situation. */ -typedef THPPointer THPObjectPtr; +using THPObjectPtr = THPPointer; diff --git a/torch/csrc/utils/python_arg_parser.cpp b/torch/csrc/utils/python_arg_parser.cpp index a295c7d645215..d36f70e460d79 100644 --- a/torch/csrc/utils/python_arg_parser.cpp +++ b/torch/csrc/utils/python_arg_parser.cpp @@ -104,9 +104,6 @@ bool FunctionParameter::check(PyObject* obj) { } case ParameterType::SCALAR: case ParameterType::DOUBLE: { - // NOTE: we don't currently accept most NumPy types as Scalars. np.float64 - // is okay because it's a subclass of PyFloat. We may want to change this - // in the future. if (THPUtils_checkDouble(obj)) { return true; } diff --git a/torch/csrc/utils/python_numbers.h b/torch/csrc/utils/python_numbers.h index 318efb1f8b707..69f952bc6ec31 100644 --- a/torch/csrc/utils/python_numbers.h +++ b/torch/csrc/utils/python_numbers.h @@ -1,9 +1,10 @@ #pragma once #include "torch/csrc/python_headers.h" -#include +#include #include #include "torch/csrc/Exceptions.h" +#include "torch/csrc/utils/tensor_numpy.h" #include "torch/csrc/jit/tracer.h" // largest integer that can be represented consecutively in a double @@ -86,10 +87,16 @@ inline int64_t THPUtils_unpackIndex(PyObject* obj) { } inline bool THPUtils_checkDouble(PyObject* obj) { + bool is_numpy_scalar; +#ifdef USE_NUMPY + is_numpy_scalar = torch::utils::is_numpy_scalar(obj); +#else + is_numpy_scalar = false; +#endif #if PY_MAJOR_VERSION == 2 - return PyFloat_Check(obj) || PyLong_Check(obj) || PyInt_Check(obj); + return PyFloat_Check(obj) || PyLong_Check(obj) || PyInt_Check(obj) || is_numpy_scalar; #else - return PyFloat_Check(obj) || PyLong_Check(obj); + return PyFloat_Check(obj) || PyLong_Check(obj) || is_numpy_scalar; #endif } diff --git a/torch/csrc/utils/python_stub.h b/torch/csrc/utils/python_stub.h index 485b3b5f9e6a0..336c530d2b1fa 100644 --- a/torch/csrc/utils/python_stub.h +++ b/torch/csrc/utils/python_stub.h @@ -1,4 +1,4 @@ #pragma once struct _object; -typedef _object PyObject; +using PyObject = _object; diff --git a/torch/csrc/utils/tensor_new.cpp b/torch/csrc/utils/tensor_new.cpp index 7dc6ff6ceb357..38fa1c8b7821b 100644 --- a/torch/csrc/utils/tensor_new.cpp +++ b/torch/csrc/utils/tensor_new.cpp @@ -72,7 +72,7 @@ Tensor new_with_sizes(const Type& type, int32_t device_index, IntList sizes) { } Tensor new_with_storage(const Type& type, Storage storage) { - auto tensor = at::empty({}, type); + auto tensor = at::empty({}, type.options()); tensor.set_(storage); return tensor; } diff --git a/torch/csrc/utils/tensor_numpy.cpp b/torch/csrc/utils/tensor_numpy.cpp index 0c9ec4cbcc764..e7170c03ae422 100644 --- a/torch/csrc/utils/tensor_numpy.cpp +++ b/torch/csrc/utils/tensor_numpy.cpp @@ -10,6 +10,9 @@ PyObject* tensor_to_numpy(const at::Tensor& tensor) { at::Tensor tensor_from_numpy(PyObject* obj) { throw std::runtime_error("PyTorch was compiled without NumPy support"); } +bool is_numpy_scalar(PyObject* obj) { + throw std::runtime_error("PyTorch was compiled without NumPy support"); +} }} #else @@ -180,6 +183,11 @@ ScalarType numpy_dtype_to_aten(int dtype) { ((PyTypeObject*)pytype.get())->tp_name); } +bool is_numpy_scalar(PyObject* obj) { + return (PyArray_IsIntegerScalar(obj) || + PyArray_IsScalar(obj, Floating)); +} + }} // namespace torch::utils #endif // USE_NUMPY diff --git a/torch/csrc/utils/tensor_numpy.h b/torch/csrc/utils/tensor_numpy.h index 4a86306cbcbf5..d3984ed853cd6 100644 --- a/torch/csrc/utils/tensor_numpy.h +++ b/torch/csrc/utils/tensor_numpy.h @@ -10,4 +10,6 @@ at::Tensor tensor_from_numpy(PyObject* obj); at::ScalarType numpy_dtype_to_aten(int dtype); +bool is_numpy_scalar(PyObject* obj); + }} // namespace torch::utils diff --git a/torch/distributed/c10d/distributed_c10d.py b/torch/distributed/c10d/distributed_c10d.py index fb9958e007114..d1a697e58288b 100644 --- a/torch/distributed/c10d/distributed_c10d.py +++ b/torch/distributed/c10d/distributed_c10d.py @@ -40,7 +40,10 @@ class GroupMember(object): NON_GROUP_MEMBER = object() -# Cached process groups, map from ProcessGroup to (DistBackend, Store) +# Cached process groups +# For NCCL and GLOO pg, it is a map from ProcessGroup to (DistBackend, Store) +# For MPI pg, it is a map from ProcessGroup to (DistBackend, Bool), where bool +# represents if the ProcessGroup objects is part of the group _pg_map = {} # Process group's names, map from ProcessGroup to str _pg_names = {} @@ -60,7 +63,15 @@ def _rank_not_in_group(group): Helper that checks if the current process's rank is not in a given group """ - return group == GroupMember.NON_GROUP_MEMBER + default_backend, _ = _pg_map[get_default_group()] + if default_backend != DistBackend.MPI: + return group == GroupMember.NON_GROUP_MEMBER + else: + if group == GroupMember.WORLD: + return False + else: + _, in_group = _pg_map[group] + return not in_group def _get_group_rank(group, rank): @@ -107,6 +118,19 @@ def _check_default_pg(): "Default process group is not initialized" +def _get_group_size(group): + """ + Helper that gets a given group's world size + + """ + if group is GroupMember.WORLD: + _check_default_pg() + return _default_pg.size() + if group not in _pg_group_ranks: + raise RuntimeError("The given group does not exist") + return len(_pg_group_ranks[group]) + + def is_mpi_available(): """ Checks if MPI is available @@ -184,7 +208,8 @@ def init_process_group(backend, raise RuntimeError("Distributed package doesn't have MPI built in") _default_pg = ProcessGroupMPI([]) - _pg_map[_default_pg] = (DistBackend.MPI, None) + _pg_map[_default_pg] = (DistBackend.MPI, True) + _pg_names[_default_pg] = group_name else: # backward compatible API if init_method != "env://" and world_size != -1 and rank != -1: @@ -212,7 +237,11 @@ def init_process_group(backend, _default_pg_init_method = init_method -def _new_process_group_helper(world_size, rank, ranks, group_name=""): +def _new_process_group_helper(world_size, + rank, + group_ranks, + in_group=True, + group_name=""): """ Create a new distributed process group. And the new process group can be used to perform collective operations. @@ -235,8 +264,8 @@ def _new_process_group_helper(world_size, rank, ranks, group_name=""): if default_backend == DistBackend.MPI: if not is_mpi_available(): raise RuntimeError("Distributed package doesn't have MPI built in") - pg = ProcessGroupMPI(ranks) - _pg_map[pg] = (DistBackend.MPI, None) + pg = ProcessGroupMPI(group_ranks) + _pg_map[pg] = (DistBackend.MPI, in_group) _pg_names[pg] = group_name else: # Create the prefix store @@ -268,18 +297,21 @@ def destroy_process_group(group=group.WORLD): groups including the default one will be destroyed. """ - if _rank_not_in_group(group): - return - global _pg_map global _pg_names global _pg_group_ranks global _default_pg global _default_pg_init_method + default_backend, _ = _pg_map[get_default_group()] + if (default_backend != DistBackend.MPI and + group == GroupMember.NON_GROUP_MEMBER): + return + if group == GroupMember.WORLD: pg = _default_pg - + else: + pg = group if _pg_map.get(pg, None) is None: raise RuntimeError("Invalid process group specified") @@ -314,11 +346,11 @@ def get_rank(group=group.WORLD): if _rank_not_in_group(group): return -1 + _check_default_pg() if group == GroupMember.WORLD: - _check_default_pg() return _default_pg.rank() - return group.rank() + return _get_group_rank(group, _default_pg.rank()) def get_world_size(group=group.WORLD): @@ -336,11 +368,7 @@ def get_world_size(group=group.WORLD): if _rank_not_in_group(group): return -1 - if group == GroupMember.WORLD: - _check_default_pg() - return _default_pg.size() - - return group.size() + return _get_group_size(group) def isend(tensor, @@ -1009,6 +1037,7 @@ def new_group(ranks=None): # checks the input ranks if ranks is not None: + input_ranks = list(ranks) group_world_size = len(ranks) if group_world_size > global_world_size: raise RuntimeError("the new group's world size should be less or " @@ -1024,19 +1053,26 @@ def new_group(ranks=None): else: group_rank = None else: - ranks = [] + input_ranks = [] + ranks = list(range(global_world_size)) group_world_size = global_world_size group_rank = global_rank if default_backend == DistBackend.MPI: - pg = _new_process_group_helper(group_world_size, group_rank, ranks) - - # Release ranks not in the group - if global_rank not in ranks: - return GroupMember.NON_GROUP_MEMBER + in_group = global_rank in ranks + pg = _new_process_group_helper(group_world_size, + group_rank, + input_ranks, + in_group) + else: + # Release ranks not in the group + if global_rank not in ranks: + return GroupMember.NON_GROUP_MEMBER - if default_backend != DistBackend.MPI: - pg = _new_process_group_helper(group_world_size, group_rank, ranks) + if default_backend != DistBackend.MPI: + pg = _new_process_group_helper(group_world_size, + group_rank, + input_ranks) # Create the global rank to group rank mapping _pg_group_ranks[pg] = {} @@ -1047,18 +1083,3 @@ def new_group(ranks=None): if rank in ranks: _pg_group_ranks[pg][rank] = ranks.index(rank) return pg - - -# TODO: delete these functions and replace DDP with public functions -DEFAULT_REDUCE_OPTIONS = AllreduceOptions() - - -def _broadcast(tensor, src, process_group): - opts = BroadcastOptions() - opts.rootRank = src - opts.rootTensor = 0 - return process_group.broadcast([tensor], opts) - - -def _all_reduce(tensor, process_group, opts=DEFAULT_REDUCE_OPTIONS): - return process_group.allreduce([tensor], opts) diff --git a/torch/distributions/bernoulli.py b/torch/distributions/bernoulli.py index 09099e07a8228..6939594035945 100644 --- a/torch/distributions/bernoulli.py +++ b/torch/distributions/bernoulli.py @@ -24,7 +24,8 @@ class Bernoulli(ExponentialFamily): probs (Number, Tensor): the probabilty of sampling `1` logits (Number, Tensor): the log-odds of sampling `1` """ - arg_constraints = {'probs': constraints.unit_interval} + arg_constraints = {'probs': constraints.unit_interval, + 'logits': constraints.real} support = constraints.boolean has_enumerate_support = True _mean_carrier_measure = 0 diff --git a/torch/distributions/binomial.py b/torch/distributions/binomial.py index da2b1062f1a20..e15138e58ea68 100644 --- a/torch/distributions/binomial.py +++ b/torch/distributions/binomial.py @@ -28,7 +28,8 @@ class Binomial(Distribution): logits (Tensor): Event log-odds """ arg_constraints = {'total_count': constraints.nonnegative_integer, - 'probs': constraints.unit_interval} + 'probs': constraints.unit_interval, + 'logits': constraints.real} has_enumerate_support = True def __init__(self, total_count=1, probs=None, logits=None, validate_args=None): diff --git a/torch/distributions/categorical.py b/torch/distributions/categorical.py index 5c42991b20606..2f3a425deaefc 100644 --- a/torch/distributions/categorical.py +++ b/torch/distributions/categorical.py @@ -37,7 +37,8 @@ class Categorical(Distribution): probs (Tensor): event probabilities logits (Tensor): event log probabilities """ - arg_constraints = {'probs': constraints.simplex} + arg_constraints = {'probs': constraints.simplex, + 'logits': constraints.real} has_enumerate_support = True def __init__(self, probs=None, logits=None, validate_args=None): diff --git a/torch/distributions/distribution.py b/torch/distributions/distribution.py index 9acb4f776e5ef..e93c15e63c780 100644 --- a/torch/distributions/distribution.py +++ b/torch/distributions/distribution.py @@ -221,7 +221,7 @@ def _validate_sample(self, value): raise ValueError('The value argument must be within the support') def __repr__(self): - param_names = [k for k, _ in self.arg_constraints.items()] + param_names = [k for k, _ in self.arg_constraints.items() if k in self.__dict__] args_string = ', '.join(['{}: {}'.format(p, self.__dict__[p] if self.__dict__[p].dim() == 0 else self.__dict__[p].size()) for p in param_names]) diff --git a/torch/distributions/geometric.py b/torch/distributions/geometric.py index c4aa9978b58d7..196a28b29e7d7 100644 --- a/torch/distributions/geometric.py +++ b/torch/distributions/geometric.py @@ -25,7 +25,8 @@ class Geometric(Distribution): probs (Number, Tensor): the probabilty of sampling `1`. Must be in range (0, 1] logits (Number, Tensor): the log-odds of sampling `1`. """ - arg_constraints = {'probs': constraints.unit_interval} + arg_constraints = {'probs': constraints.unit_interval, + 'logits': constraints.real} support = constraints.nonnegative_integer def __init__(self, probs=None, logits=None, validate_args=None): diff --git a/torch/distributions/multinomial.py b/torch/distributions/multinomial.py index ccfab11471283..dd1e20d373469 100644 --- a/torch/distributions/multinomial.py +++ b/torch/distributions/multinomial.py @@ -38,7 +38,8 @@ class Multinomial(Distribution): probs (Tensor): event probabilities logits (Tensor): event log probabilities """ - arg_constraints = {'logits': constraints.real} # Let logits be the canonical parameterization. + arg_constraints = {'probs': constraints.simplex, + 'logits': constraints.real} @property def mean(self): diff --git a/torch/distributions/negative_binomial.py b/torch/distributions/negative_binomial.py index 4f9804ead826a..de0b748f6a6fd 100644 --- a/torch/distributions/negative_binomial.py +++ b/torch/distributions/negative_binomial.py @@ -20,7 +20,8 @@ class NegativeBinomial(Distribution): logits (Tensor): Event log-odds for probabilities of success """ arg_constraints = {'total_count': constraints.greater_than_eq(0), - 'probs': constraints.half_open_interval(0., 1.)} + 'probs': constraints.half_open_interval(0., 1.), + 'logits': constraints.real} support = constraints.nonnegative_integer def __init__(self, total_count, probs=None, logits=None, validate_args=None): diff --git a/torch/distributions/one_hot_categorical.py b/torch/distributions/one_hot_categorical.py index 4213c4453225e..fe2d2d7e6f9af 100644 --- a/torch/distributions/one_hot_categorical.py +++ b/torch/distributions/one_hot_categorical.py @@ -27,7 +27,8 @@ class OneHotCategorical(Distribution): probs (Tensor): event probabilities logits (Tensor): event log probabilities """ - arg_constraints = {'probs': constraints.simplex} + arg_constraints = {'probs': constraints.simplex, + 'logits': constraints.real} support = constraints.simplex has_enumerate_support = True diff --git a/torch/distributions/relaxed_bernoulli.py b/torch/distributions/relaxed_bernoulli.py index 5923e346085c6..6b6c540ec48c0 100644 --- a/torch/distributions/relaxed_bernoulli.py +++ b/torch/distributions/relaxed_bernoulli.py @@ -25,7 +25,8 @@ class LogitRelaxedBernoulli(Distribution): [2] Categorical Reparametrization with Gumbel-Softmax (Jang et al, 2017) """ - arg_constraints = {'probs': constraints.unit_interval} + arg_constraints = {'probs': constraints.unit_interval, + 'logits': constraints.real} support = constraints.real def __init__(self, temperature, probs=None, logits=None, validate_args=None): @@ -92,7 +93,8 @@ class RelaxedBernoulli(TransformedDistribution): probs (Number, Tensor): the probabilty of sampling `1` logits (Number, Tensor): the log-odds of sampling `1` """ - arg_constraints = {'probs': constraints.unit_interval} + arg_constraints = {'probs': constraints.unit_interval, + 'logits': constraints.real} support = constraints.unit_interval has_rsample = True diff --git a/torch/distributions/relaxed_categorical.py b/torch/distributions/relaxed_categorical.py index 45b8791b4fe38..0c35defaf35bd 100644 --- a/torch/distributions/relaxed_categorical.py +++ b/torch/distributions/relaxed_categorical.py @@ -29,7 +29,8 @@ class ExpRelaxedCategorical(Distribution): [2] Categorical Reparametrization with Gumbel-Softmax (Jang et al, 2017) """ - arg_constraints = {'probs': constraints.simplex} + arg_constraints = {'probs': constraints.simplex, + 'logits': constraints.real} support = constraints.real has_rsample = True @@ -93,7 +94,8 @@ class RelaxedOneHotCategorical(TransformedDistribution): probs (Tensor): event probabilities logits (Tensor): the log probability of each event. """ - arg_constraints = {'probs': constraints.simplex} + arg_constraints = {'probs': constraints.simplex, + 'logits': constraints.real} support = constraints.simplex has_rsample = True diff --git a/torch/functional.py b/torch/functional.py index 8c78b6efe9f80..d84597f3fff43 100644 --- a/torch/functional.py +++ b/torch/functional.py @@ -18,6 +18,7 @@ 'isnan', 'split', 'stft', + 'tensordot', 'unique', ] @@ -512,6 +513,61 @@ def argmin(input, dim=None, keepdim=False): return torch._argmin(input, dim, keepdim) +def tensordot(a, b, dims=2): + """Returns a contraction of a and b over multiple dimensions. + + :attr:`tensordot` implements a generalizes the matrix product. + + Args: + a (Tensor): Left tensor to contract + b (Tensor): Right tensor to contract + dims (int or tuple of two lists of integers): number of dimensions to + contract or explicit lists of dimensions for :attr:`a` and + :attr:`b` respectively + + When called with an integer argument :attr:`dims` = :math:`d`, and the number of + dimensions of :attr:`a` and :attr:`b` is :math:`m` and :math:`n`, respectively, + it computes + + .. math:: + r_{i_0,...,i_{m-d}, i_d,...,i_n} + = \sum_{k_0,...,k_{d-1}} a_{i_0,...,i_{m-d},k_0,...,k_{d-1}} * b_{k_0,...,k_{d-1}, i_d,...,i_n}. + + When called with :attr:`dims` of the list form, the given dimensions will be contracted + in place of the last :math:`d` of :attr:`a` and the first :math:`d` of :math:`b`. The sizes + in these dimensions must match, but :attr:`tensordot` will deal with broadcasted + dimensions. + + Examples:: + + >>> a = torch.arange(60.).reshape(3, 4, 5) + >>> b = torch.arange(24.).reshape(4, 3, 2) + >>> torch.tensordot(a, b, dims=([1, 0], [0, 1])) + tensor([[4400., 4730.], + [4532., 4874.], + [4664., 5018.], + [4796., 5162.], + [4928., 5306.]]) + + >>> a = torch.randn(3, 4, 5, device='cuda') + >>> b = torch.randn(4, 5, 6, device='cuda') + >>> c = torch.tensordot(a, b, dims=2).cpu() + tensor([[ 8.3504, -2.5436, 6.2922, 2.7556, -1.0732, 3.2741], + [ 3.3161, 0.0704, 5.0187, -0.4079, -4.3126, 4.8744], + [ 0.8223, 3.9445, 3.2168, -0.2400, 3.4117, 1.7780]]) + + """ + if isinstance(dims, (list, tuple)) or \ + (isinstance(dims, torch.Tensor) and dims.numel() > 1): + dims_a, dims_b = dims + else: + if isinstance(dims, torch.Tensor): + dims = dims.item() + dims_a = list(range(-dims, 0)) + dims_b = list(range(dims)) + return torch._C._VariableFunctions.tensordot(a, b, dims_a, dims_b) + + def argsort(input, dim=None, descending=False): """Returns the indices that sort a tensor along a given dimension in ascending order by value. diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py index eab24b6f031bf..159902c81c616 100644 --- a/torch/jit/__init__.py +++ b/torch/jit/__init__.py @@ -277,16 +277,13 @@ def indent(s): class TracingCheckError(Exception): - def __init__(self, graph_diff_error, tensor_compare_error, nondeterm_warning, extra_msg=None): + def __init__(self, graph_diff_error, tensor_compare_error, extra_msg=None): self.message = 'Tracing failed sanity checks!\n' if extra_msg is not None: self.message += extra_msg + '\n' if graph_diff_error is not None: self.message += 'ERROR: Graphs differed across invocations!\n' self.message += indent(graph_diff_error) + '\n' - if nondeterm_warning is not None: - self.message += 'WARNING: ' - self.message += nondeterm_warning + '\n' if tensor_compare_error is not None: self.message += 'ERROR: Tensor-valued Constant nodes differed in value ' \ 'across invocations. This often indicates that the tracer has' \ @@ -298,6 +295,8 @@ def __init__(self, graph_diff_error, tensor_compare_error, nondeterm_warning, ex # Check the traced module against a set of user-provided validation inputs def _check_trace(check_inputs, func, executor_options, module, check_tolerance): for inputs in check_inputs: + if isinstance(inputs, torch.Tensor): + inputs = (inputs,) check_mod = torch.jit.trace(func, _clone_inputs(inputs), check_trace=False, **executor_options) def graph_diagnostic_info(): @@ -345,7 +344,7 @@ def graph_diagnostic_info(): try: torch.testing.assert_allclose(mod_tensor_val, check_tensor_val) except (RuntimeError, AssertionError) as e: - if not tensor_compare_errors: + if tensor_compare_errors is None: tensor_compare_errors = '' tensor_compare_errors += 'Node:\n' + indent(str(n_mod)) + '\n' compare_stack = n_mod.getSourceLocation() @@ -355,63 +354,66 @@ def graph_diagnostic_info(): break # For now, only print the first diverging pair - nondeterministic_ops_warning = None + return graph_diff_errors, tensor_compare_errors + + def wrap_retval(x): + return x if isinstance(x, tuple) else (x,) + + def run_mod_and_filter_tensor_outputs(mod, inputs, running_what): + try: + outs = wrap_retval(mod(*_clone_inputs(inputs))) + outs = [out for out in outs if isinstance(out, torch.Tensor)] + return outs + except Exception as e: + raise TracingCheckError(*graph_diagnostic_info(), + extra_msg='Encountered an exception while running the ' + running_what + + ' with test inputs.\nException:\n' + indent(str(e))) + + has_warned = [False] + + def maybe_warn_nondeterministic(): + if has_warned[0]: + return + has_warned[0] = True nondeterm_ops = [op for op in module.graph.nodes() if op.isNondeterministic()] if len(nondeterm_ops) > 0: nondeterministic_ops_warning = "Trace had nondeterministic nodes. Nodes:\n" - for op in nondeterm_ops: - nondeterministic_ops_warning += indent(str(op)) + nondeterministic_ops_warning += "\n".join([indent(str(op)) for op in nondeterm_ops][:20]) nondeterministic_ops_warning += "\nThis may cause errors in trace checking. To disable trace checking,"\ " pass disable_checks=True to torch.jit.trace()" - - return graph_diff_errors, tensor_compare_errors, nondeterministic_ops_warning - - def wrap_retval(x): - return x if isinstance(x, tuple) else (x,) - - def run_mod_and_filter_tensor_outputs(mod, inputs): - outs = wrap_retval(mod(*_clone_inputs(inputs))) - outs = [out for out in outs if isinstance(out, torch.Tensor)] - return outs - - try: - traced_outs = run_mod_and_filter_tensor_outputs(module, inputs) - except Exception as e: - msg = 'Encountered an exception while running trace with check inputs.\nException:\n' + indent(str(e)) - raise TracingCheckError(*graph_diagnostic_info(), extra_msg=msg) - - try: - fn_outs = run_mod_and_filter_tensor_outputs(func, inputs) - for orig, check in zip(traced_outs, fn_outs): - torch.testing.assert_allclose(orig.double(), check.double(), rtol=check_tolerance, - atol=torch.testing._get_default_tolerance(orig, check)[1]) - except (RuntimeError, AssertionError) as e: - # TODO: interpose on tracing the function again and check for - # divergence? then we can point to where in the source code - # we start diverging in python v.s. the trace - msg = 'ERROR: Traced function outputs do not match the Python function outputs.\nException: ' + str(e) - raise TracingCheckError(*graph_diagnostic_info(), - extra_msg=msg) - - try: - check_outs = run_mod_and_filter_tensor_outputs(check_mod, inputs) - except Exception as e: - msg = 'Encountered an exception while running checking trace with check inputs.\nException:\n' \ - + indent(str(e)) - raise TracingCheckError(*graph_diagnostic_info(), extra_msg=msg) - - try: - for orig, check in zip(traced_outs, check_outs): - torch.testing.assert_allclose(orig.double(), check.double(), rtol=check_tolerance, - atol=torch.testing._get_default_tolerance(orig, check)[1]) - except (RuntimeError, AssertionError) as e: - raise TracingCheckError(*graph_diagnostic_info()) + warnings.warn(nondeterministic_ops_warning, category=TracerWarning, stacklevel=5) + + def compare_outputs(original, reference, match_what): + all_ok = True + for i, (orig, ref) in enumerate(zip(original, reference)): + try: + torch.testing.assert_allclose(orig.double(), ref.double(), rtol=check_tolerance, + atol=torch.testing._get_default_tolerance(orig, ref)[1]) + except AssertionError as e: + maybe_warn_nondeterministic() + warnings.warn('Output nr ' + str(i + 1) + '. of the traced function does not match ' + 'the corresponding output of the ' + match_what + '. Detailed error:\n' + str(e), + category=TracerWarning, stacklevel=4) + all_ok = False + + return all_ok + + traced_outs = run_mod_and_filter_tensor_outputs(module, inputs, 'trace') + fn_outs = run_mod_and_filter_tensor_outputs(func, inputs, 'Python function') + if compare_outputs(traced_outs, fn_outs, 'Python function'): + check_outs = run_mod_and_filter_tensor_outputs(check_mod, inputs, 'repeated trace') + compare_outputs(traced_outs, check_outs, 'repeated trace') + + diag_info = graph_diagnostic_info() + if any(info is not None for info in diag_info): + raise TracingCheckError(*diag_info) class TracerWarning(Warning): @staticmethod def ignore_lib_warnings(): - warnings.filterwarnings('ignore', category=TracerWarning, module='torch.*') + # We ignore warnings from all submodules excluding the JIT, because we need them e.g. for _check_trace + warnings.filterwarnings('ignore', category=TracerWarning, module='torch.(?!jit)') # We ignore the tracer warnings coming form inside the library, because all our shape diff --git a/torch/jit/frontend.py b/torch/jit/frontend.py index 34d0993a9942f..659c3bc4cdd70 100644 --- a/torch/jit/frontend.py +++ b/torch/jit/frontend.py @@ -1,3 +1,4 @@ +import __future__ import torch import sys import ast @@ -124,13 +125,34 @@ def build_stmts(ctx, stmts): return list(filter(None, stmts)) +def _uses_true_division(fn): + if not PY2: + return True + if inspect.ismethod(fn): + return _uses_true_division(fn.__func__) + elif inspect.isfunction(fn): + return fn.__globals__.get('division') is __future__.division + else: + raise RuntimeError( + '_uses_true_division: expected function or method, got {}'.format(type(fn))) + + def get_jit_ast(fn, is_method): source = dedent(inspect.getsource(fn)) py_ast = ast.parse(source) if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef): raise RuntimeError("expected a single top-level function") type_line = torch.jit.annotations.get_type_line(source) - return build_def(SourceRangeFactory(source), py_ast.body[0], type_line, is_method) + ctx = SourceContext(source, _uses_true_division(fn)) + return build_def(ctx, py_ast.body[0], type_line, is_method) + + +# Thin wrapper around SourceRangeFactory to store extra metadata +# about the function-to-be-compiled. +class SourceContext(SourceRangeFactory): + def __init__(self, source, uses_true_division=True): + super(SourceContext, self).__init__(source) + self.uses_true_division = uses_true_division class Builder(object): @@ -365,6 +387,12 @@ def build_BinOp(ctx, expr): lhs = build_expr(ctx, expr.left) rhs = build_expr(ctx, expr.right) op = type(expr.op) + + if op == ast.Div and not ctx.uses_true_division: + raise RuntimeError('Division of ints in JIT script uses Python 3 true ' + 'division semantics. Please put `from __future__ ' + 'import division` at the top of your file') + op_token = ExprBuilder.binop_map.get(op) if op_token is None: err_range = ctx.make_raw_range(lhs.range().end, rhs.range().start) diff --git a/torch/lib/c10d/CUDAUtils.cpp b/torch/lib/c10d/CUDAUtils.cpp index cc3c107f127d3..2a1a7d176999b 100644 --- a/torch/lib/c10d/CUDAUtils.cpp +++ b/torch/lib/c10d/CUDAUtils.cpp @@ -13,7 +13,7 @@ CUDAEvent CUDAEvent::create(unsigned int flags) { return event; } -CUDAEvent::~CUDAEvent() { +CUDAEvent::~CUDAEvent() noexcept (false) { if (event_ != nullptr) { // cudaEventDestroy must run on the same device of the event, // otherwise it creates a context on default device as well. diff --git a/torch/lib/c10d/CUDAUtils.hpp b/torch/lib/c10d/CUDAUtils.hpp index f57f6929b38f3..c0db85e7edeea 100644 --- a/torch/lib/c10d/CUDAUtils.hpp +++ b/torch/lib/c10d/CUDAUtils.hpp @@ -16,7 +16,7 @@ class CUDAEvent { CUDAEvent() : CUDAEvent(nullptr, 0) {} - ~CUDAEvent(); + ~CUDAEvent() noexcept(false); static CUDAEvent create(unsigned int flags = cudaEventDefault); diff --git a/torch/lib/c10d/FileStore.cpp b/torch/lib/c10d/FileStore.cpp index 75d793199a609..fb842125e56cd 100644 --- a/torch/lib/c10d/FileStore.cpp +++ b/torch/lib/c10d/FileStore.cpp @@ -199,6 +199,7 @@ void FileStore::set(const std::string& key, const std::vector& value) { } std::vector FileStore::get(const std::string& key) { + const auto start = std::chrono::steady_clock::now(); while (cache_.count(key) == 0) { File file(path_, O_RDONLY); auto lock = file.lockShared(); @@ -206,6 +207,11 @@ std::vector FileStore::get(const std::string& key) { if (size == pos_) { // No new entries; release the shared lock and sleep for a bit lock.unlock(); + const auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start); + if (timeout_ != kNoTimeout && elapsed > timeout_) { + throw std::runtime_error("Timeout waiting for key: " + key); + } std::this_thread::sleep_for(std::chrono::milliseconds(10)); continue; } @@ -251,6 +257,10 @@ bool FileStore::check(const std::vector& keys) { return true; } +void FileStore::wait(const std::vector& keys) { + wait(keys, timeout_); +} + void FileStore::wait( const std::vector& keys, const std::chrono::milliseconds& timeout) { diff --git a/torch/lib/c10d/FileStore.hpp b/torch/lib/c10d/FileStore.hpp index a6753884c4629..9bbc25ef10981 100644 --- a/torch/lib/c10d/FileStore.hpp +++ b/torch/lib/c10d/FileStore.hpp @@ -22,9 +22,11 @@ class FileStore : public Store { bool check(const std::vector& keys) override; + void wait(const std::vector& keys) override; + void wait( const std::vector& keys, - const std::chrono::milliseconds& timeout = kDefaultTimeout) override; + const std::chrono::milliseconds& timeout) override; protected: std::string path_; diff --git a/torch/lib/c10d/PrefixStore.cpp b/torch/lib/c10d/PrefixStore.cpp index 5ddad65dc29e2..4535fc3d8d5c4 100644 --- a/torch/lib/c10d/PrefixStore.cpp +++ b/torch/lib/c10d/PrefixStore.cpp @@ -38,6 +38,11 @@ bool PrefixStore::check(const std::vector& keys) { return store_.check(joinedKeys); } +void PrefixStore::wait(const std::vector& keys) { + auto joinedKeys = joinKeys(keys); + store_.wait(joinedKeys); +} + void PrefixStore::wait( const std::vector& keys, const std::chrono::milliseconds& timeout) { diff --git a/torch/lib/c10d/PrefixStore.hpp b/torch/lib/c10d/PrefixStore.hpp index 8347fa213a4b4..4e572c898ec1f 100644 --- a/torch/lib/c10d/PrefixStore.hpp +++ b/torch/lib/c10d/PrefixStore.hpp @@ -18,9 +18,11 @@ class PrefixStore : public Store { bool check(const std::vector& keys) override; + void wait(const std::vector& keys) override; + void wait( const std::vector& keys, - const std::chrono::milliseconds& timeout = kDefaultTimeout) override; + const std::chrono::milliseconds& timeout) override; protected: std::string prefix_; diff --git a/torch/lib/c10d/Store.cpp b/torch/lib/c10d/Store.cpp index da7e8d8fa94bc..86bef88cdf9ec 100644 --- a/torch/lib/c10d/Store.cpp +++ b/torch/lib/c10d/Store.cpp @@ -8,4 +8,12 @@ constexpr std::chrono::milliseconds Store::kNoTimeout; // Define destructor symbol for abstract base class. Store::~Store() {} +// Set timeout function +void Store::setTimeout(const std::chrono::seconds& timeoutSec) { + if (timeoutSec.count() == 0) { + timeout_ = kNoTimeout; + } + timeout_ = timeoutSec; +} + } // namespace c10d diff --git a/torch/lib/c10d/Store.hpp b/torch/lib/c10d/Store.hpp index eadb2ecc31e51..d24fac275dcfe 100644 --- a/torch/lib/c10d/Store.hpp +++ b/torch/lib/c10d/Store.hpp @@ -15,6 +15,8 @@ class Store { static constexpr std::chrono::milliseconds kNoTimeout = std::chrono::milliseconds::zero(); + Store() : timeout_(kDefaultTimeout) {} + virtual ~Store(); virtual void set( @@ -27,9 +29,16 @@ class Store { virtual bool check(const std::vector& keys) = 0; + virtual void wait(const std::vector& keys) = 0; + virtual void wait( const std::vector& keys, - const std::chrono::milliseconds& timeout = kDefaultTimeout) = 0; + const std::chrono::milliseconds& timeout) = 0; + + void setTimeout(const std::chrono::seconds& timeoutSec); + + protected: + std::chrono::milliseconds timeout_; }; } // namespace c10d diff --git a/torch/lib/c10d/TCPStore.cpp b/torch/lib/c10d/TCPStore.cpp index da4e2f2eb98c3..25bcc0455287b 100644 --- a/torch/lib/c10d/TCPStore.cpp +++ b/torch/lib/c10d/TCPStore.cpp @@ -341,6 +341,10 @@ bool TCPStore::check(const std::vector& keys) { } } +void TCPStore::wait(const std::vector& keys) { + wait(keys, timeout_); +} + void TCPStore::wait( const std::vector& keys, const std::chrono::milliseconds& timeout) { diff --git a/torch/lib/c10d/TCPStore.hpp b/torch/lib/c10d/TCPStore.hpp index bc0c0be27535c..891f263c16ca7 100644 --- a/torch/lib/c10d/TCPStore.hpp +++ b/torch/lib/c10d/TCPStore.hpp @@ -60,9 +60,11 @@ class TCPStore : public Store { bool check(const std::vector& keys) override; + void wait(const std::vector& keys) override; + void wait( const std::vector& keys, - const std::chrono::milliseconds& timeout = kDefaultTimeout) override; + const std::chrono::milliseconds& timeout) override; protected: bool isServer_; diff --git a/torch/nn/_functions/packing.py b/torch/nn/_functions/packing.py deleted file mode 100644 index a45cb637c426c..0000000000000 --- a/torch/nn/_functions/packing.py +++ /dev/null @@ -1,56 +0,0 @@ -import torch -from torch.autograd import Function - - -class PackPadded(Function): - @staticmethod - def forward(ctx, input, lengths, batch_first): - if batch_first: - input = input.transpose(0, 1) - - if lengths[-1] <= 0: - raise ValueError("Length of all samples has to be greater than 0, " - "but found an element in 'lengths' that is <= 0") - - steps = [] - batch_sizes = [] - - # lengths is a Tensor, so we must convert to [int] before reversed() - lengths_iter = reversed(lengths.tolist()) - - batch_size = input.size(1) - - if len(lengths) != batch_size: - raise ValueError("Expected `len(lengths)` to be equal to batch_size, but got " - "{} (batch_size={}).".format(len(lengths), batch_size)) - - prev_l = 0 - for i, l in enumerate(lengths_iter): - if l > prev_l: - c_batch_size = batch_size - i - steps.append(input[prev_l:l, :c_batch_size].contiguous().view(-1, *input.size()[2:])) - batch_sizes.extend([c_batch_size] * (l - prev_l)) - prev_l = l - - elif prev_l > l: - raise ValueError("'lengths' array has to be sorted in decreasing order") - - ctx.batch_sizes = batch_sizes - ctx.batch_first = batch_first - ctx.input_size = input.size() - - return torch.cat(steps), torch.LongTensor(batch_sizes) - - @staticmethod - def backward(ctx, grad_steps, grad_batch_sizes): - grad_input = grad_steps.new(*ctx.input_size).zero_() - - offset = 0 - for i, bs in enumerate(ctx.batch_sizes): - grad_input[i, :bs] = grad_steps[offset:offset + bs] - offset += bs - - if ctx.batch_first: - grad_input = grad_input.transpose(0, 1) - - return grad_input, None, None diff --git a/torch/nn/functional.py b/torch/nn/functional.py index 76f25aeb7ef9a..47b7375bd1f45 100644 --- a/torch/nn/functional.py +++ b/torch/nn/functional.py @@ -593,7 +593,19 @@ def adaptive_avg_pool3d(input, output_size): # Activation functions -def dropout(input, p=0.5, training=False, inplace=False): +def dropout(input, p=0.5, training=True, inplace=False): + r""" + During training, randomly zeroes some of the elements of the input + tensor with probability :attr:`p` using samples from a Bernoulli + distribution. + + See :class:`~torch.nn.Dropout` for details. + + Args: + p: probability of an element to be zeroed. Default: 0.5 + training: apply dropout if is ``True``. Defualt: ``True`` + inplace: If set to ``True``, will do this operation in-place. Default: ``False`` + """ if p < 0 or p > 1: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) @@ -613,7 +625,21 @@ def alpha_dropout(input, p=0.5, training=False, inplace=False): return f(input, p, training) -def dropout2d(input, p=0.5, training=False, inplace=False): +def dropout2d(input, p=0.5, training=True, inplace=False): + r""" + Randomly zero out entire channels (a channel is a 2D feature map, + e.g., the :math:`j`-th channel of the :math:`i`-th sample in the + batched input is a 2D tensor :math:`\text{input}[i, j]`) of the input tensor). + Each channel will be zeroed out independently on every forward call. + with probability :attr:`p` using samples from a Bernoulli distribution. + + See :class:`~torch.nn.Dropout2d` for details. + + Args: + p: probability of a channel to be zeroed. Default: 0.5 + training: apply dropout if is ``True``. Defualt: ``True`` + inplace: If set to ``True``, will do this operation in-place. Default: ``False`` + """ if p < 0 or p > 1: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) @@ -621,7 +647,21 @@ def dropout2d(input, p=0.5, training=False, inplace=False): return f(input, p, training) -def dropout3d(input, p=0.5, training=False, inplace=False): +def dropout3d(input, p=0.5, training=True, inplace=False): + r""" + Randomly zero out entire channels (a channel is a 3D feature map, + e.g., the :math:`j`-th channel of the :math:`i`-th sample in the + batched input is a 3D tensor :math:`\text{input}[i, j]`) of the input tensor). + Each channel will be zeroed out independently on every forward call. + with probability :attr:`p` using samples from a Bernoulli distribution. + + See :class:`~torch.nn.Dropout3d` for details. + + Args: + p: probability of a channel to be zeroed. Default: 0.5 + training: apply dropout if is ``True``. Defualt: ``True`` + inplace: If set to ``True``, will do this operation in-place. Default: ``False`` + """ if p < 0 or p > 1: raise ValueError("dropout probability has to be between 0 and 1, " "but got {}".format(p)) diff --git a/torch/nn/modules/dropout.py b/torch/nn/modules/dropout.py index 48415f61929f6..bbc5b846947be 100644 --- a/torch/nn/modules/dropout.py +++ b/torch/nn/modules/dropout.py @@ -20,7 +20,8 @@ def extra_repr(self): class Dropout(_DropoutNd): r"""During training, randomly zeroes some of the elements of the input tensor with probability :attr:`p` using samples from a Bernoulli - distribution. The elements to zero are randomized on every forward call. + distribution. Each channel will be zeroed out independently on every forward + call. This has proven to be an effective technique for regularization and preventing the co-adaptation of neurons as described in the paper @@ -54,8 +55,11 @@ def forward(self, input): class Dropout2d(_DropoutNd): - r"""Randomly zeroes whole channels of the input tensor. - The channels to zero-out are randomized on every forward call. + r"""Randomly zero out entire channels (a channel is a 2D feature map, + e.g., the :math:`j`-th channel of the :math:`i`-th sample in the + batched input is a 2D tensor :math:`\text{input}[i, j]`) of the input tensor). + Each channel will be zeroed out independently on every forward call. + with probability :attr:`p` using samples from a Bernoulli distribution. Usually the input comes from :class:`nn.Conv2d` modules. @@ -93,8 +97,11 @@ def forward(self, input): class Dropout3d(_DropoutNd): - r"""Randomly zeroes whole channels of the input tensor. - The channels to zero are randomized on every forward call. + r"""Randomly zero out entire channels (a channel is a 3D feature map, + e.g., the :math:`j`-th channel of the :math:`i`-th sample in the + batched input is a 3D tensor :math:`\text{input}[i, j]`) of the input tensor). + Each channel will be zeroed out independently on every forward call. + with probability :attr:`p` using samples from a Bernoulli distribution. Usually the input comes from :class:`nn.Conv3d` modules. diff --git a/torch/nn/parallel/__init__.py b/torch/nn/parallel/__init__.py index 9b9eb1e3fd590..00d942abf35e0 100644 --- a/torch/nn/parallel/__init__.py +++ b/torch/nn/parallel/__init__.py @@ -5,6 +5,7 @@ from .distributed import DistributedDataParallel from .distributed_cpu import DistributedDataParallelCPU from .distributed_c10d import _DistributedDataParallelC10d +from .distributed_c10d_cpu import _DistributedDataParallelC10dCPU __all__ = ['replicate', 'scatter', 'parallel_apply', 'gather', 'data_parallel', 'DataParallel', 'DistributedDataParallel', 'DistributedDataParallelCPU'] diff --git a/torch/nn/parallel/distributed_c10d_cpu.py b/torch/nn/parallel/distributed_c10d_cpu.py new file mode 100644 index 0000000000000..d7cdeb45cae7f --- /dev/null +++ b/torch/nn/parallel/distributed_c10d_cpu.py @@ -0,0 +1,106 @@ +import torch +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +import torch.distributed.c10d as dist +from torch.nn.modules import Module +from collections import defaultdict +from torch.autograd import Variable + + +class _DistributedDataParallelC10dCPU(Module): + r"""Implements distributed data parallelism for CPU at the module level + and it will use PyTorch's new distributed package: c10d. + + This module support the ``mpi``, ``gloo``, backends. + + This container parallelizes the application of the given module by + splitting the input across the specified devices by chunking in the batch + dimension. The module is replicated on each machine, and each such replica + handles a portion of the input. During the backwards pass, gradients from + each node are averaged. + + This module could be used in conjunction with the DistributedSampler, + (see :class `torch.utils.data.distributed.DistributedSampler`) + which will load a subset of the original datset for each node with the same + batch size. So strong scaling should be configured like this: + n = 1, batch size = 128 + n = 2, batch size = 64 + n = 4, batch size = 32 + n = 8, batch size = 16 + + Creation of this class requires the distributed package to be already + initialized in the process group mode + (see :func:`torch.distributed.init_process_group`). + + .. warning:: + Constructor, forward method, and differentiation of the output (or a + function of the output of this module) is a distributed synchronization + point. Take that into account in case different node might be + executing different code. + + .. warning:: + This module assumes all parameters are registered in the model by the + time it is created. No parameters should be added nor removed later. + + .. warning:: + This module assumes all gradients are dense. + + .. warning:: + This module doesn't work with :func:`torch.autograd.grad` (i.e. it will + only work if gradients are to be accumulated in ``.grad`` attributes of + parameters). + + .. note:: + Parameters are broadcast between nodes in the __init__() function. The + module performs an all-reduce step on gradients and assumes that they + will be modified by the optimizer in all nodes in the same way. + + .. warning:: + Forward and backward hooks defined on :attr:`module` and its submodules + won't be invoked anymore, unless the hooks are initialized in the + :meth:`forward` method. + + Args: + module: module to be parallelized + + Example:: + + >>> torch.distributed.init_process_group(world_size=4, init_method='...') + >>> net = torch.nn._DistributedDataParallelC10dCPU(model) + """ + + def __init__(self, module): + super(_DistributedDataParallelC10dCPU, self).__init__() + self.module = module + self.sync_parameters() + + def allreduce_params(): + if self.needs_reduction: + self.needs_reduction = False + buckets = defaultdict(list) + for param in self.module.parameters(): + if param.requires_grad and param.grad is not None: + tp = type(param.data) + buckets[tp].append(param) + + for bucket in buckets.values(): + grads = [param.grad.data for param in bucket] + coalesced = _flatten_dense_tensors(grads) + dist.all_reduce(coalesced) + coalesced /= dist.get_world_size() + for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)): + buf.copy_(synced) + + for param in list(self.module.parameters()): + def allreduce_hook(*unused): + Variable._execution_engine.queue_callback(allreduce_params) + + if param.requires_grad: + param.register_hook(allreduce_hook) + + def sync_parameters(self): + for param in self.module.parameters(): + dist.broadcast(param.data, 0) + + def forward(self, *inputs, **kwargs): + self.needs_reduction = True + return self.module(*inputs, **kwargs) diff --git a/torch/nn/utils/rnn.py b/torch/nn/utils/rnn.py index bf8a6010898db..b61fdabd3e828 100644 --- a/torch/nn/utils/rnn.py +++ b/torch/nn/utils/rnn.py @@ -1,11 +1,9 @@ from collections import namedtuple +import warnings import torch -import torch.onnx -from .._functions.packing import PackPadded - PackedSequence_ = namedtuple('PackedSequence', ['data', 'batch_sizes']) @@ -140,45 +138,14 @@ def pack_padded_sequence(input, lengths, batch_first=False): Returns: a :class:`PackedSequence` object """ - if isinstance(lengths, list): - lengths = torch.LongTensor(lengths) - - data, batch_sizes = PackPadded.apply(input, lengths, batch_first) - - return PackedSequence(data, batch_sizes) - - -def _symbolic_pack_padded_sequence(g, input, lengths, batch_first=False, padding_value=0.0): - # There currently is no PackPadded operator in ONNX. We rely on an - # optimization pass to remove this later. It is an error if all - # PackPadded operators cannot be optimized out. - - if not isinstance(input, torch._C.Value): - raise RuntimeError("PackPadded requires `input` to be a Tensor") - if not isinstance(lengths, torch._C.Value): - raise RuntimeError("PackPadded requires `lengths` to be a Tensor") - - def _onnx_symbolic_pack_padded_sequence(g, input, lengths): - if batch_first: - input = g.op('Transpose', input, perm_i=[1, 0, 2]) - if not lengths.type().isSubtypeOf(torch._C.DynamicType.get()): - raise RuntimeError("Lengths must be a Tensor for ONNX export") - # We know it's a TensorType so this check is now safe. - if lengths.type().scalarType() != 'Int': - raise RuntimeError("ONNX export requires that the lengths passed " - "to pack_padded_sequence must be of type Int") - return g.op("prim::PackPadded", input, lengths, outputs=2) - - def pack_padded_sequence_trace_wrapper(input, lengths): - return pack_padded_sequence(input, lengths, batch_first=batch_first) - - outputs = g.wrapPyFuncWithSymbolic( - pack_padded_sequence_trace_wrapper, [input, lengths], 2, - _onnx_symbolic_pack_padded_sequence) - return tuple(o for o in outputs) - - -pack_padded_sequence = torch.onnx.symbolic_override(_symbolic_pack_padded_sequence)(pack_padded_sequence) + if torch._C._get_tracing_state() and not isinstance(lengths, torch.Tensor): + warnings.warn('pack_padded_sequence has been called with a Python list of ' + 'sequence lengths. The tracer cannot track the data flow of Python ' + 'values, and it will treat them as constants, likely rendering ' + 'the trace incorrect for any other combination of lengths.', + category=torch.jit.TracerWarning, stacklevel=2) + lengths = torch.as_tensor(lengths, dtype=torch.int64) + return PackedSequence(torch._C._VariableFunctions._pack_padded_sequence(input, lengths, batch_first)) def pad_packed_sequence(sequence, batch_first=False, padding_value=0.0, total_length=None): @@ -214,9 +181,7 @@ def pad_packed_sequence(sequence, batch_first=False, padding_value=0.0, total_le containing the list of lengths of each sequence in the batch. """ - var_data, batch_sizes = sequence - max_batch_size = int(batch_sizes[0]) - max_seq_length = batch_sizes.size(0) + max_seq_length = sequence.batch_sizes.size(0) if total_length is not None: if total_length < max_seq_length: raise ValueError("Expected total_length to be at least the length " @@ -224,56 +189,8 @@ def pad_packed_sequence(sequence, batch_first=False, padding_value=0.0, total_le "total_length={} and max sequence length being {}" .format(total_length, max_seq_length)) max_seq_length = total_length - output = var_data.data.new(max_seq_length, max_batch_size, *var_data.size()[1:]).fill_(padding_value) - - lengths = [] - data_offset = 0 - prev_batch_size = int(batch_sizes[0]) - prev_i = 0 - for i, batch_size in enumerate(batch_sizes.tolist() + [0]): - if batch_size != prev_batch_size: - l = prev_batch_size * (i - prev_i) - tmp = var_data[data_offset:data_offset + l] - output[prev_i:i, :prev_batch_size] = tmp.view(i - prev_i, prev_batch_size, *tmp.size()[1:]) - data_offset += l - prev_i = i - dec = prev_batch_size - batch_size - if dec > 0: - lengths.extend((i,) * dec) - prev_batch_size = batch_size - - lengths.reverse() - - if batch_first: - output = output.transpose(0, 1) - # This Tensor doesn't actually have any history (well, - # technically it does; it's just untracked), it is purely here to - # make ONNX export easier. That is to say, from an autodiff - # standpoint this doesn't make any sense. - return output, torch.LongTensor(lengths) - - -def _symbolic_pad_packed_sequence(g, input, batch_first=False, padding_value=0.0, total_length=None): - # Ignore total_length as it is not supported in _symbolic_pad_packed_sequence - # It is only useful/used when training using data_parallel model, so - # It shouldn't be relevant for ONNX anyway - def _onnx_symbolic_pad_packed_sequence(g, data, batch_sizes): - data, lengths = g.op("prim::PadPacked", data, batch_sizes, outputs=2) - if batch_first: - data = g.op('Transpose', data, perm_i=[1, 0, 2]) - return data, lengths - - def pad_packed_sequence_trace_wrapper(data, batch_sizes): - return pad_packed_sequence(PackedSequence(data, batch_sizes), - batch_first=batch_first, padding_value=padding_value) - - data, lengths = g.wrapPyFuncWithSymbolic( - pad_packed_sequence_trace_wrapper, [input.data, input.batch_sizes], 2, - _onnx_symbolic_pad_packed_sequence) - return data, lengths - - -pad_packed_sequence = torch.onnx.symbolic_override(_symbolic_pad_packed_sequence)(pad_packed_sequence) + return torch._C._VariableFunctions._pad_packed_sequence( + sequence.data, sequence.batch_sizes, batch_first, padding_value, max_seq_length) def pad_sequence(sequences, batch_first=False, padding_value=0): diff --git a/torch/onnx/symbolic.py b/torch/onnx/symbolic.py index e214ba26b5a04..c1a4fb7267255 100644 --- a/torch/onnx/symbolic.py +++ b/torch/onnx/symbolic.py @@ -687,6 +687,16 @@ def batch_norm(g, input, weight, bias, running_mean, running_var, training, mome # batchnorm1d accepts 2d and 3d array, but ONNX only accepts 3d input = g.op("Unsqueeze", input, axes_i=[2]) + if weight is None or weight.node().kind() == "prim::Undefined": + assert len(input_sizes) > 1 + weight_value = torch.tensor([1.] * input_sizes[1]).type( + 'torch.' + input.type().scalarType() + 'Tensor') + weight = g.op("Constant", value_t=weight_value) + if bias is None or bias.node().kind() == "prim::Undefined": + assert len(input_sizes) > 1 + bias_value = torch.tensor([0.] * input_sizes[1]).type( + 'torch.' + input.type().scalarType() + 'Tensor') + bias = g.op("Constant", value_t=bias_value) out = g.op("BatchNormalization", input, weight, bias, running_mean, running_var, epsilon_f=eps, momentum_f=1 - momentum, @@ -1163,3 +1173,31 @@ def _dim_arange(g, like, dim): def detach(g, input): # Erase aten::detach nodes because ONNX is inference only return input + + +@parse_args('v', 'v', 'i') +def _pack_padded_sequence(g, input, lengths, batch_first): + # There currently is no PackPadded operator in ONNX. We rely on an + # optimization pass to remove this later. It is an error if all + # PackPadded operators cannot be optimized out. + if batch_first: + input = g.op('Transpose', input, perm_i=[1, 0, 2]) + if not lengths.type().isSubtypeOf(torch._C.DynamicType.get()): + raise RuntimeError("Lengths must be a Tensor for ONNX export") + # We know it's a TensorType so this check is now safe. + # It's really only necessary beacuse those operators expand to something that + # only works with int32 types in Caffe2... + if lengths.type().scalarType() != 'Int': + lengths = _cast_Int(g, lengths, False) + return g.op("prim::PackPadded", input, lengths, outputs=2) + + +@parse_args('v', 'v', 'i', 't', 'i') +def _pad_packed_sequence(g, data, batch_sizes, batch_first, padding_value, total_length): + # Ignore total_length as it is not supported in _symbolic_pad_packed_sequence + # It is only useful/used when training using data_parallel model, so + # It shouldn't be relevant for ONNX anyway + data, lengths = g.op("prim::PadPacked", data, batch_sizes, outputs=2) + if batch_first: + data = g.op('Transpose', data, perm_i=[1, 0, 2]) + return data, lengths diff --git a/torch/onnx/utils.py b/torch/onnx/utils.py index 31d7cb05c7b84..963e0bc959125 100644 --- a/torch/onnx/utils.py +++ b/torch/onnx/utils.py @@ -141,6 +141,8 @@ def _optimize_graph(graph, operator_export_type): torch._C._jit_pass_peephole(graph) torch._C._jit_pass_lint(graph) + # onnx only supports tensors, but 1 / 2 = 0.5 and tensor(1) / tensor(2) = 0 + torch._C._jit_pass_prepare_division_for_onnx(graph) # onnx only supports tensors, so we turn all out number types into tensors torch._C._jit_pass_erase_number_types(graph) # onnx does not support tuples, so try to remove them diff --git a/torch/serialization.py b/torch/serialization.py index 85c1a83429f26..98e88c962157d 100644 --- a/torch/serialization.py +++ b/torch/serialization.py @@ -451,20 +451,7 @@ def persistent_load(saved_id): storage_views = pickle_module.load(f) for target_cdata, root_cdata, offset, size in storage_views: root = deserialized_objects[root_cdata] - if offset != 0 or size != root.size(): - warnings.warn("Detected storage view in legacy serialized data: " - "storage views are no longer natively supported, so we are making " - "a copy of the data instead. THIS IS A SEMANTIC CHANGE! " - "If you need aliasing, reserialize your model using " - "tensors that share storage.") - - tensor = torch._utils._rebuild_tensor(root, offset, (size,), (1,)) - obj = tensor.clone().storage() - else: - # NB: This line does not appear to be exercised by the - # test suite. - obj = root - deserialized_objects[target_cdata] = obj + deserialized_objects[target_cdata] = root[offset:offset + size] tar.extract('tensors', path=tmpdir) with open(os.path.join(tmpdir, 'tensors'), 'rb', 0) as f: