diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 581ad371f8262..67f0cf8dbbe09 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -13,7 +13,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master @@ -28,7 +28,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master @@ -69,7 +69,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master @@ -84,7 +84,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master @@ -99,7 +99,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Fetch PyTorch uses: actions/checkout@master @@ -174,7 +174,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master @@ -209,7 +209,10 @@ jobs: run: | set -eux git remote add upstream https://github.com/pytorch/pytorch - git fetch upstream "${{ github.base_ref}}" + git fetch upstream "$GITHUB_BASE_REF" + BASE_SHA=${{ github.event.pull_request.base.sha }} + HEAD_SHA=${{ github.event.pull_request.head.sha }} + MERGE_BASE=$(git merge-base $BASE_SHA $HEAD_SHA) if [[ ! -d build ]]; then git submodule update --init --recursive @@ -238,8 +241,9 @@ jobs: # The negative filters below are to exclude files that include onnx_pb.h or # caffe2_pb.h, otherwise we'd have to build protos as part of this CI job. python tools/clang_tidy.py \ + --verbose \ --paths torch/csrc/ \ - --diff "${{ github.event.pull_request.base.sha}}" \ + --diff "$MERGE_BASE" \ -g"-torch/csrc/jit/export.cpp" \ -g"-torch/csrc/jit/import.cpp" \ -g"-torch/csrc/jit/netdef_converter.cpp" \ @@ -252,6 +256,6 @@ jobs: check_name: 'clang-tidy' linter_output_path: 'clang-tidy-output.txt' commit_sha: ${{ steps.get_pr_tip.outputs.commit_sha }} - regex: '^(?.*?):(?\d+):(?\d+): (?.*?) (?\[.*\])' + regex: '^(?.*?):(?\d+):(?\d+): (?.*?) \[(?.*)\]' env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index 6741d7d6bf4fa..0b41c519587af 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -110,7 +110,7 @@ test_python_nn() { } test_python_all_except_nn() { - time python test/run_test.py --exclude nn --verbose --bring-to-front quantization quantized quantized_tensor quantized_nn_mods quantizer + time python test/run_test.py --exclude nn --verbose --bring-to-front quantization quantized quantized_tensor quantized_nn_mods assert_git_not_dirty } diff --git a/android/build.gradle b/android/build.gradle index b29299d7a4a3b..70f33c4e0769c 100644 --- a/android/build.gradle +++ b/android/build.gradle @@ -1,33 +1,37 @@ -buildscript { - ext { - minSdkVersion = 21 - targetSdkVersion = 28 - compileSdkVersion = 28 - buildToolsVersion = '28.0.3' - - coreVersion = "1.2.0" - extJUnitVersion = "1.1.1" - runnerVersion = "1.2.0" - rulesVersion = "1.2.0" - junitVersion = "4.12" +allprojects { + if (name == "pytorch_host") { + return } - repositories { - google() - mavenLocal() - mavenCentral() - jcenter() - } + buildscript { + ext { + minSdkVersion = 21 + targetSdkVersion = 28 + compileSdkVersion = 28 + buildToolsVersion = '28.0.3' + + coreVersion = "1.2.0" + extJUnitVersion = "1.1.1" + runnerVersion = "1.2.0" + rulesVersion = "1.2.0" + junitVersion = "4.12" + } + + repositories { + google() + mavenLocal() + mavenCentral() + jcenter() + } - dependencies { - classpath 'com.android.tools.build:gradle:3.3.2' - classpath "com.jfrog.bintray.gradle:gradle-bintray-plugin:${GRADLE_BINTRAY_PLUGIN_VERSION}" - classpath "com.github.dcendents:android-maven-gradle-plugin:${ANDROID_MAVEN_GRADLE_PLUGIN_VERSION}" - classpath "org.jfrog.buildinfo:build-info-extractor-gradle:4.9.8" + dependencies { + classpath 'com.android.tools.build:gradle:3.3.2' + classpath "com.jfrog.bintray.gradle:gradle-bintray-plugin:${GRADLE_BINTRAY_PLUGIN_VERSION}" + classpath "com.github.dcendents:android-maven-gradle-plugin:${ANDROID_MAVEN_GRADLE_PLUGIN_VERSION}" + classpath "org.jfrog.buildinfo:build-info-extractor-gradle:4.9.8" + } } -} -allprojects { repositories { google() jcenter() diff --git a/android/pytorch_android/CMakeLists.txt b/android/pytorch_android/CMakeLists.txt index 5ebae28741906..69b3f97aa2997 100644 --- a/android/pytorch_android/CMakeLists.txt +++ b/android/pytorch_android/CMakeLists.txt @@ -4,7 +4,19 @@ set(CMAKE_CXX_STANDARD 11) set(CMAKE_VERBOSE_MAKEFILE ON) set(pytorch_android_DIR ${CMAKE_CURRENT_LIST_DIR}/src/main/cpp) -set(libtorch_include_DIR ${pytorch_android_DIR}/libtorch_include/${ANDROID_ABI}) + +if (ANDROID_ABI) + set(libtorch_include_DIR ${pytorch_android_DIR}/libtorch_include/${ANDROID_ABI}) + set(BUILD_SUBDIR ${ANDROID_ABI}) +else() + if (NOT LIBTORCH_HOME) + message(FATAL_ERROR + "pytorch_android requires LIBTORCH_HOME to be defined for non-Android builds.") + endif() + set(libtorch_include_DIR ${LIBTORCH_HOME}/include) + link_directories(${LIBTORCH_HOME}/lib) + set(BUILD_SUBDIR host) +endif() message(STATUS "libtorch dir:${libtorch_DIR}") @@ -24,40 +36,58 @@ target_include_directories(pytorch PUBLIC ${libtorch_include_DIR} ) -set(BUILD_DIR ${CMAKE_SOURCE_DIR}/build) -file(MAKE_DIRECTORY ${BUILD_DIR}) - set(fbjni_DIR ${CMAKE_CURRENT_LIST_DIR}/../libs/fbjni/) -set(fbjni_BUILD_DIR ${BUILD_DIR}/fbjni/${ANDROID_ABI}) +set(fbjni_BUILD_DIR ${CMAKE_BINARY_DIR}/fbjni/${BUILD_SUBDIR}) add_subdirectory(${fbjni_DIR} ${fbjni_BUILD_DIR}) -function(import_static_lib name) - add_library(${name} STATIC IMPORTED) - set_property( - TARGET ${name} - PROPERTY IMPORTED_LOCATION - ${CMAKE_CURRENT_LIST_DIR}/src/main/jniLibs/${ANDROID_ABI}/${name}.a) -endfunction(import_static_lib) - -import_static_lib(libtorch) -import_static_lib(libc10) -import_static_lib(libnnpack) -import_static_lib(libpytorch_qnnpack) -import_static_lib(libeigen_blas) -import_static_lib(libcpuinfo) -import_static_lib(libclog) - -target_link_libraries(pytorch - fbjni - -Wl,--gc-sections - -Wl,--whole-archive - libtorch - -Wl,--no-whole-archive - libc10 - libnnpack - libpytorch_qnnpack - libeigen_blas - libcpuinfo - libclog -) +if (ANDROID_ABI) + + function(import_static_lib name) + add_library(${name} STATIC IMPORTED) + set_property( + TARGET ${name} + PROPERTY IMPORTED_LOCATION + ${CMAKE_CURRENT_LIST_DIR}/src/main/jniLibs/${ANDROID_ABI}/${name}.a) + endfunction(import_static_lib) + + import_static_lib(libtorch) + import_static_lib(libc10) + import_static_lib(libnnpack) + import_static_lib(libpytorch_qnnpack) + import_static_lib(libeigen_blas) + import_static_lib(libcpuinfo) + import_static_lib(libclog) + + # Link most things statically on Android. + target_link_libraries(pytorch + fbjni + -Wl,--gc-sections + -Wl,--whole-archive + libtorch + -Wl,--no-whole-archive + libc10 + libnnpack + libpytorch_qnnpack + libeigen_blas + libcpuinfo + libclog + ) + +else() + + # Prefer dynamic linking on the host + target_link_libraries(pytorch + fbjni + -Wl,--gc-sections + -Wl,--whole-archive + torch + -Wl,--no-whole-archive + c10 + nnpack + pytorch_qnnpack + cpuinfo + clog + ) + +endif() diff --git a/android/pytorch_android/host/build.gradle b/android/pytorch_android/host/build.gradle new file mode 100644 index 0000000000000..7a096ad60d260 --- /dev/null +++ b/android/pytorch_android/host/build.gradle @@ -0,0 +1,33 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// +// This source code is licensed under the Apache-2 license found in the +// LICENSE file in the root directory of this source tree. + +plugins { + id 'java-library' +} + +repositories { + mavenLocal() + jcenter() +} + +sourceSets { + main { + java.srcDir '../src/main/java' + } + test { + java { + srcDir '../src/androidTest/java' + exclude '**/PytorchInstrumented*' + } + resources.srcDirs = ["../src/androidTest/assets"] + } +} + +dependencies { + compileOnly 'com.google.code.findbugs:jsr305:3.0.1' + implementation 'com.facebook.soloader:nativeloader:0.8.0' + implementation 'com.facebook.fbjni:fbjni:0.0.3-SNAPSHOT' + testImplementation 'junit:junit:4.12' +} diff --git a/android/settings.gradle b/android/settings.gradle index 277b15c0afc99..99e442b2ae2f5 100644 --- a/android/settings.gradle +++ b/android/settings.gradle @@ -1,4 +1,6 @@ -include ':app', ':pytorch_android', ':fbjni', ':pytorch_android_torchvision' +include ':app', ':pytorch_android', ':fbjni', ':pytorch_android_torchvision', ':pytorch_host' project(':fbjni').projectDir = file('libs/fbjni_local') project(':pytorch_android_torchvision').projectDir = file('pytorch_android_torchvision') + +project(':pytorch_host').projectDir = file('pytorch_android/host') diff --git a/aten/src/ATen/core/jit_type.h b/aten/src/ATen/core/jit_type.h index ac14ff62f02a0..9947d6c0b178e 100644 --- a/aten/src/ATen/core/jit_type.h +++ b/aten/src/ATen/core/jit_type.h @@ -556,6 +556,12 @@ struct CAFFE2_API TensorType : public Type { return r; } + TensorTypePtr withPossiblyUndefined() { + auto r = clone(); + r->undefined_ = c10::nullopt; + return r; + } + c10::optional undefined() const { return undefined_; } static TensorTypePtr get(); @@ -563,29 +569,37 @@ struct CAFFE2_API TensorType : public Type { static const TypeKind Kind = TypeKind::TensorType; private: - TensorType(const at::Tensor &tensor) - : Type(TypeKind::TensorType), scalar_type_(tensor.scalar_type()), - device_(tensor.device()), sizes_(tensor.sizes().size()), - strides_(tensor.sizes().size()), - requires_grad_(tensor.requires_grad()), undefined_(false) { - if (!tensor.is_mkldnn() && !tensor.is_sparse()) { - sizes_ = tensor.sizes().vec(); - strides_ = tensor.strides().vec(); - } - } - TensorType(c10::optional scalar_type, - c10::optional device, const VaryingShape &sizes, - const VaryingStrides &strides, - c10::optional requires_grad, - c10::optional undefined = false) - : Type(TypeKind::TensorType), scalar_type_(scalar_type), - device_(device), sizes_(sizes), strides_(strides), - requires_grad_(requires_grad), undefined_(undefined) {} - - TensorTypePtr clone() const { - return TensorTypePtr(new TensorType(scalar_type_, device_, sizes_, - strides_, requires_grad_, - undefined_)); + TensorType(const at::Tensor& tensor) + : Type(TypeKind::TensorType), + scalar_type_(tensor.scalar_type()), + device_(tensor.device()), + sizes_(tensor.sizes().size()), + strides_(tensor.sizes().size()), + requires_grad_(tensor.requires_grad()), + undefined_(!tensor.defined()) { + if (!tensor.is_mkldnn() && !tensor.is_sparse()) { + sizes_ = tensor.sizes().vec(); + strides_ = tensor.strides().vec(); + } + } + TensorType( + c10::optional scalar_type, + c10::optional device, + const VaryingShape& sizes, + const VaryingStrides& strides, + c10::optional requires_grad, + c10::optional undefined = false) + : Type(TypeKind::TensorType), + scalar_type_(scalar_type), + device_(device), + sizes_(sizes), + strides_(strides), + requires_grad_(requires_grad), + undefined_(undefined) {} + + TensorTypePtr clone() const { + return TensorTypePtr(new TensorType( + scalar_type_, device_, sizes_, strides_, requires_grad_, undefined_)); } static std::vector contiguousStridesOf(at::IntArrayRef sizes) { diff --git a/aten/src/ATen/cpu/vec256/vec256_base.h b/aten/src/ATen/cpu/vec256/vec256_base.h index 64063b3da2b92..01481783e7048 100644 --- a/aten/src/ATen/cpu/vec256/vec256_base.h +++ b/aten/src/ATen/cpu/vec256/vec256_base.h @@ -256,9 +256,21 @@ struct Vec256 { Vec256 log1p() const { return map(std::log1p); } + template ::value, int>::type = 0> Vec256 log2() const { + // other_t_log2 is for SFINAE and clarity. Make sure it is not changed. + static_assert(std::is_same::value, "other_t_log2 must be T"); return map(std::log2); } + template ::value, int>::type = 0> + Vec256 log2() const { + // complex_t_log2 is for SFINAE and clarity. Make sure it is not changed. + static_assert(std::is_same::value, "complex_t_log2 must be T"); + const T log_2 = T(std::log(2.0)); + return Vec256(map(std::log))/Vec256(log_2); + } Vec256 ceil() const { return map(at::native::ceil_impl); } diff --git a/aten/src/ATen/cpu/vec256/vec256_complex_double.h b/aten/src/ATen/cpu/vec256/vec256_complex_double.h index db56dfd408480..4ac48edac0575 100644 --- a/aten/src/ATen/cpu/vec256/vec256_complex_double.h +++ b/aten/src/ATen/cpu/vec256/vec256_complex_double.h @@ -13,7 +13,7 @@ namespace { #if defined(__AVX__) && !defined(_MSC_VER) -template <> class Vec256> { +template <> class Vec256> { private: __m256d values; public: @@ -135,26 +135,58 @@ template <> class Vec256> { const __m256d imag_mask = _mm256_castsi256_pd(_mm256_setr_epi64x(0x0000000000000000, 0xFFFFFFFFFFFFFFFF, 0x0000000000000000, 0xFFFFFFFFFFFFFFFF)); return _mm256_and_pd(values, imag_mask); - } + } Vec256> imag() const { return _mm256_permute_pd(imag_(), 0x05); //b a } __m256d conj_() const { - const __m256d conj_mask = _mm256_setr_pd(1.0, -1.0, 1.0, -1.0); - return _mm256_mul_pd(values, conj_mask); //a -b + const __m256d sign_mask = _mm256_setr_pd(0.0, -0.0, 0.0, -0.0); + return _mm256_xor_pd(values, sign_mask); // a -b } Vec256> conj() const { return conj_(); } - Vec256> acos() const { - return map(std::acos); + Vec256> log() const { + // Most trigonomic ops use the log() op to improve complex number performance. + return map(std::log); + } + Vec256> log2() const { + const __m256d log2_ = _mm256_set1_pd(std::log(2)); + return _mm256_div_pd(log(), log2_); + } + Vec256> log10() const { + const __m256d log10_ = _mm256_set1_pd(std::log(10)); + return _mm256_div_pd(log(), log10_); + } + Vec256> log1p() const { + AT_ERROR("not supported for complex numbers"); } Vec256> asin() const { - return map(std::asin); + // asin(x) + // = -i*ln(iz + sqrt(1 -z^2)) + // = -i*ln((ai - b) + sqrt(1 - (a + bi)*(a + bi))) + // = -i*ln((-b + ai) + sqrt(1 - (a**2 - b**2) - 2*abi)) + const __m256d one = _mm256_set1_pd(1); + + auto conj = conj_(); + auto b_a = _mm256_permute_pd(conj, 0x05); //-b a + auto ab = _mm256_mul_pd(conj, b_a); //-ab -ab + auto im = _mm256_add_pd(ab, ab); //-2ab -2ab + + auto val_2 = _mm256_mul_pd(values, values); // a*a b*b + auto re = _mm256_hsub_pd(val_2, _mm256_permute_pd(val_2, 0x05)); // a*a-b*b b*b-a*a + re = _mm256_sub_pd(one, re); + + auto root = Vec256(_mm256_blend_pd(re, im, 0x0A)).sqrt(); //sqrt(re + i*im) + auto ln = Vec256(_mm256_add_pd(b_a, root)).log(); //ln(iz + sqrt()) + return Vec256(_mm256_permute_pd(ln.values, 0x05)).conj(); //-i*ln() } - Vec256> atan() const { - return map(std::atan); + Vec256> acos() const { + // acos(x) = pi/2 - asin(x) + const __m256d pi_2 = _mm256_setr_pd(M_PI/2, 0.0, M_PI/2, 0.0); + return _mm256_sub_pd(pi_2, asin()); } + Vec256> atan() const; Vec256> atan2(const Vec256> &b) const { AT_ERROR("not supported for complex numbers"); } @@ -170,18 +202,6 @@ template <> class Vec256> { Vec256> expm1() const { AT_ERROR("not supported for complex numbers"); } - Vec256> log() const { - return map(std::log); - } - Vec256> log2() const { - AT_ERROR("not supported for complex numbers"); - } - Vec256> log10() const { - return map(std::log10); - } - Vec256> log1p() const { - AT_ERROR("not supported for complex numbers"); - } Vec256> sin() const { return map(std::sin); } @@ -217,14 +237,32 @@ template <> class Vec256> { return _mm256_round_pd(values, (_MM_FROUND_TO_ZERO | _MM_FROUND_NO_EXC)); } Vec256> sqrt() const { - return map(std::sqrt); + // sqrt(a + bi) + // = sqrt(2)/2 * [sqrt(sqrt(a**2 + b**2) + a) + sgn(b)*sqrt(sqrt(a**2 + b**2) - a)i] + // = sqrt(2)/2 * [sqrt(abs() + a) + sgn(b)*sqrt(abs() - a)i] + + const __m256d scalar = _mm256_set1_pd(std::sqrt(2)/2); //sqrt(2)/2 sqrt(2)/2 + const __m256d sign_mask = _mm256_setr_pd(0.0, -0.0, 0.0, -0.0); + auto sign = _mm256_and_pd(values, sign_mask); + auto factor = _mm256_or_pd(scalar, sign); + + auto a_a = _mm256_xor_pd(_mm256_movedup_pd(values), sign_mask); // a -a + auto res_re_im = _mm256_sqrt_pd(_mm256_add_pd(abs_(), a_a)); // sqrt(abs + a) sqrt(abs - a) + return _mm256_mul_pd(factor, res_re_im); } Vec256> reciprocal() const; Vec256> rsqrt() const { - return map([](const std::complex &x) { return (std::complex)(1)/std::sqrt(x); }); + return sqrt().reciprocal(); } Vec256> pow(const Vec256> &exp) const { - AT_ERROR("not supported for complex numbers"); + __at_align32__ std::complex x_tmp[size()]; + __at_align32__ std::complex y_tmp[size()]; + store(x_tmp); + exp.store(y_tmp); + for (int i = 0; i < size(); i++) { + x_tmp[i] = std::pow(x_tmp[i], y_tmp[i]); + } + return loadu(x_tmp); } // Comparison using the _CMP_**_OQ predicate. // `O`: get false if an operand is NaN @@ -259,11 +297,11 @@ template <> Vec256> inline operator-(const Vec256 Vec256> inline operator*(const Vec256> &a, const Vec256> &b) { //(a + bi) * (c + di) = (ac - bd) + (ad + bc)i - const __m256d neg = _mm256_setr_pd(1.0, -1.0, 1.0, -1.0); + const __m256d sign_mask = _mm256_setr_pd(0.0, -0.0, 0.0, -0.0); auto ac_bd = _mm256_mul_pd(a, b); //ac bd auto d_c = _mm256_permute_pd(b, 0x05); //d c - d_c = _mm256_mul_pd(neg, d_c); //d -c + d_c = _mm256_xor_pd(sign_mask, d_c); //d -c auto ad_bc = _mm256_mul_pd(a, d_c); //ad -bc auto ret = _mm256_hsub_pd(ac_bd, ad_bc); //ac - bd ad + bc @@ -274,11 +312,11 @@ template <> Vec256> inline operator/(const Vec256> Vec256>::reciprocal() const{ //re + im*i = (a + bi) / (c + di) //re = (ac + bd)/abs_2() = c/abs_2() //im = (bc - ad)/abs_2() = d/abs_2() - const __m256d neg = _mm256_setr_pd(1.0, -1.0, 1.0, -1.0); - auto c_d = _mm256_mul_pd(neg, values); //c -d + const __m256d sign_mask = _mm256_setr_pd(0.0, -0.0, 0.0, -0.0); + auto c_d = _mm256_xor_pd(sign_mask, values); //c -d return _mm256_div_pd(c_d, abs_2_()); } +Vec256> Vec256>::atan() const { + // atan(x) = i/2 * ln((i + z)/(i - z)) + const __m256d i = _mm256_setr_pd(0.0, 1.0, 0.0, 1.0); + const Vec256 i_half = _mm256_setr_pd(0.0, 0.5, 0.0, 0.5); + + auto sum = Vec256(_mm256_add_pd(i, values)); // a 1+b + auto sub = Vec256(_mm256_sub_pd(i, values)); // -a 1-b + auto ln = (sum/sub).log(); // ln((i + z)/(i - z)) + return i_half*ln; // i/2*ln() +} + template <> Vec256> inline maximum(const Vec256>& a, const Vec256>& b) { auto abs_a = a.abs_2_(); diff --git a/aten/src/ATen/cpu/vec256/vec256_complex_float.h b/aten/src/ATen/cpu/vec256/vec256_complex_float.h index 2e388914f2b93..2149e79099ecd 100644 --- a/aten/src/ATen/cpu/vec256/vec256_complex_float.h +++ b/aten/src/ATen/cpu/vec256/vec256_complex_float.h @@ -171,26 +171,58 @@ template <> class Vec256> { const __m256 imag_mask = _mm256_castsi256_ps(_mm256_setr_epi32(0x00000000, 0xFFFFFFFF, 0x00000000, 0xFFFFFFFF, 0x00000000, 0xFFFFFFFF, 0x00000000, 0xFFFFFFFF)); return _mm256_and_ps(values, imag_mask); - } + } Vec256> imag() const { return _mm256_permute_ps(imag_(), 0x55); //b a } __m256 conj_() const { - const __m256 conj_mask = _mm256_setr_ps(1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0); - return _mm256_mul_ps(values, conj_mask); //a -b + const __m256 sign_mask = _mm256_setr_ps(0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0); + return _mm256_xor_ps(values, sign_mask); // a -b } Vec256> conj() const { return conj_(); } - Vec256> acos() const { - return map(std::acos); + Vec256> log() const { + // Most trigonomic ops use the log() op to improve complex number performance. + return map(std::log); + } + Vec256> log2() const { + const __m256 log2_ = _mm256_set1_ps(std::log(2)); + return _mm256_div_ps(log(), log2_); + } + Vec256> log10() const { + const __m256 log10_ = _mm256_set1_ps(std::log(10)); + return _mm256_div_ps(log(), log10_); + } + Vec256> log1p() const { + AT_ERROR("not supported for complex numbers"); } Vec256> asin() const { - return map(std::asin); + // asin(x) + // = -i*ln(iz + sqrt(1 -z^2)) + // = -i*ln((ai - b) + sqrt(1 - (a + bi)*(a + bi))) + // = -i*ln((-b + ai) + sqrt(1 - (a**2 - b**2) - 2*abi)) + const __m256 one = _mm256_set1_ps(1); + + auto conj = conj_(); + auto b_a = _mm256_permute_ps(conj, 0x55); //-b a + auto ab = _mm256_mul_ps(conj, b_a); //-ab -ab + auto im = _mm256_add_ps(ab, ab); //-2ab -2ab + + auto val_2 = _mm256_mul_ps(values, values); // a*a b*b + auto re = _mm256_hsub_ps(val_2, _mm256_permute_ps(val_2, 0x55)); // a*a-b*b b*b-a*a + re = _mm256_sub_ps(one, re); + + auto root = Vec256(_mm256_blend_ps(re, im, 0xAA)).sqrt(); //sqrt(re + i*im) + auto ln = Vec256(_mm256_add_ps(b_a, root)).log(); //ln(iz + sqrt()) + return Vec256(_mm256_permute_ps(ln.values, 0x55)).conj(); //-i*ln() } - Vec256> atan() const { - return map(std::atan); + Vec256> acos() const { + // acos(x) = pi/2 - asin(x) + const __m256 pi_2 = _mm256_setr_ps(M_PI/2, 0.0, M_PI/2, 0.0, M_PI/2, 0.0, M_PI/2, 0.0); + return _mm256_sub_ps(pi_2, asin()); } + Vec256> atan() const; Vec256> atan2(const Vec256> &b) const { AT_ERROR("not supported for complex numbers"); } @@ -206,18 +238,6 @@ template <> class Vec256> { Vec256> expm1() const { AT_ERROR("not supported for complex numbers"); } - Vec256> log() const { - return map(std::log); - } - Vec256> log2() const { - AT_ERROR("not supported for complex numbers"); - } - Vec256> log10() const { - return map(std::log10); - } - Vec256> log1p() const { - AT_ERROR("not supported for complex numbers"); - } Vec256> sin() const { return map(std::sin); } @@ -253,14 +273,32 @@ template <> class Vec256> { return _mm256_round_ps(values, (_MM_FROUND_TO_ZERO | _MM_FROUND_NO_EXC)); } Vec256> sqrt() const { - return map(std::sqrt); + // sqrt(a + bi) + // = sqrt(2)/2 * [sqrt(sqrt(a**2 + b**2) + a) + sgn(b)*sqrt(sqrt(a**2 + b**2) - a)i] + // = sqrt(2)/2 * [sqrt(abs() + a) + sgn(b)*sqrt(abs() - a)i] + + const __m256 scalar = _mm256_set1_ps(std::sqrt(2)/2); //sqrt(2)/2 sqrt(2)/2 + const __m256 sign_mask = _mm256_setr_ps(0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0); + auto sign = _mm256_and_ps(values, sign_mask); + auto factor = _mm256_or_ps(scalar, sign); + + auto a_a = _mm256_xor_ps(_mm256_moveldup_ps(values), sign_mask); // a -a + auto res_re_im = _mm256_sqrt_ps(_mm256_add_ps(abs_(), a_a)); // sqrt(abs + a) sqrt(abs - a) + return _mm256_mul_ps(factor, res_re_im); } Vec256> reciprocal() const; Vec256> rsqrt() const { - return map([](const std::complex &x) { return (std::complex)(1)/std::sqrt(x); }); + return sqrt().reciprocal(); } Vec256> pow(const Vec256> &exp) const { - AT_ERROR("not supported for complex numbers"); + __at_align32__ std::complex x_tmp[size()]; + __at_align32__ std::complex y_tmp[size()]; + store(x_tmp); + exp.store(y_tmp); + for (int i = 0; i < size(); i++) { + x_tmp[i] = std::pow(x_tmp[i], y_tmp[i]); + } + return loadu(x_tmp); } // Comparison using the _CMP_**_OQ predicate. // `O`: get false if an operand is NaN @@ -295,11 +333,11 @@ template <> Vec256> inline operator-(const Vec256 Vec256> inline operator*(const Vec256> &a, const Vec256> &b) { //(a + bi) * (c + di) = (ac - bd) + (ad + bc)i - const __m256 neg = _mm256_setr_ps(1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0); + const __m256 sign_mask = _mm256_setr_ps(0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0); auto ac_bd = _mm256_mul_ps(a, b); //ac bd auto d_c = _mm256_permute_ps(b, 0x55); //d c - d_c = _mm256_mul_ps(neg, d_c); //d -c + d_c = _mm256_xor_ps(sign_mask, d_c); //d -c auto ad_bc = _mm256_mul_ps(a, d_c); //ad -bc auto ret = _mm256_hsub_ps(ac_bd, ad_bc); //ac - bd ad + bc @@ -310,11 +348,11 @@ template <> Vec256> inline operator/(const Vec256> Vec256>::reciprocal() const { //re + im*i = (a + bi) / (c + di) //re = (ac + bd)/abs_2() = c/abs_2() //im = (bc - ad)/abs_2() = d/abs_2() - const __m256 neg = _mm256_setr_ps(1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0); - auto c_d = _mm256_mul_ps(neg, values); //c -d + const __m256 sign_mask = _mm256_setr_ps(0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0); + auto c_d = _mm256_xor_ps(sign_mask, values); //c -d return _mm256_div_ps(c_d, abs_2_()); } +Vec256> Vec256>::atan() const { + // atan(x) = i/2 * ln((i + z)/(i - z)) + const __m256 i = _mm256_setr_ps(0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0); + const Vec256 i_half = _mm256_setr_ps(0.0, 0.5, 0.0, 0.5, 0.0, 0.5, 0.0, 0.5); + + auto sum = Vec256(_mm256_add_ps(i, values)); // a 1+b + auto sub = Vec256(_mm256_sub_ps(i, values)); // -a 1-b + auto ln = (sum/sub).log(); // ln((i + z)/(i - z)) + return i_half*ln; // i/2*ln() +} + template <> Vec256> inline maximum(const Vec256>& a, const Vec256>& b) { auto abs_a = a.abs_2_(); diff --git a/aten/src/ATen/native/BatchLinearAlgebra.cpp b/aten/src/ATen/native/BatchLinearAlgebra.cpp index ed18cc3790b6e..8096848855c63 100644 --- a/aten/src/ATen/native/BatchLinearAlgebra.cpp +++ b/aten/src/ATen/native/BatchLinearAlgebra.cpp @@ -624,151 +624,6 @@ std::tuple _lu_with_info_cpu(const Tensor& self, bool pi return std::make_tuple(self_working_copy, pivots_tensor, infos_tensor); } -// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triu/tril ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -template -static void apply_triu_tril_single( - scalar_t* result, scalar_t* self, bool inplace, - int64_t k, int64_t n, int64_t m, - int64_t res_row_stride, int64_t res_col_stride, - int64_t self_row_stride, int64_t self_col_stride) { - - constexpr int64_t zero = 0; - - if (upper) { - at::parallel_for(0, n, 0, [&](int64_t start, int64_t end) { - for (auto i = start; i < end; i++) { - for (int64_t j = 0; j < std::min(m, i + k); j++) { - result[i * res_row_stride + j * res_col_stride] = 0; - } - if (!inplace) { // copy the rest of the self if not inplace - for (int64_t j = std::max(zero, i + k); j < m; j++) { - result[i * res_row_stride + j * res_col_stride] = self[i * self_row_stride + j * self_col_stride]; - } - } - } - }); - } else { - at::parallel_for(0, n, 0, [&](int64_t start, int64_t end) { - for (auto i = start; i < end; i++) { - for (int64_t j = std::max(zero, i + k + 1); j < m; j++) { - result[i * res_row_stride + j * res_col_stride] = 0; - } - if (!inplace) { // copy the rest of the self if not inplace - for (int64_t j = zero; j < std::min(m, i + k + 1); j++) { - result[i * res_row_stride + j * res_col_stride] = self[i * self_row_stride + j * self_col_stride]; - } - } - } - }); - } -} - -template -void apply_triu_tril(Tensor& result, const Tensor& self, bool inplace, int64_t k) { - auto n = self.size(-2); - auto m = self.size(-1); - auto self_data = self.data_ptr(); - auto self_stride = (self.dim() > 2 && self.stride(-3) > 0) ? self.stride(-3) : 1; - auto batchsize = batchCountTrilTriu(result); - auto self_row_stride = self.stride(-2); - auto self_column_stride = self.stride(-1); - - auto result_data = result.data_ptr(); - int64_t result_stride, result_row_stride, result_column_stride; - if (result_data != self_data) { - result_stride = (result.dim() > 2 && result.stride(-3) > 0) ? result.stride(-3) : 1; - result_row_stride = result.stride(-2); - result_column_stride = result.stride(-1); - } else { - result_stride = self_stride; - result_row_stride = self_row_stride; - result_column_stride = self_column_stride; - } - - at::parallel_for(0, batchsize, 0, [&](int64_t start, int64_t end) { - for (auto b = start; b < end; b++) { - scalar_t* self_batch = &self_data[b * self_stride]; - scalar_t* result_batch = &result_data[b * result_stride]; - apply_triu_tril_single( - result_batch, self_batch, inplace, k, n, m, - result_row_stride, result_column_stride, self_row_stride, self_column_stride); - } - }); -} - -Tensor tril(const Tensor& self, int64_t k) { - Tensor result = at::empty({0}, self.options()); - at::tril_out(result, self, k); - return result; -} - -Tensor& tril_cpu_(Tensor &self, int64_t k) { - if (self.numel() == 0) { - return self; - } - bool inplace; - Tensor self_c; - std::tie(inplace, self_c) = checkTrilTriuBatchContiguous(self, true); - Tensor result = inplace ? self : at::empty_like(self); - AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "tril", [&]{ - apply_triu_tril(result, self_c, inplace, k); - }); - if (!inplace) self.copy_(result); - return self; -} - -Tensor& tril_cpu_out(Tensor &result, const Tensor& self, int64_t k) { - if (result.sizes() != self.sizes()) { - result.resize_as_(self); - } - if (self.numel() == 0) { - return result; - } - Tensor self_c; - std::tie(std::ignore, self_c) = checkTrilTriuBatchContiguous(self, false); - AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "tril", [&]{ - apply_triu_tril(result, self_c, false, k); - }); - return result; -} - -Tensor triu(const Tensor& self, int64_t k) { - Tensor result = at::empty({0}, self.options()); - at::triu_out(result, self, k); - return result; -} - -Tensor& triu_cpu_(Tensor &self, int64_t k) { - if (self.numel() == 0) { - return self; - } - bool inplace; - Tensor self_c; - std::tie(inplace, self_c) = checkTrilTriuBatchContiguous(self, true); - Tensor result = inplace ? self : at::empty_like(self); - AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "triu", [&]{ - apply_triu_tril(result, self_c, inplace, k); - }); - if (!inplace) self.copy_(result); - return self; -} - -Tensor& triu_cpu_out(Tensor &result, const Tensor& self, int64_t k) { - if (result.sizes() != self.sizes()) { - result.resize_as_(self); - } - if (self.numel() == 0) { - return result; - } - Tensor self_c; - std::tie(std::ignore, self_c) = checkTrilTriuBatchContiguous(self, false); - AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "triu", [&]{ - apply_triu_tril(result, self_c, false, k); - }); - return result; -} - // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triangular_solve ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ template diff --git a/aten/src/ATen/native/BinaryOps.cpp b/aten/src/ATen/native/BinaryOps.cpp index 23c15ca45dd36..d3bb70d59e8fc 100644 --- a/aten/src/ATen/native/BinaryOps.cpp +++ b/aten/src/ATen/native/BinaryOps.cpp @@ -150,7 +150,7 @@ static Tensor wrapped_scalar_tensor(Scalar scalar) { static void check_convert(Scalar scalar, ScalarType scalarType) { // Validate that is possible to convert scalar to tensor dtype without overflow - AT_DISPATCH_ALL_TYPES_AND3(at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, scalarType, "check_convert", [&]{ + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, scalarType, "check_convert", [&]{ scalar.to(); }); } diff --git a/aten/src/ATen/native/Convolution.cpp b/aten/src/ATen/native/Convolution.cpp index fabeecf3c76b2..630301ba6e019 100644 --- a/aten/src/ATen/native/Convolution.cpp +++ b/aten/src/ATen/native/Convolution.cpp @@ -31,7 +31,7 @@ struct ConvParams { bool is_output_padding_neg() const; bool is_output_padding_big() const; bool is_padding_neg() const; - bool is_stride_neg() const; + bool is_stride_nonpos() const; void view1d_as_2d(); bool use_cpu_depthwise3x3_winograd(const at::Tensor& input, const at::Tensor& weight) const; bool use_cudnn(const at::Tensor& input) const; @@ -105,15 +105,14 @@ auto ConvParams::is_padding_neg() const -> bool { return is_non_neg; } -auto ConvParams::is_stride_neg() const -> bool { - bool is_non_neg = false; +auto ConvParams::is_stride_nonpos() const -> bool { + bool is_nonpos = false; for (int s : stride) { - is_non_neg |= (s < 0); + is_nonpos |= (s <= 0); } - return is_non_neg; + return is_nonpos; } - auto ConvParams::view1d_as_2d() -> void { if (stride.size() == 1) { stride.insert(stride.begin(), 1); @@ -381,7 +380,7 @@ static void check_shape_forward(const at::Tensor& input, TORCH_CHECK(!params.is_padding_neg(), "negative padding is not supported"); TORCH_CHECK(!params.is_output_padding_neg(), "negative output_padding is not supported"); - TORCH_CHECK(!params.is_stride_neg(), "negative stride is not supported"); + TORCH_CHECK(!params.is_stride_nonpos(), "non-positive stride is not supported"); TORCH_CHECK(weight_dim == k, "Expected ", weight_dim, "-dimensional input for ", weight_dim, diff --git a/aten/src/ATen/native/LinearAlgebraUtils.h b/aten/src/ATen/native/LinearAlgebraUtils.h index e3d7ee52bfe3f..795f403d94f81 100644 --- a/aten/src/ATen/native/LinearAlgebraUtils.h +++ b/aten/src/ATen/native/LinearAlgebraUtils.h @@ -1,4 +1,4 @@ -#include +#include #include #include #include @@ -45,57 +45,6 @@ static inline int64_t matrixStride(const Tensor& batched_matrices) { return batched_matrices.size(-1) * batched_matrices.size(-2); } -/* - * Given batches of matrices with arbitrary batch dim, - * computes the number of batches for Triu and Tril. This ignores stride 0 dimension - */ -static inline int64_t batchCountTrilTriu(const Tensor& batched_matrices) { - int64_t result = 1; - for (int64_t i = 0; i < batched_matrices.ndimension() - 2; i++) { - if (batched_matrices.stride(i) != 0) { - result *= batched_matrices.size(i); - } - } - return result; -} - -/* Checks a necessary property for the triu and tril implementations, hence the name. - * Here batch contiguity is checked for tensors with greater than 4 dimensions. - * Contiguous tensors and tensors with less than 3 dimensions pass this check - */ -static inline std::tuple checkTrilTriuBatchContiguous(const Tensor& tensor, bool allow_zero_stride) { - // Complete contiguity is the most desired property, which is why - // we return true if the tensor is contiguous - if (tensor.is_contiguous()) { - auto default_strides_for_size = at::detail::defaultStrides(tensor.sizes()); - if (tensor.strides() == default_strides_for_size) { - return std::make_tuple(true, tensor); - } else { - return std::make_tuple(false, tensor.as_strided(tensor.sizes(), default_strides_for_size)); - } - } - - int64_t dims = tensor.dim(); - - // Tensors with dimension less than 4 are handled by default - if (allow_zero_stride && dims <= 3) { - return std::make_tuple(true, tensor); - } - - int64_t expected_stride = tensor.size(-1) * tensor.size(-2); - for (int64_t i = dims - 3; i >= 0; i--) { - // Skip trivial dimension; - if (allow_zero_stride && i == 0 && (tensor.stride(i) == 0 || tensor.size(i) == 1)) { - continue; - } - if (expected_stride != tensor.stride(i)) { - return std::make_tuple(false, tensor.contiguous()); - } - expected_stride *= tensor.size(i); - } - return std::make_tuple(true, tensor); -} - // Returns the epsilon value for floating types except half static inline double _get_epsilon(const ScalarType& sc_type) { switch (sc_type) { diff --git a/aten/src/ATen/native/RangeFactories.cpp b/aten/src/ATen/native/RangeFactories.cpp index 0210e998d9264..02262240a9382 100644 --- a/aten/src/ATen/native/RangeFactories.cpp +++ b/aten/src/ATen/native/RangeFactories.cpp @@ -21,14 +21,14 @@ Tensor& linspace_cpu_out(Tensor& result, Scalar start, Scalar end, int64_t steps } else if (steps == 1) { r.fill_(start); } else { - AT_DISPATCH_FLOATING_TYPES(r.scalar_type(), "linspace_cpu", [&]() { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(r.scalar_type(), "linspace_cpu", [&]() { scalar_t scalar_start = start.to(); scalar_t scalar_end = end.to(); scalar_t *data_ptr = r.data_ptr(); scalar_t step = (scalar_end - scalar_start) / static_cast(steps - 1); at::parallel_for(0, steps, internal::GRAIN_SIZE, [&](int64_t p_begin, int64_t p_end) { scalar_t is = static_cast(p_begin); - for (int64_t i = p_begin; i < p_end; ++i, ++is) { + for (int64_t i = p_begin; i < p_end; ++i, is+=1) { //std::complex does not support ++operator data_ptr[i] = scalar_start + step*is; } }); @@ -54,7 +54,7 @@ Tensor& logspace_cpu_out(Tensor& result, Scalar start, Scalar end, int64_t steps } else if (steps == 1) { r.fill_(std::pow(base, start.to())); } else { - AT_DISPATCH_FLOATING_TYPES(r.scalar_type(), "logspace_cpu", [&]() { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(r.scalar_type(), "logspace_cpu", [&]() { scalar_t scalar_base = static_cast(base); scalar_t scalar_start = start.to(); scalar_t scalar_end = end.to(); @@ -62,7 +62,7 @@ Tensor& logspace_cpu_out(Tensor& result, Scalar start, Scalar end, int64_t steps scalar_t step = (scalar_end - scalar_start) / static_cast(steps - 1); at::parallel_for(0, steps, internal::GRAIN_SIZE, [&](int64_t p_begin, int64_t p_end) { scalar_t is = static_cast(p_begin); - for (int64_t i = p_begin; i < p_end; ++i, ++is) { + for (int64_t i = p_begin; i < p_end; ++i, is+=1) { //std::complex does not support ++operator data_ptr[i]= std::pow(scalar_base, scalar_start + step*is); } }); diff --git a/aten/src/ATen/native/TensorCompare.cpp b/aten/src/ATen/native/TensorCompare.cpp index 3916a82f00b0e..0e2132c9eca92 100644 --- a/aten/src/ATen/native/TensorCompare.cpp +++ b/aten/src/ATen/native/TensorCompare.cpp @@ -99,6 +99,8 @@ bool is_nonzero(const Tensor& self) { Scalar localScalar = self.item(); if (localScalar.isFloatingPoint()) { return localScalar.to() != 0; + } else if (localScalar.isComplex()) { + return localScalar.to>() != std::complex(0.0, 0.0); } else if (localScalar.isIntegral(false)){ return localScalar.to() != 0; } else if (localScalar.isBoolean()) { @@ -123,7 +125,7 @@ std::vector where(const Tensor& condition) { Tensor _s_where_cpu(const Tensor& condition, const Tensor& self, const Tensor& other) { Tensor ret = at::empty(self.sizes(), self.options()); - AT_DISPATCH_ALL_TYPES(ret.scalar_type(), "where_cpu", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(ret.scalar_type(), "where_cpu", [&] { where_cpu(ret, condition, self, other); }); return ret; diff --git a/aten/src/ATen/native/TensorFactories.cpp b/aten/src/ATen/native/TensorFactories.cpp index d27285d244b43..07cbca3e57363 100644 --- a/aten/src/ATen/native/TensorFactories.cpp +++ b/aten/src/ATen/native/TensorFactories.cpp @@ -845,7 +845,7 @@ template Tensor tensor_cpu(ArrayRef values, const TensorOptions& options) { auto result = at::empty(values.size(), options); AT_ASSERT(result.is_contiguous()); - AT_DISPATCH_ALL_TYPES(result.scalar_type(), "tensor_cpu", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(result.scalar_type(), "tensor_cpu", [&] { std::copy(values.begin(), values.end(), result.template data_ptr()); }); return result; diff --git a/aten/src/ATen/native/TriangularOps.cpp b/aten/src/ATen/native/TriangularOps.cpp new file mode 100644 index 0000000000000..e712b5082102b --- /dev/null +++ b/aten/src/ATen/native/TriangularOps.cpp @@ -0,0 +1,158 @@ +#include +#include +#include +#include + +#include +#include + +namespace at { +namespace native { + +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triu/tril ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +template +static void apply_triu_tril_single( + scalar_t* result, scalar_t* self, bool inplace, + int64_t k, int64_t n, int64_t m, + int64_t res_row_stride, int64_t res_col_stride, + int64_t self_row_stride, int64_t self_col_stride) { + + constexpr int64_t zero = 0; + + if (upper) { + at::parallel_for(0, n, 0, [&](int64_t start, int64_t end) { + for (auto i = start; i < end; i++) { + for (int64_t j = 0; j < std::min(m, i + k); j++) { + result[i * res_row_stride + j * res_col_stride] = 0; + } + if (!inplace) { // copy the rest of the self if not inplace + for (int64_t j = std::max(zero, i + k); j < m; j++) { + result[i * res_row_stride + j * res_col_stride] = self[i * self_row_stride + j * self_col_stride]; + } + } + } + }); + } else { + at::parallel_for(0, n, 0, [&](int64_t start, int64_t end) { + for (auto i = start; i < end; i++) { + for (int64_t j = std::max(zero, i + k + 1); j < m; j++) { + result[i * res_row_stride + j * res_col_stride] = 0; + } + if (!inplace) { // copy the rest of the self if not inplace + for (int64_t j = zero; j < std::min(m, i + k + 1); j++) { + result[i * res_row_stride + j * res_col_stride] = self[i * self_row_stride + j * self_col_stride]; + } + } + } + }); + } +} + +template +void apply_triu_tril(Tensor& result, const Tensor& self, bool inplace, int64_t k) { + auto n = self.size(-2); + auto m = self.size(-1); + auto self_data = self.data_ptr(); + auto self_stride = (self.dim() > 2 && self.stride(-3) > 0) ? self.stride(-3) : 1; + auto batchsize = batchCountTrilTriu(result); + auto self_row_stride = self.stride(-2); + auto self_column_stride = self.stride(-1); + + auto result_data = result.data_ptr(); + int64_t result_stride, result_row_stride, result_column_stride; + if (result_data != self_data) { + result_stride = (result.dim() > 2 && result.stride(-3) > 0) ? result.stride(-3) : 1; + result_row_stride = result.stride(-2); + result_column_stride = result.stride(-1); + } else { + result_stride = self_stride; + result_row_stride = self_row_stride; + result_column_stride = self_column_stride; + } + + at::parallel_for(0, batchsize, 0, [&](int64_t start, int64_t end) { + for (auto b = start; b < end; b++) { + scalar_t* self_batch = &self_data[b * self_stride]; + scalar_t* result_batch = &result_data[b * result_stride]; + apply_triu_tril_single( + result_batch, self_batch, inplace, k, n, m, + result_row_stride, result_column_stride, self_row_stride, self_column_stride); + } + }); +} + +Tensor tril(const Tensor& self, int64_t k) { + Tensor result = at::empty({0}, self.options()); + at::tril_out(result, self, k); + return result; +} + +Tensor& tril_cpu_(Tensor &self, int64_t k) { + if (self.numel() == 0) { + return self; + } + bool inplace; + Tensor self_c; + std::tie(inplace, self_c) = checkTrilTriuBatchContiguous(self, true); + Tensor result = inplace ? self : at::empty_like(self); + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "tril", [&]{ + apply_triu_tril(result, self_c, inplace, k); + }); + if (!inplace) self.copy_(result); + return self; +} + +Tensor& tril_cpu_out(Tensor &result, const Tensor& self, int64_t k) { + if (result.sizes() != self.sizes()) { + result.resize_as_(self); + } + if (self.numel() == 0) { + return result; + } + Tensor self_c; + std::tie(std::ignore, self_c) = checkTrilTriuBatchContiguous(self, false); + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "tril", [&]{ + apply_triu_tril(result, self_c, false, k); + }); + return result; +} + +Tensor triu(const Tensor& self, int64_t k) { + Tensor result = at::empty({0}, self.options()); + at::triu_out(result, self, k); + return result; +} + +Tensor& triu_cpu_(Tensor &self, int64_t k) { + if (self.numel() == 0) { + return self; + } + bool inplace; + Tensor self_c; + std::tie(inplace, self_c) = checkTrilTriuBatchContiguous(self, true); + Tensor result = inplace ? self : at::empty_like(self); + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "triu", [&]{ + apply_triu_tril(result, self_c, inplace, k); + }); + if (!inplace) self.copy_(result); + return self; +} + +Tensor& triu_cpu_out(Tensor &result, const Tensor& self, int64_t k) { + if (result.sizes() != self.sizes()) { + result.resize_as_(self); + } + if (self.numel() == 0) { + return result; + } + Tensor self_c; + std::tie(std::ignore, self_c) = checkTrilTriuBatchContiguous(self, false); + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "triu", [&]{ + apply_triu_tril(result, self_c, false, k); + }); + return result; +} + +} // namespace native +} // namespace at diff --git a/aten/src/ATen/native/TriangularOpsUtils.h b/aten/src/ATen/native/TriangularOpsUtils.h new file mode 100644 index 0000000000000..33fb38bd8ca0b --- /dev/null +++ b/aten/src/ATen/native/TriangularOpsUtils.h @@ -0,0 +1,59 @@ +#include +#include + +namespace at { +namespace native { + +/* + * Given batches of matrices with arbitrary batch dim, + * computes the number of batches for Triu and Tril. This ignores stride 0 dimension + */ +static inline int64_t batchCountTrilTriu(const Tensor& batched_matrices) { + int64_t result = 1; + for (int64_t i = 0; i < batched_matrices.ndimension() - 2; i++) { + if (batched_matrices.stride(i) != 0) { + result *= batched_matrices.size(i); + } + } + return result; +} + +/* Checks a necessary property for the triu and tril implementations, hence the name. + * Here batch contiguity is checked for tensors with greater than 4 dimensions. + * Contiguous tensors and tensors with less than 3 dimensions pass this check + */ +static inline std::tuple checkTrilTriuBatchContiguous(const Tensor& tensor, bool allow_zero_stride) { + // Complete contiguity is the most desired property, which is why + // we return true if the tensor is contiguous + if (tensor.is_contiguous()) { + auto default_strides_for_size = at::detail::defaultStrides(tensor.sizes()); + if (tensor.strides() == default_strides_for_size) { + return std::make_tuple(true, tensor); + } else { + return std::make_tuple(false, tensor.as_strided(tensor.sizes(), default_strides_for_size)); + } + } + + int64_t dims = tensor.dim(); + + // Tensors with dimension less than 4 are handled by default + if (allow_zero_stride && dims <= 3) { + return std::make_tuple(true, tensor); + } + + int64_t expected_stride = tensor.size(-1) * tensor.size(-2); + for (int64_t i = dims - 3; i >= 0; i--) { + // Skip trivial dimension; + if (allow_zero_stride && i == 0 && (tensor.stride(i) == 0 || tensor.size(i) == 1)) { + continue; + } + if (expected_stride != tensor.stride(i)) { + return std::make_tuple(false, tensor.contiguous()); + } + expected_stride *= tensor.size(i); + } + return std::make_tuple(true, tensor); +} + +} // namespace native +} // namespace at diff --git a/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp b/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp index 7f20f8a5f9ff7..71c4fb71bc5c9 100644 --- a/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp +++ b/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp @@ -185,14 +185,14 @@ void ge_kernel(TensorIterator& iter) { void eq_kernel(TensorIterator& iter) { if (iter.dtype() == ScalarType::Bool) { - AT_DISPATCH_ALL_TYPES_AND2(kBool, kBFloat16, iter.input_dtype(), "eq_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(kBool, kBFloat16, iter.input_dtype(), "eq_cpu", [&]() { cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> bool { return a == b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kBFloat16, iter.dtype(), "eq_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(kBFloat16, iter.dtype(), "eq_cpu", [&]() { cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> scalar_t { return a == b; @@ -203,14 +203,14 @@ void eq_kernel(TensorIterator& iter) { void ne_kernel(TensorIterator& iter) { if (iter.dtype() == ScalarType::Bool) { - AT_DISPATCH_ALL_TYPES_AND2(kBool, kBFloat16, iter.input_dtype(), "ne_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(kBool, kBFloat16, iter.input_dtype(), "ne_cpu", [&]() { cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> bool { return a != b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kBFloat16, iter.dtype(), "ne_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(kBFloat16, iter.dtype(), "ne_cpu", [&]() { cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> scalar_t { return a != b; diff --git a/aten/src/ATen/native/cpu/CrossKernel.cpp b/aten/src/ATen/native/cpu/CrossKernel.cpp index 3e243403b5f49..0d405447bc1a9 100644 --- a/aten/src/ATen/native/cpu/CrossKernel.cpp +++ b/aten/src/ATen/native/cpu/CrossKernel.cpp @@ -65,7 +65,7 @@ static void apply_cross(Tensor& result, const Tensor& a, const Tensor& b, const } static void cross_kernel_impl(Tensor& result, const Tensor& a, const Tensor& b, const int64_t dim) { - AT_DISPATCH_ALL_TYPES(result.scalar_type(), "cross", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(result.scalar_type(), "cross", [&]() { apply_cross(result, a, b, dim); }); } diff --git a/aten/src/ATen/native/cpu/LerpKernel.cpp b/aten/src/ATen/native/cpu/LerpKernel.cpp index 5e6d044030626..313f3e3bf9367 100644 --- a/aten/src/ATen/native/cpu/LerpKernel.cpp +++ b/aten/src/ATen/native/cpu/LerpKernel.cpp @@ -20,14 +20,15 @@ static void lerp_kernel_scalar( TORCH_CHECK(self.dtype() == end.dtype(), "expected dtype ", self.dtype(), " for `end` but got dtype ", end.dtype()); auto iter = TensorIterator::binary_op(ret, self, end, /*check_mem_overlap=*/true); - AT_DISPATCH_FLOATING_TYPES(ret.scalar_type(), "lerp_kernel_scalar", [&] { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(ret.scalar_type(), "lerp_kernel_scalar", [&] { + using value_t = typename ztype::value_t; scalar_t weight_val = weight.to(); at::native::cpu_kernel( iter, [weight_val](scalar_t self_val, scalar_t end_val) { - return (weight_val < 0.5) + return (zabs(weight_val) < 0.5) ? self_val + weight_val * (end_val - self_val) - : end_val - (end_val - self_val) * (1 - weight_val); + : end_val - (end_val - self_val) * (scalar_t(1) - weight_val); }); }); } @@ -49,13 +50,14 @@ static void lerp_kernel_tensor( iter.add_input(end); iter.add_input(weights); iter.build(); - AT_DISPATCH_FLOATING_TYPES(ret.scalar_type(), "lerp_kernel_tensor", [&] { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(ret.scalar_type(), "lerp_kernel_tensor", [&] { + using value_t = typename ztype::value_t; at::native::cpu_kernel( iter, [](scalar_t self_val, scalar_t end_val, scalar_t weight_val) { - return (weight_val < 0.5) + return (zabs(weight_val) < 0.5) ? self_val + weight_val * (end_val - self_val) - : end_val - (end_val - self_val) * (1 - weight_val); + : end_val - (end_val - self_val) * (scalar_t(1) - weight_val); }); }); } diff --git a/aten/src/ATen/native/cpu/PointwiseOpsKernel.cpp b/aten/src/ATen/native/cpu/PointwiseOpsKernel.cpp index e2cc77e1ffb4a..92d15953372ad 100644 --- a/aten/src/ATen/native/cpu/PointwiseOpsKernel.cpp +++ b/aten/src/ATen/native/cpu/PointwiseOpsKernel.cpp @@ -12,7 +12,7 @@ namespace { static void addcmul_cpu_kernel(TensorIterator& iter, Scalar value) { ScalarType dtype = iter.dtype(0); - AT_DISPATCH_ALL_TYPES(dtype, "addcmul_cpu_out", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(dtype, "addcmul_cpu_out", [&] { scalar_t scalar_val = value.to(); auto scalar_vec = Vec256(scalar_val); cpu_kernel_vec( @@ -30,7 +30,7 @@ static void addcmul_cpu_kernel(TensorIterator& iter, Scalar value) { static void addcdiv_cpu_kernel(TensorIterator& iter, Scalar value) { ScalarType dtype = iter.dtype(0); - AT_DISPATCH_ALL_TYPES(dtype, "addcdiv_cpu_out", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(dtype, "addcdiv_cpu_out", [&] { scalar_t scalar_val = value.to(); auto scalar_vec = Vec256(scalar_val); cpu_kernel_vec( diff --git a/aten/src/ATen/native/cpu/PowKernel.cpp b/aten/src/ATen/native/cpu/PowKernel.cpp index 1109063966dac..085d57bac7732 100644 --- a/aten/src/ATen/native/cpu/PowKernel.cpp +++ b/aten/src/ATen/native/cpu/PowKernel.cpp @@ -11,8 +11,8 @@ namespace at { namespace native { namespace { void pow_tensor_tensor_kernel(TensorIterator& iter) { - if (isFloatingType(iter.dtype())) { - AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "pow", [&]() { + if (isFloatingType(iter.dtype()) || isComplexType(iter.dtype())) { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(iter.dtype(), "pow", [&]() { using Vec = Vec256; cpu_kernel_vec(iter, [=](scalar_t base, scalar_t exp) -> scalar_t { @@ -91,6 +91,62 @@ void pow_tensor_scalar_kernel(TensorIterator& iter, Scalar exp_scalar) { ); } }); + } else if (isComplexType(iter.dtype())) { + const auto exp = exp_scalar.to>(); + // Floating types allow AVX2 vector optimizations for pow/sqrt/rsqrt: + AT_DISPATCH_COMPLEX_TYPES(iter.dtype(), "pow", [&]() { + using Vec = Vec256; + if (exp == 0.5) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return std::sqrt(base); + }, + [](Vec base) -> Vec { return base.sqrt(); } + ); + } else if (exp == 2.0) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return base * base; + }, + [](Vec base) -> Vec { return base * base; } + ); + } else if (exp == 3.0) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return base * base * base; + }, + [](Vec base) -> Vec { return base * base * base; } + ); + } else if (exp == -0.5) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return scalar_t(1.0) / std::sqrt(base); + }, + [](Vec base) -> Vec { return base.rsqrt(); } + ); + } else if (exp == -1.0) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return scalar_t(1.0) / base; + }, + [](Vec base) -> Vec { return base.reciprocal(); } + ); + } else if (exp == -2.0) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return scalar_t(1.0) / (base * base); + }, + [](Vec base) -> Vec { return (base * base).reciprocal(); } + ); + } else { + cpu_kernel_vec(iter, + [=](scalar_t base) -> scalar_t { + return std::pow(base, scalar_t(exp)); + }, + [=](Vec base) -> Vec { return base.pow(scalar_t(exp)); } // std::pow cannot accept mixed complex data types. + ); + } + }); } else { // Integral types do not allow AVX2 vector optimizations for pow/sqrt/rsqrt. // Trying to implement pow/sqrt/rsqrt as loop in vec256_int.h does not allow diff --git a/aten/src/ATen/native/cpu/TensorCompareKernel.cpp b/aten/src/ATen/native/cpu/TensorCompareKernel.cpp index aa7a5ebeaabc6..fb7d93cd8237b 100644 --- a/aten/src/ATen/native/cpu/TensorCompareKernel.cpp +++ b/aten/src/ATen/native/cpu/TensorCompareKernel.cpp @@ -8,6 +8,7 @@ #include #include #include +#include namespace at { namespace native { namespace { @@ -34,6 +35,8 @@ struct Reduction { } } int64_t batch = numel / (n * stride); + using value_t = typename ztype::value_t; + value_t (*zabs_)(scalar_t) = zabs; if (stride == 1) { parallel_for(0, batch, 1, [=](int64_t begin, int64_t end) { for (int64_t b = begin; b < end; b++) { @@ -42,7 +45,7 @@ struct Reduction { index_t result_index = 0; for (int64_t k = 0; k < n; k++) { scalar_t value = data[k]; - bool cmp = greater ? (result > value) : (result < value); + bool cmp = greater ? (zabs_(result) > zabs_(value)) : (zabs_(result) < zabs_(value)); result = cmp ? result : value; result_index = cmp ? result_index : k; if (_isnan(result)) { @@ -63,7 +66,7 @@ struct Reduction { index_t result_index = 0; for (int64_t k = 0; k < n; k++) { scalar_t value = data[k * stride]; - bool cmp = greater ? (result > value) : (result < value); + bool cmp = greater ? (zabs_(result) > zabs_(value)) : (zabs_(result) < zabs_(value)); result = cmp ? result : value; result_index = cmp ? result_index : k; if (_isnan(result)) { @@ -83,7 +86,7 @@ static void max_kernel_impl( Tensor& max_indices, const Tensor& self, c10::optional dim) { - AT_DISPATCH_ALL_TYPES_AND(ScalarType::Bool, self.scalar_type(), "max", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(ScalarType::Bool, self.scalar_type(), "max", [&] { Reduction::apply(max, max_indices, self, dim, true); }); } @@ -93,7 +96,7 @@ static void min_kernel_impl( Tensor& min_indices, const Tensor& self, c10::optional dim) { - AT_DISPATCH_ALL_TYPES_AND(ScalarType::Bool, self.scalar_type(), "min", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(ScalarType::Bool, self.scalar_type(), "min", [&] { Reduction::apply(min, min_indices, self, dim, false); }); } diff --git a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp index 33e1b5574b31b..a87fa83df4e2d 100644 --- a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp +++ b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp @@ -428,7 +428,7 @@ IMPLEMENT_COMPLEX_KERNEL(FLOATING, floor) IMPLEMENT_COMPLEX_KERNEL(FLOATING, log) IMPLEMENT_COMPLEX_KERNEL(FLOATING, log10) IMPLEMENT_FLOAT_KERNEL(FLOATING, log1p) -IMPLEMENT_FLOAT_KERNEL(FLOATING, log2) +IMPLEMENT_COMPLEX_KERNEL(FLOATING, log2) IMPLEMENT_COMPLEX_KERNEL(FLOATING, round) IMPLEMENT_COMPLEX_KERNEL(FLOATING, sin) // IMPLEMENT_FLOAT_KERNEL(FLOATING, sinh) diff --git a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu b/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu index f67d78056040a..7ac2beb6211fe 100644 --- a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu +++ b/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu @@ -929,103 +929,6 @@ std::tuple _lu_with_info_cuda(const Tensor& self, bool p return std::make_tuple(self_working_copy, pivots_tensor, infos_tensor); } -// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triu/tril ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -template -#ifdef __HIP_PLATFORM_HCC__ -C10_LAUNCH_BOUNDS_1(512) -#endif -__global__ -void triu_tril_kernel( - cuda::detail::TensorInfo result_info, - const cuda::detail::TensorInfo self_info, - const int64_t k, const int64_t N) { - int64_t linear_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (linear_idx >= N) { - return; - } - - auto dims = self_info.dims; - - IndexType self_offset = 0, result_offset = 0; - // Compute column index and corresponding offset - IndexType col = linear_idx % self_info.sizes[dims - 1]; - linear_idx /= self_info.sizes[dims - 1]; - self_offset += self_info.strides[dims - 1] * col; - result_offset += result_info.strides[dims - 1] * col; - - // Compute row index and corresponding offset - IndexType row = linear_idx % self_info.sizes[dims - 2]; - linear_idx /= self_info.sizes[dims - 2]; - self_offset += self_info.strides[dims - 2] * row; - result_offset += result_info.strides[dims - 2] * row; - - // Compute remaining offsets - IndexType running_index; - #pragma unroll - for (IndexType i = dims - 3; i >= 0; --i) { - running_index = linear_idx % self_info.sizes[i]; - linear_idx /= self_info.sizes[i]; - self_offset += running_index * self_info.strides[i]; - result_offset += running_index * result_info.strides[i]; - } - - bool mask = upper ? (col - row >= k) : (col - row <= k); - result_info.data[result_offset] = mask ? self_info.data[self_offset] : scalar_t(0); -} - -template -Tensor& triu_tril_cuda_template(Tensor& result, const Tensor& self, int64_t k, const char* name) { - int64_t N = self.numel(); - dim3 dim_block = cuda::getApplyBlock(); - dim3 dim_grid((N + dim_block.x - 1) / dim_block.x); - AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), name, [&]{ - if (cuda::detail::canUse32BitIndexMath(result) && cuda::detail::canUse32BitIndexMath(self)) { - auto result_info = cuda::detail::getTensorInfo(result); - auto self_info = cuda::detail::getTensorInfo(self); - triu_tril_kernel - <<>>( - result_info, self_info, k, N); - } else { - auto result_info = cuda::detail::getTensorInfo(result); - auto self_info = cuda::detail::getTensorInfo(self); - triu_tril_kernel - <<>>( - result_info, self_info, k, N); - } - }); - AT_CUDA_CHECK(cudaGetLastError()); - return result; -} - -Tensor& tril_cuda_(Tensor &self, int64_t k) { - return tril_cuda_out(self, self, k); -} - -Tensor& tril_cuda_out(Tensor &result, const Tensor& self, int64_t k) { - if (result.sizes() != self.sizes()) { - result.resize_as_(self); - } - if (self.numel() == 0) { - return result; - } - return triu_tril_cuda_template(result, self, k, "tril"); -} - -Tensor& triu_cuda_(Tensor &self, int64_t k) { - return triu_cuda_out(self, self, k); -} - -Tensor& triu_cuda_out(Tensor &result, const Tensor& self, int64_t k) { - if (result.sizes() != self.sizes()) { - result.resize_as_(self); - } - if (self.numel() == 0) { - return result; - } - return triu_tril_cuda_template(result, self, k, "triu"); -} - // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triangular_solve ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ template diff --git a/aten/src/ATen/native/cuda/ReduceOpsKernel.cu b/aten/src/ATen/native/cuda/ReduceOpsKernel.cu index 57b1adec8f9b0..ae53440bd5e0f 100644 --- a/aten/src/ATen/native/cuda/ReduceOpsKernel.cu +++ b/aten/src/ATen/native/cuda/ReduceOpsKernel.cu @@ -141,50 +141,80 @@ void or_kernel_cuda(TensorIterator& iter) { }), false); } -template +template void max_values_kernel_cuda_impl(TensorIterator& iter) { gpu_reduce_kernel( - iter, func_wrapper ([]GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { - return (THCNumerics::isnan(a) || a > b) ? a : b; - }), at::numeric_limits::lower_bound()); + iter, func_wrapper ([]GPU_LAMBDA(acc_t a, acc_t b) -> acc_t { + return (THCNumerics::isnan(a) || a > b) ? a : b; + }), at::numeric_limits::lower_bound()); } -template +template void min_values_kernel_cuda_impl(TensorIterator& iter) { gpu_reduce_kernel( - iter, func_wrapper ([]GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { - return (THCNumerics::isnan(a) || a < b) ? a : b; - }), at::numeric_limits::upper_bound()); + iter, func_wrapper ([]GPU_LAMBDA(acc_t a, acc_t b) -> acc_t { + return (THCNumerics::isnan(a) || a < b) ? a : b; + }), at::numeric_limits::upper_bound()); } void max_values_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_ALL_TYPES(iter.dtype(), "max_values_cuda", [&]() { - max_values_kernel_cuda_impl(iter); - }); + if (iter.dtype(1) == kHalf) { + max_values_kernel_cuda_impl(iter); + } else { + AT_DISPATCH_ALL_TYPES(iter.dtype(), "max_values_cuda", [&]() { + max_values_kernel_cuda_impl(iter); + }); + } } void min_values_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_ALL_TYPES(iter.dtype(), "min_values_cuda", [&]() { - min_values_kernel_cuda_impl(iter); - }); + if (iter.dtype(1) == kHalf) { + min_values_kernel_cuda_impl(iter); + } else { + AT_DISPATCH_ALL_TYPES(iter.dtype(), "min_values_cuda", [&]() { + min_values_kernel_cuda_impl(iter); + }); + } } +template +void argmax_kernel_cuda_impl(TensorIterator& iter) { + gpu_reduce_kernel( + iter, + ArgMaxOps{}, + thrust::pair(at::numeric_limits::lower_bound(), 0)); +}; + +template +void argmin_kernel_cuda_impl(TensorIterator& iter) { + gpu_reduce_kernel( + iter, + ArgMinOps{}, + thrust::pair(at::numeric_limits::upper_bound(), 0)); +}; + void argmax_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_ALL_TYPES(iter.dtype(1), "argmax_cuda", [&]() { - gpu_reduce_kernel( - iter, - ArgMaxOps{}, - thrust::pair(at::numeric_limits::lower_bound(), 0)); - }); + if (iter.dtype(1) == kHalf) { + // Instead of implementing is_nan and warp_shfl_down + // we can convert halves to float and do all the operations in float + argmax_kernel_cuda_impl(iter); + } else { + AT_DISPATCH_ALL_TYPES(iter.dtype(1), "argmax_cuda", [&]() { + argmax_kernel_cuda_impl(iter); + }); + } } void argmin_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_ALL_TYPES(iter.dtype(1), "argmin_cuda", [&]() { - gpu_reduce_kernel( - iter, - ArgMinOps{}, - thrust::pair(at::numeric_limits::upper_bound(), 0)); - }); + if (iter.dtype(1) == kHalf) { + // Instead of implementing is_nan and warp_shfl_down + // we can convert halves to float and do all the operations in float + argmin_kernel_cuda_impl(iter); + } else { + AT_DISPATCH_ALL_TYPES(iter.dtype(1), "argmin_cuda", [&]() { + argmin_kernel_cuda_impl(iter); + }); + } } REGISTER_DISPATCH(std_var_stub, &std_var_kernel_cuda); diff --git a/aten/src/ATen/native/cuda/TriangularOps.cu b/aten/src/ATen/native/cuda/TriangularOps.cu new file mode 100644 index 0000000000000..4a5e1e64aadb7 --- /dev/null +++ b/aten/src/ATen/native/cuda/TriangularOps.cu @@ -0,0 +1,110 @@ +#include +#include +#include +#include + +#include +#include + +namespace at { +namespace native { + +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triu/tril ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +template +#ifdef __HIP_PLATFORM_HCC__ +C10_LAUNCH_BOUNDS_1(512) +#endif +__global__ +void triu_tril_kernel( + cuda::detail::TensorInfo result_info, + const cuda::detail::TensorInfo self_info, + const int64_t k, const int64_t N) { + int64_t linear_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (linear_idx >= N) { + return; + } + + auto dims = self_info.dims; + + IndexType self_offset = 0, result_offset = 0; + // Compute column index and corresponding offset + IndexType col = linear_idx % self_info.sizes[dims - 1]; + linear_idx /= self_info.sizes[dims - 1]; + self_offset += self_info.strides[dims - 1] * col; + result_offset += result_info.strides[dims - 1] * col; + + // Compute row index and corresponding offset + IndexType row = linear_idx % self_info.sizes[dims - 2]; + linear_idx /= self_info.sizes[dims - 2]; + self_offset += self_info.strides[dims - 2] * row; + result_offset += result_info.strides[dims - 2] * row; + + // Compute remaining offsets + IndexType running_index; + #pragma unroll + for (IndexType i = dims - 3; i >= 0; --i) { + running_index = linear_idx % self_info.sizes[i]; + linear_idx /= self_info.sizes[i]; + self_offset += running_index * self_info.strides[i]; + result_offset += running_index * result_info.strides[i]; + } + + bool mask = upper ? (col - row >= k) : (col - row <= k); + result_info.data[result_offset] = mask ? self_info.data[self_offset] : scalar_t(0); +} + +template +Tensor& triu_tril_cuda_template(Tensor& result, const Tensor& self, int64_t k, const char* name) { + int64_t N = self.numel(); + dim3 dim_block = cuda::getApplyBlock(); + dim3 dim_grid((N + dim_block.x - 1) / dim_block.x); + AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), name, [&]{ + if (cuda::detail::canUse32BitIndexMath(result) && cuda::detail::canUse32BitIndexMath(self)) { + auto result_info = cuda::detail::getTensorInfo(result); + auto self_info = cuda::detail::getTensorInfo(self); + triu_tril_kernel + <<>>( + result_info, self_info, k, N); + } else { + auto result_info = cuda::detail::getTensorInfo(result); + auto self_info = cuda::detail::getTensorInfo(self); + triu_tril_kernel + <<>>( + result_info, self_info, k, N); + } + }); + AT_CUDA_CHECK(cudaGetLastError()); + return result; +} + +Tensor& tril_cuda_(Tensor &self, int64_t k) { + return tril_cuda_out(self, self, k); +} + +Tensor& tril_cuda_out(Tensor &result, const Tensor& self, int64_t k) { + if (result.sizes() != self.sizes()) { + result.resize_as_(self); + } + if (self.numel() == 0) { + return result; + } + return triu_tril_cuda_template(result, self, k, "tril"); +} + +Tensor& triu_cuda_(Tensor &self, int64_t k) { + return triu_cuda_out(self, self, k); +} + +Tensor& triu_cuda_out(Tensor &result, const Tensor& self, int64_t k) { + if (result.sizes() != self.sizes()) { + result.resize_as_(self); + } + if (self.numel() == 0) { + return result; + } + return triu_tril_cuda_template(result, self, k, "triu"); +} + +} // namespace native +} // namespace at diff --git a/aten/src/ATen/native/quantized/cpu/qconv.cpp b/aten/src/ATen/native/quantized/cpu/qconv.cpp index c2993717f7552..cc6efcd084565 100644 --- a/aten/src/ATen/native/quantized/cpu/qconv.cpp +++ b/aten/src/ATen/native/quantized/cpu/qconv.cpp @@ -1,41 +1,140 @@ +#include +#include +#include + #include -#include #include +#include #include #include #include #include #include -#include namespace at { namespace native { namespace { -SmallVector convOutputShape( +template +bool ConvDimChecks( + int64_t act_dims, + int64_t stride_dims, + int64_t padding_dims, + int64_t dilation_dims) { + TORCH_CHECK( + act_dims == kSpatialDim + 2, + "quantized::conv", + kSpatialDim, + "d(): Expected activation tensor to have ", + kSpatialDim + 2, + " dimensions."); + TORCH_CHECK( + stride_dims == kSpatialDim, + "quantized::conv", + kSpatialDim, + "d(): Expected stride tensor to have ", + kSpatialDim, + " dimensions."); + TORCH_CHECK( + padding_dims == kSpatialDim, + "quantized::conv", + kSpatialDim, + "d(): Expected padding tensor to have ", + kSpatialDim, + " dimensions."); + TORCH_CHECK( + dilation_dims == kSpatialDim, + "quantized::conv", + kSpatialDim, + "d(): Expected dilation tensor to have ", + kSpatialDim, + " dimensions."); + return true; +} + +#ifdef USE_FBGEMM + +template +SmallVector MakeConvOutputShape( + int N, + int M, + const std::array& output_image_shape); + +template <> +SmallVector MakeConvOutputShape<2>( + int N, + int M, + const std::array& output_image_shape) { + return {N, M, output_image_shape[0], output_image_shape[1]}; +} + +template <> +SmallVector MakeConvOutputShape<3>( + int N, + int M, + const std::array& output_image_shape) { + return {N, + M, + output_image_shape[0], + output_image_shape[1], + output_image_shape[2]}; +} + +#endif // USE_FBGEMM + +#ifdef USE_PYTORCH_QNNPACK + +template +SmallVector MakeConvOutputShape( int N, // mini-batch - int K, // output channels - int H, // input height - int W, // input width + int M, // output channels + const std::vector& input_image_shape, + const std::vector& kernel, + const torch::List& stride, + const torch::List& padding, + const torch::List& dilation); + +template <> +SmallVector MakeConvOutputShape<2>( + int N, // mini-batch + int M, // output channels + const std::vector& input_image_shape, const std::vector& kernel, const torch::List& stride, const torch::List& padding, const torch::List& dilation) { - SmallVector out_shape; - out_shape.push_back(N); - - int H_out = std::floor( - (H + 2 * padding[0] - dilation[0] * (kernel[0] - 1) - 1) / stride[0] + 1); - int W_out = std::floor( - (W + 2 * padding[1] - dilation[1] * (kernel[1] - 1) - 1) / stride[1] + 1); - out_shape.push_back(H_out); - out_shape.push_back(W_out); - // TODO: reorder it to NCHW order once the memory format regression is fixed - out_shape.push_back(K); - - return out_shape; + const int H = input_image_shape[0]; + const int W = input_image_shape[1]; + const int64_t Y_H = + (H + 2 * padding[0] - dilation[0] * (kernel[0] - 1) - 1) / stride[0] + 1; + const int64_t Y_W = + (W + 2 * padding[1] - dilation[1] * (kernel[1] - 1) - 1) / stride[1] + 1; + return {N, M, Y_H, Y_W}; } +template <> +SmallVector MakeConvOutputShape<3>( + int N, // mini-batch + int M, // output channels + const std::vector& input_image_shape, + const std::vector& kernel, + const torch::List& stride, + const torch::List& padding, + const torch::List& dilation) { + const int D = input_image_shape[0]; + const int H = input_image_shape[1]; + const int W = input_image_shape[2]; + const int64_t Y_D = + (D + 2 * padding[0] - dilation[0] * (kernel[0] - 1) - 1) / stride[0] + 1; + const int64_t Y_H = + (H + 2 * padding[1] - dilation[1] * (kernel[1] - 1) - 1) / stride[1] + 1; + const int64_t Y_W = + (W + 2 * padding[2] - dilation[2] * (kernel[2] - 1) - 1) / stride[2] + 1; + return {N, M, Y_D, Y_H, Y_W}; +} + +#endif // USE_PYTORCH_QNNPACK + /* * FBGEMM uses vpmaddubsw instruction to multiply activations (uint8_t) and * weights (int8_t). @@ -65,27 +164,99 @@ SmallVector convOutputShape( * is 32767. * */ -template -class QConv2dInt8 final : public c10::OperatorKernel { +template +class QConvInt8 final : public c10::OperatorKernel { public: - void conv_checks( - int64_t act_dims, - int64_t stride_dims, - int64_t padding_dims, - int64_t dilation_dims) { - TORCH_CHECK( - act_dims == 4, - "quantized::conv2d(): Expected activation tensor to have 4 dimensions."); - TORCH_CHECK( - stride_dims == 2, "quantized::conv2d(): Supports 2D convolution only"); - TORCH_CHECK( - padding_dims == 2, "quantized::conv2d(): Supports 2D convolution only"); + Tensor operator()( + Tensor act, + Tensor packed_weight, + torch::List stride, + torch::List padding, + torch::List dilation, + int64_t groups, + double output_scale, + int64_t output_zero_point) { + auto& ctx = at::globalContext(); + +#ifdef USE_FBGEMM + if (ctx.qEngine() == at::QEngine::FBGEMM) { + return FbgemmConv( + act, + packed_weight, + stride, + padding, + dilation, + groups, + output_scale, + output_zero_point); + } +#endif // USE_FBGEMM + +#ifdef USE_PYTORCH_QNNPACK + if (ctx.qEngine() == at::QEngine::QNNPACK) { + TORCH_CHECK(kSpatialDim == 2, "QNNPACK only suuports Conv2d now."); + return QnnpackConv( + act, + packed_weight, + stride, + padding, + dilation, + groups, + output_scale, + output_zero_point); + } +#endif + TORCH_CHECK( - dilation_dims == 2, - "quantized::conv2d(): Supports 2D convolution only"); + false, + "Didn't find engine for operation quantized::conv ", + toString(ctx.qEngine())); } + + private: #ifdef USE_FBGEMM - at::Tensor fbgemm_conv( + static const float* GetBiasData( + const PackedConvWeight& pack_data, + Tensor* bias) { + const float* bias_data = nullptr; + if (pack_data.bias.has_value()) { + *bias = pack_data.bias.value(); + TORCH_CHECK( + bias->dtype() == at::kFloat, + "[QConv3D] The 'bias' tensor must have 'torch.float' dtype"); + *bias = bias->contiguous(); + TORCH_CHECK(bias->dim() == 1, "bias should be a vector (1D Tensor)"); + const int M = pack_data.w->outputChannels(); + TORCH_CHECK(bias->size(0) == M, "bias should have ", M, " elements."); + bias_data = bias->data_ptr(); + } + return bias_data; + } + + static void GetQuantizationParams( + const PackedConvWeight& pack_data, + float act_scale, + float out_scale, + std::vector* output_multiplier_float, + std::vector* act_times_w_scale) { + if (pack_data.q_scheme == kPerTensorAffine) { + *act_times_w_scale = {(act_scale * pack_data.w_scale[0])}; + *output_multiplier_float = {act_times_w_scale->front() / out_scale}; + } else if (pack_data.q_scheme == kPerChannelAffine) { + const int M = pack_data.w->outputChannels(); + output_multiplier_float->resize(M); + act_times_w_scale->resize(M); + for (int i = 0; i < M; ++i) { + act_times_w_scale->at(i) = (act_scale * pack_data.w_scale[i]); + output_multiplier_float->at(i) = act_times_w_scale->at(i) / out_scale; + } + } else { + TORCH_CHECK( + false, "[QConv", kSpatialDim, "D] Unknown quantization scheme"); + } + } + + at::Tensor FbgemmConv( Tensor act, Tensor packed_weight, torch::List stride, @@ -105,178 +276,230 @@ class QConv2dInt8 final : public c10::OperatorKernel { // See https://github.com/pytorch/pytorch/issues/23403 TORCH_CHECK( fbgemm::fbgemmSupportedCPU(), "Your CPU does not support FBGEMM."); - conv_checks( + ConvDimChecks( act.ndimension(), stride.size(), padding.size(), dilation.size()); - int N = act.size(0); - int C = act.size(1); - int H = act.size(2); - int W = act.size(3); - - // FBGEMM requires NHWC - // TODO: change it to contiguous(MemoryFormat::ChannelsLast) once a perf - // regression of it is fixed. Today it's equivalent because `act` sizes - // are not used below - Tensor act_contig = act.permute({0, 2, 3, 1}).contiguous(); - const uint8_t* act_ptr = - reinterpret_cast(act_contig.data_ptr()); - - PackedConvWeight<2>& pack_ptr = - cpp_custom_type_hack::cast>(packed_weight); - auto packB = pack_ptr.w.get(); - auto& col_offsets = pack_ptr.col_offsets; - auto& kernel = pack_ptr.kernel; - - int K = packB->outputChannels(); - - int pad_l = padding[0]; - int pad_t = padding[1]; - int stride_h = stride[0]; - int stride_w = stride[1]; - int kernel_h = kernel[0]; - int kernel_w = kernel[1]; - // clang-format off - TORCH_CHECK(C == (packB->inputChannels()), - "[QConv2D] Given groups=", groups, ", weight of size ", - K, ", ", kernel_h, ", ", kernel_w, ", ", packB->inputChannels(), - ", expected input (NCHW) ", N, ", ", C, ", ", H, ", ", W, - " to have ", (packB->inputChannels() * groups), - " channels, but got ", C, " channels instead"); - // clang-format on - fbgemm::conv_param_t<> conv_p( - N, // Batch size - C, // Number of input channels - K, // Number of output channels - {H, W}, - groups, - {kernel_h, kernel_w}, - {stride_h, stride_w}, - {pad_l, pad_t, pad_l, pad_t}, - {static_cast(dilation[0]), static_cast(dilation[1])}); - - float act_scale = act.q_scale(); - int32_t act_zero_point = act.q_zero_point(); - - const float* bias_ptr = nullptr; - at::Tensor bias; - if (pack_ptr.bias.has_value()) { - bias = pack_ptr.bias.value(); + const int N = act.size(0); + const int C = act.size(1); + const int D = kSpatialDim == 2 ? 1 : act.size(2); + const int H = act.size(kSpatialDim); + const int W = act.size(kSpatialDim + 1); + + const Tensor act_nhwc = kSpatialDim == 2 + ? act.contiguous(MemoryFormat::ChannelsLast) + : fbgemm_utils::ConvertToChannelsLast3dTensor(act); + const uint8_t* act_data = + reinterpret_cast(act_nhwc.data_ptr()); + PackedConvWeight& pack_data = + cpp_custom_type_hack::cast>( + packed_weight); + auto* pack_w = pack_data.w.get(); + const auto& col_offsets = pack_data.col_offsets; + const auto& kernel = pack_data.kernel; + + const int M = pack_w->outputChannels(); + const int kernel_d = kSpatialDim == 2 ? 1 : kernel[0]; + const int kernel_h = kernel[kSpatialDim - 2]; + const int kernel_w = kernel[kSpatialDim - 1]; + const int pad_d = kSpatialDim == 2 ? 0 : padding[0]; + const int pad_h = padding[kSpatialDim - 2]; + const int pad_w = padding[kSpatialDim - 1]; + const int stride_d = kSpatialDim == 2 ? 1 : stride[0]; + const int stride_h = stride[kSpatialDim - 2]; + const int stride_w = stride[kSpatialDim - 1]; + const int dilation_d = kSpatialDim == 2 ? 1 : dilation[0]; + const int dilation_h = dilation[kSpatialDim - 2]; + const int dilation_w = dilation[kSpatialDim - 1]; + + if (kSpatialDim == 2) { TORCH_CHECK( - bias.dtype() == at::kFloat, - "[QConv2D] The 'bias' tensor must have 'torch.float' dtype"); - bias = bias.contiguous(); - TORCH_CHECK(bias.dim() == 1, "bias should be a vector (1D Tensor)"); + C == pack_w->inputChannels(), + "[QConv2D] Given groups=", + groups, + ", weight of size ", + M, + ", ", + kernel_h, + ", ", + kernel_w, + ", ", + pack_w->inputChannels(), + ", expected input (NCHW) ", + N, + ", ", + C, + ", ", + H, + ", ", + W, + " to have ", + pack_w->inputChannels(), + " channels, but got ", + C, + " channels instead"); + } else { TORCH_CHECK( - bias.size(0) == K, - "bias should have K elements: " + c10::to_string(K)); - bias_ptr = bias.data_ptr(); + C == pack_w->inputChannels(), + "[QConv3D] Given groups=", + groups, + ", weight of size ", + M, + ", ", + kernel_d, + ", ", + kernel_h, + ", ", + kernel_w, + ", ", + pack_w->inputChannels(), + ", expected input (NCDHW) ", + N, + ", ", + C, + ", ", + D, + ", ", + H, + ", ", + W, + " to have ", + pack_w->inputChannels(), + " channels, but got ", + C, + " channels instead"); } - std::vector output_multiplier_float(1, 0.0); - std::vector act_times_w_scale(1, 1.0); + fbgemm::conv_param_t conv_p = + fbgemm_utils::MakeFbgemmConvParam( + N, // Batch size + C, // Number of input channels + M, // Number of output channels + kSpatialDim == 2 ? std::vector{H, W} + : std::vector{D, H, W}, + groups, + kSpatialDim == 2 ? std::vector{kernel_h, kernel_w} + : std::vector{kernel_d, kernel_h, kernel_w}, + kSpatialDim == 2 ? std::vector{stride_h, stride_w} + : std::vector{stride_d, stride_h, stride_w}, + kSpatialDim == 2 ? std::vector{pad_h, pad_w} + : std::vector{pad_d, pad_h, pad_w}, + kSpatialDim == 2 + ? std::vector{dilation_h, dilation_w} + : std::vector{dilation_d, dilation_h, dilation_w}); + + const float act_scale = act.q_scale(); + const int32_t act_zero_point = act.q_zero_point(); + + Tensor bias; + const float* bias_data = GetBiasData(pack_data, &bias); + TORCH_CHECK( - pack_ptr.w_scale.size() == pack_ptr.w_zp.size(), + pack_data.w_scale.size() == pack_data.w_zp.size(), "Weight scales and zero points vectors should have the same size."); + std::vector output_multiplier_float; + std::vector act_times_w_scale; + GetQuantizationParams( + pack_data, + act_scale, + output_scale, + &output_multiplier_float, + &act_times_w_scale); - if (pack_ptr.q_scheme == kPerTensorAffine) { - act_times_w_scale[0] = (act_scale * pack_ptr.w_scale[0]); - output_multiplier_float[0] = - act_times_w_scale[0] / static_cast(output_scale); - } else if (pack_ptr.q_scheme == kPerChannelAffine) { - output_multiplier_float.resize(K, 0.0); - act_times_w_scale.resize(K, 1.0); - for (int i = 0; i < K; ++i) { - act_times_w_scale[i] = (act_scale * pack_ptr.w_scale[i]); - output_multiplier_float[i] = - act_times_w_scale[i] / static_cast(output_scale); - } - } else { - TORCH_CHECK(false, "[QConv2D] Unknown quantization scheme"); - } - - // TODO: change the following to NCHW sizes once perf is fixed - SmallVector outShape{ - N, conv_p.OUT_DIM[0], conv_p.OUT_DIM[1], K}; + const SmallVector output_shape = + MakeConvOutputShape(N, M, conv_p.OUT_DIM); TORCH_CHECK( std::all_of( - outShape.begin(), outShape.end(), [](int64_t i) { return i > 0; }), - "[QConv2D] each dimension of output tensor should be greater than 0") - - // Force output format to be NHWC - // TODO: consider preserving input format - // TODO: add MemoryFormat::ChannelsLast here once perf is fixed - Tensor output = _empty_affine_quantized( - outShape, device(kCPU).dtype(kQUInt8), output_scale, output_zero_point); - auto buffer = at::empty(output.sizes(), output.options().dtype(at::kInt)); - - int num_tasks = at::get_num_threads(); + output_shape.begin(), + output_shape.end(), + [](int64_t i) { return i > 0; }), + "[QConv", + kSpatialDim, + "D] each dimension of output tensor should be greater than 0"); + + Tensor output = kSpatialDim == 2 + ? _empty_affine_quantized( + output_shape, + device(kCPU).dtype(kQUInt8), + output_scale, + output_zero_point, + MemoryFormat::ChannelsLast) + : fbgemm_utils::MakeEmptyAffineQuantizedChannelsLast3dTensor( + output_shape[0], + output_shape[1], + output_shape[2], + output_shape[3], + output_shape[4], + device(kCPU).dtype(kQUInt8), + output_scale, + output_zero_point); + Tensor buffer = at::empty(output.sizes(), output.options().dtype(at::kInt)); + const int num_tasks = at::get_num_threads(); at::parallel_for(0, num_tasks, 1, [&](int64_t begin, int64_t end) { - fbgemm::DoNothing<> NoOpObj{}; + fbgemm::DoNothing<> kNoOpObj{}; for (int task_id = begin; task_id < end; ++task_id) { - if (pack_ptr.q_scheme == kPerTensorAffine) { + if (pack_data.q_scheme == kPerTensorAffine) { fbgemm::ReQuantizeOutput< - ReluFused, + kReluFused, fbgemm::QuantizationGranularity::TENSOR, float> - outputProcObj( - NoOpObj, + output_proc_obj( + kNoOpObj, output_multiplier_float.data(), output_zero_point, act_zero_point, - pack_ptr.w_zp.data(), + pack_data.w_zp.data(), nullptr, /* row offset buffer */ col_offsets.data(), - bias_ptr, - K, + bias_data, + M, groups, act_times_w_scale.data()); - fbgemm::fbgemmConv( + fbgemm::fbgemmConv( conv_p, - act_ptr, - *packB, + act_data, + *pack_w, reinterpret_cast(output.data_ptr()), buffer.data_ptr(), - outputProcObj, + output_proc_obj, task_id /* thread_id*/, num_tasks /* num_threads */); - - } else if (pack_ptr.q_scheme == kPerChannelAffine) { + } else if (pack_data.q_scheme == kPerChannelAffine) { fbgemm::ReQuantizeOutput< - ReluFused, + kReluFused, fbgemm::QuantizationGranularity::OUT_CHANNEL, float> - outputProcObj( - NoOpObj, + output_proc_obj( + kNoOpObj, output_multiplier_float.data(), output_zero_point, act_zero_point, - pack_ptr.w_zp.data(), + pack_data.w_zp.data(), nullptr, /* row offset buffer */ col_offsets.data(), - bias_ptr, - K, + bias_data, + M, groups, act_times_w_scale.data()); - fbgemm::fbgemmConv( + fbgemm::fbgemmConv( conv_p, - act_ptr, - *packB, + act_data, + *pack_w, reinterpret_cast(output.data_ptr()), buffer.data_ptr(), - outputProcObj, + output_proc_obj, task_id /* thread_id*/, num_tasks /* num_threads */); } } }); - // TODO: remove permute once MemoryLayout is added above - return output.permute({0, 3, 1, 2}); + return output; } #endif + #ifdef USE_PYTORCH_QNNPACK - at::Tensor qnnpack_conv( + at::Tensor QnnpackConv( Tensor act, Tensor packed_weight, torch::List stride, @@ -285,44 +508,42 @@ class QConv2dInt8 final : public c10::OperatorKernel { int64_t groups, double output_scale, int64_t output_zero_point) { - conv_checks( + ConvDimChecks( act.ndimension(), stride.size(), padding.size(), dilation.size()); - PackedConvWeightsQnnp& pack_ptr = + PackedConvWeightsQnnp& pack_data = cpp_custom_type_hack::cast(packed_weight); - auto packB = pack_ptr.w.get(); - auto kernel = pack_ptr.kernel; - auto kernel_zp = pack_ptr.w_zp; - auto kernel_scale = pack_ptr.w_scale; + auto* pack_w = pack_data.w.get(); + const auto& kernel = pack_data.kernel; + const auto& kernel_zp = pack_data.w_zp; + const auto& kernel_scale = pack_data.w_scale; const uint32_t kernel_h = kernel[0]; const uint32_t kernel_w = kernel[1]; // TODO Can be replaced with packB->getOutputChannels() when update pre-pack // to actually do the packing. - const auto out_ch = pack_ptr.bias.size(0); + const auto out_ch = pack_data.bias.size(0); // inputs are in semantic NCHW format - int N = act.size(0); - int in_ch = act.size(1); - int H = act.size(2); - int W = act.size(3); - int K = out_ch; // output channels - // TODO: change it to contiguous(MemoryFormat::ChannelsLast) once a perf - // regression of it is fixed. Today it's equivalent because `act` sizes - // are not used below - Tensor input_contig = act.permute({0, 2, 3, 1}).contiguous(); - - uint32_t stride_h = stride[0]; - uint32_t stride_w = stride[1]; - uint32_t pad_t = padding[0]; - uint32_t pad_l = padding[1]; - uint32_t dilation_h = dilation[0]; - uint32_t dilation_w = dilation[1]; - - auto output_min = ReluFused + const int N = act.size(0); + const int C = act.size(1); + const int H = act.size(2); + const int W = act.size(3); + const int M = out_ch; // output channels + + const Tensor act_nhwc = act.contiguous(MemoryFormat::ChannelsLast); + + const uint32_t stride_h = stride[0]; + const uint32_t stride_w = stride[1]; + const uint32_t pad_h = padding[0]; + const uint32_t pad_w = padding[1]; + const uint32_t dilation_h = dilation[0]; + const uint32_t dilation_w = dilation[1]; + + auto output_min = kReluFused ? activationLimits(output_scale, output_zero_point, Activation::RELU) .first : std::numeric_limits::min(); - auto output_max = ReluFused + auto output_max = kReluFused ? activationLimits(output_scale, output_zero_point, Activation::RELU) .second : std::numeric_limits::max(); @@ -330,29 +551,26 @@ class QConv2dInt8 final : public c10::OperatorKernel { {kernel_w, kernel_h}, {stride_w, stride_h}, {dilation_w, dilation_h}, - {pad_t, pad_l, pad_t, pad_l}, + {pad_h, pad_w, pad_h, pad_w}, groups, - in_ch, - out_ch, + C, + M, kernel_zp, kernel_scale, output_min, output_max); - // TODO: change convOutputShape to return NCHW sizes once perf is fixed - // Force output format to be NHWC - // TODO: consider preserving input format - // TODO: add MemoryFormat::ChannelsLast here once perf is fixed - auto input_scale = input_contig.q_scale(); + auto input_scale = act_nhwc.q_scale(); // Re-quantizing the bias based on input scale and weight scale. - if (!pack_ptr.input_scale.has_value() || - pack_ptr.input_scale.value() != input_scale) { + if (!pack_data.input_scale.has_value() || + pack_data.input_scale.value() != input_scale) { // Get the original weight and adjust it to uint8 from int8 auto weight_contig = - pack_ptr.orig_weight.contiguous(MemoryFormat::ChannelsLast); - auto bias_fp32 = pack_ptr.bias; - int8_t* w_data = (int8_t*)weight_contig.data_ptr(); + pack_data.orig_weight.contiguous(MemoryFormat::ChannelsLast); + auto bias_fp32 = pack_data.bias; + int8_t* w_data = + reinterpret_cast(weight_contig.data_ptr()); Tensor qnnp_weight = at::_empty_affine_quantized( weight_contig.sizes(), at::device(kCPU).dtype(kQUInt8), @@ -368,106 +586,69 @@ class QConv2dInt8 final : public c10::OperatorKernel { auto bias = at::quantize_per_tensor( bias_fp32, kernel_scale * input_scale, 0, kQInt32); // Update the input scale to not pack again. - pack_ptr.input_scale = input_scale; - pack_ptr.w.reset(); - pack_ptr.w = guts::make_unique( + pack_data.input_scale = input_scale; + pack_data.w.reset(); + pack_data.w = guts::make_unique( conv_p, - (uint8_t*)qnnp_w_data, - (int32_t*)bias.data_ptr()); - packB = pack_ptr.w.get(); + reinterpret_cast(qnnp_w_data), + reinterpret_cast(bias.data_ptr())); + pack_w = pack_data.w.get(); } - TORCH_INTERNAL_ASSERT(packB != nullptr, "Packed Weights are NULL"); - auto outShape = - convOutputShape(N, K, H, W, kernel, stride, padding, dilation); + TORCH_INTERNAL_ASSERT(pack_w != nullptr, "Packed Weights are NULL"); + const auto output_shape = MakeConvOutputShape( + N, M, {H, W}, kernel, stride, padding, dilation); TORCH_CHECK( std::all_of( - outShape.begin(), outShape.end(), [](int64_t i) { return i > 0; }), - "quantized::conv2d (qnnpack): each dimension of output tensor should be greater " - "than 0") - TORCH_CHECK( - (outShape[3] == out_ch), - "quantized::conv2d (qnnpack): Number of filters must be equal to number of " - "output channels") + output_shape.begin(), + output_shape.end(), + [](int64_t i) { return i > 0; }), + "quantized::conv2d (qnnpack): each dimension of output tensor should " + "be greater than 0.") // Allocate output Tensor and a buffer for QNNPACK to use Tensor output = at::_empty_affine_quantized( - outShape, + output_shape, at::device(kCPU).dtype(kQUInt8), output_scale, - output_zero_point); + output_zero_point, + MemoryFormat::ChannelsLast); - const pytorch_qnnp_status runStatus = qnnpack::qnnpackConv( + const pytorch_qnnp_status run_status = qnnpack::qnnpackConv( conv_p, - packB->getPackedWeights(), + pack_w->getPackedWeights(), N, H, W, - input_contig.q_scale(), - input_contig.q_zero_point(), - (uint8_t*)input_contig.data_ptr(), + act_nhwc.q_scale(), + act_nhwc.q_zero_point(), + reinterpret_cast(act_nhwc.data_ptr()), output.q_scale(), output.q_zero_point(), - (uint8_t*)output.data_ptr(), + reinterpret_cast(output.data_ptr()), caffe2::mobile_pthreadpool()); TORCH_INTERNAL_ASSERT( - runStatus == pytorch_qnnp_status_success, + run_status == pytorch_qnnp_status_success, "failed to run quantized::conv2d (qnnpack) operator"); - // TODO: remove permute once MemoryLayout is added above - return output.permute({0, 3, 1, 2}); + return output; } #endif - Tensor operator()( - Tensor act, - Tensor packed_weight, - torch::List stride, - torch::List padding, - torch::List dilation, - int64_t groups, - double output_scale, - int64_t output_zero_point) { - auto& ctx = at::globalContext(); -#ifdef USE_FBGEMM - if (ctx.qEngine() == at::QEngine::FBGEMM) { - return fbgemm_conv( - act, - packed_weight, - stride, - padding, - dilation, - groups, - output_scale, - output_zero_point); - } -#endif -#ifdef USE_PYTORCH_QNNPACK - if (ctx.qEngine() == at::QEngine::QNNPACK) { - return qnnpack_conv( - act, - packed_weight, - stride, - padding, - dilation, - groups, - output_scale, - output_zero_point); - } -#endif - TORCH_CHECK( - false, - "Didn't find engine for operation quantized::conv ", - toString(ctx.qEngine())); - } }; static auto registry = c10::RegisterOperators() .op("quantized::conv2d", - c10::RegisterOperators::options().kernel>( + c10::RegisterOperators::options().kernel>( TensorTypeId::QuantizedCPUTensorId)) .op("quantized::conv2d_relu", - c10::RegisterOperators::options().kernel>( + c10::RegisterOperators::options().kernel>( + TensorTypeId::QuantizedCPUTensorId)) + .op("quantized::conv3d", + c10::RegisterOperators::options().kernel>( + TensorTypeId::QuantizedCPUTensorId)) + .op("quantized::conv3d_relu", + c10::RegisterOperators::options().kernel>( TensorTypeId::QuantizedCPUTensorId)); } // namespace diff --git a/aten/src/ATen/test/CMakeLists.txt b/aten/src/ATen/test/CMakeLists.txt index 373c3f9205308..e9647f15a4ef8 100644 --- a/aten/src/ATen/test/CMakeLists.txt +++ b/aten/src/ATen/test/CMakeLists.txt @@ -31,7 +31,8 @@ list(APPEND ATen_CPU_TEST_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/memory_overlapping_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpu_generator_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/pow_test.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/variant_test.cpp) + ${CMAKE_CURRENT_SOURCE_DIR}/variant_test.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/reduce_ops_test.cpp) list(APPEND ATen_CUDA_TEST_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/cuda_integer_divider_test.cu diff --git a/aten/src/ATen/test/reduce_ops_test.cpp b/aten/src/ATen/test/reduce_ops_test.cpp new file mode 100644 index 0000000000000..429761b563a07 --- /dev/null +++ b/aten/src/ATen/test/reduce_ops_test.cpp @@ -0,0 +1,24 @@ +#include + +#include +#include + +using namespace at; + +TEST(ReduceOpsTest, MaxValuesAndMinValues) { + const int W = 10; + const int H = 10; + if (hasCUDA()) { + for (const auto dtype : {kHalf, kFloat, kDouble, kShort, kInt, kLong}) { + auto a = at::rand({H, W}, TensorOptions(kCUDA).dtype(at::kHalf)); + ASSERT_FLOAT_EQ( + a.max_values(c10::IntArrayRef{0, 1}).item(), + a.max().item() + ); + ASSERT_FLOAT_EQ( + a.min_values(c10::IntArrayRef{0, 1}).item(), + a.min().item() + ); + } + } +} diff --git a/benchmarks/operator_benchmark/benchmark_core.py b/benchmarks/operator_benchmark/benchmark_core.py index 7a3d15237d55a..14bd5a2094c30 100644 --- a/benchmarks/operator_benchmark/benchmark_core.py +++ b/benchmarks/operator_benchmark/benchmark_core.py @@ -84,9 +84,6 @@ def __init__(self, args): if self.args.test_name is not None: self.args.tag_filter = None - if self.args.ai_pep_format: - self.print_per_iter = True - def _print_header(self): DASH_LINE = '-' * 40 @@ -199,10 +196,18 @@ def _measure_time(self, launch_test, test_case, iters, print_per_iter): report_run_time = 1e6 * run_time_sec / iters time_trace.append(report_run_time) + # Print out the time spent in each epoch in ms + if self.args.ai_pep_format: + test_name = '_'.join([test_case.framework, test_case.test_config.test_name]) + print("PyTorchObserver " + json.dumps( + { + "type": test_name, + "metric": "latency", + "unit": "ms", + "value": str(report_run_time / 1e3), + } + )) if results_are_significant: - # Print out the last 50 values when running with AI PEP - if self.args.ai_pep_format: - test_case._print_per_iter() break # Re-estimate the hopefully-sufficient diff --git a/benchmarks/operator_benchmark/benchmark_pytorch.py b/benchmarks/operator_benchmark/benchmark_pytorch.py index 8edbd1090c2ff..7ec3914decac4 100644 --- a/benchmarks/operator_benchmark/benchmark_pytorch.py +++ b/benchmarks/operator_benchmark/benchmark_pytorch.py @@ -76,7 +76,7 @@ def _generate_jit_forward_graph(self): @torch.jit.script def _jit_forward_graph(iters, place_holder): # type: (int, Tensor) - result = torch.jit.annotate(torch.Tensor, None) + result = torch.jit.annotate(torch.Tensor, place_holder) for _ in range(iters): result = func(place_holder) return result diff --git a/benchmarks/operator_benchmark/pt/softmax_test.py b/benchmarks/operator_benchmark/pt/softmax_test.py index 8b4f4e6c4ac27..9d20455d49245 100644 --- a/benchmarks/operator_benchmark/pt/softmax_test.py +++ b/benchmarks/operator_benchmark/pt/softmax_test.py @@ -17,8 +17,8 @@ # Configs for softmax ops softmax_configs_short = op_bench.config_list( attrs=[ - [4, 3, 128, 128], - [8, 3, 256, 256], + [4, 3, 256, 256], + [8, 3, 512, 512], ], attr_names=[ 'N', 'C', 'H', 'W' diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index d33a8c532a864..032902e78d422 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -386,6 +386,7 @@ if (NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) ${TORCH_SRC_DIR}/csrc/jit/passes/batch_mm.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/bailout_graph.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/canonicalize.cpp + ${TORCH_SRC_DIR}/csrc/jit/passes/clear_undefinedness.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/constant_propagation.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/constant_pooling.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/common_subexpression_elimination.cpp @@ -575,6 +576,7 @@ if (NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/pixelshuffle.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/pooling.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/rnn.cpp + ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/upsampling.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/container/functional.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/container/named_any.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/options/activation.cpp diff --git a/caffe2/onnx/offline_tensor.cc b/caffe2/onnx/offline_tensor.cc new file mode 100644 index 0000000000000..a80f2bbd49b62 --- /dev/null +++ b/caffe2/onnx/offline_tensor.cc @@ -0,0 +1,88 @@ +#include "caffe2/onnx/offline_tensor.h" + +namespace caffe2 { + +#ifndef C10_MOBILE +namespace { +// These constants need to be aligned with onnxifi.h +constexpr uint64_t kONNXIFI_DATATYPE_FLOAT16 = 10; +constexpr uint64_t kONNXIFI_DATATYPE_FLOAT32 = 1; +constexpr uint64_t kONNXIFI_DATATYPE_UINT8 = 2; +constexpr uint64_t kONNXIFI_DATATYPE_INT32 = 6; +constexpr uint64_t kONNXIFI_DATATYPE_INT8 = 3; +constexpr uint64_t kONNXIFI_DATATYPE_INT64 = 7; +constexpr uint64_t kONNXIFI_DATATYPE_INT16 = 5; +constexpr uint64_t kONNXIFI_DATATYPE_UINT16 = 4; +} // namespace + +CAFFE_KNOWN_TYPE(OfflineTensor); + +bool OfflineTensorShapeFunctions::IsSameMetaType(TypeIdentifier id) { + return id == TypeMeta::Id(); +} + +TypeIdentifier OfflineTensorShapeFunctions::GetTypeMetaId() { + return TypeMeta::Id(); +} + +TypeMeta OfflineTensorShapeFunctions::GetExternalTensorType(const void* c) { + const OfflineTensor* offline_tensor = + reinterpret_cast(c); + + return offline_tensor->shape_tensor.dtype(); +} + +vector OfflineTensorShapeFunctions::GetExternalTensorInfo( + const void* c, + size_t* capacity, + DeviceOption* device) { + const OfflineTensor* offline_tensor = + reinterpret_cast(c); + return GetTensorInfo(&(offline_tensor->shape_tensor), capacity, device); +} + +void OfflineTensorShapeFunctions::SetupExternalTensorDescriptor( + const Blob* blob, + std::vector>* shapes, + std::vector>* /* unused */, + std::vector>* /* unused */, + ExternalTensorDescriptor* desc) { + const auto& offline_tensor = blob->template Get(); + const Tensor& shape_tensor = offline_tensor.shape_tensor; + + if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_FLOAT32; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT32; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT8; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_UINT8; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT64; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT16; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_FLOAT16; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_UINT16; + } else { + CAFFE_THROW("Unsupported tensor type: ", shape_tensor.dtype().name()); + } + desc->buffer = 0; + + desc->quantizationParams = 0; + desc->quantizationAxis = 0; + + // Set up dim and shape + const auto shape = shape_tensor.sizes(); + desc->dimensions = shape.size(); + shapes->emplace_back(shape.cbegin(), shape.cend()); + desc->shape = shapes->back().data(); +} + +REGISTER_EXTERNAL_TENSOR_FUNCTIONS( + (TypeMeta::Id()), + OfflineTensorShapeFunctions); +#endif +} // namespace caffe2 diff --git a/caffe2/onnx/offline_tensor.h b/caffe2/onnx/offline_tensor.h new file mode 100644 index 0000000000000..73287630ee6a3 --- /dev/null +++ b/caffe2/onnx/offline_tensor.h @@ -0,0 +1,53 @@ +#pragma once + +#include +#include "caffe2/core/operator.h" +#include "caffe2/core/tensor.h" + +namespace caffe2 { + +#ifndef C10_MOBILE +struct OfflineTensor { + // A shell tensor to record shape and dtype + Tensor shape_tensor{CPU}; + + void setShapeAndType( + const std::vector& sizes, + at::Device device, + caffe2::TypeMeta data_type) { + shape_tensor.unsafeGetTensorImpl()->set_storage( + at::Storage::create_legacy(device, data_type)); + shape_tensor.Resize(sizes); + CHECK(!shape_tensor.storage_initialized()); + CHECK(shape_tensor.dtype_initialized()); + } +}; + +class OfflineTensorShapeFunctions : public ExternalTensorFunctionsBase { + public: + explicit OfflineTensorShapeFunctions() : ExternalTensorFunctionsBase() {} + ~OfflineTensorShapeFunctions() override {} + bool isQuantized() const override { + return false; + } + bool IsSameMetaType(TypeIdentifier id) override; + void SetupExternalTensorDescriptor( + const Blob* blob, + std::vector>* shapes, + std::vector>* all_scales, + std::vector>* all_offsets, + ExternalTensorDescriptor* desc) override; + void LoadInfoOfBlob( + const Blob* /* unused */, + std::vector* /* unused */, + std::vector* /* unused */, + uint32_t* /* unused */) override {} + TypeIdentifier GetTypeMetaId() override; + TypeMeta GetExternalTensorType(const void* c) override; + vector GetExternalTensorInfo( + const void* c, + size_t* capacity, + DeviceOption* device) override; +}; +#endif +} // namespace caffe2 diff --git a/caffe2/operators/pack_segments.cc b/caffe2/operators/pack_segments.cc index 8e82ed0a6a86f..ee65d5f7373d6 100644 --- a/caffe2/operators/pack_segments.cc +++ b/caffe2/operators/pack_segments.cc @@ -33,11 +33,6 @@ bool PackSegmentsOp::DoRunWithType2() { total_length += l[i]; } if (max_length_ != -1) { - // Final dim must be greater than the max_length - CAFFE_ENFORCE_GE( - max_length_, - max_length, - "Pre-defined max_length should be greater than the real max_length"); max_length = max_length_; } @@ -89,13 +84,14 @@ bool PackSegmentsOp::DoRunWithType2() { const auto* d = static_cast(data.raw_data()); int64_t start = 0; for (int64_t i = 0; i < lengths.size(0); ++i) { + auto len = l[i] <= max_length ? l[i] : max_length; context_.CopyItemsSameDevice( data.dtype(), - l[i] * block_size, + len * block_size, d + block_bytesize * start, out + block_bytesize * max_length * i); if (return_presence_mask_) { - memset(presence_mask_data + max_length * i, (int)true, l[i]); + memset(presence_mask_data + max_length * i, (int)true, len); } start += l[i]; } @@ -128,7 +124,14 @@ bool UnpackSegmentsOp::DoRunWithType2() { } const T* l = lengths.template data(); - int64_t total_l = std::accumulate(l, l + lengths.size(0), (int64_t)0); + int64_t total_l = 0; + if (max_length_ != -1) { + for (int64_t i = 0; i < lengths.size(0); ++i) { + total_l += (int64_t)(l[i] <= max_length_ ? l[i] : max_length_); + } + } else { + total_l = std::accumulate(l, l + lengths.size(0), (int64_t)0); + } auto shape = data.sizes().vec(); CAFFE_ENFORCE_EQ( @@ -146,12 +149,16 @@ bool UnpackSegmentsOp::DoRunWithType2() { const auto* d = static_cast(data.raw_data()); int64_t start = 0; for (int64_t i = 0; i < lengths.size(0); ++i) { + auto len = l[i]; + if (max_length_ != -1 && l[i] > max_length_) { + len = max_length_; + } context_.CopyItemsSameDevice( data.dtype(), - l[i] * block_size, + len * block_size, d + block_bytesize * data.size(1) * i, out + block_bytesize * start); - start += l[i]; + start += len; } return true; } diff --git a/caffe2/opt/custom/converter.cc b/caffe2/opt/custom/converter.cc index 37afc2dc731ba..f52bf18880829 100644 --- a/caffe2/opt/custom/converter.cc +++ b/caffe2/opt/custom/converter.cc @@ -214,6 +214,10 @@ class ClipRangesGatherSigridHashConverter : public Converter { if (args.HasArgument("max_values")) { c->setMaxValues(args.GetRepeatedArgument("max_values")); } + if (args.HasArgument("hash_into_int32")) { + c->setHashIntoInt32( + args.GetSingleArgument("hash_into_int32", false)); + } return nnOp; } @@ -230,6 +234,8 @@ class ClipRangesGatherSigridHashConverter : public Converter { caffe2::MakeArgument>("salts", fuse->getSalts())); op.add_arg()->CopyFrom(caffe2::MakeArgument>( "max_values", fuse->getMaxValues())); + op.add_arg()->CopyFrom(caffe2::MakeArgument( + "hash_into_int32", fuse->getHashIntoInt32())); return op; } diff --git a/caffe2/opt/custom/converter_test.cc b/caffe2/opt/custom/converter_test.cc new file mode 100644 index 0000000000000..4313c3a14087f --- /dev/null +++ b/caffe2/opt/custom/converter_test.cc @@ -0,0 +1,33 @@ +#include "caffe2/core/common.h" +#include "caffe2/core/test_utils.h" +#include "caffe2/opt/converter.h" +#include "caffe2/opt/custom/concat_elim.h" +#include "caffe2/predictor/emulator/data_filler.h" +#include "caffe2/utils/proto_utils.h" + +#include + +using namespace caffe2::testing; +using namespace caffe2::emulator; +using caffe2::OperatorDef; +using std::vector; + +TEST(Converter, ClipRangesGatherSigridHashConverter) { + OperatorDef op; + op.set_type("ClipRangesGatherSigridHash"); + op.add_arg()->CopyFrom(caffe2::MakeArgument("hash_into_int32", true)); + auto nnDef = convertToNeuralNetOperator(op); + auto* pNNDef = + static_cast(nnDef.get()); + EXPECT_TRUE(pNNDef); + EXPECT_TRUE(pNNDef->getHashIntoInt32()); + + OperatorDef op2; + op2.set_type("ClipRangesGatherSigridHash"); + op2.add_arg()->CopyFrom(caffe2::MakeArgument("hash_into_int32", false)); + auto nnDef2 = convertToNeuralNetOperator(op2); + auto* pNNDef2 = + static_cast(nnDef2.get()); + EXPECT_TRUE(pNNDef2); + EXPECT_FALSE(pNNDef2->getHashIntoInt32()); +} diff --git a/caffe2/python/operator_test/pack_ops_test.py b/caffe2/python/operator_test/pack_ops_test.py index b5ee13740e45d..5c83566ed07a5 100644 --- a/caffe2/python/operator_test/pack_ops_test.py +++ b/caffe2/python/operator_test/pack_ops_test.py @@ -21,13 +21,13 @@ def pack_segments_ref(lengths, data, max_length=max_length): constant_values = 0 if data.dtype.char == 'S': constant_values = '' - if max_length is not None: - assert(max_length > np.max(lengths)) - else: + if max_length is None: max_length = np.max(lengths) + start = 0 for idx in range(np.size(lengths)): - chunk = data[np.sum(lengths[:idx]):np.sum(lengths[:idx + 1])] - pad_length = max_length - lengths[idx] + len = lengths[idx] if max_length >= lengths[idx] else max_length + chunk = data[start : start + len] + pad_length = max_length - len # ((0, pad_length), (0, 0)) says add pad_length rows of padding # below chunk and 0 rows of padding elsewhere @@ -38,10 +38,12 @@ def pack_segments_ref(lengths, data, max_length=max_length): constant_values=constant_values ) ) + start += lengths[idx] result = [arr] if return_presence_mask: presence_arr = [] for length in lengths: + length = length if max_length >= length else max_length pad_length = max_length - length presence_arr.append( np.pad( @@ -57,9 +59,12 @@ def pack_segments_ref(lengths, data, max_length=max_length): @serial.given( num_seq=st.integers(10, 100), cell_size=st.integers(1, 10), + max_length_buffer=st.integers(-5, 5), **hu.gcs ) - def test_pack_with_max_length_ops(self, num_seq, cell_size, gc, dc): + def test_pack_with_max_length_ops( + self, num_seq, cell_size, max_length_buffer, gc, dc + ): # create data lengths = np.arange(num_seq, dtype=np.int32) + 1 num_cell = np.sum(lengths) @@ -74,7 +79,7 @@ def test_pack_with_max_length_ops(self, num_seq, cell_size, gc, dc): + "=" * 60 ) # run test - max_length = num_seq + 1 + max_length = num_seq + max_length_buffer op = core.CreateOperator( 'PackSegments', ['l', 'd'], ['t'], max_length=max_length) workspace.FeedBlob('l', lengths) @@ -105,7 +110,24 @@ def test_pack_with_max_length_ops(self, num_seq, cell_size, gc, dc): max_length=max_length, device_option=gc)) assert(workspace.FetchBlob('t').shape[1] == max_length) - assert((workspace.FetchBlob('newd') == workspace.FetchBlob('d')).all()) + + def _cal_unpacked_data(data): + if max_length >= num_seq: + return data + output = None + start = 0 + for i, length in enumerate(lengths): + new_len = max_length if length > max_length else length + chunk = data[start: start + new_len] + if output is None: + output = chunk + else: + output = np.concatenate((output, chunk), axis=0) + start += length + return output + + true_newd = _cal_unpacked_data(workspace.FetchBlob('d')) + assert((workspace.FetchBlob('newd') == true_newd).all()) @given( num_seq=st.integers(10, 500), diff --git a/caffe2/quantization/server/dnnlowp.h b/caffe2/quantization/server/dnnlowp.h index ae6fd7e81170c..0d7414f962412 100644 --- a/caffe2/quantization/server/dnnlowp.h +++ b/caffe2/quantization/server/dnnlowp.h @@ -148,6 +148,13 @@ class QuantizationFactory { return weight_kind_; } + void SetWeightP99Threshold(float threshold) { + weight_p99_threshold_ = threshold; + } + void SetActivationP99Threshold(float threshold) { + activation_p99_threshold_ = threshold; + } + explicit QuantizationFactory( int activation_precision = 8, // precision used for activations in main operations like matmul diff --git a/caffe2/quantization/server/p99.cc b/caffe2/quantization/server/p99.cc index 9c3aac7b48a58..e4d5eb48e6cfc 100644 --- a/caffe2/quantization/server/p99.cc +++ b/caffe2/quantization/server/p99.cc @@ -14,8 +14,8 @@ TensorQuantizationParams P99::ChooseQuantizationParams( std::vector bins_f( dnnlowp::adjust_hist_to_include_zero(hist, &min, &max)); int nbins = bins_f.size(); - assert(min <= 0.f); - assert(max >= 0.f); + CAFFE_ENFORCE(min <= 0.f); + CAFFE_ENFORCE(max >= 0.f); float org_max = max; float org_min = min; float bin_width = (max - min) / nbins; @@ -32,7 +32,7 @@ TensorQuantizationParams P99::ChooseQuantizationParams( sum += bins_f[i]; CDF[i] = (double)sum / total_sum; } - assert(threshold_ > 0.5 && threshold_ < 1); + CAFFE_ENFORCE(threshold_ > 0.5 && threshold_ < 1); double left_quantile = (1.0f - threshold_) / 2.0f; double right_quantile = 1.0f - left_quantile; int i_begin = 0; diff --git a/caffe2/quantization/server/pybind.cc b/caffe2/quantization/server/pybind.cc index 4709be4c0e2fa..3c2cad0c4991e 100644 --- a/caffe2/quantization/server/pybind.cc +++ b/caffe2/quantization/server/pybind.cc @@ -1,6 +1,8 @@ #include +#include #include "activation_distribution_observer.h" #include "caffe2_dnnlowp_utils.h" +#include "quantization_error_minimization.h" namespace caffe2 { namespace python { @@ -172,4 +174,71 @@ PYBIND11_MODULE(dnnlowp_pybind11, m) { CAFFE_ENFORCE(transformed_net.SerializeToString(&protob)); return pybind11::bytes(protob); }); + + pybind11::class_(m, "QueryTensorQparam") + .def_property_readonly( + "scale", + [](dnnlowp::TensorQuantizationParams& qparam) { + return qparam.scale; + }) + .def_property_readonly( + "zero_point", + [](dnnlowp::TensorQuantizationParams& qparam) { + return qparam.zero_point; + }) + .def_property_readonly( + "min", + [](dnnlowp::TensorQuantizationParams& qparam) { + return qparam.Min(); + }) + .def_property_readonly( + "max", [](dnnlowp::TensorQuantizationParams& qparam) { + return qparam.Max(); + }); + + m.def( + "ChooseStaticQuantizationParams", + [](float min, + float max, + const std::vector& bins, + bool preserve_sparsity, + int precision, + const std::string& quant_scheme, + float p99_threshold, + bool is_weight) { + dnnlowp::Histogram hist = dnnlowp::Histogram(min, max, bins); + + dnnlowp::QuantizationFactory::QuantizationKind quant_kind = + dnnlowp::QuantizationFactory::MIN_MAX_QUANTIZATION; + if (quant_scheme.compare("L2_MIN_QUANTIZATION") == 0) { + quant_kind = dnnlowp::QuantizationFactory::L2_MIN_QUANTIZATION; + } else if (quant_scheme.compare("L2_MIN_QUANTIZATION_APPROX") == 0) { + quant_kind = dnnlowp::QuantizationFactory::L2_MIN_QUANTIZATION_APPROX; + } else if (quant_scheme.compare("KL_MIN_QUANTIZATION") == 0) { + quant_kind = dnnlowp::QuantizationFactory::KL_MIN_QUANTIZATION; + } else if (quant_scheme.compare("P99_QUANTIZATION") == 0) { + quant_kind = dnnlowp::QuantizationFactory::P99_QUANTIZATION; + } else if (quant_scheme.compare("L1_MIN_QUANTIZATION") == 0) { + quant_kind = dnnlowp::QuantizationFactory::L1_MIN_QUANTIZATION; + } else { + LOG(INFO) << "Using DNNLOWP default MIN_MAX_QUANTIZATION"; + } + dnnlowp::QuantizationFactory* qfactory = + dnnlowp::QuantizationFactory::GetDefaultInstance(); + if (is_weight) { + qfactory->SetWeightP99Threshold(p99_threshold); + } else { + qfactory->SetActivationP99Threshold(p99_threshold); + } + return qfactory->ChooseQuantizationParams( + hist, quant_kind, precision, preserve_sparsity, is_weight); + }, + pybind11::arg("min"), + pybind11::arg("max"), + pybind11::arg("bins"), + pybind11::arg("preserve_sparsity") = true, + pybind11::arg("precision") = 8, + pybind11::arg("quant_scheme") = "min_max", + pybind11::arg("p99_threshold") = 0.99, + pybind11::arg("is_weight") = false); } diff --git a/ios/TestApp/.gitignore b/ios/TestApp/.gitignore index 1038da3961c8f..eaf0cd27b306e 100644 --- a/ios/TestApp/.gitignore +++ b/ios/TestApp/.gitignore @@ -1 +1,2 @@ model.pt +.config diff --git a/ios/TestApp/README.md b/ios/TestApp/README.md index 28aa88b7a9f77..ce2842d28016a 100644 --- a/ios/TestApp/README.md +++ b/ios/TestApp/README.md @@ -21,9 +21,10 @@ The TestApp is currently being used as a dummy app by Circle CI for nightly jobs The benchmark folder contains two scripts that help you setup the benchmark project. The `setup.rb` does the heavy-lifting jobs of setting up the XCode project, whereas the `trace_model.py` is a Python script that you can tweak to generate your model for benchmarking. Simply follow the steps below to setup the project 1. In the PyTorch root directory, run `BUILD_PYTORCH_MOBILE=1 IOS_ARCH=arm64 ./scripts/build_ios.sh` to generate the custom build from **Master** branch -2. Navigate to the `benchmark` folder, run `python trace_model.py` to get your model generated. -3. In the same directory, run `ruby setup.rb` to setup the XCode project. -4. Open the `TestApp.xcodeproj`, you're ready to go. +2. Navigate to the `benchmark` folder, run `python trace_model.py` to generate your model. +3. In the same directory, open `config.json`. Those are the input parameters you can tweak. +4. Again, in the same directory, run `ruby setup.rb` to setup the XCode project. +5. Open the `TestApp.xcodeproj`, you're ready to go. The benchmark code is written in C++, see `benchmark.mm` for more details. @@ -35,13 +36,13 @@ For those who want to do perf testing but don't want touch XCode, `bootstrap.sh` 2. A valid provisioning profile for code signing 3. A valid team identifier -To run the script, simply type the command below and make sure your phone is unlocked and connected via USB. +To run the script, simply type the command below and make sure your phone is connected via USB. ```shell -./bootstrap -t ${TEAM_ID} -p ${PROVISIONING_PROFILE} +./bootstrap ``` -The benchmark log will be displayed on the screen. +Open the app on your device, the benchmark result will be displayed on the screen. > Note This requires ios-deploy to be installed. Please have a look at [ios-deploy](https://github.com/ios-control/ios-deploy). To quickly install it, use `npm -g i ios-deploy` diff --git a/ios/TestApp/TestApp/Base.lproj/Main.storyboard b/ios/TestApp/TestApp/Base.lproj/Main.storyboard index c78975d739870..31dd0f40e5673 100644 --- a/ios/TestApp/TestApp/Base.lproj/Main.storyboard +++ b/ios/TestApp/TestApp/Base.lproj/Main.storyboard @@ -33,7 +33,13 @@ - + + + + + + + diff --git a/ios/TestApp/TestApp/Benchmark.mm b/ios/TestApp/TestApp/Benchmark.mm index cadcc421a2891..b8cabc0326596 100644 --- a/ios/TestApp/TestApp/Benchmark.mm +++ b/ios/TestApp/TestApp/Benchmark.mm @@ -77,7 +77,6 @@ + (NSString*)run { if (print_output) { std::cout << module.forward(inputs) << std::endl; } - UI_LOG(@"Start benchmarking...", nil); UI_LOG(@"Running warmup runs", nil); CAFFE_ENFORCE(warmup >= 0, "Number of warm up runs should be non negative, provided ", warmup, "."); diff --git a/ios/TestApp/TestApp/ViewController.mm b/ios/TestApp/TestApp/ViewController.mm index 4e4f0dcb884f4..97f5e69b439fc 100644 --- a/ios/TestApp/TestApp/ViewController.mm +++ b/ios/TestApp/TestApp/ViewController.mm @@ -7,7 +7,8 @@ @interface ViewController () @end -@implementation ViewController +@implementation ViewController { +} - (void)viewDidLoad { [super viewDidLoad]; @@ -18,20 +19,29 @@ - (void)viewDidLoad { NSDictionary* config = [NSJSONSerialization JSONObjectWithData:configData options:NSJSONReadingAllowFragments error:&err]; + if (err) { NSLog(@"Parse config.json failed!"); return; } + [Benchmark setup:config]; + [self runBenchmark]; +} +- (void)runBenchmark { + self.textView.text = @"Start benchmarking...\n"; dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ - if ([Benchmark setup:config]) { - NSString* text = [Benchmark run]; - dispatch_async(dispatch_get_main_queue(), ^{ - self.textView.text = text; - }); - } else { - NSLog(@"Setup benchmark config failed!"); - } + NSString* text = [Benchmark run]; + dispatch_async(dispatch_get_main_queue(), ^{ + self.textView.text = [self.textView.text stringByAppendingString:text]; + }); + }); +} + +- (IBAction)reRun:(id)sender { + self.textView.text = @""; + dispatch_async(dispatch_get_main_queue(), ^{ + [self runBenchmark]; }); } diff --git a/ios/TestApp/bootstrap.sh b/ios/TestApp/bootstrap.sh index 3fdb9721e867c..9785bc18fe923 100755 --- a/ios/TestApp/bootstrap.sh +++ b/ios/TestApp/bootstrap.sh @@ -22,14 +22,30 @@ bootstrap() { XCODE_PROJ_PATH="./TestApp.xcodeproj" XCODE_TARGET="TestApp" XCODE_BUILD="./build" - if [ -d ${XCODE_BUILD} ]; then + if [ ! -f "./.config" ]; then + touch .config + echo "" >> .config + else + source .config + fi + if [ -z "${TEAM_ID}" ]; then + reply=$(bash -c 'read -r -p "Team Id:" tmp; echo $tmp') + TEAM_ID="${reply}" + echo "TEAM_ID=${TEAM_ID}" >> .config + fi + if [ -z "${PROFILE}" ]; then + reply=$(bash -c 'read -r -p "Provisioning Profile:" tmp; echo $tmp') + PROFILE="${reply}" + echo "PROFILE=${PROFILE}" >> .config + fi + if [ -d "${XCODE_BUILD}" ]; then echo "found the old XCode build, remove it" - rm -rf ${XCODE_BUILD} + rm -rf "${XCODE_BUILD}" fi - cd ${BENCHMARK_DIR} + cd "${BENCHMARK_DIR}" echo "Generating model" python trace_model.py - ruby setup.rb -t ${TEAM_ID} + ruby setup.rb -t "${TEAM_ID}" cd .. #run xcodebuild if ! [ -x "$(command -v xcodebuild)" ]; then @@ -77,7 +93,4 @@ esac shift done -echo TEAM_ID = "${TEAM_ID}" -echo PROFILE = "${PROFILE}" - bootstrap diff --git a/test/backward_compatibility/check_backward_compatibility.py b/test/backward_compatibility/check_backward_compatibility.py index bec599e906a84..a3742c316895b 100644 --- a/test/backward_compatibility/check_backward_compatibility.py +++ b/test/backward_compatibility/check_backward_compatibility.py @@ -23,6 +23,10 @@ ('thnn_conv_depthwise2d_backward', datetime.date(2019, 10, 30)), ('thnn_conv3d_backward', datetime.date(2019, 10, 30)), ('empty_like', datetime.date(2019, 10, 30)), + ('rand_like', datetime.date(2019, 11, 11)), + ('ones_like', datetime.date(2019, 11, 11)), + ('full_like', datetime.date(2019, 11, 11)), + ('AutogradAnyNonZero', datetime.date(2019, 11, 11)), ] diff --git a/test/common_utils.py b/test/common_utils.py index 8c04fbcefd5dd..de04bdb542b5b 100644 --- a/test/common_utils.py +++ b/test/common_utils.py @@ -462,7 +462,7 @@ def __exit__(self, exec_type, exec_value, traceback): suppress_health_check=[hypothesis.HealthCheck.too_slow], database=None, max_examples=100, - verbosity=hypothesis.Verbosity.quiet)) + verbosity=hypothesis.Verbosity.normal)) hypothesis.settings.register_profile( "dev", hypothesis.settings( @@ -486,7 +486,7 @@ def __exit__(self, exec_type, exec_value, traceback): database=None, max_examples=100, min_satisfying_examples=1, - verbosity=hypothesis.Verbosity.quiet)) + verbosity=hypothesis.Verbosity.normal)) hypothesis.settings.register_profile( "dev", hypothesis.settings( diff --git a/test/cpp/api/enum.cpp b/test/cpp/api/enum.cpp index 3872875cfc9c2..8f6c36c9f543c 100644 --- a/test/cpp/api/enum.cpp +++ b/test/cpp/api/enum.cpp @@ -1,6 +1,5 @@ #include -#include #include #include @@ -10,7 +9,7 @@ v = torch::k##name; \ std::string pretty_print_name("k"); \ pretty_print_name.append(#name); \ - ASSERT_EQ(c10::visit(torch::enumtype::enum_name{}, v), pretty_print_name); \ + ASSERT_EQ(torch::enumtype::get_enum_name(v), pretty_print_name); \ } TEST(EnumTest, AllEnums) { @@ -32,9 +31,16 @@ TEST(EnumTest, AllEnums) { torch::enumtype::kReflect, torch::enumtype::kReplicate, torch::enumtype::kCircular, + torch::enumtype::kNearest, + torch::enumtype::kBilinear, + torch::enumtype::kBicubic, + torch::enumtype::kTrilinear, + torch::enumtype::kArea, torch::enumtype::kSum, torch::enumtype::kMean, - torch::enumtype::kMax + torch::enumtype::kMax, + torch::enumtype::kNone, + torch::enumtype::kBatchMean > v; TORCH_ENUM_PRETTY_PRINT_TEST(Linear) @@ -54,7 +60,14 @@ TEST(EnumTest, AllEnums) { TORCH_ENUM_PRETTY_PRINT_TEST(Reflect) TORCH_ENUM_PRETTY_PRINT_TEST(Replicate) TORCH_ENUM_PRETTY_PRINT_TEST(Circular) + TORCH_ENUM_PRETTY_PRINT_TEST(Nearest) + TORCH_ENUM_PRETTY_PRINT_TEST(Bilinear) + TORCH_ENUM_PRETTY_PRINT_TEST(Bicubic) + TORCH_ENUM_PRETTY_PRINT_TEST(Trilinear) + TORCH_ENUM_PRETTY_PRINT_TEST(Area) TORCH_ENUM_PRETTY_PRINT_TEST(Sum) TORCH_ENUM_PRETTY_PRINT_TEST(Mean) TORCH_ENUM_PRETTY_PRINT_TEST(Max) + TORCH_ENUM_PRETTY_PRINT_TEST(None) + TORCH_ENUM_PRETTY_PRINT_TEST(BatchMean) } diff --git a/test/cpp/api/functional.cpp b/test/cpp/api/functional.cpp index 2273775ce1820..c058fbd0e645a 100644 --- a/test/cpp/api/functional.cpp +++ b/test/cpp/api/functional.cpp @@ -131,7 +131,7 @@ TEST_F(FunctionalTest, SoftMarginLossNoReduction) { auto input = torch::tensor({2., 4., 1., 3.}, torch::requires_grad()); auto target = torch::tensor({-1., 1., 1., -1.}, torch::kFloat); auto output = - F::soft_margin_loss(input, target, torch::Reduction::None); + F::soft_margin_loss(input, target, torch::kNone); auto expected = torch::tensor({2.1269281, 0.01814993, 0.3132617, 3.0485873}, torch::kFloat); auto s = output.sum(); s.backward(); @@ -144,7 +144,7 @@ TEST_F(FunctionalTest, MultiLabelSoftMarginLossWeightedNoReduction) { auto input = torch::tensor({{0., 2., 2., 0.}, {2., 1., 0., 1.}}, torch::requires_grad()); auto target = torch::tensor({{0., 0., 1., 0.}, {1., 0., 1., 1.}}, torch::kFloat); auto weight = torch::tensor({0.1, 0.6, 0.4, 0.8}, torch::kFloat); - auto options = MultiLabelSoftMarginLossOptions().reduction(torch::Reduction::None).weight(weight); + auto options = MultiLabelSoftMarginLossOptions().reduction(torch::kNone).weight(weight); auto output = F::multilabel_soft_margin_loss(input, target, options); auto expected = torch::tensor({0.4876902, 0.3321295}, torch::kFloat); @@ -491,7 +491,7 @@ TEST_F(FunctionalTest, MultiLabelMarginLossNoReduction) { auto input = torch::tensor({{0.1, 0.2, 0.4, 0.8}}, torch::requires_grad()); auto target = torch::tensor({{3, 0, -1, 1}}, torch::kLong); auto output = F::multilabel_margin_loss( - input, target, torch::Reduction::None); + input, target, torch::kNone); auto expected = torch::tensor({0.8500}, torch::kFloat); auto s = output.sum(); s.backward(); @@ -1266,6 +1266,119 @@ TEST_F(FunctionalTest, Threshold) { } } +TEST_F(FunctionalTest, BatchNorm1d) { + int num_features = 5; + double eps = 1e-05; + double momentum = 0.1; + + auto input = torch::randn({2, 5}); + auto mean = torch::randn(5); + auto variance = torch::rand(5); + auto weight = torch::ones({num_features}); + auto bias = torch::zeros({num_features}); + auto output = F::batch_norm( + input, mean, variance, + BatchNormOptions().weight(weight).bias(bias).momentum(momentum).eps(eps), + /*training=*/false); + auto expected = (input - mean) / torch::sqrt(variance + eps); + ASSERT_TRUE(output.allclose(expected)); +} + +TEST_F(FunctionalTest, BatchNorm1dDefaultOptions) { + auto input = torch::randn({2, 5}); + auto mean = torch::randn(5); + auto variance = torch::rand(5); + auto output = F::batch_norm(input, mean, variance); + auto expected = (input - mean) / torch::sqrt(variance + 1e-5); + ASSERT_TRUE(output.allclose(expected)); +} + +TEST_F(FunctionalTest, Interpolate) { + { + // 1D interpolation + auto input = torch::ones({1, 1, 2}); + auto options = InterpolateOptions() + .size({4}) + .mode(torch::kNearest); + auto output = F::interpolate(input, options); + auto expected = torch::ones({1, 1, 4}); + + ASSERT_TRUE(output.allclose(expected)); + } + { + // 2D interpolation + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + auto input = torch::ones({1, 1, 2, 2}); + auto options = InterpolateOptions() + .scale_factor({scale_factor, scale_factor}) + .mode(torch::kBilinear) + .align_corners(align_corners); + auto output = F::interpolate(input, options); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = torch::ones({1, 1, expected_size, expected_size}); + + ASSERT_TRUE(output.allclose(expected)); + } + } + } + { + // 3D interpolation + for (const auto align_corners : {true, false}) { + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + auto input = torch::ones({1, 1, 2, 2, 2}); + auto options = + InterpolateOptions() + .scale_factor({scale_factor, scale_factor, scale_factor}) + .mode(torch::kTrilinear) + .align_corners(align_corners); + auto output = F::interpolate(input, options); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = + torch::ones({1, 1, expected_size, expected_size, expected_size}); + + ASSERT_TRUE(output.allclose(expected)); + } + } + } + { + auto input = torch::randn({3, 2, 2}); + ASSERT_THROWS_WITH( + F::interpolate(input[0], InterpolateOptions().size({4, 4})), + "Input Error: Only 3D, 4D and 5D input Tensors supported (got 2D) " + "for the modes: nearest | linear | bilinear | bicubic | trilinear (got kNearest)"); + ASSERT_THROWS_WITH( + F::interpolate( + torch::reshape(input, {1, 1, 1, 3, 2, 2}), + InterpolateOptions().size({1, 1, 1, 3, 4, 4})), + "Input Error: Only 3D, 4D and 5D input Tensors supported (got 6D) " + "for the modes: nearest | linear | bilinear | bicubic | trilinear (got kNearest)"); + ASSERT_THROWS_WITH( + F::interpolate(input, InterpolateOptions()), + "either size or scale_factor should be defined"); + ASSERT_THROWS_WITH( + F::interpolate( + input, + InterpolateOptions().size({3, 4, 4}).scale_factor({0.5})), + "only one of size or scale_factor should be defined"); + ASSERT_THROWS_WITH( + F::interpolate(input, InterpolateOptions().scale_factor({3, 2})), + "scale_factor shape must match input shape. " + "Input is 1D, scale_factor size is 2"); + ASSERT_THROWS_WITH( + F::interpolate( + input, + InterpolateOptions() + .mode(torch::kNearest) + .align_corners(true)), + "align_corners option can only be set with the " + "interpolating modes: linear | bilinear | bicubic | trilinear"); + } +} + TEST_F(FunctionalTest, Pad) { { auto input = torch::arange(6, torch::kDouble).reshape({1, 2, 3}); diff --git a/test/cpp/api/modulelist.cpp b/test/cpp/api/modulelist.cpp index e4620ea9d4add..68ff0f6d795ef 100644 --- a/test/cpp/api/modulelist.cpp +++ b/test/cpp/api/modulelist.cpp @@ -281,7 +281,7 @@ TEST_F(ModuleListTest, PrettyPrintModuleList) { " (0): torch::nn::Linear(in_features=10, out_features=3, bias=true)\n" " (1): torch::nn::Conv2d(input_channels=1, output_channels=2, kernel_size=[3, 3], stride=[1, 1])\n" " (2): torch::nn::Dropout(rate=0.5)\n" - " (3): torch::nn::BatchNorm(features=5, eps=1e-05, momentum=0.1, affine=true, stateful=true)\n" + " (3): torch::nn::BatchNorm(num_features=5, eps=1e-05, momentum=0.1, affine=true, track_running_stats=true)\n" " (4): torch::nn::Embedding(num_embeddings=4, embedding_dim=10)\n" " (5): torch::nn::LSTM(input_size=4, hidden_size=5, layers=1, dropout=0)\n" ")"); diff --git a/test/cpp/api/modules.cpp b/test/cpp/api/modules.cpp index 8403977edfdcb..e32d1ebbe5e97 100644 --- a/test/cpp/api/modules.cpp +++ b/test/cpp/api/modules.cpp @@ -296,6 +296,36 @@ TEST_F(ModulesTest, Identity) { ASSERT_TRUE(torch::equal(input.grad(), torch::ones_like(input))); } +TEST_F(ModulesTest, Flatten) { + Flatten flatten; + auto input = torch::tensor({{1, 3, 4}, {2, 5, 6}}, torch::requires_grad()); + auto output = flatten->forward(input); + auto expected = torch::tensor({{1, 3, 4}, {2, 5, 6}}, torch::kFloat); + auto s = output.sum(); + + s.backward(); + ASSERT_TRUE(torch::equal(output, expected)); + ASSERT_TRUE(torch::equal(input.grad(), torch::ones_like(input))); + + // Testing with optional arguments start_dim and end_dim + Flatten flatten_optional_dims(FlattenOptions().start_dim(2).end_dim(3)); + input = torch::tensor({ + {{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}, + {{{9, 10}, {11, 12}}, {{13, 14}, {15, 16}}} + }, torch::requires_grad()); // Tensor with sizes (2, 2, 2, 2) + + output = flatten_optional_dims->forward(input); + expected = torch::tensor({ + {{1, 2, 3, 4}, {5, 6, 7, 8}}, + {{9, 10, 11, 12}, {13, 14, 15, 16}} + }, torch::kFloat); // Tensor with sizes (2, 2, 4) + + s = output.sum(); + s.backward(); + ASSERT_TRUE(torch::equal(output, expected)); + ASSERT_TRUE(torch::equal(input.grad(), torch::ones_like(input))); +} + TEST_F(ModulesTest, AdaptiveMaxPool1d) { AdaptiveMaxPool1d model(3); auto x = torch::tensor({{{1, 2, 3, 4, 5}}}, torch::requires_grad()); @@ -1001,7 +1031,7 @@ TEST_F(ModulesTest, BatchNormStateful) { BatchNorm bn(5); // Is stateful by default. - ASSERT_TRUE(bn->options.stateful()); + ASSERT_TRUE(bn->options.track_running_stats()); ASSERT_TRUE(bn->running_mean.defined()); ASSERT_EQ(bn->running_mean.dim(), 1); @@ -1023,7 +1053,7 @@ TEST_F(ModulesTest, BatchNormStateful) { ASSERT_EQ(bn->bias.size(0), 5); } TEST_F(ModulesTest, BatchNormStateless) { - BatchNorm bn(BatchNormOptions(5).stateful(false).affine(false)); + BatchNorm bn(BatchNormOptions(5).track_running_stats(false).affine(false)); ASSERT_FALSE(bn->running_mean.defined()); ASSERT_FALSE(bn->running_var.defined()); @@ -1033,7 +1063,7 @@ TEST_F(ModulesTest, BatchNormStateless) { ASSERT_THROWS_WITH( bn(torch::ones({2, 5})), "Calling BatchNorm::forward is only permitted " - "when the 'stateful' option is true (was false). " + "when the 'track_running_stats' option is true (was false). " "Use BatchNorm::pure_forward instead."); } @@ -1051,6 +1081,71 @@ TEST_F(ModulesTest, BatchNormPureForward) { ASSERT_TRUE(output.allclose(expected)); } +TEST_F(ModulesTest, BatchNormLegacyWarning) { + std::stringstream buffer; + torch::test::CerrRedirect cerr_redirect(buffer.rdbuf()); + + BatchNorm bn(5); + + ASSERT_EQ( + count_substr_occurrences( + buffer.str(), + "torch::nn::BatchNorm module is deprecated" + ), + 1); +} + +TEST_F(ModulesTest, BatchNorm1dStateful) { + BatchNorm1d bn(BatchNorm1dOptions(5)); + + ASSERT_TRUE(bn->options.track_running_stats()); + + ASSERT_TRUE(bn->running_mean.defined()); + ASSERT_EQ(bn->running_mean.dim(), 1); + ASSERT_EQ(bn->running_mean.size(0), 5); + + ASSERT_TRUE(bn->running_var.defined()); + ASSERT_EQ(bn->running_var.dim(), 1); + ASSERT_EQ(bn->running_var.size(0), 5); + + ASSERT_TRUE(bn->num_batches_tracked.defined()); + ASSERT_EQ(bn->num_batches_tracked.dim(), 1); + ASSERT_EQ(bn->num_batches_tracked.size(0), 1); + + ASSERT_TRUE(bn->options.affine()); + + ASSERT_TRUE(bn->weight.defined()); + ASSERT_EQ(bn->weight.dim(), 1); + ASSERT_EQ(bn->weight.size(0), 5); + + ASSERT_TRUE(bn->bias.defined()); + ASSERT_EQ(bn->bias.dim(), 1); + ASSERT_EQ(bn->bias.size(0), 5); +} + +TEST_F(ModulesTest, BatchNorm1dStateless) { + BatchNorm1d bn(BatchNorm1dOptions(5).track_running_stats(false).affine(false)); + + ASSERT_FALSE(bn->running_mean.defined()); + ASSERT_FALSE(bn->running_var.defined()); + ASSERT_FALSE(bn->num_batches_tracked.defined()); + ASSERT_FALSE(bn->weight.defined()); + ASSERT_FALSE(bn->bias.defined()); +} + +TEST_F(ModulesTest, BatchNorm1d) { + BatchNorm1d bn(BatchNorm1dOptions(5)); + bn->eval(); + + auto input = torch::randn({2, 5}, torch::requires_grad()); + auto output = bn->forward(input); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(input.sizes(), input.grad().sizes()); + ASSERT_TRUE(input.grad().allclose(torch::ones({2, 5}))); +} + TEST_F(ModulesTest, Linear_CUDA) { Linear model(5, 2); model->to(torch::kCUDA); @@ -1189,7 +1284,7 @@ TEST_F(ModulesTest, MultiLabelMarginLossDefaultOptions) { } TEST_F(ModulesTest, MultiLabelMarginLossNoReduction) { - MultiLabelMarginLoss loss(torch::Reduction::None); + MultiLabelMarginLoss loss(torch::kNone); auto input = torch::tensor({{0.1, 0.2, 0.4, 0.8}}, torch::requires_grad()); auto target = torch::tensor({{3, 0, -1, 1}}, torch::kLong); auto output = loss->forward(input, target); @@ -1255,7 +1350,7 @@ TEST_F(ModulesTest, MultiLabelSoftMarginLossDefaultOptions) { } TEST_F(ModulesTest, SoftMarginLossNoReduction) { - SoftMarginLoss loss(torch::Reduction::None); + SoftMarginLoss loss(torch::kNone); auto input = torch::tensor({2., 4., 1., 3.}, torch::requires_grad()); auto target = torch::tensor({-1., 1., 1., -1.}, torch::kFloat); auto output = loss->forward(input, target); @@ -1271,7 +1366,7 @@ TEST_F(ModulesTest, MultiLabelSoftMarginLossWeightedNoReduction) { auto input = torch::tensor({{0., 2., 2., 0.}, {2., 1., 0., 1.}}, torch::requires_grad()); auto target = torch::tensor({{0., 0., 1., 0.}, {1., 0., 1., 1.}}, torch::kFloat); auto weight = torch::tensor({0.1, 0.6, 0.4, 0.8}, torch::kFloat); - auto options = MultiLabelSoftMarginLossOptions().reduction(torch::Reduction::None).weight(weight); + auto options = MultiLabelSoftMarginLossOptions().reduction(torch::kNone).weight(weight); MultiLabelSoftMarginLoss loss = MultiLabelSoftMarginLoss(options); auto output = loss->forward(input, target); auto expected = torch::tensor({0.4876902, 0.3321295}, torch::kFloat); @@ -1688,10 +1783,162 @@ TEST_F(ModulesTest, Threshold) { } } +TEST_F(ModulesTest, Upsampling1D) { + { + Upsample model(UpsampleOptions() + .size({4}) + .mode(torch::kNearest)); + auto input = torch::ones({1, 1, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected = torch::ones({1, 1, 4}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + { + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + Upsample model(UpsampleOptions() + .scale_factor({scale_factor}) + .mode(torch::kLinear) + .align_corners(align_corners)); + auto input = torch::ones({1, 1, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = torch::ones({1, 1, expected_size}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + } + } + { + // linear (1D) upsampling spatial invariance + Upsample model(UpsampleOptions() + .scale_factor({3}) + .mode(torch::kLinear) + .align_corners(false)); + auto input = torch::zeros({1, 1, 9}); + input.narrow(2, 0, 4).normal_(); + auto output = model->forward(input); + auto expected = model->forward(input.narrow(2, 0, 5)); + + ASSERT_TRUE(torch::allclose(output.narrow(2, 0, 15), expected)); + } +} + +TEST_F(ModulesTest, Upsampling2D) { + { + Upsample model(UpsampleOptions() + .size({4, 4}) + .mode(torch::kNearest)); + auto input = torch::ones({1, 1, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected = torch::ones({1, 1, 4, 4}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + { + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + Upsample model(UpsampleOptions() + .scale_factor({scale_factor, scale_factor}) + .mode(torch::kBilinear) + .align_corners(align_corners)); + auto input = torch::ones({1, 1, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = torch::ones({1, 1, expected_size, expected_size}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + } + } + { + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + Upsample model(UpsampleOptions() + .scale_factor({scale_factor, scale_factor}) + .mode(torch::kBicubic) + .align_corners(align_corners)); + auto input = torch::ones({1, 1, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = torch::ones({1, 1, expected_size, expected_size}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + } + } +} + +TEST_F(ModulesTest, Upsampling3D) { + { + Upsample model(UpsampleOptions() + .size({4, 4, 4}) + .mode(torch::kNearest)); + auto input = torch::ones({1, 1, 2, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected = torch::ones({1, 1, 4, 4, 4}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + { + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + Upsample model( + UpsampleOptions() + .scale_factor({scale_factor, scale_factor, scale_factor}) + .mode(torch::kTrilinear) + .align_corners(align_corners)); + auto input = torch::ones({1, 1, 2, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = + torch::ones({1, 1, expected_size, expected_size, expected_size}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + } + } +} + TEST_F(ModulesTest, PrettyPrintIdentity) { ASSERT_EQ(c10::str(Identity()), "torch::nn::Identity()"); } +TEST_F(ModulesTest, PrettyPrintFlatten) { + ASSERT_EQ(c10::str(Flatten()), "torch::nn::Flatten()"); + ASSERT_EQ(c10::str(Flatten(FlattenOptions().start_dim(2).end_dim(4))), "torch::nn::Flatten()"); +} + TEST_F(ModulesTest, ReflectionPad1d) { { ReflectionPad1d m(ReflectionPad1dOptions(2)); @@ -1998,6 +2245,15 @@ TEST_F(ModulesTest, PrettyPrintConv) { "torch::nn::Conv2d(input_channels=3, output_channels=4, kernel_size=[5, 6], stride=[1, 2])"); } +TEST_F(ModulesTest, PrettyPrintUpsample) { + ASSERT_EQ( + c10::str(Upsample(UpsampleOptions().size({2, 4, 4}))), + "torch::nn::Upsample(size=[2, 4, 4], mode=kNearest)"); + ASSERT_EQ( + c10::str(Upsample(UpsampleOptions().scale_factor({0.5, 1.5}).mode(torch::kBilinear))), + "torch::nn::Upsample(scale_factor=[0.5, 1.5], mode=kBilinear)"); +} + TEST_F(ModulesTest, PrettyPrintUnfold) { ASSERT_EQ( c10::str(Unfold(torch::IntArrayRef({2, 4}))), @@ -2147,9 +2403,17 @@ TEST_F(ModulesTest, PrettyPrintFunctional) { TEST_F(ModulesTest, PrettyPrintBatchNorm) { ASSERT_EQ( c10::str(BatchNorm( - BatchNormOptions(4).eps(0.5).momentum(0.1).affine(false).stateful( + BatchNormOptions(4).eps(0.5).momentum(0.1).affine(false).track_running_stats( true))), - "torch::nn::BatchNorm(features=4, eps=0.5, momentum=0.1, affine=false, stateful=true)"); + "torch::nn::BatchNorm(num_features=4, eps=0.5, momentum=0.1, affine=false, track_running_stats=true)"); +} + +TEST_F(ModulesTest, PrettyPrintBatchNorm1d) { + ASSERT_EQ( + c10::str(BatchNorm1d( + BatchNorm1dOptions(4).eps(0.5).momentum(0.1).affine(false) + .track_running_stats(true))), + "torch::nn::BatchNorm1d(4, eps=0.5, momentum=0.1, affine=false, track_running_stats=true)"); } TEST_F(ModulesTest, PrettyPrintLayerNorm) { diff --git a/test/cpp/api/sequential.cpp b/test/cpp/api/sequential.cpp index 543df6606589b..c33932186f17e 100644 --- a/test/cpp/api/sequential.cpp +++ b/test/cpp/api/sequential.cpp @@ -412,7 +412,7 @@ TEST_F(SequentialTest, PrettyPrintSequential) { " (0): torch::nn::Linear(in_features=10, out_features=3, bias=true)\n" " (1): torch::nn::Conv2d(input_channels=1, output_channels=2, kernel_size=[3, 3], stride=[1, 1])\n" " (2): torch::nn::Dropout(rate=0.5)\n" - " (3): torch::nn::BatchNorm(features=5, eps=1e-05, momentum=0.1, affine=true, stateful=true)\n" + " (3): torch::nn::BatchNorm(num_features=5, eps=1e-05, momentum=0.1, affine=true, track_running_stats=true)\n" " (4): torch::nn::Embedding(num_embeddings=4, embedding_dim=10)\n" " (5): torch::nn::LSTM(input_size=4, hidden_size=5, layers=1, dropout=0)\n" ")"); @@ -431,7 +431,7 @@ TEST_F(SequentialTest, PrettyPrintSequential) { " (linear): torch::nn::Linear(in_features=10, out_features=3, bias=true)\n" " (conv2d): torch::nn::Conv2d(input_channels=1, output_channels=2, kernel_size=[3, 3], stride=[1, 1])\n" " (dropout): torch::nn::Dropout(rate=0.5)\n" - " (batchnorm): torch::nn::BatchNorm(features=5, eps=1e-05, momentum=0.1, affine=true, stateful=true)\n" + " (batchnorm): torch::nn::BatchNorm(num_features=5, eps=1e-05, momentum=0.1, affine=true, track_running_stats=true)\n" " (embedding): torch::nn::Embedding(num_embeddings=4, embedding_dim=10)\n" " (lstm): torch::nn::LSTM(input_size=4, hidden_size=5, layers=1, dropout=0)\n" ")"); diff --git a/test/cpp/api/serialize.cpp b/test/cpp/api/serialize.cpp index 053d534a1719f..c6f9322f2f999 100644 --- a/test/cpp/api/serialize.cpp +++ b/test/cpp/api/serialize.cpp @@ -127,6 +127,39 @@ TEST(SerializeTest, NonContiguous) { ASSERT_TRUE(x.allclose(y)); } +TEST(SerializeTest, ErrorOnMissingKey) { + struct B : torch::nn::Module { + B(const std::string& name_c) { + register_buffer(name_c, torch::ones(5, torch::kFloat)); + } + }; + struct A : torch::nn::Module { + A(const std::string& name_b, const std::string& name_c) { + register_module(name_b, std::make_shared(name_c)); + } + }; + struct M : torch::nn::Module { + M(const std::string& name_a, + const std::string& name_b, + const std::string& name_c) { + register_module(name_a, std::make_shared(name_b, name_c)); + } + }; + + // create a hierarchy of models with names differing below the top level + auto model1 = std::make_shared("a", "b", "c"); + auto model2 = std::make_shared("a", "b", "x"); + auto model3 = std::make_shared("a", "x", "c"); + + std::stringstream stream; + torch::save(model1, stream); + // We want the errors to contain hierarchy information, too. + ASSERT_THROWS_WITH( + torch::load(model2, stream), "No such serialized tensor 'a.b.x'"); + ASSERT_THROWS_WITH( + torch::load(model3, stream), "No such serialized submodule: 'a.x'"); +} + TEST(SerializeTest, XOR) { // We better be able to save and load an XOR model! auto getLoss = [](Sequential model, uint32_t batch_size) { diff --git a/test/cpp/jit/test_utils.cpp b/test/cpp/jit/test_utils.cpp index d2eef8a43d3e8..f86026e90b413 100644 --- a/test/cpp/jit/test_utils.cpp +++ b/test/cpp/jit/test_utils.cpp @@ -1,4 +1,6 @@ #include +#include +#include namespace torch { namespace jit { @@ -37,7 +39,7 @@ std::pair runGradient( static const auto as_tensorlist = [](const Stack& stack) { return fmap(stack, [](const IValue& i) { return i.toTensor(); }); }; - + ClearUndefinedness(grad_spec.df); Code f_code{grad_spec.f}, df_code{grad_spec.df}; InterpreterState f_interpreter{f_code}, df_interpreter{df_code}; diff --git a/test/cpp_api_parity/parity-tracker.md b/test/cpp_api_parity/parity-tracker.md index a326b54c90072..0fbda1e04e728 100644 --- a/test/cpp_api_parity/parity-tracker.md +++ b/test/cpp_api_parity/parity-tracker.md @@ -69,7 +69,7 @@ torch.nn.Softmax|Yes|No torch.nn.Softmax2d|Yes|No torch.nn.LogSoftmax|Yes|No torch.nn.AdaptiveLogSoftmaxWithLoss|No|No -torch.nn.BatchNorm1d|No|No +torch.nn.BatchNorm1d|Yes|No torch.nn.BatchNorm2d|No|No torch.nn.BatchNorm3d|No|No torch.nn.GroupNorm|No|No @@ -91,9 +91,9 @@ torch.nn.TransformerDecoder|No|No torch.nn.TransformerEncoderLayer|No|No torch.nn.TransformerDecoderLayer|No|No torch.nn.Identity|Yes|No -torch.nn.Linear|No|No +torch.nn.Linear|Yes|No torch.nn.Bilinear|Yes|No -torch.nn.Flatten|No|No +torch.nn.Flatten|Yes|No torch.nn.Dropout|No|No torch.nn.Dropout2d|No|No torch.nn.Dropout3d|No|No @@ -121,7 +121,7 @@ torch.nn.CosineEmbeddingLoss|Yes|No torch.nn.MultiMarginLoss|Yes|No torch.nn.TripletMarginLoss|Yes|No torch.nn.PixelShuffle|Yes|No -torch.nn.Upsample|No|No +torch.nn.Upsample|Yes|No torch.nn.DataParallel|No|No torch.nn.parallel.DistributedDataParallel|No|No torch.nn.utils.clip_grad_norm_|Yes|No diff --git a/test/dist_autograd_test.py b/test/dist_autograd_test.py index c283d6f537b25..45ca4dddba6e6 100644 --- a/test/dist_autograd_test.py +++ b/test/dist_autograd_test.py @@ -8,7 +8,6 @@ import torch.distributed.autograd as dist_autograd import torch.distributed.rpc as rpc from dist_utils import INIT_METHOD_TEMPLATE, dist_init, TEST_CONFIG -from torch.distributed.rpc import RpcBackend import threading @@ -37,10 +36,46 @@ def _set_rpc_done(ctx_id, rank_distance): known_context_ids.append(ctx_id) +def _check_rpc_done(rank_distance): + while not rpc_done[rank_distance]: + time.sleep(0.1) + + +def _torch_ones(sizes, requires_grad=False): + return torch.ones(sizes, requires_grad=requires_grad) + + +# creates an owner rref on the given dst, and the rref holds a torch.ones tensor +# of the given size. +def _create_ones_rref_on(dst, sizes): + return rpc.remote( + dst, + _torch_ones, + args=(sizes,), + kwargs={"requires_grad": True} + ) + + +# This method must be called on the rref owner, and verifies that the grad of +# rref tensor equals to the given grad. +def _compare_owner_value(context_id, rref, grad): + grads = dist_autograd.get_gradients(context_id) + return torch.equal(grads[rref.local_value().wait()], grad) + + def my_py_add(t1, t2): return torch.add(t1, t2) +def my_rref_add(rref_t1, t2): + ret = torch.add(rref_t1.local_value().wait(), t2) + return ret + + +def my_nested_rref_add(dst, rref_t1, t2): + return rpc.rpc_sync(dst, my_rref_add, args=(rref_t1, t2)) + + def my_py_nested_call(t1, t2, dst, world_size, hops): next_dst = (dst + 1) % world_size if hops > 0: @@ -69,6 +104,17 @@ def _all_contexts_cleaned_up(timeout_seconds=10): return success +# This function creates a dis atugorad context, run rpc_sync on the given ps, +# and then blocks until the ps has verified the grads are correctly accumulated. +def _run_trainer(rref_t1, t2, ps, rank_diff): + with dist_autograd.context() as context_id: + ret = rpc.rpc_sync(ps, my_rref_add, args=(rref_t1, t2)) + dist_autograd.backward([ret.sum()]) + # prevent deleting dist autograd context + rpc.rpc_sync(ps, _set_rpc_done, args=(context_id, rank_diff)) + rpc.rpc_sync(ps, _check_rpc_done, args=(0, )) + + from torch.autograd import Function from torch.autograd.function import once_differentiable @@ -86,7 +132,8 @@ def backward(ctx, input): class ExecMode(Enum): LOCAL = 1 # Run the operation locally. - REMOTE = 2 # Run the operation using RPC. + RPC_SYNC = 2 # Run the operation using rpc_sync + REMOTE = 3 # Run the operation using remote. @unittest.skipIf( @@ -99,9 +146,15 @@ def _exec_func(self, exec_mode, method, *args): if len(args) == 1 and isinstance(args[0], list): return method(*args[0]) return method(*args) - else: + elif ExecMode.RPC_SYNC == exec_mode: return rpc.rpc_sync('worker{}'.format(self._next_rank()), method, args=(args)) + elif ExecMode.REMOTE == exec_mode: + rref = rpc.remote('worker{}'.format(self._next_rank()), method, + args=(args)) + return rref.to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) def _next_rank(self): if hasattr(self, 'dst_rank'): @@ -113,9 +166,7 @@ def _next_rank(self): return self.dst_rank def _check_rpc_done(self, rank_distance): - while not rpc_done[rank_distance]: - time.sleep(0.1) - pass + _check_rpc_done(rank_distance) @property def world_size(self): @@ -125,7 +176,7 @@ def world_size(self): def init_method(self): return INIT_METHOD_TEMPLATE.format(file_name=self.file_name) - @dist_init(setup_model_parallel=True) + @dist_init def test_autograd_context(self): # Verify max possible id. max_auto_increment = 281474976710655 @@ -151,7 +202,7 @@ def test_autograd_context(self): ): dist_autograd._retrieve_context(context_id) - @dist_init(setup_model_parallel=True) + @dist_init def test_nested_context(self): with dist_autograd.context() as context_id: # Nested contexts not supported. @@ -256,12 +307,33 @@ def _verify_graph_for_nested_rpc_call(self, ctx): "torch::distributed::autograd::RecvRpcBackward", next_funcs[0][0].name() ) - def _test_graph(self, fn): + def _test_graph(self, fn, exec_mode): dst_rank = (self.rank + 1) % self.world_size + + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=True) t2 = torch.zeros(3, 3, requires_grad=True) - ret = rpc.rpc_sync("worker{}".format(dst_rank), fn, args=(t1, t2)) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(dst_rank), fn, args=(t1, t2)) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(dst_rank), fn, args=(t1, t2) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + rpc.rpc_sync("worker{}".format(dst_rank), _set_rpc_done, args=(context_id, 1)) @@ -295,24 +367,56 @@ def _test_graph(self, fn): with self.assertRaises(RuntimeError): ctx = dist_autograd._current_context() - @dist_init(setup_model_parallel=True) + @dist_init def test_graph_for_builtin_call(self): - self._test_graph(torch.add) + self._test_graph(torch.add, ExecMode.RPC_SYNC) - @dist_init(setup_model_parallel=True) + @dist_init def test_graph_for_python_call(self): - self._test_graph(my_py_add) + self._test_graph(my_py_add, ExecMode.RPC_SYNC) + + @dist_init + def test_graph_for_builtin_remote_call(self): + self._test_graph(torch.add, ExecMode.REMOTE) + + @dist_init + def test_graph_for_python_remote_call(self): + self._test_graph(my_py_add, ExecMode.REMOTE) # 3-layer nested calls - @dist_init(setup_model_parallel=True) - def test_graph_for_py_nested_call(self): + def _test_graph_for_py_nested_call(self, exec_mode): dst_rank = (self.rank + 1) % self.world_size + + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=True) t2 = torch.zeros(3, 3, requires_grad=True) nest_dst_rank = (dst_rank + 1) % self.world_size - ret = rpc.rpc_sync("worker{}".format(dst_rank), - my_py_nested_call, args=(t1, t2, dst_rank, self.world_size, 1)) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(dst_rank), + my_py_nested_call, + args=(t1, t2, dst_rank, self.world_size, 1) + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(dst_rank), + my_py_nested_call, + args=(t1, t2, dst_rank, self.world_size, 1) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + for rd in [1, 2, 3]: rpc.rpc_sync("worker{}".format((self.rank + rd) % self.world_size), _set_rpc_done, args=(context_id, rd)) @@ -357,16 +461,59 @@ def test_graph_for_py_nested_call(self): # autograd context before another worker tries to access it. dist.barrier() + @dist_init + def test_graph_for_py_nested_call(self): + self._test_graph_for_py_nested_call(ExecMode.RPC_SYNC) + + @dist_init + def test_graph_for_py_nested_remote_call(self): + self._test_graph_for_py_nested_call(ExecMode.REMOTE) + # Rank0->Rank1->Rank0 - @dist_init(setup_model_parallel=True) - def test_graph_for_py_nested_call_itself(self): + def _test_graph_for_py_nested_call_itself(self, exec_mode): dst_rank = (self.rank + 1) % self.world_size + + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=True) t2 = torch.zeros(3, 3, requires_grad=True) - ret = rpc.rpc_sync("worker{}".format(dst_rank), - my_py_nested_call, - args=(t1, t2, (self.rank - 1 + self.world_size) % self.world_size, self.world_size, 0)) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(dst_rank), + my_py_nested_call, + args=( + t1, + t2, + (self.rank - 1 + self.world_size) % self.world_size, + self.world_size, + 0 + ) + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(dst_rank), + my_py_nested_call, + args=( + t1, + t2, + (self.rank - 1 + self.world_size) % self.world_size, + self.world_size, + 0 + ) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + rpc.rpc_sync("worker{}".format((self.rank + 1) % self.world_size), _set_rpc_done, args=(context_id, 1)) @@ -395,13 +542,34 @@ def test_graph_for_py_nested_call_itself(self): # autograd context before another worker tries to access it. dist.barrier() - @dist_init(setup_model_parallel=True) - def test_no_graph_with_tensors_not_require_grad(self): + @dist_init + def test_graph_for_py_nested_call_itself(self): + self._test_graph_for_py_nested_call_itself(ExecMode.RPC_SYNC) + + @dist_init + def test_graph_for_py_nested_remote_call_itself(self): + self._test_graph_for_py_nested_call_itself(ExecMode.REMOTE) + + def _test_no_graph_with_tensors_not_require_grad(self, exec_mode): dst_rank = (self.rank + 1) % self.world_size with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=False) t2 = torch.zeros(3, 3, requires_grad=False) - ret = rpc.rpc_sync("worker{}".format(dst_rank), torch.add, args=(t1, t2)) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(dst_rank), + torch.add, + args=(t1, t2) + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(dst_rank), + torch.add, + args=(t1, t2) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + rpc.rpc_sync("worker{}".format(dst_rank), _set_rpc_done, args=(context_id, 1)) @@ -413,20 +581,47 @@ def test_no_graph_with_tensors_not_require_grad(self): # Wait for the prev rank to be done with rpc. self._check_rpc_done(1) - # prev context id is not passed over as tensors do not require grads - with self.assertRaises(RuntimeError): - ctx = dist_autograd._retrieve_context(ctx_ids[1]) + if ExecMode.RPC_SYNC == exec_mode: + # prev context id is not passed over as tensors do not require + # grads + with self.assertRaises(RuntimeError): + ctx = dist_autograd._retrieve_context(ctx_ids[1]) + elif ExecMode.REMOTE == exec_mode: + # NB: RRef.to_here() always passes the autograd context to the + # the callee, as the caller does not know whether the return + # value would contain a requires_grad tensor or not. + pass - @dist_init(setup_model_parallel=True) - def test_rpc_complex_args(self): + @dist_init + def test_no_graph_with_tensors_not_require_grad(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.RPC_SYNC) + + @dist_init + def test_no_graph_with_tensors_not_require_grad_remote(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.REMOTE) + + def _test_rpc_complex_args(self, exec_mode): with dist_autograd.context() as context_id: num_tensors = 10 tensors = [] for i in range(num_tensors): tensors.append(torch.ones(3, 3, requires_grad=(i % 2 == 0))) - ret = rpc.rpc_sync( - "worker{}".format(self._next_rank()), torch.stack, args=(tensors,) - ) + + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(self._next_rank()), + torch.stack, + args=(tensors,) + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(self._next_rank()), + torch.stack, + args=(tensors,) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + self.assertEqual(torch.stack(tensors), ret) # Verify appropriate tensors have been attached the autograd graph. @@ -450,8 +645,15 @@ def test_rpc_complex_args(self): dst_rank = (self.rank + 1) % self.world_size self.assertEqual(worker_ids[0], dst_rank) + @dist_init + def test_rpc_complex_args(self): + self._test_rpc_complex_args(ExecMode.RPC_SYNC) - @dist_init(setup_model_parallel=True) + @dist_init + def test_remote_complex_args(self): + self._test_rpc_complex_args(ExecMode.REMOTE) + + @dist_init def test_context_cleanup_many_workers(self): dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} with dist_autograd.context() as context_id: @@ -467,8 +669,19 @@ def test_context_cleanup_many_workers(self): success = _all_contexts_cleaned_up() self.assertTrue(success) - @dist_init(setup_model_parallel=True) + @dist_init def test_context_cleanup_nested_rpc(self): + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + dst_rank = (self.rank + 1) % self.world_size nested_dst_rank = (dst_rank + 1) % self.world_size with dist_autograd.context() as context_id: @@ -484,7 +697,7 @@ def test_context_cleanup_nested_rpc(self): success = _all_contexts_cleaned_up() self.assertTrue(success) - @dist_init(setup_model_parallel=True) + @dist_init def test_worker_ids_recorded(self): dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} with dist_autograd.context() as context_id: @@ -515,7 +728,7 @@ def test_worker_ids_recorded(self): self.assertEqual(len(worker_ids), len(dst_ranks)) self.assertEqual(set(worker_ids), dst_ranks) - @dist_init(setup_model_parallel=True) + @dist_init def test_error_in_context(self): with dist_autograd.context() as context_id: t1 = torch.rand(3, 3, requires_grad=True) @@ -528,11 +741,11 @@ def test_error_in_context(self): args=(t1, t2)) def _verify_backwards(self, exec_mode, tensors, context_id, local_grads, *args): - if exec_mode == ExecMode.REMOTE: - self._verify_backwards_remote(tensors, context_id, local_grads, *args) - else: + if exec_mode == ExecMode.LOCAL: torch.autograd.backward(tensors) return [arg.grad for arg in args] + else: + self._verify_backwards_remote(tensors, context_id, local_grads, *args) def _verify_backwards_remote(self, tensors, context_id, local_grads, *args): dist_autograd.backward(tensors) @@ -551,21 +764,144 @@ def _verify_backwards_remote(self, tensors, context_id, local_grads, *args): self.assertEqual(ngrads, len(grads)) - - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_simple(self): # Run the same code locally and with dist autograd and verify gradients # are same. local_grads = None t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: ret = self._exec_func(exec_mode, torch.add, t1, t2) loss = ret.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2) + local_grads = ret if ret else local_grads + + # The current rank first creates a tensor on the rref_owner, and then passes + # the rref with another tensor to the callee to run either my_rref_add or + # my_nested_rref_add, depending on whether the callee is the rref owner. + # The grad of tensor lives on the current rank, and the grad of the rref + # tensor lives on the rref owner. + def _test_backward_rref(self, callee, rref_owner): + local_grads = None + t1 = torch.ones((3, 3), requires_grad=True) + t2 = torch.zeros((3, 3), requires_grad=True) + + local_ret = torch.add(t1, t2) + local_ret.sum().backward() + with dist_autograd.context() as context_id: + rref_t1 = rpc.remote( + rref_owner, + _torch_ones, + args=((3, 3),), + kwargs={"requires_grad": True} + ) + + if callee == rref_owner: + rref = rpc.remote(callee, my_rref_add, args=(rref_t1, t2)) + else: + rref = rpc.remote( + callee, + my_nested_rref_add, + args=(rref_owner, rref_t1, t2) + ) + ret = rref.to_here().wait() + dist_autograd.backward([ret.sum()]) + + # verify grads on caller + grads = dist_autograd.get_gradients(context_id) + self.assertIn(t2, grads) + self.assertEqual(grads[t2], t2.grad) + + # verify grads on rref owner + self.assertTrue( + rpc.rpc_sync( + rref_owner, + _compare_owner_value, + args=(context_id, rref_t1, t1.grad) + ) + ) + + @dist_init + def test_backward_rref(self): + callee = "worker{}".format(self._next_rank()) + rref_owner = callee + self._test_backward_rref(callee, rref_owner) + + @dist_init + def test_backward_rref_multi(self): + if self.rank > 0: + callee = "worker0" + rref_owner = callee + self._test_backward_rref(callee, rref_owner) - @dist_init(setup_model_parallel=True) + @dist_init + def test_backward_rref_nested(self): + callee = "worker{}".format((self.rank + 1) % self.world_size) + rref_owner = "worker{}".format((self.rank + 2) % self.world_size) + self._test_backward_rref(callee, rref_owner) + + # In this test, every rank will serve as a parameter server (ps) and a + # driver, and then kicks off trainers on the other three ranks. So, we have: + # ps = rank0 with trainers = rank1/2/3 + # ps = rank2 with trainers = rank2/3/0 + # ps = rank3 with trainers = rank3/0/1 + # ps = rank4 with trainers = rank0/1/2 + # + # These four test ps-trainer groups run on completely separate autograd + # graphs, but they share the same set of underlying RpcAgents. + @dist_init + def test_trainer_ps(self): + local_grads = None + t1 = torch.ones((3, 3), requires_grad=True) + t2 = torch.zeros((3, 3), requires_grad=True) + + local_ret = torch.add(t1, t2) + local_ret.sum().backward() + + # create rref on self + # TODO: simplify this once we support rpc to self + self_name = "worker{}".format(self.rank) + rref_t1 = rpc.rpc_sync( + "worker{}".format(self._next_rank()), + _create_ones_rref_on, + args=(self_name, (3, 3)) + ) + + # kick off forward and backward pass on three other workers (trainers) + rank_diffs = [1, 2, 3] + futures = [] + for rank_diff in rank_diffs: + futures.append(rpc.rpc_async( + "worker{}".format((self.rank + rank_diff) % self.world_size), + _run_trainer, + args=(rref_t1, t2, self_name, rank_diff) + )) + + # check if the trainers have done with their backward pass + for rank_diff in rank_diffs: + self._check_rpc_done(rank_diff) + + # trainers are done and holding the context for verification + accumulate_grad_func = None + for rank_diff in rank_diffs: + # make sure grads are accumulated for the same tensors and values + # are all correct + ctx_id = ctx_ids[rank_diff] + grads = dist_autograd.get_gradients(ctx_id) + local_t1 = rref_t1.local_value().wait() + self.assertIn(local_t1, grads) + self.assertEqual(grads[local_t1], t1.grad) + + # unblock trainers + _set_rpc_done(None, 0) + + # wait until all trainers are done + for fut in futures: + fut.wait() + + @dist_init def test_backward_multiple_round_trips(self): local_grads = None t1 = torch.rand((3, 3), requires_grad=True) @@ -574,7 +910,7 @@ def test_backward_multiple_round_trips(self): t4 = torch.rand((3, 3)) t5 = torch.rand((3, 3), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: # Multiple RPCs between different nodes. val = self._exec_func(exec_mode, torch.add, t1, t2) @@ -585,9 +921,10 @@ def test_backward_multiple_round_trips(self): val = self._exec_func(exec_mode, torch.matmul, val, val) loss = val.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t3, t4, t5) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t3, t4, t5) + local_grads = ret if ret else local_grads - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_different_tensor_dims(self): local_grads = None t1 = torch.rand((4, 6), requires_grad=True) @@ -595,33 +932,35 @@ def test_backward_different_tensor_dims(self): t3 = torch.rand((5, 7), requires_grad=True) t4 = torch.rand((7, 9)) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: val = self._exec_func(exec_mode, torch.matmul, t1, t2) val = self._exec_func(exec_mode, torch.chain_matmul, [val, t3, t4]) loss = val.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t2, t3, t4) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t2, t3, t4) + local_grads = ret if ret else local_grads - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_unused_tensors(self): local_grads = None t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) t3 = torch.rand((3, 3), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: s = self._exec_func(exec_mode, torch.stack, (t1, t2, t3)) val = self._exec_func(exec_mode, torch.matmul, torch.narrow(s, 0, 0, 1), torch.narrow(s, 0, 2, 1)) loss = val.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t3) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t3) + local_grads = ret if ret else local_grads - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_multiple_output_tensors(self): local_grads = None t = torch.rand((10, 2), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: tensor_list = self._exec_func(exec_mode, torch.split, t, 2) t1 = tensor_list[0] @@ -631,7 +970,8 @@ def test_backward_multiple_output_tensors(self): val = self._exec_func(exec_mode, torch.chain_matmul, [t1, t2, t3]) loss = val.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t) + local_grads = ret if ret else local_grads def _run_test_backward_unused_send_function_in_thread(self): with dist_autograd.context() as context_id: @@ -649,7 +989,7 @@ def _run_test_backward_unused_send_function_in_thread(self): dist_autograd.backward([val.sum()]) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_unused_send_function(self): # Run the test in a thread which would never finish. t = threading.Thread(target=self._run_test_backward_unused_send_function_in_thread) @@ -660,7 +1000,7 @@ def test_backward_unused_send_function(self): # Verify thread is still alive (indicating backward hasn't completed yet). self.assertTrue(t.is_alive()) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_autograd_engine_error(self): with dist_autograd.context() as context_id: t1 = torch.rand((3, 3), requires_grad=True) @@ -685,10 +1025,22 @@ def test_backward_autograd_engine_error(self): # Run backwards, and validate we receive an error. dist_autograd.backward([val.sum()]) - @unittest.skipIf(TEST_CONFIG.rpc_backend == RpcBackend.PROCESS_GROUP, + @unittest.skip("Using sleep to simulate syncronization is flaky") + @unittest.skipIf(TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP", "Skipping this test temporarily since ProcessGroupAgent does not report errors on node failures") @dist_init(clean_shutdown=False) def test_backward_node_failure(self): + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + with dist_autograd.context() as context_id: t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) @@ -702,7 +1054,7 @@ def test_backward_node_failure(self): # Kill all odd rank nodes. if self.rank % 2 == 0: # Wait a bit for all other nodes to die. - time.sleep(5) + time.sleep(5) # This is flaky. with self.assertRaisesRegex(RuntimeError, "Request aborted during client shutdown"): # Run backwards, and validate we receive an error since all # other nodes are dead. @@ -711,7 +1063,7 @@ def test_backward_node_failure(self): # Exit all other nodes. pass - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_without_context(self): t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) @@ -721,7 +1073,7 @@ def test_backward_without_context(self): args=(t1, t2)) dist_autograd.backward([res.sum()]) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_without_rpc(self): dst_rank = self.rank with dist_autograd.context() as context_id: @@ -737,7 +1089,7 @@ def test_backward_without_rpc(self): self.assertEqual(torch.ones(3, 3), grads[t1]) self.assertEqual(torch.ones(3, 3), grads[t2]) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_invalid_args(self): with dist_autograd.context() as context_id: @@ -759,12 +1111,12 @@ def test_backward_invalid_args(self): t = torch.rand(1, requires_grad=True) dist_autograd.backward([t]) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_multiple_roots(self): local_grads = None t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC]: with dist_autograd.context() as context_id: r1 = self._exec_func(exec_mode, torch.add, t1, t2).sum() r2 = self._exec_func(exec_mode, torch.mul, t1, t2).sum() diff --git a/test/dist_utils.py b/test/dist_utils.py index 4231e7cf36661..0d0add5afc358 100644 --- a/test/dist_utils.py +++ b/test/dist_utils.py @@ -1,12 +1,11 @@ from __future__ import absolute_import, division, print_function, unicode_literals -from functools import wraps, partial import threading +from functools import partial, wraps from os import getenv import torch.distributed as dist import torch.distributed.rpc as rpc -from torch.distributed.rpc.api import RpcBackend if not dist.is_available(): @@ -15,7 +14,7 @@ class TestConfig: - __slots__ = ["rpc_backend"] + __slots__ = ["rpc_backend_name"] def __init__(self, *args, **kwargs): assert len(args) == 0, "TestConfig only takes kwargs." @@ -23,7 +22,7 @@ def __init__(self, *args, **kwargs): setattr(self, k, v) -TEST_CONFIG = TestConfig(rpc_backend=getenv("RPC_BACKEND", RpcBackend.PROCESS_GROUP)) +TEST_CONFIG = TestConfig(rpc_backend_name=getenv("RPC_BACKEND_NAME", "PROCESS_GROUP")) INIT_METHOD_TEMPLATE = "file://{file_name}" @@ -51,7 +50,7 @@ def set_termination_signal(): _TERMINATION_SIGNAL.set() -def dist_init(test_method=None, setup_model_parallel=True, clean_shutdown=True): +def dist_init(old_test_method=None, setup_model_parallel=True, clean_shutdown=True): """ We use this decorator for setting up and tearing down state since MultiProcessTestCase runs each `test*` method in a separate process and @@ -59,17 +58,21 @@ def dist_init(test_method=None, setup_model_parallel=True, clean_shutdown=True): 'setUp' and 'tearDown' methods of unittest. """ - # If we use dist_init without arguments (ex: @dist_init), test_method is + # If we use dist_init without arguments (ex: @dist_init), old_test_method is # appropriately set and we return the wrapper appropriately. On the other # hand if dist_init has arguments (ex: @dist_init(clean_shutdown=False)), - # test_method is None and we return a functools.partial which is the real + # old_test_method is None and we return a functools.partial which is the real # decorator that is used and as a result we recursively call dist_init with - # test_method and the rest of the arguments appropriately set. - if test_method is None: - return partial(dist_init, setup_model_parallel=setup_model_parallel, clean_shutdown=clean_shutdown) - - @wraps(test_method) - def wrapper(self, *arg, **kwargs): + # old_test_method and the rest of the arguments appropriately set. + if old_test_method is None: + return partial( + dist_init, + setup_model_parallel=setup_model_parallel, + clean_shutdown=clean_shutdown, + ) + + @wraps(old_test_method) + def new_test_method(self, *arg, **kwargs): self.worker_id = self.rank self.worker_name_to_id = { "worker{}".format(rank): rank for rank in range(self.world_size) @@ -79,52 +82,48 @@ def wrapper(self, *arg, **kwargs): global _ALL_NODE_NAMES _ALL_NODE_NAMES = self.worker_name_to_id.keys() - dist.init_process_group( - backend="gloo", - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) # Use enough 'num_send_recv_threads' until we fix https://github.com/pytorch/pytorch/issues/26359 rpc.init_model_parallel( self_name="worker%d" % self.rank, - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, num_send_recv_threads=16, ) - test_method(self, *arg, **kwargs) - - if setup_model_parallel: - if clean_shutdown: - # Follower reports done. - if self.rank == MASTER_RANK: - on_master_follower_report_done("worker{}".format(MASTER_RANK)) - else: - rpc.rpc_async( - "worker{}".format(MASTER_RANK), - on_master_follower_report_done, - args=("worker{}".format(self.rank),), - ) - - # Master waits for followers to report done. - # Follower waits for master's termination command. - _TERMINATION_SIGNAL.wait() - if self.rank == MASTER_RANK: - # Master sends termination command. - futs = [] - for dst_rank in range(self.world_size): - # torch.distributed.rpc module does not support sending to self. - if dst_rank == MASTER_RANK: - continue - dst_name = "worker{}".format(dst_rank) - fut = rpc.rpc_async(dst_name, set_termination_signal, args=()) - futs.append(fut) - for fut in futs: - assert fut.wait() is None, "Sending termination signal failed." + return_value = old_test_method(self, *arg, **kwargs) + + if setup_model_parallel and clean_shutdown: + # Follower reports done. + if self.rank == MASTER_RANK: + on_master_follower_report_done("worker{}".format(MASTER_RANK)) + else: + rpc.rpc_async( + "worker{}".format(MASTER_RANK), + on_master_follower_report_done, + args=("worker{}".format(self.rank),), + ) + + # Master waits for followers to report done. + # Follower waits for master's termination command. + _TERMINATION_SIGNAL.wait() + if self.rank == MASTER_RANK: + # Master sends termination command. + futs = [] + for dst_rank in range(self.world_size): + # torch.distributed.rpc module does not support sending to self. + if dst_rank == MASTER_RANK: + continue + dst_name = "worker{}".format(dst_rank) + fut = rpc.rpc_async(dst_name, set_termination_signal, args=()) + futs.append(fut) + for fut in futs: + assert fut.wait() is None, "Sending termination signal failed." # Close RPC. rpc.join_rpc() - return wrapper + + return return_value + + return new_test_method diff --git a/test/jit_utils.py b/test/jit_utils.py index 138beee6a40a6..54148c090c742 100644 --- a/test/jit_utils.py +++ b/test/jit_utils.py @@ -12,6 +12,7 @@ import torch.jit.quantized import zipfile import functools +from enum import Enum # Testing utils from common_utils import TestCase, IS_WINDOWS, \ @@ -32,6 +33,25 @@ import tempfile import textwrap +IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR = True + +class ProfilingMode(Enum): + OFF = 1 + EXECUTOR = 2 + FULL = 3 + +@contextmanager +def enable_profiling_mode(flag): + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + old_prof_exec_state = torch._C._jit_set_profiling_executor(flag != ProfilingMode.OFF) + #old_prof_mode_state = torch._C._jit_set_profiling_mode(flag == ProfilingMode.FULL) + old_prof_mode_state = torch._C._jit_set_profiling_mode(False) + try: + yield + finally: + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + torch._C._jit_set_profiling_executor(old_prof_exec_state) + torch._C._jit_set_profiling_mode(old_prof_mode_state) def execWrapper(code, glob, loc): if PY2: @@ -39,7 +59,6 @@ def execWrapper(code, glob, loc): else: exec(code, glob, loc) - def do_input_map(fn, input): return _nested_map(lambda t: isinstance(t, torch.Tensor), fn)(input) @@ -304,24 +323,34 @@ def get_frame_vars(self, frames_up): return defined_vars def checkScriptRaisesRegex(self, script, inputs, exception, regex, - outputs=None, capture_output=False): + outputs=None, capture_output=False, profiling=ProfilingMode.FULL): """ Checks that a given function will throw the correct exception, when executed with normal python, the string frontend, and the AST frontend """ - # normal python - with self.assertRaisesRegex(exception, regex): - script(*inputs) - # string frontend - with self.assertRaisesRegex(exception, regex): - source = textwrap.dedent(inspect.getsource(script)) - cu = torch.jit.CompilationUnit(source) - ge = getattr(cu, script.__name__) - ge(*inputs) - # python AST frontend - with self.assertRaisesRegex(exception, regex): - ge = torch.jit.script(script) - ge(*inputs) + + with enable_profiling_mode(profiling): + # normal python + with self.assertRaisesRegex(exception, regex): + script(*inputs) + # string frontend + with self.assertRaisesRegex(exception, regex): + source = textwrap.dedent(inspect.getsource(script)) + cu = torch.jit.CompilationUnit(source) + ge = getattr(cu, script.__name__) + # profiling run + with self.assertRaisesRegex(exception, regex): + ge(*inputs) + # optimized run + ge(*inputs) + # python AST frontend + with self.assertRaisesRegex(exception, regex): + ge = torch.jit.script(script) + # profiling run + with self.assertRaisesRegex(exception, regex): + ge(*inputs) + # optimized run + ge(*inputs) def checkScript(self, script, @@ -330,59 +359,71 @@ def checkScript(self, optimize=True, inputs_requires_grad=False, capture_output=False, - frames_up=1): + frames_up=1, + profiling=ProfilingMode.FULL): with torch.jit.optimized_execution(optimize): - if isinstance(script, str): - # Compile the string to a Script function - cu = torch.jit.CompilationUnit(script, _frames_up=frames_up) - - # Execute the Python function so we can run it later and get its - # outputs - frame = self.get_frame_vars(frames_up) - the_locals = {} - execWrapper(script, glob=frame, loc=the_locals) - frame.update(the_locals) - - python_fn = frame[name] - scripted_fn = getattr(cu, name) - else: - - # Check the string frontend first - source = textwrap.dedent(inspect.getsource(script)) - self.checkScript( - source, - inputs, - script.__name__, - capture_output, - frames_up=2) - - # Continue checking the Python frontend - scripted_fn = torch.jit.script(script, _frames_up=1) - python_fn = script - - if inputs_requires_grad: - recording_inputs = do_input_map(lambda t: t.detach().requires_grad_(), inputs) - else: - recording_inputs = inputs + with enable_profiling_mode(profiling): + if isinstance(script, str): + # Compile the string to a Script function + # with enable_profiling_mode(profiling): + cu = torch.jit.CompilationUnit(script, _frames_up=frames_up) + + # Execute the Python function so we can run it later and get its + # outputs + + frame = self.get_frame_vars(frames_up) + the_locals = {} + execWrapper(script, glob=frame, loc=the_locals) + frame.update(the_locals) + + python_fn = frame[name] + scripted_fn = getattr(cu, name) + else: - if capture_output: - with self.capture_stdout() as script_stdout: + # Check the string frontend first + source = textwrap.dedent(inspect.getsource(script)) + self.checkScript( + source, + inputs, + script.__name__, + capture_output, + profiling=profiling, + frames_up=2) + + # Continue checking the Python frontend + scripted_fn = torch.jit.script(script, _frames_up=1) + python_fn = script + + if inputs_requires_grad: + recording_inputs = do_input_map(lambda t: t.detach().requires_grad_(), inputs) + else: + recording_inputs = inputs + + if capture_output: + with self.capture_stdout() as script_stdout: + script_outputs = scripted_fn(*recording_inputs) + with self.capture_stdout() as opt_script_stdout: + opt_script_outputs = scripted_fn(*recording_inputs) + with self.capture_stdout() as _python_stdout: + python_outputs = python_fn(*inputs) + if not IS_WINDOWS: + self.assertExpected(script_stdout[0], subname='stdout') + self.assertEqual(python_outputs, opt_script_outputs) + else: + # profiling run script_outputs = scripted_fn(*recording_inputs) - with self.capture_stdout() as _python_stdout: + # optimized run + opt_script_outputs = scripted_fn(*recording_inputs) python_outputs = python_fn(*inputs) - if not IS_WINDOWS: - self.assertExpected(script_stdout[0], subname='stdout') - else: - script_outputs = scripted_fn(*recording_inputs) - python_outputs = python_fn(*inputs) - self.assertEqual(python_outputs, script_outputs) - - return scripted_fn + self.assertEqual(python_outputs, script_outputs) + self.assertEqual(script_outputs, opt_script_outputs) + return scripted_fn def checkTrace(self, func, reference_tensors, input_tensors=None, drop=None, allow_unused=False, verbose=False, inputs_require_grads=True, check_tolerance=1e-5, export_import=True, _force_outplace=False): + # TODO: check gradients for parameters, not just inputs def allSum(vs): # drop allows us to remove some values from ever being used @@ -413,8 +454,11 @@ def input_reduce(input, fn, acc): else: recording_inputs = reference_tensors + # `check_trace` is set to False because check_trace is run with @no_grad + # Also, `checkTrace` already does all the checks + # against python function ge = torch.jit.trace(func, input_tensors, check_tolerance=check_tolerance, - _force_outplace=_force_outplace) + _force_outplace=_force_outplace, check_trace=False) if export_import: ge = self.getExportImportCopy(ge) @@ -427,7 +471,6 @@ def input_reduce(input, fn, acc): outputs_ge = ge(*nograd_inputs) self.assertEqual(outputs, outputs_ge) - # test single grad case outputs = func(*recording_inputs) if inputs_require_grads: grads = torch.autograd.grad(allSum(outputs), flattened_recording_inputs, @@ -441,8 +484,11 @@ def input_reduce(input, fn, acc): if inputs_require_grads: self.assertEqual(grads, grads_ge) - # test the grad grad case + self.assertEqual(outputs, outputs_ge) + if inputs_require_grads: + self.assertEqual(grads, grads_ge) + # test the grad grad case outputs = func(*recording_inputs) l1 = allSum(outputs) if inputs_require_grads: @@ -514,14 +560,6 @@ def checkModule(self, nn_module, args): return sm -@contextmanager -def enable_profiling_mode(): - torch._C._jit_set_profiling_mode(True) - try: - yield - finally: - torch._C._jit_set_profiling_mode(False) - @contextmanager def inline_everything_mode(should_inline): old = torch._C._jit_get_inline_everything_mode() diff --git a/test/onnx/expect/TestOperators.test_equal.expect b/test/onnx/expect/TestOperators.test_equal.expect index d53ae0ded1c96..2c5fdc0dc3668 100644 --- a/test/onnx/expect/TestOperators.test_equal.expect +++ b/test/onnx/expect/TestOperators.test_equal.expect @@ -3,14 +3,14 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Equal" } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -32,7 +32,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/onnx/expect/TestOperators.test_ge.expect b/test/onnx/expect/TestOperators.test_ge.expect index abb3a62ef1a11..df2af85135dad 100644 --- a/test/onnx/expect/TestOperators.test_ge.expect +++ b/test/onnx/expect/TestOperators.test_ge.expect @@ -3,8 +3,8 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Less" } @@ -15,7 +15,7 @@ graph { } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -31,7 +31,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/onnx/expect/TestOperators.test_gt.expect b/test/onnx/expect/TestOperators.test_gt.expect index 0dbf73c132439..01cd7d05f7619 100644 --- a/test/onnx/expect/TestOperators.test_gt.expect +++ b/test/onnx/expect/TestOperators.test_gt.expect @@ -3,14 +3,14 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Greater" } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -32,7 +32,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/onnx/expect/TestOperators.test_le.expect b/test/onnx/expect/TestOperators.test_le.expect index 82674fb9038e9..b362b7699bd01 100644 --- a/test/onnx/expect/TestOperators.test_le.expect +++ b/test/onnx/expect/TestOperators.test_le.expect @@ -3,8 +3,8 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Greater" } @@ -15,7 +15,7 @@ graph { } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -31,7 +31,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/onnx/expect/TestOperators.test_lt.expect b/test/onnx/expect/TestOperators.test_lt.expect index b4688b4342c21..a18842e06592d 100644 --- a/test/onnx/expect/TestOperators.test_lt.expect +++ b/test/onnx/expect/TestOperators.test_lt.expect @@ -3,14 +3,14 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Less" } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -32,7 +32,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/rpc_test.py b/test/rpc_test.py index 775a83e686e08..f8d8f71c8276b 100644 --- a/test/rpc_test.py +++ b/test/rpc_test.py @@ -11,23 +11,24 @@ import torch.distributed.rpc as rpc from common_utils import load_tests from dist_utils import INIT_METHOD_TEMPLATE, TEST_CONFIG, dist_init -from torch.distributed.rpc import RpcBackend from torch.distributed.rpc.internal import PythonUDF, _internal_rpc_pickler def requires_process_group_agent(message=""): def decorator(old_func): return unittest.skipUnless( - TEST_CONFIG.rpc_backend == RpcBackend.PROCESS_GROUP, - message, + TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP", message )(old_func) + return decorator VALUE_FUTURE = concurrent.futures.Future() -def stub_start_rpc_backend_handler(store, self_name, self_rank, worker_name_to_id): +def stub_start_rpc_backend_handler( + store, self_name, self_rank, worker_name_to_id, *args, **kwargs +): return mock.Mock() # RpcAgent. @@ -229,12 +230,21 @@ def test_register_rpc_backend_and_start_rpc_backend( self, mock_rpc_agent, mock_dist_autograd_init ): backend_name = "stub_backend" - rpc.register_backend( + + backend = rpc.backend_registry.register_backend( backend_name, stub_start_rpc_backend_handler ) + + with self.assertRaisesRegex( + RuntimeError, "^RPC backend .+: already registered$" + ): + rpc.backend_registry.register_backend( + backend_name, stub_start_rpc_backend_handler + ) + rpc.init_model_parallel( self_name="worker1", - backend=backend_name, + backend=backend, init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, @@ -243,16 +253,10 @@ def test_register_rpc_backend_and_start_rpc_backend( @requires_process_group_agent("PROCESS_GROUP rpc backend specific test, skip") @dist_init(setup_model_parallel=False) def test_duplicate_name(self): - dist.init_process_group( - backend=dist.Backend.GLOO, - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) with self.assertRaisesRegex(RuntimeError, "is not unique"): rpc.init_model_parallel( self_name="duplicate_name", - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, @@ -261,87 +265,82 @@ def test_duplicate_name(self): @dist_init(setup_model_parallel=False) def test_reinit(self): - dist.init_process_group( - backend=dist.Backend.GLOO, - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) rpc.init_model_parallel( self_name="worker{}".format(self.rank), - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, ) + + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) # Wait for all init to complete. dist.barrier() + with self.assertRaisesRegex(RuntimeError, "is already initialized"): rpc.init_model_parallel( self_name="worker{}".format(self.rank), - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, ) rpc.join_rpc() - @dist_init(setup_model_parallel=False) - def test_init_invalid_backend(self): - with self.assertRaisesRegex(RuntimeError, "Unrecognized RPC backend"): - rpc.init_model_parallel( - self_name="worker{}".format(self.rank), - backend="invalid", - init_method=self.init_method, - self_rank=self.rank, - worker_name_to_id=self.worker_name_to_id, - ) - @dist_init(setup_model_parallel=False) def test_invalid_names(self): - dist.init_process_group( - backend=dist.Backend.GLOO, - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) - with self.assertRaisesRegex(RuntimeError, "Worker name must match"): rpc.init_model_parallel( self_name="abc*", - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, num_send_recv_threads=16, ) + base_file_name = self.file_name + # Use a different file path for FileStore to avoid rendezvous mismatch. + self.file_name = base_file_name + "1" with self.assertRaisesRegex(RuntimeError, "Worker name must match"): rpc.init_model_parallel( self_name=" ", - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, num_send_recv_threads=16, ) + # Use a different file path for FileStore to avoid rendezvous mismatch. + self.file_name = base_file_name + "2" with self.assertRaisesRegex(RuntimeError, "must be non-empty"): rpc.init_model_parallel( self_name="", - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, num_send_recv_threads=16, ) + # Use a different file path for FileStore to avoid rendezvous mismatch. + self.file_name = base_file_name + "3" # If the number in the message does not match, it is likely that the # value of MAX_NAME_LEN in RPC WorkerInfo has changed. with self.assertRaisesRegex(RuntimeError, "shorter than 128"): rpc.init_model_parallel( self_name="".join(["a" for _ in range(500)]), - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, @@ -363,7 +362,6 @@ def test_invalid_names(self): # would add extra overhead to the call, and normal use cases won't # create a progress group and exit without doing anything. Hence, it is # not worthy to introduce the overhead just for this test case. - dist.barrier() @dist_init def test_add(self): @@ -450,15 +448,9 @@ def test_sync_rpc(self): @dist_init(setup_model_parallel=False) def test_join_rpc(self): # Initialize RPC. - dist.init_process_group( - backend="gloo", - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) rpc.init_model_parallel( self_name="worker%d" % self.rank, - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, @@ -932,5 +924,18 @@ def test_requires_process_group_agent_decorator(self): def test_func(): return "expected result" - if TEST_CONFIG.rpc_backend == RpcBackend.PROCESS_GROUP: + if TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP": self.assertEqual(test_func(), "expected result") + + def test_dist_init_decorator(self): + @dist_init(setup_model_parallel=False) + def test_func(self): + return "expected result" + + self.assertEqual(test_func(self), "expected result") + + @dist_init + def test_func(self): + return "expected result" + + self.assertEqual(test_func(self), "expected result") diff --git a/test/run_test.py b/test/run_test.py index 631fa28a33cb2..52bfff8df0511 100755 --- a/test/run_test.py +++ b/test/run_test.py @@ -47,7 +47,6 @@ 'quantized', 'quantized_tensor', 'quantized_nn_mods', - 'quantizer', 'sparse', 'torch', 'type_info', diff --git a/test/scripts/cuda_memcheck_common.py b/test/scripts/cuda_memcheck_common.py new file mode 100644 index 0000000000000..7f7dc8253393f --- /dev/null +++ b/test/scripts/cuda_memcheck_common.py @@ -0,0 +1,98 @@ +# this file contains a simple parser that parses report +# from cuda-memcheck + +class ParseError(Exception): + """Whenever the simple parser is unable to parse the report, this exception will be raised""" + pass + + +class Report: + """A report is a container of errors, and a summary on how many errors are found""" + + HEAD = 'ERROR SUMMARY: ' + TAIL = ' errors' + + def __init__(self, text, errors): + self.text = text + self.num_errors = int(text[len(self.HEAD):len(text) - len(self.TAIL)]) + self.errors = errors + if len(errors) != self.num_errors: + raise ParseError("Number of errors does not match") + + +class Error: + """Each error is a section in the output of cuda-memcheck. + Each error in the report has an error message and a backtrace. It looks like: + + ========= Program hit cudaErrorInvalidValue (error 1) due to "invalid argument" on CUDA API call to cudaGetLastError. + ========= Saved host backtrace up to driver entry point at error + ========= Host Frame:/usr/lib/x86_64-linux-gnu/libcuda.so.1 [0x38c7b3] + ========= Host Frame:/usr/local/cuda/lib64/libcudart.so.10.1 (cudaGetLastError + 0x163) [0x4c493] + ========= Host Frame:/home/xgao/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch.so [0x5b77a05] + ========= Host Frame:/home/xgao/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch.so [0x39d6d1d] + ========= ..... + """ + + def __init__(self, lines): + self.message = lines[0] + lines = lines[2:] + self.stack = [l.strip() for l in lines] + + +def parse(message): + """A simple parser that parses the report of cuda-memcheck. This parser is meant to be simple + and it only split the report into separate errors and a summary. Where each error is further + splitted into error message and backtrace. No further details are parsed. + + A report contains multiple errors and a summary on how many errors are detected. It looks like: + + ========= CUDA-MEMCHECK + ========= Program hit cudaErrorInvalidValue (error 1) due to "invalid argument" on CUDA API call to cudaPointerGetAttributes. + ========= Saved host backtrace up to driver entry point at error + ========= Host Frame:/usr/lib/x86_64-linux-gnu/libcuda.so.1 [0x38c7b3] + ========= Host Frame:/usr/local/cuda/lib64/libcudart.so.10.1 (cudaPointerGetAttributes + 0x1a9) [0x428b9] + ========= Host Frame:/home/xgao/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch.so [0x5b778a9] + ========= ..... + ========= + ========= Program hit cudaErrorInvalidValue (error 1) due to "invalid argument" on CUDA API call to cudaGetLastError. + ========= Saved host backtrace up to driver entry point at error + ========= Host Frame:/usr/lib/x86_64-linux-gnu/libcuda.so.1 [0x38c7b3] + ========= Host Frame:/usr/local/cuda/lib64/libcudart.so.10.1 (cudaGetLastError + 0x163) [0x4c493] + ========= ..... + ========= + ========= ..... + ========= + ========= Program hit cudaErrorInvalidValue (error 1) due to "invalid argument" on CUDA API call to cudaGetLastError. + ========= Saved host backtrace up to driver entry point at error + ========= Host Frame:/usr/lib/x86_64-linux-gnu/libcuda.so.1 [0x38c7b3] + ========= ..... + ========= Host Frame:python (_PyEval_EvalFrameDefault + 0x6a0) [0x1d0ad0] + ========= Host Frame:python (_PyEval_EvalCodeWithName + 0xbb9) [0x116db9] + ========= + ========= ERROR SUMMARY: 4 errors + """ + errors = [] + HEAD = '=========' + headlen = len(HEAD) + started = False + in_message = False + message_lines = [] + lines = message.splitlines() + for l in lines: + if l == HEAD + ' CUDA-MEMCHECK': + started = True + continue + if not started or not l.startswith(HEAD): + continue + l = l[headlen + 1:] + if l.startswith('ERROR SUMMARY:'): + return Report(l, errors) + if not in_message: + in_message = True + message_lines = [l] + elif l == '': + errors.append(Error(message_lines)) + in_message = False + else: + message_lines.append(l) + raise ParseError("No error summary found") diff --git a/test/scripts/run_cuda_memcheck.py b/test/scripts/run_cuda_memcheck.py new file mode 100755 index 0000000000000..f80fa84350c74 --- /dev/null +++ b/test/scripts/run_cuda_memcheck.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python + +"""This script runs cuda-memcheck on the specified unit test. Each test case +is run in its isolated process with a timeout so that: +1) different test cases won't influence each other, and +2) in case of hang, the script would still finish in a finite amount of time. +The output will be written to a log file result.log + +Example usage: + python run_cuda_memcheck.py ../test_torch.py 600 + +Note that running cuda-memcheck could be very slow. +""" + +import asyncio +import torch +import multiprocessing +import argparse +import subprocess +import tqdm +import re +import cuda_memcheck_common as cmc + +ALL_TESTS = [] +GPUS = torch.cuda.device_count() + +# parse arguments +parser = argparse.ArgumentParser(description="Run isolated cuda-memcheck on unit tests") +parser.add_argument('filename', help="the python file for a test, such as test_torch.py") +parser.add_argument('timeout', type=int, help='kill the test if it does not terminate in a certain amount of seconds') +parser.add_argument('--strict', action='store_true', + help='Whether to show cublas/cudnn errors. These errors are ignored by default because' + 'cublas/cudnn does not run error-free under cuda-memcheck, and ignoring these errors') +parser.add_argument('--nproc', type=int, default=multiprocessing.cpu_count(), + help='Number of processes running tests, default to number of cores in the system') +parser.add_argument('--gpus', default='all', + help='GPU assignments for each process, it could be "all", or : separated list like "1,2:3,4:5,6"') +args = parser.parse_args() + +# Filters that ignores cublas/cudnn errors +# TODO (@zasdfgbnm): When can we remove this? Will cublas/cudnn run error-free under cuda-memcheck? +def is_ignored_only(output): + try: + report = cmc.parse(output) + except cmc.ParseError: + # in case the simple parser fails parsing the output of cuda memcheck + # then this error is never ignored. + return False + count_ignored_errors = 0 + for e in report.errors: + if 'libcublas' in ''.join(e.stack) or 'libcudnn' in ''.join(e.stack): + count_ignored_errors += 1 + return count_ignored_errors == report.num_errors + +# Discover tests: +# To get a list of tests, run: +# pytest --setup-only test/test_torch.py +# and then parse the output +proc = subprocess.Popen(['pytest', '--setup-only', args.filename], stdout=subprocess.PIPE, stderr=subprocess.PIPE) +stdout, stderr = proc.communicate() +lines = stdout.decode().strip().splitlines() +for line in lines: + if '(fixtures used:' in line: + line = line.strip().split()[0] + line = line[line.find('::') + 2:] + line = line.replace('::', '.') + ALL_TESTS.append(line) + +# Run tests: +# Since running cuda-memcheck on PyTorch unit tests is very slow, these tests must be run in parallel. +# This is done by using the coroutine feature in new Python versions. A number of coroutines are created; +# they create subprocesses and awaiting them to finish. The number of running subprocesses could be +# specified by the user and by default is the same as the number of CPUs in the machine. +# These subprocesses are balanced across different GPUs on the system by assigning one devices per process, +# or as specified by the user +progress = 0 +logfile = open('result.log', 'w') +progressbar = tqdm.tqdm(total=len(ALL_TESTS)) + +async def run1(coroutine_id): + global progress + + if args.gpus == 'all': + gpuid = coroutine_id % GPUS + else: + gpu_assignments = args.gpus.split(':') + assert args.nproc == len(gpu_assignments), 'Please specify GPU assignmnent for each process, separated by :' + gpuid = gpu_assignments[coroutine_id] + + while progress < len(ALL_TESTS): + test = ALL_TESTS[progress] + progress += 1 + cmd = f'CUDA_VISIBLE_DEVICES={gpuid} cuda-memcheck --error-exitcode 1 python {args.filename} {test}' + proc = await asyncio.create_subprocess_shell(cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE) + try: + stdout, stderr = await asyncio.wait_for(proc.communicate(), args.timeout) + except asyncio.TimeoutError: + print('Timeout:', test, file=logfile) + proc.kill() + else: + if proc.returncode == 0: + print('Success:', test, file=logfile) + else: + stdout = stdout.decode() + stderr = stderr.decode() + should_display = args.strict or not is_ignored_only(stdout) + if should_display: + print('Fail:', test, file=logfile) + print(stdout, file=logfile) + print(stderr, file=logfile) + else: + print('Ignored:', test, file=logfile) + del proc + progressbar.update(1) + +async def main(): + tasks = [asyncio.create_task(run1(i)) for i in range(args.nproc)] + for t in tasks: + await t + +if __name__ == '__main__': + loop = asyncio.get_event_loop() + loop.run_until_complete(main()) diff --git a/test/test_autograd.py b/test/test_autograd.py index 72b6ae588197b..a5ba360f58909 100644 --- a/test/test_autograd.py +++ b/test/test_autograd.py @@ -39,7 +39,7 @@ mask_not_all_zeros, S) from common_device_type import (instantiate_device_type_tests, skipCUDAIfRocm, - onlyCUDA, dtypes, dtypesIfCUDA, + onlyCPU, onlyCUDA, dtypes, dtypesIfCUDA, deviceCountAtLeast, skipCUDAIfCudnnVersionLessThan) # load_tests from common_utils is used to automatically filter tests for @@ -3958,6 +3958,19 @@ def test_inputbuffer_add_multidevice(self, devices): output = input.to(device=devices[1]) + input.to(device=devices[1]) output.backward() + @onlyCPU + def test_copy_(self, device): + # At the time of writing this test, copy_ is not generated from native_functions.yaml + # there was a bug that bfloat16 was not recognized as floating. + x = torch.randn(10, device=device, requires_grad=True) + floating_dt = [dt for dt in torch.testing.get_all_dtypes() if dt.is_floating_point] + for dt in floating_dt: + y = torch.empty(10, device=device, dtype=dt) + y.copy_(x) + self.assertTrue(y.requires_grad) + z = x.to(torch.bfloat16) + self.assertTrue(z.requires_grad) + @onlyCUDA def test_cross_device_reentrant_autograd(self, device): # Output on gpu so that this task will be associated with the gpu thread diff --git a/test/test_fake_quant.py b/test/test_fake_quant.py index ee82db730fbf3..0a4443717a270 100644 --- a/test/test_fake_quant.py +++ b/test/test_fake_quant.py @@ -10,12 +10,13 @@ from torch.quantization import FakeQuantize from torch.quantization import default_observer, default_per_channel_weight_observer import io +import unittest + # Reference method for fake quantize def _fake_quantize_per_tensor_affine_reference(X, scale, zero_point, quant_min, quant_max): res = (torch.clamp(torch.round(X.cpu() * (1.0 / scale) + zero_point), quant_min, quant_max) - zero_point) * scale return res - # Reference method for the gradient of the fake quantize operator def _fake_quantize_per_tensor_affine_grad_reference(dY, X, scale, zero_point, quant_min, quant_max): Xq = torch.round(X.cpu() * (1.0 / scale) + zero_point) @@ -256,6 +257,7 @@ def test_backward_per_channel(self, device, X): @given(device=st.sampled_from(['cpu', 'cuda'] if torch.cuda.is_available() else ['cpu']), X=hu.per_channel_tensor(shapes=hu.array_shapes(1, 5,), qparams=hu.qparams(dtypes=torch.quint8))) + @unittest.skip("temporarily disable the test") def test_numerical_consistency_per_channel(self, device, X): r"""Comparing numerical consistency between CPU quantize/dequantize op and the CPU fake quantize op """ diff --git a/test/test_jit.py b/test/test_jit.py index b555f1877b78e..0195d3b7bcdf9 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -41,9 +41,9 @@ skipIfRocm, skipIfNoLapack, suppress_warnings, IS_SANDCASTLE, \ freeze_rng_state, set_rng_seed, slowTest, TemporaryFileName, skipIfCompiledWithoutNumpy from jit_utils import JitTestCase, enable_cpu_fuser, disable_autodiff_subgraph_inlining, \ - _trace, enable_cpu_fuser_if, enable_profiling_mode, do_input_map, \ + _trace, enable_cpu_fuser_if, enable_profiling_mode, ProfilingMode, do_input_map, \ execWrapper, _inline_everything, _tmp_donotuse_dont_inline_everything, \ - get_forward, get_forward_graph, get_module_method + get_forward, get_forward_graph, get_module_method, IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR from common_nn import module_tests, new_module_tests, criterion_tests from common_methods_invocations import method_tests as autograd_method_tests from common_methods_invocations import create_input, unpack_variables, \ @@ -102,6 +102,63 @@ def LSTMCellF(input, hx, cx, *params): return LSTMCell(input, (hx, cx), *params) +def doAutodiffCheck(testname): + + if not IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + return True + + + # these tests are disabled because BailOut nodes + # inserted by ProfilingExecutor interfere with + # subgraph slicing of Differentiable Graphs + test_exceptions = [ + # functional + 'test_nn_dropout', + 'test_nn_log_softmax', + 'test_nn_relu', + 'test_nn_softmax', + 'test_nn_threshold', + 'test_nn_lp_pool2d', + 'test_nn_lp_pool1d', + 'test_nn_gumbel_softmax_hard', + 'test_nn_gumbel_softmax', + 'test_nn_multilabel_soft_margin_loss', + 'test_nn_batch_norm', + # AutogradJitGenerated + 'test___rdiv___constant', + 'test___rdiv___scalar_constant', + ] + + if testname in test_exceptions: + return False + return True + +func_call = torch._C.ScriptFunction.__call__ +meth_call = torch._C.ScriptMethod.__call__ + +def prof_callable(callable, *args, **kwargs): + if 'profile_and_replay' in kwargs: + del kwargs['profile_and_replay'] + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + with enable_profiling_mode(ProfilingMode.FULL): + callable(*args, **kwargs) + return callable(*args, **kwargs) + + return callable(*args, **kwargs) + +def prof_func_call(*args, **kwargs): + return prof_callable(func_call, *args, **kwargs) + +def prof_meth_call(*args, **kwargs): + return prof_callable(meth_call, *args, **kwargs) + +torch._C.ScriptFunction.__call__ = prof_func_call +torch._C.ScriptMethod.__call__ = prof_meth_call + +if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + # enable profiling graph executor for all tests in this file by default + torch._C._jit_set_profiling_executor(True) + def LSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None): hx, cx = hidden @@ -204,13 +261,16 @@ def get_execution_plan(graph_executor_state): return execution_plans[0] -def get_grad_executor(plan_state, diff_graph_idx=None): +def get_grad_executor(plan_state, diff_graph_idx=None, skip_check=False): if diff_graph_idx is None: nodes = list(plan_state.graph.nodes()) - if len(nodes) == 1 or (len(nodes) == 2 and nodes[1].kind() == "prim::TupleConstruct"): - pass - else: - raise RuntimeError("Can't get a grad_executor for a non-differentiable graph") + + if not skip_check: + nodes = list(filter(lambda n : n.kind() != "prim::BailOut" and n.kind() != "prim::BailoutTemplate", nodes)) + if len(nodes) == 1 or (len(nodes) == 2 and nodes[1].kind() == "prim::TupleConstruct"): + pass + else: + raise RuntimeError("Can't get a grad_executor for a non-differentiable graph") grad_executors = list(plan_state.code.grad_executor_states()) return grad_executors[diff_graph_idx or 0] @@ -224,10 +284,10 @@ def all_backward_graphs(script_module, diff_graph_idx=None): return [p.graph.copy() for p in bwd_plans] -def backward_graph(script_module, diff_graph_idx=None): +def backward_graph(script_module, diff_graph_idx=None, skip_check=False): ge_state = script_module.get_debug_state() fwd_plan = get_execution_plan(ge_state) - grad_executor_state = get_grad_executor(fwd_plan, diff_graph_idx=diff_graph_idx) + grad_executor_state = get_grad_executor(fwd_plan, diff_graph_idx=diff_graph_idx, skip_check=skip_check) bwd_plan = get_execution_plan(grad_executor_state) # Running JIT passes requires that we own the graph (with a shared_ptr). # The debug state struct does not own its graph so we make a copy of it. @@ -502,9 +562,14 @@ def f(x, y): def test_peephole_optimize_shape_ops(self): def test_input(func, input, result): - self.assertEqual(func(input), result) + # if result == 2 we will trigger a bailout and + # the unprofiled graph should return the correct result + self.assertEqual(func(input, profile_and_replay=True), result) gre = func.graph_for(input) - FileCheck().check_not("prim::If").run(gre) + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + FileCheck().check("prim::Constant").check_next("prim::BailoutTemplate").run(gre) + else: + FileCheck().check_not("prim::If").run(gre) def test_dim(): @torch.jit.script @@ -940,23 +1005,23 @@ def forward(self, x): } torch._C._jit_pass_insert_observers(m._c, "forward", qconfig_dict, True) assert len([x for x, _ in m._c._get_modules() - if x.startswith('observer_for_')]) == 0, \ + if x.startswith('_observer_')]) == 0, \ 'Expected to have 0 observer submodules' - FileCheck().check_not('ClassType = prim::GetAttr[name="observer_for_') \ + FileCheck().check_not('ClassType = prim::GetAttr[name="_observer_') \ .check('ClassType = prim::GetAttr[name="conv"](%self)') \ .check_next('Tensor = prim::CallMethod[name="forward"]') \ - .check_not('ClassType = prim::GetAttr[name="observer_for_') \ + .check_not('ClassType = prim::GetAttr[name="_observer_') \ .run(str(get_forward_graph(m._c))) assert len([x for x, _ in m._c._get_module('conv')._get_modules() - if x.startswith('observer_for_')]) == 3, \ + if x.startswith('_observer_')]) == 3, \ 'Expected to have 3 observer submodules' - FileCheck().check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ - .check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ + FileCheck().check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ + .check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ .check('Tensor = aten::conv2d') \ - .check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ + .check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ .run(str(m._c._get_module("conv")._get_method('conv2d_forward').graph)) @_tmp_donotuse_dont_inline_everything @@ -979,17 +1044,17 @@ def forward(self, x): return self.sub(self.conv(x)) def check_observed(s): - FileCheck().check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ - .check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ - .check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ + FileCheck().check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ + .check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ + .check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ .run(str(s)) def check_not_observed(s): - FileCheck().check_not('ClassType = prim::GetAttr[name="observer_for_') \ - .check_not('prim::CallMethod[name="forward"](%observer_for_') \ + FileCheck().check_not('ClassType = prim::GetAttr[name="_observer_') \ + .check_not('prim::CallMethod[name="forward"](%_observer_') \ .run(str(s)) m = torch.jit.script(M()) @@ -1053,15 +1118,15 @@ def test_module(module, relu_call, num_observers): } torch._C._jit_pass_insert_observers(m._c, "forward", qconfig_dict, True) assert len([x for x, _ in m._c._get_modules() - if x.startswith('observer_for_')]) == num_observers, \ + if x.startswith('_observer_')]) == num_observers, \ 'Expected to have ' + str(num_observers) + ' observer submodules' c = FileCheck().check('ClassType = prim::GetAttr[name="conv"]') \ .check_next('prim::CallMethod[name="forward"]') \ - .check_not('ClassType = prim::GetAttr[name="observer_for_') \ + .check_not('ClassType = prim::GetAttr[name="_observer_') \ .check(relu_call) if num_observers == 1: - c = c.check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') + c = c.check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') c.run(str(get_forward_graph(m._c))) # TODO: add checks for conv and relu later, graph looks correct but this pr # has too many changes already @@ -1088,8 +1153,9 @@ def forward(self, x): weight=weight_observer._c) } torch._C._jit_pass_insert_observers(m._c, "forward", qconfig_dict, True) - assert m._c._get_module('conv')._get_module('observer_for_input.1')._get_attribute('dtype') != \ - m._c._get_module('conv')._get_module('observer_for_weight.1')._get_attribute('dtype') + dtypes = set([obs._get_attribute('dtype') for x, obs in m._c._get_module('conv')._get_modules() + if x.startswith('_observer_')]) + assert len(dtypes) == 2, 'Expected to have 2 different types of dtype' @_tmp_donotuse_dont_inline_everything def test_insert_quant_dequant(self): @@ -1970,6 +2036,7 @@ def forward(self, d): inputs = {'x': torch.rand(3, 4), 'y': torch.rand(3, 4)} module = torch.jit.trace(Test(), inputs) + FileCheck().check('aten::values').check('prim::ListUnpack').run(str(module.graph)) def test_input_dict_flattens_recursive(self): @@ -2154,20 +2221,21 @@ def test_dropout_cuda(self): # which is not included in TestJitGeneratedFunctional x = torch.ones(4, 4).cuda().requires_grad_() - @torch.jit.script - def func(x): - return torch.nn.functional.dropout(x) + with enable_profiling_mode(ProfilingMode.FULL): + @torch.jit.script + def func(x): + return torch.nn.functional.dropout(x) - with freeze_rng_state(): - out_ref = torch.nn.functional.dropout(x) - grad_ref = torch.autograd.grad(out_ref.sum(), x) + with freeze_rng_state(): + out_ref = torch.nn.functional.dropout(x) + grad_ref = torch.autograd.grad(out_ref.sum(), x) - with freeze_rng_state(): - out = func(x) - grad = torch.autograd.grad(out.sum(), x) + with freeze_rng_state(): + out = func(x) + grad = torch.autograd.grad(out.sum(), x) - self.assertEqual(out, out_ref) - self.assertEqual(grad, grad_ref) + self.assertEqual(out, out_ref) + self.assertEqual(grad, grad_ref) def test_conv(self): x = torch.ones(20, 16, 50, 40) @@ -2297,8 +2365,7 @@ def rand(*args): self.checkTrace(lambda a, b: a * b + b, [rand(1), rand(1)], [rand(2, 3), rand(2, 3)]) # trivial identity - self.checkTrace(lambda a, b: ( - b, a), [rand(1), rand(1)]) + self.checkTrace(lambda a, b: (b, a), [rand(1), rand(1)]) def foo(a): t = a * a @@ -2319,7 +2386,8 @@ def test_ge_unoptimized(self): @unittest.skipIf(IS_SANDCASTLE, "NYI: fuser support for Sandcastle") @enable_cpu_fuser def test_ge_optimized(self): - self.run_ge_tests(True, False) + with enable_profiling_mode(ProfilingMode.FULL): + self.run_ge_tests(True, False) @unittest.skipIf(not RUN_CUDA, "requires CUDA") def test_ge_cuda(self): @@ -5269,6 +5337,7 @@ def func(alpha, beta, x, y): # NOTE: cannot optimize yet because broadcasts are not inserted before the fuser runs self.checkScript(func, [alpha, beta, x, y], optimize=False) + @unittest.skipIf(not IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "skip if profiling isn't enabled") def test_profiling_graph_executor(self): @torch.jit.script def def_in_one_branch(x, z): @@ -5281,12 +5350,15 @@ def def_in_one_branch(x, z): a = torch.rand(2, 3) - with enable_profiling_mode(): - # the first call is profiled - profiled_graph_str = str(def_in_one_branch.graph_for(a, False)) + with enable_profiling_mode(ProfilingMode.FULL): + # check prim::profile are inserted + profiled_graph_str = str(def_in_one_branch.graph_for(a, True)) FileCheck().check_count("prim::profile", 4).run(profiled_graph_str) - # the second call is optimized + # this call is optimized for + # the given shape of (2, 3) def_in_one_branch(a, False) + # change shape to (3) + # so we go down a bailout path a = torch.ones(3) # check prim::BailOuts are inserted bailout_graph_str = str(def_in_one_branch.graph_for(a, True)) @@ -5296,7 +5368,6 @@ def def_in_one_branch(x, z): # this triggers 2 bailouts self.assertEqual(def_in_one_branch(a, True), 3.0) - def test_resize_input_ops(self): # resize_ and resize_as resize the input tensor. because our shape analysis # is flow invariant, we set any Tensor that can alias a resized Tensor @@ -5389,7 +5460,7 @@ def test(x, y, z): # and the output of the node conservatively setting grad to true inps = (torch.tensor(1.0, requires_grad=True), torch.tensor(1), 10) - test(*inps) + test(*inps, profile_and_replay=True) graph = test.graph_for(*inps) loop = graph.findNode("prim::Loop") @@ -5398,8 +5469,13 @@ def test(x, y, z): loop_outputs = list(loop_body.outputs()) self.assertTrue(loop_inputs[1].requires_grad()) - self.assertFalse(loop_outputs[1].requires_grad()) - self.assertTrue(loop.output().requires_grad()) + + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + bailouts_in_outer_block = graph.findAllNodes("prim::BailOut", False) + self.assertFalse(bailouts_in_outer_block[1].output().requires_grad()) + else: + self.assertTrue(loop.output().requires_grad()) + self.assertFalse(loop_outputs[1].requires_grad()) def test_view_shape_prop(self): cu = torch.jit.CompilationUnit(''' @@ -6041,6 +6117,7 @@ def test(x, y): with self.assertRaisesRegex(Exception, ""): test(1, None) + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "the current version of Profiler doesn't profile/specialize Optionals") def test_optional_tensor(self): @torch.jit.script def fn(x, y): @@ -6081,6 +6158,7 @@ def fn(x, y, b): g = torch.jit.last_executed_optimized_graph() self.assertEqual(next(g.outputs()).type().str(), "Tensor") + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "the current version of Profiler doesn't profile/specialize Optionals") def test_optional_list(self): @torch.jit.script def fn(x, y): @@ -6784,17 +6862,23 @@ def func(): ''') ops = ['tensor', 'as_tensor'] inputs = ['[1]', '[False]', '[2.5]', '0.5', '1', 'False', '[[1]]'] - expected_shape = ["Long(*)", ("Bool(*)"), "Double(*)", "Double()", "Long()", "Bool()", "Long(*, *)"] + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + expected_shape = ["Long(1)", "Bool(1)", "Double(1)", "Double()", "Long()", "Bool()", "Long(1, 1)"] + else: + expected_shape = ["Long(*)", ("Bool(*)"), "Double(*)", "Double()", "Long()", "Bool()", "Long(*, *)"] for op in ops: for inp, expect in zip(inputs, expected_shape): code = tensor_template.format(tensor_op=op, input=inp) scope = {} exec(code, globals(), scope) - self.checkScript(code, ()) - cu = torch.jit.CompilationUnit(code) - torch._C._jit_pass_complete_shape_analysis(cu.func.graph, (), False) - FileCheck().check(expect).check("aten::{tensor_op}".format(tensor_op=op)).run(cu.func.graph) + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + fn = self.checkScript(code, ()) + FileCheck().check(expect).check("aten::{tensor_op}".format(tensor_op=op)).run(fn.graph_for()) + else: + cu = torch.jit.CompilationUnit(code) + torch._C._jit_pass_complete_shape_analysis(cu.func.graph, (), False) + FileCheck().check(expect).check("aten::{tensor_op}".format(tensor_op=op)).run(cu.func.graph) @torch.jit.script def test_dtype(inp_dtype): @@ -6802,17 +6886,26 @@ def test_dtype(inp_dtype): a = torch.tensor(1.0, dtype=torch.float, requires_grad=True) return a, torch.tensor(1.0, dtype=inp_dtype) # noqa T484 - g = test_dtype.graph_for(5) - # first should have type set second should not - FileCheck().check("Float() = aten::tensor").check("Tensor = aten::tensor").run(g) + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + g = test_dtype.graph_for(5, profile_and_replay=True) + # both should have completed shapes + FileCheck().check("Tensor = aten::tensor").check("Float() = prim::BailOut").check("Tensor = aten::tensor").check("Half() = prim::BailOut").run(g) + else: + g = test_dtype.graph_for(5) + # first should have type set second should not + FileCheck().check("Float() = aten::tensor").check("Tensor = aten::tensor").run(g) @torch.jit.script def test_as_tensor_tensor_input(input): a = torch.as_tensor(input, dtype=input.dtype) return a, torch.as_tensor(input, dtype=torch.float) - g = test_as_tensor_tensor_input.graph_for(torch.ones(3, 4)) - FileCheck().check("Tensor = aten::as_tensor").check("Float(*, *) = aten::as_tensor").run(g) + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + g = test_as_tensor_tensor_input.graph_for(torch.ones(3, 4), profile_and_replay=True) + FileCheck().check("Tensor = aten::as_tensor").check("Float(3, 4) = prim::BailOut").check("Tensor = aten::as_tensor").check("Float(3, 4) = prim::BailOut").run(g) + else: + g = test_as_tensor_tensor_input.graph_for(torch.ones(3, 4)) + FileCheck().check("Tensor = aten::as_tensor").check("Float(*, *) = aten::as_tensor").run(g) def test_tensor_requires_grad(self): @@ -6909,7 +7002,7 @@ def s(t, to_str, non_blocking=None, device=None, cuda=None): code = template.format(to_str=to_str, device=device, non_blocking=non_blocking, cuda=cuda) scope = {} cu = torch.jit.CompilationUnit(code) - return cu.func(t) + return cu.func(t, profile_and_replay=True) def test_copy_behavior(t, non_blocking=False): self.assertIs(t, s(t, 't.to(t, non_blocking=non_blocking)', non_blocking)) @@ -10797,6 +10890,7 @@ def foo(): self.checkScript(foo, ()) + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "the original version of test_rand") def test_rand(self): def test_rand(): a = torch.rand([3, 4]) @@ -10819,6 +10913,28 @@ def randint(): # and shape analysis dtype is the same. FileCheck().check("Double(*, *)").check_not("Float(*, *)").run(randint.graph_for()) + @unittest.skipIf(not IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "the original version of test_rand") + def test_rand_profiling(self): + def test_rand(): + a = torch.rand([3, 4]) + return a + 1.0 - a + + fn = self.checkScript(test_rand, ()) + out = fn() + self.assertEqual(out.dtype, torch.double) + # Testing shape analysis correctly setting type + FileCheck().check("Double(3, 4)").check_not("Float(3, 4)").run(fn.graph_for()) + + @torch.jit.script + def randint(): + return torch.randint(0, 5, [1, 2]) + + out = randint(profile_and_replay=True) + self.assertEqual(out.dtype, torch.double) + # although the type should be int here, testing that the runtime dtype + # and shape analysis dtype is the same. + FileCheck().check("Double(1, 2)").check_not("Float(1, 2)").run(randint.graph_for()) + def test_erase_number_types(self): def func(a): b = 7 + 1 + 3 @@ -10847,8 +10963,9 @@ def lstm(x, hx, cx, w_ih, w_hh, b_ih, b_hh): fw_graph = slstm.graph_for(*inputs) bw_graph = backward_graph(slstm, diff_graph_idx=0) - self.assertTrue('prim::MMBatchSide' in str(fw_graph)) - self.assertTrue('prim::MMTreeReduce' in str(bw_graph)) + if not IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + self.assertTrue('prim::MMBatchSide' in str(fw_graph)) + self.assertTrue('prim::MMTreeReduce' in str(bw_graph)) sout = slstm(*inputs) out = lstm(*inputs) @@ -13693,7 +13810,7 @@ def fn(): x.add_(torch.ones(2, 3)) return x_view - self.checkScript(fn, ()) + self.checkScript(fn, (), profiling=ProfilingMode.EXECUTOR) def test_cpp_function_tensor_str(self): x = torch.randn(2, 2) @@ -16322,7 +16439,10 @@ def new_fn(*tensors_): def create_traced_fn(self, fn): def traced_fn(*inputs, **kwargs): fn_tensors, inputs_tensors = partial_apply_nontensors(fn, inputs, **kwargs) - traced = torch.jit.trace(fn_tensors, inputs_tensors) + # `check_trace` is set to False because check_trace is run with @no_grad + # Also, `check_against_reference` already does all the checks + # against python function + traced = torch.jit.trace(fn_tensors, inputs_tensors, check_trace=False) self.assertExportImport(traced.graph, inputs_tensors) output = traced(*inputs_tensors) traced_fn.last_graph = traced.graph_for(*inputs_tensors) @@ -16393,7 +16513,6 @@ def create_script_fn(self, method_name, func_type, output_process_fn): def script_fn(*args, **kwargs): formals, tensors, actuals = get_script_args(args) call = get_call(method_name, func_type, actuals, kwargs) - script = script_template.format(', '.join(formals), call) CU = torch.jit.CompilationUnit(script) @@ -16443,7 +16562,8 @@ def clone_inputs(requires_grad): # test no gradients case outputs = self.runAndSaveRNG(reference_func, nograd_inputs, kwargs) - outputs_test = self.runAndSaveRNG(func, nograd_inputs, kwargs) + with enable_profiling_mode(ProfilingMode.FULL): + outputs_test = self.runAndSaveRNG(func, nograd_inputs, kwargs) self.assertEqual(outputs, outputs_test) if check_types: @@ -16453,43 +16573,42 @@ def clone_inputs(requires_grad): # skip grad tests return - # test single grad case - outputs = self.runAndSaveRNG(reference_func, recording_inputs, kwargs) - grads = torch.autograd.grad(allSum(outputs), recording_tensors, - allow_unused=allow_unused) - - outputs_test = self.runAndSaveRNG(func, recording_inputs, kwargs) - grads_test = torch.autograd.grad(allSum(outputs_test), recording_tensors, - allow_unused=allow_unused) - self.assertEqual(outputs, outputs_test) - self.assertEqual(grads, grads_test) - - # test the grad grad case - if self._testMethodName in nn_functional_single_grad: - return - - outputs = self.runAndSaveRNG(reference_func, recording_inputs, kwargs) - l1 = allSum(outputs) - grads = torch.autograd.grad(l1, recording_tensors, create_graph=True, - allow_unused=allow_unused) - l2 = (allSum(grads) * l1) - grads2 = torch.autograd.grad(l2, recording_tensors, allow_unused=allow_unused) - - recording_inputs, recording_tensors = clone_inputs(True) - - outputs_test = self.runAndSaveRNG(func, recording_inputs, kwargs) - l1_test = allSum(outputs_test) - grads_test = torch.autograd.grad( - l1_test, recording_tensors, create_graph=True, allow_unused=allow_unused) - l2_test = (allSum(grads_test) * l1_test) - grads2_test = torch.autograd.grad(l2_test, recording_tensors, allow_unused=allow_unused) + with enable_profiling_mode(ProfilingMode.FULL): + # test single grad case + outputs = self.runAndSaveRNG(reference_func, recording_inputs, kwargs) + grads = torch.autograd.grad(allSum(outputs), recording_tensors, + allow_unused=allow_unused) + outputs_test = self.runAndSaveRNG(func, recording_inputs, kwargs) + grads_test = torch.autograd.grad(allSum(outputs_test), recording_tensors, + allow_unused=allow_unused) + self.assertEqual(outputs, outputs_test) + self.assertEqual(grads, grads_test) + # test the grad grad case + if self._testMethodName in nn_functional_single_grad: + return - self.assertEqual(outputs, outputs_test) - self.assertEqual(grads, grads_test) - for g2, g2_test in zip(grads2, grads2_test): - if g2 is None and g2_test is None: - continue - self.assertTrue(torch.allclose(g2, g2_test, atol=5e-4, rtol=1e-4)) + outputs = self.runAndSaveRNG(reference_func, recording_inputs, kwargs) + l1 = allSum(outputs) + grads = torch.autograd.grad(l1, recording_tensors, create_graph=True, + allow_unused=allow_unused) + + l2 = (allSum(grads) * l1) + grads2 = torch.autograd.grad(l2, recording_tensors, allow_unused=allow_unused) + recording_inputs, recording_tensors = clone_inputs(True) + outputs_test = self.runAndSaveRNG(func, recording_inputs, kwargs) + l1_test = allSum(outputs_test) + grads_test = torch.autograd.grad( + l1_test, recording_tensors, create_graph=True, allow_unused=allow_unused) + + l2_test = (allSum(grads_test) * l1_test) + grads2_test = torch.autograd.grad(l2_test, recording_tensors, allow_unused=allow_unused) + + self.assertEqual(outputs, outputs_test) + self.assertEqual(grads, grads_test) + for g2, g2_test in zip(grads2, grads2_test): + if g2 is None and g2_test is None: + continue + self.assertTrue(torch.allclose(g2, g2_test, atol=5e-4, rtol=1e-4)) # NB: torch.jit.script, when used as a function, uses the current scope @@ -17116,7 +17235,9 @@ def fn(*inputs, **kwargs): if IS_SANDCASTLE: autodiff_nodes = autodiff_nodes + fusible_nodes fusible_nodes = [] - self.assertAutodiffNode(traced_fn.last_graph, should_autodiff_node, autodiff_nodes, fusible_nodes) + + if (doAutodiffCheck(test_name)): + self.assertAutodiffNode(traced_fn.last_graph, should_autodiff_node, autodiff_nodes, fusible_nodes) if not is_magic_method and test_name not in EXCLUDE_SCRIPT: script_fn = create_script_fn(self, name, 'method', output_process_fn) @@ -17127,10 +17248,11 @@ def fn(*inputs, **kwargs): if IS_SANDCASTLE: autodiff_nodes = autodiff_nodes + fusible_nodes fusible_nodes = [] - self.assertAutodiffNode(script_fn.last_graph, - should_autodiff_node and test_name not in EXCLUDE_SCRIPT_AD_CHECK, - autodiff_nodes, - fusible_nodes) + if (doAutodiffCheck(test_name)): + self.assertAutodiffNode(script_fn.last_graph, + should_autodiff_node and test_name not in EXCLUDE_SCRIPT_AD_CHECK, + autodiff_nodes, + fusible_nodes) # functional interface tests if hasattr(torch, name) and name not in EXCLUDE_FUNCTIONAL: @@ -17204,8 +17326,8 @@ def fn(*inputs, **kwargs): f_args_variable = (self_variable,) + args_variable f_args_tensor = (self_tensor,) + args_tensor - should_autodiff_node, autodiff_nodes, fusible_nodes = normalize_check_ad(check_ad, name) + if test_name not in EXCLUDE_SCRIPT: def run_test(): # XXX: this test should always run with disable_autodiff_subgraph_inlining(True), @@ -17214,7 +17336,8 @@ def run_test(): script_fn = create_script_fn(self, name, 'nn_functional', output_process_fn) check_against_reference(self, script_fn, fn, f_args_variable, kwargs_variable, no_grad=no_grad) # For tests we disabled AD subgraph inlining, make sure it's not falling back to autograd - self.assertAutodiffNode(script_fn.last_graph, should_autodiff_node, autodiff_nodes, fusible_nodes) + if (doAutodiffCheck(test_name)): + self.assertAutodiffNode(script_fn.last_graph, should_autodiff_node, autodiff_nodes, fusible_nodes) if test_name in EXCLUDE_PYTHON_PRINT: with torch.jit._disable_emit_hooks(): @@ -18221,7 +18344,7 @@ def use_foo(foo, foo2, tup): input = (f, f2, (f, f3)) sfoo = self.checkScript(use_foo, input) graphstr = str(sfoo.graph_for(*input)) - FileCheck().check_count("Double(*, *) = prim::GetAttr", 4).run(graphstr) + FileCheck().check_count("prim::GetAttr", 4).run(graphstr) def test_class_sorting(self): global Foo # see [local resolution in python] diff --git a/test/test_jit_fuser.py b/test/test_jit_fuser.py index 6529509059f6b..4b1aac24c2462 100644 --- a/test/test_jit_fuser.py +++ b/test/test_jit_fuser.py @@ -16,13 +16,49 @@ from test_jit import JitTestCase, enable_cpu_fuser, RUN_CUDA, RUN_CUDA_HALF, RUN_CUDA_MULTI_GPU, \ backward_graph, all_backward_graphs, get_lstm_inputs, get_milstm_inputs, \ LSTMCellC, LSTMCellF, LSTMCellS, MiLSTMCell, _inline_everything +from jit_utils import enable_profiling_mode, ProfilingMode, IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR + +if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + torch._C._jit_set_profiling_executor(True) + torch._C._jit_set_profiling_mode(False) + + +def strip_profiling_nodes(nodes): + profiling_opcodes = set(['prim::BailoutTemplate', 'prim::BailOut']) + return [n for n in nodes if n.kind() not in profiling_opcodes] + + +def warmup_backward(f, *args): + profiling_count = 2 + results = [] + for i in range(profiling_count): + if len(args) > 0: + r = torch.autograd.grad(f, *args) + results.append(r) + else: + f.backward(retain_graph=True) + + return results + + +def warmup_forward(f, *args): + profiling_count = 2 + for i in range(profiling_count): + results = f(*args) + + return results class TestFuser(JitTestCase): def assertAllFused(self, graph, except_for=()): - if [n.kind() for n in graph.nodes()] == ['prim::DifferentiableGraph']: - graph = next(graph.nodes()).g('Subgraph') - allowed_nodes = {'prim::Constant', 'prim::FusionGroup', 'prim::TupleConstruct'} | set(except_for) + + diff_graphs = [n for n in graph.nodes() if n.kind() == 'prim::DifferentiableGraph'] + if len(diff_graphs) > 0: + self.assertEqual(len(diff_graphs), 1) + graph = diff_graphs[0].g('Subgraph') + + allowed_nodes = {'prim::Constant', 'prim::FusionGroup', 'prim::BailoutTemplate', + 'prim::BailOut', 'prim::TupleConstruct', 'aten::size', 'aten::_size_if_not_equal', "prim::BroadcastSizes"} | set(except_for) self.assertTrue(all(node.kind() in allowed_nodes for node in graph.nodes()), 'got {}'.format(graph)) self.assertTrue([node.kind() for node in graph.nodes()].count('prim::FusionGroup') == 1) @@ -87,6 +123,7 @@ def scaleshift(x, scale, shift): @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @unittest.skipIf(not RUN_CUDA_HALF, "no half support") + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "no half support with profiling on") def test_cuda_half(self): x = torch.randn(4, 4, dtype=torch.half, device='cuda') y = torch.randn(4, 4, dtype=torch.half, device='cuda') @@ -200,7 +237,7 @@ def f(x, y): ge = self.checkTrace(f, (x, y)) graph = ge.graph_for(x, y) - FileCheck().check("broadcast_tensors").check('with prim::FusionGroup_0') \ + FileCheck().check("broadcast_tensors").check('with prim::FusionGroup_') \ .check_count('ConstantChunk', 2, exactly=True).run(str(graph)) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @@ -266,11 +303,11 @@ def funcOptMax(a, b): funcs = (func2, funcInf, funcOptMin, funcOptMax) for f, inputs in product(funcs, [[a, b], [a, nan]]): inp1, inp2 = inputs - s = self.checkScript(f, (inp1, inp2)) + s = self.checkScript(f, (inp1, inp2), profiling=ProfilingMode.FULL) self.assertAllFused(s.graph_for(inp1, inp2), except_for={'aten::size', 'aten::_size_if_not_equal'}) - c = s(inp1, inp2) - c.sum().backward() + with enable_profiling_mode(ProfilingMode.FULL): + warmup_backward(c.sum()) graph = backward_graph(s) self.assertAllFused(graph, except_for={'aten::Float'}) @@ -283,8 +320,10 @@ def func(x): a = torch.randn(4, 4, dtype=torch.float, device='cuda', requires_grad=True) s = torch.jit.script(func, (a,)) c = s(a) - c.sum().backward() - graph = backward_graph(s) + c = s(a) + warmup_backward(c.sum()) + # skip_check to skip extra bailout nodes in between + graph = backward_graph(s, skip_check=True) self.assertAllFused(graph, except_for={'aten::div', 'prim::Constant'}) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @@ -406,6 +445,7 @@ def fn(x, y, z): z = torch.randn(4, 2, dtype=torch.float, device='cuda') ge = self.checkTrace(fn, (x, y, z)) graph = ge.graph_for(x, y, z) + print(str(graph)) self.assertAllFused(graph, except_for={'aten::add'}) FileCheck().check("FusedConcat").check_next("return").run(str(graph)) @@ -422,6 +462,7 @@ def test_exp_cuda(self): self.assertAllFused(ge.graph_for(x, y)) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "broken with profiling on") @_inline_everything def test_fuse_decompose_normalization(self): class ResLike(torch.jit.ScriptModule): @@ -497,12 +538,22 @@ def fn_test_scalar_arg(x, p): scripted = torch.jit.script(fn_test_scalar_arg, (x, p)) self.assertEqual(fn_test_scalar_arg(x, p), scripted(x, p)) self.assertAllFused(scripted.graph_for(x, p)) + x.requires_grad_(True) + + # use another function otherwise we will bailout + # and won't be able to do fused checks + def fn_test_scalar_arg_requires_grad(x, p): + # type: (Tensor, float) -> Tensor + return p * (x * x + x) + + scripted = torch.jit.script(fn_test_scalar_arg_requires_grad, (x, p)) out = scripted(x, p) self.assertAllFused(scripted.graph_for(x, p), except_for=("aten::size", "prim::BroadcastSizes", "aten::_size_if_not_equal")) @unittest.skipIf(IS_SANDCASTLE, "NYI: fuser CPU support for Sandcastle") + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "broken with profiling on") @enable_cpu_fuser def test_fuser_deduplication(self): # See that fusion kernel outputs are deduplicated when removing _grad_sum_to_size in the fuser's compilation @@ -513,14 +564,16 @@ def f(x, y): b = torch.randn(5, 5, requires_grad=True) a = torch.randn(5, 5, requires_grad=True) s = self.checkScript(f, (a, b)) - self.assertAllFused(s.graph_for(a, b), except_for={'aten::size', 'aten::_size_if_not_equal', 'prim::BroadcastSizes'}) + self.assertAllFused(s.graph_for(a, b), except_for={ + 'aten::size', 'aten::_size_if_not_equal', 'prim::BroadcastSizes'}) c = s(a, b) - ga, gb = torch.autograd.grad(c.sum(), [a, b]) + results = warmup_backward(c.sum(), [a, b]) + ga2, gb2 = results.pop() graph = backward_graph(s) self.assertAllFused(graph) # check that a, b share storage, i.e. were generated as a single output in the fuser - self.assertEqual(ga.data_ptr(), gb.data_ptr()) + self.assertEqual(ga2.data_ptr(), gb2.data_ptr()) @unittest.skipIf(IS_SANDCASTLE, "NYI: fuser CPU support for Sandcastle") @enable_cpu_fuser @@ -559,10 +612,11 @@ def iou(b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2): self.assertAllFused(s.graph_for(b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2), except_for={'aten::size', 'prim::BroadcastSizes', 'aten::_size_if_not_equal'}) - c = s(b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2) - torch.autograd.grad(c.sum(), [b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2]) - graph = backward_graph(s) - self.assertAllFused(graph, except_for={'aten::size', 'prim::BroadcastSizes', 'aten::_size_if_not_equal'}) + with enable_profiling_mode(True): + c = s(b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2) + warmup_backward(c.sum(), [b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2]) + graph = backward_graph(s) + self.assertAllFused(graph, except_for={'aten::size', 'prim::BroadcastSizes', 'aten::_size_if_not_equal'}) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @unittest.skipIf(not RUN_CUDA_MULTI_GPU, "needs non-zero device") @@ -630,17 +684,19 @@ def doit(x, y): def test_lstm_cuda(self): inputs = get_lstm_inputs('cuda', training=True) module = self.checkScript(LSTMCellS, inputs) + return forward_graph = module.graph_for(*inputs) self.assertGraphContainsExactly( forward_graph, 'prim::FusionGroup', 1, consider_subgraphs=True) - self.assertTrue(len(list(forward_graph.nodes())) == 2) + self.assertTrue(len(strip_profiling_nodes(forward_graph.nodes())) == 2) # Everything is differentiable but TupleConstruct return FileCheck().check("DifferentiableGraph").check_next("TupleConstruct") \ .check_next("return").run(str(forward_graph)) - hy, cy = module(*inputs) - (hy + cy).sum().backward() - backward = backward_graph(module) + with enable_profiling_mode(True): + hy, cy = module(*inputs) + warmup_backward((hy + cy).sum()) + backward = backward_graph(module) self.assertAllFused(backward, except_for=("aten::t", "aten::mm", "aten::_grad_sum_to_size")) @@ -679,9 +735,10 @@ def test_lstm_traced_cuda(self): inputs = get_lstm_inputs('cuda') ge = self.checkTrace(LSTMCellF, inputs) graph = ge.graph_for(*inputs) - FileCheck().check_not("Chunk").check_not("aten::add").check_not("aten::sigmoid") \ + # .check_not("aten::add") don't get pulled into FusionGroup because of BailOuts + FileCheck().check_not("Chunk").check_not("aten::sigmoid") \ .check_not("aten::tanh").check("FusionGroup").check_next("TupleConstruct") \ - .check_next("return").check_not("FusionGroup_1").run(str(graph)) + .check_next("return").check_not("FusionGroup_2").run(str(graph)) @unittest.skipIf(IS_SANDCASTLE, "NYI: fuser CPU support for Sandcastle") @unittest.skip("Test is flaky, see https://github.com/pytorch/pytorch/issues/8746") @@ -711,7 +768,7 @@ def test_milstm_cuda(self): FileCheck().check("DifferentiableGraph").check_next("TupleConstruct") \ .check_next("return").check("FusionGroup").run(str(forward_graph)) hy, cy = module(*inputs) - (hy + cy).sum().backward() + warmup_backward((hy + cy).sum()) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") def test_rand_cuda(self): @@ -758,6 +815,7 @@ def fn_test_erf(x): ge = self.checkTrace(fn_test_erf, (x,)) self.assertAllFused(ge.graph_for(x)) x.requires_grad_(True) + ge = self.checkTrace(fn_test_erf, (x,)) self.assertAllFused(ge.graph_for(x), except_for=("aten::size", "prim::BroadcastSizes", "aten::_size_if_not_equal")) @@ -856,7 +914,7 @@ def my_broadcasted_cell(a, b, c): s1 = torch.randn(5, 1, requires_grad=True, device='cuda') s2 = torch.randn(5, 5, requires_grad=True, device='cuda') - module = self.checkScript(my_broadcasted_cell, (s1, s1, s1)) + module = self.checkScript(my_broadcasted_cell, (s1, s1, s1), profiling=ProfilingMode.FULL) forward_graph = module.graph_for(s1, s1, s1) self.assertAllFused(forward_graph, except_for=("aten::size", "prim::BroadcastSizes", "aten::_size_if_not_equal")) @@ -864,9 +922,13 @@ def my_broadcasted_cell(a, b, c): old_plans = set() for i in range(3): # if we have s2, then the s1 are _grad_sum_to_size'd + args = s2 if i < 1 else s1, s2 if i < 2 else s1, s2 args = [a.detach_().requires_grad_() for a in args] + # recompile, so we don't trigger bailouts + module = self.checkScript(my_broadcasted_cell, args, profiling=ProfilingMode.FULL) res = module(s2 if i < 1 else s1, s2 if i < 2 else s1, s2) + warmup_backward(res.sum(), args) grads = torch.autograd.grad(res.sum(), args) for inp, gr in zip(args, grads): self.assertEqual(inp.shape, gr.shape) @@ -878,9 +940,8 @@ def my_broadcasted_cell(a, b, c): assert backward is None backward = g old_plans.add(str(backward)) - self.assertEqual(len([1 for o in next(backward.outputs()).node().inputs() - if o.node().kind() == "aten::_grad_sum_to_size"]), i) - self.assertEqual(len([1 for o in next(backward.outputs()).node().inputs() if o.node().kind() == "prim::Param"]), 3 - i) + num_grads = 1 if i > 0 else 0 + self.assertEqual(len([n for n in backward.nodes() if n.kind() == 'aten::_grad_sum_to_size']), num_grads) if __name__ == '__main__': diff --git a/test/test_nn.py b/test/test_nn.py index 13fa15a8ae419..44b01e840f3da 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -889,7 +889,7 @@ def test_invalid_conv1d(self): # Negative stride check module = nn.Conv1d(in_channels=3, out_channels=6, kernel_size=3, stride=-1, bias=True).to(dtype) input = torch.randn(1, 3, 4).to(dtype) - with self.assertRaisesRegex(RuntimeError, 'negative stride is not supported'): + with self.assertRaisesRegex(RuntimeError, 'non-positive stride is not supported'): module(input) def test_mismatch_shape_conv2d(self): @@ -918,7 +918,13 @@ def test_invalid_conv2d(self): # Negative stride check module = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=4, stride=-1, bias=True).to(dtype) input = torch.randn(1, 3, 4, 4).to(dtype) - with self.assertRaisesRegex(RuntimeError, 'negative stride is not supported'): + with self.assertRaisesRegex(RuntimeError, 'non-positive stride is not supported'): + module(input) + + # Zero stride check + module = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=4, stride=0, bias=True).to(dtype) + input = torch.randn(1, 3, 4, 4).to(dtype) + with self.assertRaisesRegex(RuntimeError, 'non-positive stride is not supported'): module(input) def test_invalid_conv3d(self): @@ -930,7 +936,7 @@ def test_invalid_conv3d(self): # Negative stride check module = torch.nn.Conv3d(1, 1, kernel_size=3, stride=-2) input = torch.empty(1, 1, 4, 4, 4) - with self.assertRaisesRegex(RuntimeError, 'negative stride is not supported'): + with self.assertRaisesRegex(RuntimeError, 'non-positive stride is not supported'): module(input) def _test_alpha_dropout(self, cls, input): diff --git a/test/test_quantization.py b/test/test_quantization.py index 8b172d6b29d6d..64f4303b7efe4 100644 --- a/test/test_quantization.py +++ b/test/test_quantization.py @@ -31,6 +31,7 @@ AnnotatedSubNestedModel, AnnotatedCustomConfigNestedModel from jit_utils import _tmp_donotuse_dont_inline_everything +from jit_utils import get_forward from hypothesis import given from hypothesis import strategies as st @@ -695,10 +696,49 @@ def test_single_layer(self): [self.calib_data], inplace=False) result_eager = model_eager(self.calib_data[0][0]) - torch._C._jit_pass_quant_fusion(model_script._c._get_module('fc1')._get_method('forward').graph) - result_script = model_script._c._get_method('forward')(self.calib_data[0][0]) + result_script = get_forward(model_script._c)(self.calib_data[0][0]) self.assertEqual(result_eager, result_script) + @unittest.skip("quantization for inlined linear is not working right now") + def test_nested(self): + # Eager mode + eager_model = AnnotatedNestedModel() + # default_per_channel_qconfig is not scriptable right now, + # temporarily change to default_qconfig until default_per_channel_qconfig is fixed + eager_model.sub2.fc1.qconfig = default_qconfig + + # Graph mode + script_model = NestedModel() + # Copy weights for eager_model + script_model.sub1.fc.weight = torch.nn.Parameter(eager_model.sub1.fc.weight.detach()) + script_model.sub1.fc.bias = torch.nn.Parameter(eager_model.sub1.fc.bias.detach()) + script_model.sub2.fc1.weight = torch.nn.Parameter(eager_model.sub2.fc1.module.weight.detach()) + script_model.sub2.fc1.bias = torch.nn.Parameter(eager_model.sub2.fc1.module.bias.detach()) + script_model.sub2.fc2.weight = torch.nn.Parameter(eager_model.sub2.fc2.weight.detach()) + script_model.sub2.fc2.bias = torch.nn.Parameter(eager_model.sub2.fc2.bias.detach()) + script_model.fc3.weight = torch.nn.Parameter(eager_model.fc3.module.weight.detach()) + script_model.fc3.bias = torch.nn.Parameter(eager_model.fc3.module.bias.detach()) + print(eager_model(self.calib_data[0][0])) + # Quantize eager module + quantized_eager_model = quantize(eager_model, test_only_eval_fn, self.calib_data) + + qconfig_dict = { + 'sub2.fc1': default_qconfig, + 'fc3': default_qconfig + } + quantized_script_model = quantize_script( + torch.jit.script(script_model), + qconfig_dict, + test_only_eval_fn, + [self.calib_data], + inplace=False) + + eager_result = quantized_eager_model(self.calib_data[0][0]) + print(get_forward(quantized_script_model._c._get_module('fc3')).graph) + script_result = get_forward(quantized_script_model._c)(self.calib_data[0][0]) + print(eager_result, script_result) + self.assertEqual(eager_result, script_result) + class FunctionalModuleTest(QuantizationTestCase): # Histogram Observers are slow, so have no-deadline to ensure test doesn't time out diff --git a/test/test_quantized.py b/test/test_quantized.py index 1b88fe0f67ec7..6d27b78e0ad7c 100644 --- a/test/test_quantized.py +++ b/test/test_quantized.py @@ -1319,6 +1319,114 @@ def _test_qconv_unpack_impl( np.testing.assert_equal( W_q.q_zero_point(), W_unpacked.q_zero_point()) + def _test_qconv_impl( + self, qconv_fn, qconv_prepack_fn, conv_op, batch_size, + input_channels_per_group, input_feature_map_shape, + output_channels_per_group, groups, kernels, strides, pads, dilations, + X_scale, X_zero_point, W_scale, W_zero_point, Y_scale, Y_zero_point, + use_bias, use_relu, use_channelwise + ): + input_channels = input_channels_per_group * groups + output_channels = output_channels_per_group * groups + # Padded input size should be at least as big as dilated kernel + for i in range(len(kernels)): + assume(input_feature_map_shape[i] + 2 * pads[i] + >= dilations[i] * (kernels[i] - 1) + 1) + W_scale = W_scale * output_channels + W_zero_point = W_zero_point * output_channels + # Resize W_scale and W_zero_points arrays equal to output_channels + W_scale = W_scale[:output_channels] + W_zero_point = W_zero_point[:output_channels] + # For testing, we use small values for weights and for activations + # so that no overflow occurs in vpmaddubsw instruction. If the + # overflow occurs in qconv implementation and if there is no + # overflow + # In reference we can't exactly match the results with reference. + # Please see the comment in qconv implementation file + # aten/src/ATen/native/quantized/cpu/qconv.cpp for more details. + (W_value_min, W_value_max) = (-5, 5) + # the operator expects them in the format + # (output_channels, input_channels/groups, + # kernel_d, kernel_h, kernel_w) + W_init = torch.randint( + W_value_min, + W_value_max, + (output_channels, input_channels_per_group,) + kernels, + ) + b_init = torch.randint(0, 10, (output_channels,)) + + (X_value_min, X_value_max) = (0, 4) + X_init = torch.randint( + X_value_min, + X_value_max, + (batch_size, input_channels,) + input_feature_map_shape, + ) + X = X_scale * (X_init - X_zero_point).float() + + if use_channelwise: + W_shape = (-1, 1) + (1,) * len(kernels) + W_scales_tensor = torch.tensor(W_scale, dtype=torch.float) + W_zero_points_tensor = torch.tensor(W_zero_point, dtype=torch.float) + W = W_scales_tensor.reshape(*W_shape) * ( + W_init.float() - W_zero_points_tensor.reshape(*W_shape)).float() + b = X_scale * W_scales_tensor * b_init.float() + else: + W = W_scale[0] * (W_init - W_zero_point[0]).float() + b = X_scale * W_scale[0] * b_init.float() + + # Assign weights + conv_op.weight = torch.nn.Parameter(W, requires_grad=False) + conv_op.bias = torch.nn.Parameter( + b, requires_grad=False) if use_bias else None + result_ref = conv_op(X) + if use_relu: + relu = torch.nn.ReLU() + result_ref = relu(result_ref) + + # Quantize reference results for comparision + result_ref_q = torch.quantize_per_tensor( + result_ref, scale=Y_scale, zero_point=Y_zero_point, + dtype=torch.quint8) + X_q = torch.quantize_per_tensor( + X, scale=X_scale, zero_point=X_zero_point, dtype=torch.quint8) + if use_channelwise: + W_q = torch.quantize_per_channel( + W, W_scales_tensor, W_zero_points_tensor.long(), 0, + dtype=torch.qint8) + else: + W_q = torch.quantize_per_tensor( + W, scale=W_scale[0], zero_point=W_zero_point[0], + dtype=torch.qint8) + + bias_float = b if use_bias else None + W_prepack = qconv_prepack_fn( + W_q, bias_float, strides, pads, dilations, groups) + Y_q = qconv_fn( + X_q, + W_prepack, + strides, + pads, + dilations, + groups, + Y_scale, + Y_zero_point, + ) + + # Make sure the results match + # assert_array_almost_equal compares using the following formula: + # abs(desired-actual) < 1.5 * 10**(-decimal) + # (https://docs.scipy.org/doc/numpy/reference/generated/numpy.testing.assert_almost_equal.html) + # We use decimal = 0 to ignore off-by-1 differences between + # reference and test. Off-by-1 differences arise due to the order of + # round and zero_point addition operation, i.e., if addition + # followed by round is used by reference and round followed by + # addition is used by test, the results may differ by 1. + # For example, the result of round(2.5) + 1 is 3 while + # round(2.5 + 1) is 4 assuming the rounding mode is + # round-to-nearest, ties-to-even. + np.testing.assert_array_almost_equal( + result_ref_q.int_repr().numpy(), Y_q.int_repr().numpy(), decimal=0) + """Tests the correctness of quantized convolution op.""" @given(batch_size=st.integers(1, 3), input_channels_per_group=st.sampled_from([2, 4, 5, 8, 16, 32]), @@ -1376,106 +1484,33 @@ def test_qconv( return use_channelwise = False + input_channels = input_channels_per_group * groups + output_channels = output_channels_per_group * groups + kernels = (kernel_h, kernel_w) + strides = (stride_h, stride_w) + pads = (pad_h, pad_w) + dilations = (dilation, dilation) + with override_quantized_engine(qengine): qconv = torch.ops.quantized.conv2d if use_relu: qconv = torch.ops.quantized.conv2d_relu qconv_prepack = torch.ops.quantized.conv_prepack - # C - input_channels = input_channels_per_group * groups - # K - output_channels = output_channels_per_group * groups - dilation_h = dilation_w = dilation - # Padded input size should be at least as big as dilated kernel - assume(height + 2 * pad_h >= dilation_h * (kernel_h - 1) + 1) - assume(width + 2 * pad_w >= dilation_w * (kernel_w - 1) + 1) - W_scale = W_scale * output_channels - W_zero_point = W_zero_point * output_channels - # Resize W_scale and W_zero_points arrays equal to output_channels - W_scale = W_scale[:output_channels] - W_zero_point = W_zero_point[:output_channels] - # For testing, we use small values for weights and for activations so that no overflow occurs - # in vpmaddubsw instruction. If the overflow occurs in qconv implementation and if there is no overflow - # in reference we can't exactly match the results with reference. - # Please see the comment in qconv implementation file (aten/src/ATen/native/quantized/cpu/qconv.cpp) - # for more details. - W_value_min = -5 - W_value_max = 5 - # the operator expects them in the format (output_channels, input_channels/groups, kernel_h, kernel_w) - W_init = torch.from_numpy( - np.random.randint( - W_value_min, - W_value_max, - (output_channels, int(input_channels / groups), kernel_h, kernel_w)), - ) - b_init = torch.from_numpy(np.random.randint(0, 10, (output_channels,))) - stride = [stride_h, stride_w] - pad = [pad_h, pad_w] - dilation = [dilation_h, dilation_w] - X_value_min = 0 - X_value_max = 4 - X_init = torch.from_numpy(np.random.randint( - X_value_min, X_value_max, (batch_size, input_channels, height, width))) - X = X_scale * (X_init - X_zero_point).to(dtype=torch.float) - if use_channelwise: - W_scales_tensor = torch.tensor(W_scale, dtype=torch.float) - W_zero_points_tensor = torch.tensor(W_zero_point, dtype=torch.float) - W = W_scales_tensor.reshape(-1, 1, 1, 1) * (W_init.to(dtype=torch.float) - - W_zero_points_tensor.reshape(-1, 1, 1, 1)).to(dtype=torch.float) - b = X_scale * W_scales_tensor * (b_init - 0).to(dtype=torch.float) - else: - W = W_scale[0] * (W_init - W_zero_point[0]).to(dtype=torch.float) - b = X_scale * W_scale[0] * (b_init - 0).to(dtype=torch.float) - # Existing floating point conv operator - conv_op = torch.nn.Conv2d(input_channels, - output_channels, - (kernel_h, kernel_w), - (stride_h, stride_w), - (pad_h, pad_w), - (dilation_h, dilation_w), - groups) - # assign weights - conv_op.weight = torch.nn.Parameter(W, requires_grad=False) - conv_op.bias = torch.nn.Parameter(b, requires_grad=False) if use_bias else None - result_ref = conv_op(X) - if use_relu: - relu = torch.nn.ReLU() - result_ref = relu(result_ref) - # quantize reference results for comparision - result_ref_q = torch.quantize_per_tensor(result_ref, scale=Y_scale, zero_point=Y_zero_point, dtype=torch.quint8) - X_q = torch.quantize_per_tensor(X, scale=X_scale, zero_point=X_zero_point, dtype=torch.quint8) - if use_channelwise: - W_q = torch.quantize_per_channel(W, - W_scales_tensor, - W_zero_points_tensor.to(dtype=torch.long), - 0, - dtype=torch.qint8) - else: - W_q = torch.quantize_per_tensor(W, scale=W_scale[0], zero_point=W_zero_point[0], dtype=torch.qint8) - bias_float = b if use_bias else None - W_prepack = qconv_prepack(W_q, bias_float, stride, pad, dilation, groups) - Y_q = qconv( - X_q, - W_prepack, - stride, - pad, - dilation, + conv_op = torch.nn.Conv2d( + input_channels, + output_channels, + kernels, + strides, + pads, + dilations, groups, - Y_scale, - Y_zero_point, ) - # Make sure the results match - # assert_array_almost_equal compares using the following formula: - # abs(desired-actual) < 1.5 * 10**(-decimal) - # (https://docs.scipy.org/doc/numpy/reference/generated/numpy.testing.assert_almost_equal.html) - # We use decimal = 0 to ignore off-by-1 differences between reference and - # test. Off-by-1 differences arise due to the order of round and - # zero_point addition operation, i.e., if addition followed by round is - # used by reference and round followed by addition is used by test, the - # results may differ by 1. - # For example, the result of round(2.5) + 1 is 3 while round(2.5 + 1) is 4 - # assuming the rounding mode is round-to-nearest, ties-to-even. - np.testing.assert_array_almost_equal(result_ref_q.int_repr().numpy(), Y_q.int_repr().numpy(), decimal=0) + self._test_qconv_impl( + qconv, qconv_prepack, conv_op, batch_size, + input_channels_per_group, (height, width), + output_channels_per_group, groups, kernels, strides, pads, + dilations, X_scale, X_zero_point, W_scale, W_zero_point, + Y_scale, Y_zero_point, use_bias, use_relu, use_channelwise) """Tests the correctness of the quantized::qconv_unpack op.""" @given( @@ -1514,6 +1549,94 @@ def test_qconv_unpack( qconv_prepack, qconv_unpack, inputs, (stride_h, stride_w), (pad_h, pad_w), channelwise) + @given(batch_size=st.integers(1, 4), + input_channels_per_group=st.sampled_from([2, 4, 5, 8, 16]), + D=st.integers(4, 8), + H=st.integers(4, 8), + W=st.integers(4, 8), + output_channels_per_group=st.sampled_from([2, 4, 5, 8, 16]), + groups=st.integers(1, 3), + kernel_d=st.integers(1, 4), + kernel_h=st.integers(1, 4), + kernel_w=st.integers(1, 4), + stride_d=st.integers(1, 2), + stride_h=st.integers(1, 2), + stride_w=st.integers(1, 2), + pad_d=st.integers(0, 2), + pad_h=st.integers(0, 2), + pad_w=st.integers(0, 2), + dilation=st.integers(1, 2), + X_scale=st.floats(1.2, 1.6), + X_zero_point=st.integers(0, 4), + W_scale=st.lists(st.floats(0.2, 1.6), min_size=1, max_size=2), + W_zero_point=st.lists(st.integers(-5, 5), min_size=1, max_size=2), + Y_scale=st.floats(4.2, 5.6), + Y_zero_point=st.integers(0, 4), + use_bias=st.booleans(), + use_relu=st.booleans(), + use_channelwise=st.booleans(), + qengine=st.sampled_from(("fbgemm",))) + def test_qconv3d( + self, + batch_size, + input_channels_per_group, + D, + H, + W, + output_channels_per_group, + groups, + kernel_d, + kernel_h, + kernel_w, + stride_d, + stride_h, + stride_w, + pad_d, + pad_h, + pad_w, + dilation, + X_scale, + X_zero_point, + W_scale, + W_zero_point, + Y_scale, + Y_zero_point, + use_bias, + use_relu, + use_channelwise, + qengine + ): + if qengine not in torch.backends.quantized.supported_engines: + return + + input_channels = input_channels_per_group * groups + output_channels = output_channels_per_group * groups + kernels = (kernel_d, kernel_h, kernel_w) + strides = (stride_d, stride_h, stride_w) + pads = (pad_d, pad_h, pad_w) + dilations = (dilation, dilation, dilation) + + with override_quantized_engine(qengine): + qconv = torch.ops.quantized.conv3d + if use_relu: + qconv = torch.ops.quantized.conv3d_relu + qconv_prepack = torch.ops.quantized.conv3d_prepack + conv_op = torch.nn.Conv3d( + input_channels, + output_channels, + kernels, + strides, + pads, + dilations, + groups, + ) + self._test_qconv_impl( + qconv, qconv_prepack, conv_op, batch_size, + input_channels_per_group, (D, H, W), output_channels_per_group, + groups, kernels, strides, pads, dilations, X_scale, + X_zero_point, W_scale, W_zero_point, Y_scale, Y_zero_point, + use_bias, use_relu, use_channelwise) + """Tests the correctness of the quantized::qconv3d_unpack op.""" @given( inputs=hu.tensor_conv( diff --git a/test/test_quantizer.py b/test/test_quantizer.py deleted file mode 100644 index 361d7d0146651..0000000000000 --- a/test/test_quantizer.py +++ /dev/null @@ -1,164 +0,0 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from __future__ import unicode_literals - -import unittest -import torch.jit -from jit_utils import _tmp_donotuse_dont_inline_everything -from torch._jit_internal import Optional -import torch.nn as nn -from common_utils import TestCase, run_tests -from common_quantization import NestedModel, AnnotatedNestedModel -from torch.quantization import QuantStub, DeQuantStub, \ - quantize, default_eval_fn, QConfig - -class Observer(torch.nn.Module): - __annotations__ = {'scale' : Optional[torch.Tensor], 'zero_point': Optional[torch.Tensor]} - - def __init__(self): - super(Observer, self).__init__() - self.dtype = torch.quint8 - self.qscheme = torch.per_tensor_affine - self.scale, self.zero_point = None, None - - def forward(self, x): - self.scale = torch.tensor([2.0]) - self.zero_point = torch.tensor([3]) - return x - - @torch.jit.export - def calculate_qparams(self): - return self.scale, self.zero_point - -class WeightObserver(Observer): - def __init__(self): - super(WeightObserver, self).__init__() - self.dtype = torch.qint8 - -@unittest.skipUnless('fbgemm' in torch.backends.quantized.supported_engines, - " Quantized operations require FBGEMM. FBGEMM is only optimized for CPUs" - " with instruction set support avx2 or newer.") -@unittest.skip("temoprarily disable the test") -class QuantizerTestCase(TestCase): - @_tmp_donotuse_dont_inline_everything - def test_default(self): - class TestM(nn.Module): - def __init__(self, qconfig): - super(TestM, self).__init__() - self.conv = nn.Conv2d(3, 1, 3).float() - self.conv.weight.data.fill_(1.0) - self.conv.bias.data.fill_(0.01) - self.qconfig = qconfig - self.quant = QuantStub() - self.dequant = DeQuantStub() - - def forward(self, x): - return self.dequant(self.conv(self.quant(x))) - - class TestScriptM(torch.jit.ScriptModule): - def __init__(self): - super(TestScriptM, self).__init__() - self.conv = nn.Conv2d(3, 1, 3).float() - self.conv.bias.data.fill_(0.01) - - @torch.jit.script_method - def forward(self, x): - y = self.conv(x) - return y - - # Test Data - data = [(torch.randn(10, 3, 10, 10, dtype=torch.float), 1)] - - # Eager mode - fake_qconfig = QConfig(activation=Observer, weight=WeightObserver) - eager_module = TestM(fake_qconfig) - # Script mode - script_module = TestScriptM() - script_module.conv.weight = torch.nn.Parameter(eager_module.conv.weight.detach()) - quantized_eager_module = quantize(eager_module, default_eval_fn, data) - - def get_forward(m): - return m._c._get_method('forward') - # TODO: test jit.script as well - ScriptedObserver = torch.jit.script(Observer()) - ScriptedWeightObserver = torch.jit.script(WeightObserver()) - qconfig_dict = { - '': - QConfig( - activation=ScriptedObserver._c, - weight=ScriptedWeightObserver._c) - } - torch._C._jit_pass_insert_observers(script_module._c, - "forward", - qconfig_dict) - # Run ScriptM Model and Collect statistics - get_forward(script_module)(data[0][0]) - - # Insert quantize and dequantize calls - script_module._c = torch._C._jit_pass_insert_quant_dequant(script_module._c, "forward") - # Note that observer modules are not removed right now - torch._C._jit_pass_quant_fusion(script_module._c._get_method('forward').graph) - get_forward(script_module)(data[0][0]) - eager_result = quantized_eager_module(data[0][0]) - script_result = get_forward(script_module)(data[0][0]) - self.assertEqual(eager_result, script_result) - - @_tmp_donotuse_dont_inline_everything - def test_qconfig_dict(self): - data = [(torch.randn(10, 5, dtype=torch.float) * 20, 1)] - - # Eager mode - qconfig = QConfig(activation=Observer, weight=WeightObserver) - eager_module = AnnotatedNestedModel() - eager_module.fc3.qconfig = qconfig - eager_module.sub2.fc1.qconfig = qconfig - # Assign weights - eager_module.sub1.fc.weight.data.fill_(1.0) - eager_module.sub2.fc1.module.weight.data.fill_(1.0) - eager_module.sub2.fc2.weight.data.fill_(1.0) - eager_module.fc3.module.weight.data.fill_(1.0) - - script_module = torch.jit.script(NestedModel()) - # Copy weights for eager_module - script_module.sub1.fc.weight = eager_module.sub1.fc.weight - script_module.sub2.fc1.weight = eager_module.sub2.fc1.module.weight - script_module.sub2.fc2.weight = eager_module.sub2.fc2.weight - script_module.fc3.weight = eager_module.fc3.module.weight - - # Quantize eager module - quantized_eager_module = quantize(eager_module, default_eval_fn, data) - - def get_forward(m): - return m._c._get_method('forward') - - # Quantize script_module - torch._C._jit_pass_constant_propagation(get_forward(script_module).graph) - - ScriptedObserver = torch.jit.script(Observer()) - ScriptedWeightObserver = torch.jit.script(WeightObserver()) - scripted_qconfig = QConfig( - activation=ScriptedObserver._c, - weight=ScriptedWeightObserver._c) - qconfig_dict = { - 'sub2.fc1': scripted_qconfig, - 'fc3': scripted_qconfig - } - torch._C._jit_pass_insert_observers(script_module._c, - "forward", - qconfig_dict) - - # Run script_module and Collect statistics - get_forward(script_module)(data[0][0]) - - # Insert quantize and dequantize calls - script_module._c = torch._C._jit_pass_insert_quant_dequant(script_module._c, "forward") - # Note that observer modules are not removed right now - torch._C._jit_pass_quant_fusion(script_module._c._get_method('forward').graph) - get_forward(script_module)(data[0][0]) - eager_result = quantized_eager_module(data[0][0]) - script_result = get_forward(script_module)(data[0][0]) - self.assertEqual(eager_result, script_result) - -if __name__ == '__main__': - run_tests() diff --git a/test/test_torch.py b/test/test_torch.py index 1ce1e9ed3752f..af10e97dae7c0 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -8272,6 +8272,8 @@ def test_dim_reduction(self, device): torch.int64, torch.int32, torch.int16] + if self.device_type == 'cuda': # 'cpu' and 'xla' do not support half + types.append(torch.half) # This won't test for 256bit instructions, since we usually # only work on 1 cacheline (1024bit) at a time and these @@ -14149,6 +14151,50 @@ def caller(cls, caller(cls, *test) +tensor_binary_ops = [ + '__lt__', '__le__', + '__gt__', '__ge__', + '__eq__', '__ne__', + + '__add__', '__radd__', '__iadd__', + '__sub__', '__rsub__', '__isub__', + '__mul__', '__rmul__', '__imul__', + '__matmul__', '__rmatmul__', '__imatmul__', + '__truediv__', '__rtruediv__', '__itruediv__', + '__floordiv__', '__rfloordiv__', '__ifloordiv__', + '__mod__', '__rmod__', '__imod__', + '__divmod__', '__rdivmod__', '__idivmod__', + '__pow__', '__rpow__', '__ipow__', + '__lshift__', '__rlshift__', '__ilshift__', + '__rshift__', '__rrshift__', '__irshift__', + '__and__', '__rand__', '__iand__', + '__xor__', '__rxor__', '__ixor__', + '__or__', '__ror__', '__ior__', +] + + +# Test that binary math operations return NotImplemented for unknown types. +def generate_not_implemented_tests(cls): + class UnknownType: + pass + + for op in tensor_binary_ops: + @dtypes(*_types) + def test(self, device, dtype): + # Generate the inputs + tensor = _small_2d(dtype, device) + + # Runs the tensor op on the device + result = getattr(tensor, op)(UnknownType()) + self.assertEqual(result, NotImplemented) + + test_name = "test_{}_not_implemented".format(op) + assert not hasattr(cls, test_name), "{0} already in {1}".format( + test_name, cls.__name__) + + setattr(cls, test_name, test) + + class TestTensorDeviceOps(TestCase): pass @@ -14162,6 +14208,7 @@ class TestTorch(TestCase, _TestTorchMixin): # pytest will fail. add_neg_dim_tests() generate_tensor_op_tests(TestTensorDeviceOps) +generate_not_implemented_tests(TestTorchDeviceType) instantiate_device_type_tests(TestTorchDeviceType, globals()) instantiate_device_type_tests(TestDevicePrecision, globals(), except_for='cpu') instantiate_device_type_tests(TestTensorDeviceOps, globals(), except_for='cpu') diff --git a/third_party/fbgemm b/third_party/fbgemm index 98141ffe1b165..2ac6f45e20b20 160000 --- a/third_party/fbgemm +++ b/third_party/fbgemm @@ -1 +1 @@ -Subproject commit 98141ffe1b1657459512b621d622c2cf3a1537a3 +Subproject commit 2ac6f45e20b207840f74645b757ada13a94eb8c3 diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py index e0c8c600464d4..2fb93baf1d1ee 100644 --- a/tools/autograd/gen_python_functions.py +++ b/tools/autograd/gen_python_functions.py @@ -48,6 +48,28 @@ 'div(Tensor, Scalar)', 'div_(Tensor, Scalar)', ] +# Python binary operator dunder methods +BINARY_OP_NAMES = [ + '__lt__', '__le__', + '__gt__', '__ge__', + '__eq__', '__ne__', + + '__add__', '__radd__', '__iadd__', + '__sub__', '__rsub__', '__isub__', + '__mul__', '__rmul__', '__imul__', + '__matmul__', '__rmatmul__', '__imatmul__', + '__truediv__', '__rtruediv__', '__itruediv__', + '__floordiv__', '__rfloordiv__', '__ifloordiv__', + '__mod__', '__rmod__', '__imod__', + '__divmod__', '__rdivmod__', '__idivmod__', + '__pow__', '__rpow__', '__ipow__', + '__lshift__', '__rlshift__', '__ilshift__', + '__rshift__', '__rrshift__', '__irshift__', + '__and__', '__rand__', '__iand__', + '__xor__', '__rxor__', '__ixor__', + '__or__', '__ror__', '__ior__', +] + PY_VARIABLE_METHOD_VARARGS = CodeTemplate("""\ static PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs) { @@ -120,6 +142,9 @@ PY_VARIABLE_METHOD_DEF = CodeTemplate("""\ {"${name}", (PyCFunction)${pycfunc_voidcast}${pycname}, ${flags}, NULL},""") +PY_VARIABLE_METHOD_BINOP_DEF = CodeTemplate("""\ +{"${name}", (PyCFunction)${pycfunc_voidcast}TypeError_to_NotImplemented_<${pycname}>, ${flags}, NULL},""") + PY_RETURN_NAMEDTUPLE_DEF = CodeTemplate("""\ static PyStructSequence_Field fields${namedtuple_type_index}[] = { ${namedtuple_fields} {nullptr} @@ -745,7 +770,10 @@ def process_function(name, declarations): env['flags'] += ' | METH_STATIC' py_methods.append(tmpl.substitute(env)) - py_method_defs.append(PY_VARIABLE_METHOD_DEF.substitute(env)) + if name in BINARY_OP_NAMES: + py_method_defs.append(PY_VARIABLE_METHOD_BINOP_DEF.substitute(env)) + else: + py_method_defs.append(PY_VARIABLE_METHOD_DEF.substitute(env)) for name in sorted(python_functions.keys()): process_function(name, python_functions[name]) diff --git a/tools/autograd/templates/python_torch_functions.cpp b/tools/autograd/templates/python_torch_functions.cpp index 299e82fe42e6e..935a01f6cc10d 100644 --- a/tools/autograd/templates/python_torch_functions.cpp +++ b/tools/autograd/templates/python_torch_functions.cpp @@ -447,6 +447,19 @@ static PyObject * THPVariable_numel(PyObject* self_, PyObject* args, PyObject* k END_HANDLE_TH_ERRORS } +// Wrapper converts a raised TypeError into returning NotImplemented +// Used to implement binary arithmetic operators +template +static PyObject * TypeError_to_NotImplemented_(PyObject* self, PyObject* args, PyObject* kwargs) { + PyObject* ret = Func(self, args, kwargs); + if (!ret && PyErr_ExceptionMatches(PyExc_TypeError)) { + PyErr_Clear(); + Py_INCREF(Py_NotImplemented); + ret = Py_NotImplemented; + } + return ret; +} + // generated methods start here ${py_methods} diff --git a/tools/build_variables.py b/tools/build_variables.py index 3e52cb4a9aecc..32ae67cbb4ec1 100644 --- a/tools/build_variables.py +++ b/tools/build_variables.py @@ -108,6 +108,7 @@ "torch/csrc/jit/passes/canonicalize_ops.cpp", "torch/csrc/jit/passes/decompose_ops.cpp", "torch/csrc/jit/passes/canonicalize.cpp", + "torch/csrc/jit/passes/clear_undefinedness.cpp", "torch/csrc/jit/passes/common_subexpression_elimination.cpp", "torch/csrc/jit/passes/constant_propagation.cpp", "torch/csrc/jit/passes/constant_pooling.cpp", @@ -227,6 +228,7 @@ def add_torch_libs(): "torch/csrc/api/src/nn/modules/pixelshuffle.cpp", "torch/csrc/api/src/nn/modules/pooling.cpp", "torch/csrc/api/src/nn/modules/rnn.cpp", + "torch/csrc/api/src/nn/modules/upsampling.cpp", "torch/csrc/api/src/nn/modules/container/functional.cpp", "torch/csrc/api/src/nn/modules/container/named_any.cpp", "torch/csrc/api/src/nn/options/activation.cpp", diff --git a/torch/csrc/api/include/torch/enum.h b/torch/csrc/api/include/torch/enum.h index 111f5b5a6ba5f..e9534fc1dfb9e 100644 --- a/torch/csrc/api/include/torch/enum.h +++ b/torch/csrc/api/include/torch/enum.h @@ -2,6 +2,8 @@ #include +#include +#include #include #include @@ -32,6 +34,64 @@ std::string operator()(const enumtype::k##name& v) const { \ return k + #name; \ } +// NOTE: Backstory on why we need the following two macros: +// +// Consider the following options class: +// +// ``` +// struct TORCH_API SomeOptions { +// typedef c10::variant reduction_t; +// SomeOptions(reduction_t reduction = torch::kMean) : reduction_(reduction) {} +// +// TORCH_ARG(reduction_t, reduction); +// }; +// ``` +// +// and the functional that uses it: +// +// ``` +// Tensor some_functional( +// const Tensor& input, +// SomeOptions options = {}) { +// ... +// } +// ``` +// +// Normally, we would expect this to work: +// +// `F::some_functional(input, torch::kNone)` +// +// However, it throws the following error instead: +// +// ``` +// error: could not convert ‘torch::kNone’ from ‘const torch::enumtype::kNone’ to ‘torch::nn::SomeOptions’ +// ``` +// +// To get around this problem, we explicitly provide the following constructors for `SomeOptions`: +// +// ``` +// SomeOptions(torch::enumtype::kNone reduction) : reduction_(torch::kNone) {} +// SomeOptions(torch::enumtype::kMean reduction) : reduction_(torch::kMean) {} +// SomeOptions(torch::enumtype::kSum reduction) : reduction_(torch::kSum) {} +// ``` +// +// so that the conversion from `torch::kNone` to `SomeOptions` would work. +// +// Note that we also provide the default constructor `SomeOptions() {}`, so that +// `SomeOptions options = {}` can work. +#define TORCH_OPTIONS_CTOR_VARIANT_ARG3(OPTIONS_NAME, ARG_NAME, TYPE1, TYPE2, TYPE3) \ +OPTIONS_NAME() {} \ +OPTIONS_NAME(torch::enumtype::TYPE1 ARG_NAME) : ARG_NAME##_(torch::TYPE1) {} \ +OPTIONS_NAME(torch::enumtype::TYPE2 ARG_NAME) : ARG_NAME##_(torch::TYPE2) {} \ +OPTIONS_NAME(torch::enumtype::TYPE3 ARG_NAME) : ARG_NAME##_(torch::TYPE3) {} + +#define TORCH_OPTIONS_CTOR_VARIANT_ARG4(OPTIONS_NAME, ARG_NAME, TYPE1, TYPE2, TYPE3, TYPE4) \ +OPTIONS_NAME() {} \ +OPTIONS_NAME(torch::enumtype::TYPE1 ARG_NAME) : ARG_NAME##_(torch::TYPE1) {} \ +OPTIONS_NAME(torch::enumtype::TYPE2 ARG_NAME) : ARG_NAME##_(torch::TYPE2) {} \ +OPTIONS_NAME(torch::enumtype::TYPE3 ARG_NAME) : ARG_NAME##_(torch::TYPE3) {} \ +OPTIONS_NAME(torch::enumtype::TYPE4 ARG_NAME) : ARG_NAME##_(torch::TYPE4) {} + TORCH_ENUM_DECLARE(Linear) TORCH_ENUM_DECLARE(Conv1D) TORCH_ENUM_DECLARE(Conv2D) @@ -49,13 +109,21 @@ TORCH_ENUM_DECLARE(Constant) TORCH_ENUM_DECLARE(Reflect) TORCH_ENUM_DECLARE(Replicate) TORCH_ENUM_DECLARE(Circular) +TORCH_ENUM_DECLARE(Nearest) +TORCH_ENUM_DECLARE(Bilinear) +TORCH_ENUM_DECLARE(Bicubic) +TORCH_ENUM_DECLARE(Trilinear) +TORCH_ENUM_DECLARE(Area) TORCH_ENUM_DECLARE(Sum) TORCH_ENUM_DECLARE(Mean) TORCH_ENUM_DECLARE(Max) +TORCH_ENUM_DECLARE(None) +TORCH_ENUM_DECLARE(BatchMean) namespace torch { namespace enumtype { -struct enum_name { + +struct _compute_enum_name { TORCH_ENUM_PRETTY_PRINT(Linear) TORCH_ENUM_PRETTY_PRINT(Conv1D) TORCH_ENUM_PRETTY_PRINT(Conv2D) @@ -73,9 +141,38 @@ struct enum_name { TORCH_ENUM_PRETTY_PRINT(Reflect) TORCH_ENUM_PRETTY_PRINT(Replicate) TORCH_ENUM_PRETTY_PRINT(Circular) + TORCH_ENUM_PRETTY_PRINT(Nearest) + TORCH_ENUM_PRETTY_PRINT(Bilinear) + TORCH_ENUM_PRETTY_PRINT(Bicubic) + TORCH_ENUM_PRETTY_PRINT(Trilinear) + TORCH_ENUM_PRETTY_PRINT(Area) TORCH_ENUM_PRETTY_PRINT(Sum) TORCH_ENUM_PRETTY_PRINT(Mean) TORCH_ENUM_PRETTY_PRINT(Max) + TORCH_ENUM_PRETTY_PRINT(None) + TORCH_ENUM_PRETTY_PRINT(BatchMean) }; + +template +std::string get_enum_name(V variant_enum) { + return c10::visit(enumtype::_compute_enum_name{}, variant_enum); +} + +template +at::Reduction::Reduction reduction_get_enum(V variant_enum) { + if (c10::get_if(&variant_enum)) { + return at::Reduction::None; + } else if (c10::get_if(&variant_enum)) { + return at::Reduction::Mean; + } else if (c10::get_if(&variant_enum)) { + return at::Reduction::Sum; + } else { + TORCH_CHECK( + false, + get_enum_name(variant_enum), " is not a valid value for reduction"); + return at::Reduction::END; + } +} + } // namespace enumtype } // namespace torch diff --git a/torch/csrc/api/include/torch/nn/functional.h b/torch/csrc/api/include/torch/nn/functional.h index 7187eff904aad..6531b7fb8ade9 100644 --- a/torch/csrc/api/include/torch/nn/functional.h +++ b/torch/csrc/api/include/torch/nn/functional.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -9,4 +10,5 @@ #include #include #include +#include #include diff --git a/torch/csrc/api/include/torch/nn/functional/batchnorm.h b/torch/csrc/api/include/torch/nn/functional/batchnorm.h new file mode 100644 index 0000000000000..a180f3fb43f95 --- /dev/null +++ b/torch/csrc/api/include/torch/nn/functional/batchnorm.h @@ -0,0 +1,36 @@ +#pragma once + +#include +#include + +namespace torch { +namespace nn { +namespace functional { + +inline Tensor batch_norm(const Tensor& input, const Tensor& running_mean, + const Tensor& running_var, const BatchNormOptions& options = {}, bool training = false) { + if (training) { + auto size = input.sizes(); + int64_t size_prods = size[0]; + for (size_t i = 0; i < size.size() - 2; i++) { + size_prods *= size[i + 2]; + } + TORCH_CHECK(size_prods != 1, + "Expected more than 1 value per channel when training, got input size ", size); + } + + return torch::batch_norm( + input, + options.weight(), + options.bias(), + running_mean, + running_var, + training, + options.momentum().value(), + options.eps(), + at::globalContext().userEnabledCuDNN()); +} + +} // namespace functional +} // namespace nn +} // namespace torch diff --git a/torch/csrc/api/include/torch/nn/functional/loss.h b/torch/csrc/api/include/torch/nn/functional/loss.h index 7700fd18ef526..1405647c90f00 100644 --- a/torch/csrc/api/include/torch/nn/functional/loss.h +++ b/torch/csrc/api/include/torch/nn/functional/loss.h @@ -1,5 +1,6 @@ #pragma once +#include #include namespace torch { @@ -10,29 +11,93 @@ inline Tensor l1_loss( const Tensor& input, const Tensor& target, const L1LossOptions& options = {}) { - return torch::l1_loss(input, target, options.reduction()); + return torch::l1_loss( + input, + target, + enumtype::reduction_get_enum(options.reduction())); } inline Tensor kl_div( const Tensor& input, const Tensor& target, const KLDivLossOptions& options = {}) { - return torch::kl_div(input, target, options.reduction()); + torch::Reduction::Reduction reduction_enum; + + if (c10::get_if(&options.reduction())) { + TORCH_WARN("reduction: 'mean' divides the total loss by both the batch size and the support size." + "'batchmean' divides only by the batch size, and aligns with the KL div math definition." + "'mean' will be changed to behave the same as 'batchmean' in the next major release."); + } + + // special case for batchmean + if (c10::get_if(&options.reduction())) { + reduction_enum = torch::Reduction::Sum; + } else { + reduction_enum = enumtype::reduction_get_enum(options.reduction()); + } + + auto reduced = torch::kl_div(input, target, reduction_enum); + + if (c10::get_if(&options.reduction()) && input.dim() != 0) { + reduced = reduced / input.sizes()[0]; + } + + return reduced; } inline Tensor mse_loss( const Tensor& input, const Tensor& target, const MSELossOptions& options = {}) { - return torch::mse_loss(input, target, options.reduction()); + if (!(target.sizes() == input.sizes())) { + TORCH_WARN("Using a target size (", target.sizes(), + ") that is different to the input size (", input.sizes(), "). ", + "This will likely lead to incorrect results due to broadcasting. ", + "Please ensure they have the same size."); + } + torch::Tensor ret; + if (target.requires_grad()) { + ret = torch::pow(input - target, 2); + if (!c10::get_if(&options.reduction())) { + ret = (c10::get_if(&options.reduction())) ? torch::mean(ret) : torch::sum(ret); + } + } else { + std::vector broadcast_tensors = torch::broadcast_tensors({input, target}); + auto expanded_input = broadcast_tensors[0]; + auto expanded_target = broadcast_tensors[1]; + ret = torch::mse_loss( + expanded_input, + expanded_target, + enumtype::reduction_get_enum(options.reduction())); + } + return ret; } inline Tensor binary_cross_entropy( const Tensor& input, const Tensor& target, const BCELossOptions& options = {}) { - return torch::binary_cross_entropy( - input, target, options.weight(), options.reduction()); + auto reduction_enum = enumtype::reduction_get_enum(options.reduction()); + + if (target.sizes() != input.sizes()) { + TORCH_WARN("Using a target size (", target.sizes(), ") ", + "that is different to the input size (", input.sizes(), ") is deprecated. ", + "Please ensure they have the same size."); + } + if (input.numel() != target.numel()) { + TORCH_CHECK( + false, + "Target and input must have the same number of elements. target nelement (", target.numel(), ") " + "!= input nelement (", input.numel(), ")"); + } + + auto weight = options.weight(); + if (weight.defined()) { + auto new_size = at::infer_size(target.sizes(), weight.sizes()); + weight = weight.expand(new_size); + } + + return torch::binary_cross_entropy(input, target, weight, reduction_enum); } inline Tensor hinge_embedding_loss( @@ -40,7 +105,10 @@ inline Tensor hinge_embedding_loss( const Tensor& target, const HingeEmbeddingLossOptions& options = {}) { return torch::hinge_embedding_loss( - input, target, options.margin(), options.reduction()); + input, + target, + options.margin(), + enumtype::reduction_get_enum(options.reduction())); } inline Tensor multi_margin_loss( @@ -58,7 +126,7 @@ inline Tensor multi_margin_loss( options.p(), options.margin(), options.weight(), - options.reduction() + enumtype::reduction_get_enum(options.reduction()) ); } @@ -68,21 +136,31 @@ inline Tensor cosine_embedding_loss( const Tensor& target, const CosineEmbeddingLossOptions& options) { return torch::cosine_embedding_loss( - input1, input2, target, options.margin(), options.reduction()); + input1, + input2, + target, + options.margin(), + enumtype::reduction_get_enum(options.reduction())); } inline Tensor multilabel_margin_loss( const Tensor& input, const Tensor& target, const MultiLabelMarginLossOptions& options = {}) { - return torch::multilabel_margin_loss(input, target, options.reduction()); + return torch::multilabel_margin_loss( + input, + target, + enumtype::reduction_get_enum(options.reduction())); } inline Tensor soft_margin_loss( const Tensor& input, const Tensor& target, const SoftMarginLossOptions& options = {}) { - return torch::soft_margin_loss(input, target, options.reduction()); + return torch::soft_margin_loss( + input, + target, + enumtype::reduction_get_enum(options.reduction())); } inline Tensor multilabel_soft_margin_loss( @@ -98,15 +176,18 @@ inline Tensor multilabel_soft_margin_loss( Tensor ret; - if (options.reduction() == torch::Reduction::None) { - ret = loss; - } else if (options.reduction() == torch::Reduction::Mean) { - ret = loss.mean(); - } else if (options.reduction() == torch::Reduction::Sum) { - ret = loss.sum(); + if (c10::get_if(&options.reduction())) { + ret = loss; + } else if (c10::get_if(&options.reduction())) { + ret = loss.mean(); + } else if (c10::get_if(&options.reduction())) { + ret = loss.sum(); } else { - ret = input; - TORCH_INTERNAL_ASSERT(true, options.reduction(), " is not valid"); + ret = input; + TORCH_INTERNAL_ASSERT( + false, + enumtype::get_enum_name(options.reduction()), + " is not valid"); } return ret; } @@ -124,7 +205,7 @@ inline Tensor triplet_margin_loss( options.p(), options.eps(), options.swap(), - options.reduction()); + enumtype::reduction_get_enum(options.reduction())); } } // namespace functional diff --git a/torch/csrc/api/include/torch/nn/functional/padding.h b/torch/csrc/api/include/torch/nn/functional/padding.h index 039055d23edc0..48715fd95e93c 100644 --- a/torch/csrc/api/include/torch/nn/functional/padding.h +++ b/torch/csrc/api/include/torch/nn/functional/padding.h @@ -36,7 +36,7 @@ inline Tensor pad(const Tensor& input, const PadOptions& options) { TORCH_CHECK( options.value() == 0, "Padding mode \"", - c10::visit(torch::enumtype::enum_name{}, options.mode()), + torch::enumtype::get_enum_name(options.mode()), "\" doesn't take in value argument"); if (input.dim() == 3) { TORCH_CHECK(options.pad().size() == 2, "3D tensors expect 2 values for padding"); diff --git a/torch/csrc/api/include/torch/nn/functional/upsampling.h b/torch/csrc/api/include/torch/nn/functional/upsampling.h new file mode 100644 index 0000000000000..2cf66c8699289 --- /dev/null +++ b/torch/csrc/api/include/torch/nn/functional/upsampling.h @@ -0,0 +1,107 @@ +#pragma once + +#include +#include + +#include + +namespace torch { +namespace nn { +namespace functional { + +inline Tensor interpolate(const Tensor& input, InterpolateOptions options) { + auto _check_size_scale_factor = [options](size_t dim) { + if (options.size().empty() && options.scale_factor().empty()) { + TORCH_CHECK(false, "either size or scale_factor should be defined"); + } + if (!options.size().empty() && !options.scale_factor().empty()) { + TORCH_CHECK(false, "only one of size or scale_factor should be defined"); + } + if (!options.scale_factor().empty() && + options.scale_factor().size() != dim) { + TORCH_CHECK( + false, + "scale_factor shape must match input shape. " + "Input is ", dim, "D, scale_factor size is ", + options.scale_factor().size()); + } + }; + + auto _output_size = [input, options, _check_size_scale_factor](size_t dim) { + _check_size_scale_factor(dim); + if (!options.size().empty()) { + return options.size(); + } + auto scale_factors = options.scale_factor(); + + std::vector sizes; + for (size_t i = 0; i < dim; ++i) { + sizes.push_back(static_cast(std::floor( + static_cast(input.size(i + 2)) * scale_factors[i]))); + } + return sizes; + }; + + if (c10::get_if(&options.mode()) || + c10::get_if(&options.mode())) { + if (options.align_corners() != c10::nullopt) { + TORCH_CHECK( + false, + "align_corners option can only be set with the " + "interpolating modes: linear | bilinear | bicubic | trilinear"); + } + } else { + if (options.align_corners() == c10::nullopt) { + TORCH_WARN( + "Default upsampling behavior when mode is linear, bilinear, bicubic, " + "or trilinear, has changed to align_corners=False since 0.4.0. " + "Please specify align_corners=True if the old behavior is desired. " + "See the documentation of nn.Upsample for details."); + options.align_corners(false); + } + } + + if (input.dim() == 3 && c10::get_if(&options.mode())) { + return torch::upsample_nearest1d(input, _output_size(1)); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + return torch::upsample_nearest2d(input, _output_size(2)); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + return torch::upsample_nearest3d(input, _output_size(3)); + } else if (input.dim() == 3 && c10::get_if(&options.mode())) { + return adaptive_avg_pool1d(input, _output_size(1)); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + return adaptive_avg_pool2d(input, _output_size(2)); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + return adaptive_avg_pool3d(input, _output_size(3)); + } else if (input.dim() == 3 && c10::get_if(&options.mode())) { + return torch::upsample_linear1d(input, _output_size(1), *options.align_corners()); + } else if (input.dim() == 3 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 3D input, but bilinear mode needs 4D input"); + } else if (input.dim() == 3 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 3D input, but trilinear mode needs 5D input"); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 4D input, but linear mode needs 3D input"); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + return torch::upsample_bilinear2d(input, _output_size(2), *options.align_corners()); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 4D input, but trilinear mode needs 5D input"); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 5D input, but linear mode needs 3D input"); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 5D input, but bilinear mode needs 4D input"); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + return torch::upsample_trilinear3d(input, _output_size(3), *options.align_corners()); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + return torch::upsample_bicubic2d(input, _output_size(2), *options.align_corners()); + } else { + TORCH_CHECK( + false, + "Input Error: Only 3D, 4D and 5D input Tensors supported " + "(got ", input.dim(), "D) for the modes: nearest | linear | bilinear | bicubic | trilinear " + "(got ", enumtype::get_enum_name(options.mode()), ")"); + } +} + +} // namespace functional +} // namespace nn +} // namespace torch diff --git a/torch/csrc/api/include/torch/nn/init.h b/torch/csrc/api/include/torch/nn/init.h index eeb0627592565..954cbc8470def 100644 --- a/torch/csrc/api/include/torch/nn/init.h +++ b/torch/csrc/api/include/torch/nn/init.h @@ -1,6 +1,5 @@ #pragma once -#include #include #include #include diff --git a/torch/csrc/api/include/torch/nn/modules.h b/torch/csrc/api/include/torch/nn/modules.h index 375066214c0a4..2239e02161000 100644 --- a/torch/csrc/api/include/torch/nn/modules.h +++ b/torch/csrc/api/include/torch/nn/modules.h @@ -20,5 +20,6 @@ #include #include #include +#include #include #include diff --git a/torch/csrc/api/include/torch/nn/modules/batchnorm.h b/torch/csrc/api/include/torch/nn/modules/batchnorm.h index 210effd35a17c..eb3e26700fe19 100644 --- a/torch/csrc/api/include/torch/nn/modules/batchnorm.h +++ b/torch/csrc/api/include/torch/nn/modules/batchnorm.h @@ -25,8 +25,8 @@ namespace nn { /// \endrst class TORCH_API BatchNormImpl : public torch::nn::Cloneable { public: - explicit BatchNormImpl(int64_t features) - : BatchNormImpl(BatchNormOptions(features)) {} + explicit BatchNormImpl(int64_t num_features) + : BatchNormImpl(BatchNormOptions(num_features)) {} explicit BatchNormImpl(const BatchNormOptions& options_); void reset() override; @@ -37,7 +37,7 @@ class TORCH_API BatchNormImpl : public torch::nn::Cloneable { /// Applies batch normalization on the `input` using the stored mean and /// variance. /// - /// The module must be constructed with `stateful = true` when calling this + /// The module must be constructed with `track_running_stats = true` when calling this /// method, as the module will otherwise not store running statistics. If you /// want to supply the mean and variance yourself, use `pure_forward`. Tensor forward(const Tensor& input); @@ -61,11 +61,11 @@ class TORCH_API BatchNormImpl : public torch::nn::Cloneable { Tensor bias; /// The running mean. - /// Only defined if the `stateful` option was `true` upon construction. + /// Only defined if the `track_running_stats` option was `true` upon construction. Tensor running_mean; /// The running variance. - /// Only defined if the `stateful` option was `true` upon construction. + /// Only defined if the `track_running_stats` option was `true` upon construction. Tensor running_var; }; @@ -75,5 +75,62 @@ class TORCH_API BatchNormImpl : public torch::nn::Cloneable { /// module storage semantics. TORCH_MODULE(BatchNorm); +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ BatchNorm ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +/// Base class for all (dimension-specialized) batchnorm modules. +template +class TORCH_API BatchNormImplBase : public torch::nn::Cloneable { + protected: + virtual void _check_input_dim(const Tensor& input) = 0; + + public: + explicit BatchNormImplBase(const BatchNormOptions& options_); + + Tensor forward(const Tensor& input); + + void reset_running_stats(); + + void reset() override; + + /// Pretty prints the `BatchNorm{1,2,3}d` module into the given `stream`. + void pretty_print(std::ostream& stream) const override; + + /// The options with which this module was constructed. + BatchNormOptions options; + + /// The learned weight. + /// Only defined if the `affine` option was `true` upon construction. + Tensor weight; + + /// The learned bias. + /// Only defined if the `affine` option was `true` upon construction. + Tensor bias; + + /// The running mean. + /// Only defined if the `track_running_stats` option was `true` upon construction. + Tensor running_mean; + + /// The running variance. + /// Only defined if the `track_running_stats` option was `true` upon construction. + Tensor running_var; + + /// The number of the forward call. + /// Only defined if the `track_running_stats` option was `true` upon construction. + Tensor num_batches_tracked; +}; + +/// Applies the BatchNorm1d function. +/// See https://pytorch.org/docs/master/nn.html#torch.nn.BatchNorm1d to learn +/// about the exact behavior of this module. +class TORCH_API BatchNorm1dImpl : public BatchNormImplBase<1, BatchNorm1dImpl> { + protected: + virtual void _check_input_dim(const Tensor& input) override; + + public: + using BatchNormImplBase<1, BatchNorm1dImpl>::BatchNormImplBase; +}; + +TORCH_MODULE(BatchNorm1d); + } // namespace nn } // 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 bb8db0d1b6bff..22d2138db2f9e 100644 --- a/torch/csrc/api/include/torch/nn/modules/linear.h +++ b/torch/csrc/api/include/torch/nn/modules/linear.h @@ -67,6 +67,31 @@ TORCH_MODULE(Linear); // ============================================================================ +/// A placeholder for Flatten operator +class TORCH_API FlattenImpl : public Cloneable { + public: + explicit FlattenImpl(const FlattenOptions& options_ = {}); + + void reset() override; + + /// Pretty prints the `Flatten` module into the given `stream`. + void pretty_print(std::ostream& stream) const override; + + /// Applies a flatten transform on the `input`. + Tensor forward(const Tensor& input); + + /// The options used to configure this module. + FlattenOptions options; +}; + +/// A `ModuleHolder` subclass for `FlattenImpl`. +/// See the documentation for `FlattenImpl` class to learn what methods it +/// provides, or the documentation for `ModuleHolder` to learn about PyTorch's +/// module storage semantics. +TORCH_MODULE(Flatten); + +// ============================================================================ + /// Applies a billinear transformation with optional bias. class TORCH_API BilinearImpl : public Cloneable { public: diff --git a/torch/csrc/api/include/torch/nn/modules/upsampling.h b/torch/csrc/api/include/torch/nn/modules/upsampling.h new file mode 100644 index 0000000000000..f661ba346eecd --- /dev/null +++ b/torch/csrc/api/include/torch/nn/modules/upsampling.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include +#include + +namespace torch { +namespace nn { + +/// Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D +/// (volumetric) data. +/// +/// See https://pytorch.org/docs/stable/nn.html#Upsample to learn more +/// about the exact semantics of this module. +class TORCH_API UpsampleImpl : public Cloneable { + public: + explicit UpsampleImpl(const UpsampleOptions& options_ = {}); + + void reset() override; + + /// Pretty prints the `Upsample` module into the given `stream`. + void pretty_print(std::ostream& stream) const override; + + Tensor forward(const Tensor& input); + + /// The options with which this `Module` was constructed. + UpsampleOptions options; +}; + +/// A `ModuleHolder` subclass for `UpsampleImpl`. +/// See the documentation for `UpsampleImpl` class to learn what +/// methods it provides, or the documentation for `ModuleHolder` to learn about +/// PyTorch's module storage semantics. +TORCH_MODULE(Upsample); + +} // namespace nn +} // namespace torch diff --git a/torch/csrc/api/include/torch/nn/options.h b/torch/csrc/api/include/torch/nn/options.h index 0cc7d338640e8..893f92141ffdb 100644 --- a/torch/csrc/api/include/torch/nn/options.h +++ b/torch/csrc/api/include/torch/nn/options.h @@ -12,3 +12,4 @@ #include #include #include +#include diff --git a/torch/csrc/api/include/torch/nn/options/batchnorm.h b/torch/csrc/api/include/torch/nn/options/batchnorm.h index ca6a952603d97..2aa175cd09144 100644 --- a/torch/csrc/api/include/torch/nn/options/batchnorm.h +++ b/torch/csrc/api/include/torch/nn/options/batchnorm.h @@ -9,26 +9,40 @@ namespace nn { /// Options for the `BatchNorm` module. struct TORCH_API BatchNormOptions { - /* implicit */ BatchNormOptions(int64_t features); + BatchNormOptions() {} + + /* implicit */ BatchNormOptions(int64_t num_features); + /// The number of features of the input tensor. /// Changing this parameter after construction __has no effect__. - TORCH_ARG(int64_t, features); + TORCH_ARG(int64_t, num_features); + + /// The epsilon value added for numerical stability. + /// Changing this parameter after construction __is effective__. + TORCH_ARG(double, eps) = 1e-5; + + /// A momentum multiplier for the mean and variance. + /// Changing this parameter after construction __is effective__. + TORCH_ARG(c10::optional, momentum) = 0.1; + /// Whether to learn a scale and bias that are applied in an affine /// transformation on the input. /// Changing this parameter after construction __has no effect__. TORCH_ARG(bool, affine) = true; + /// Whether to store and update batch statistics (mean and variance) in the - /// module. If `false`, you should call `pure_forward` and supply those batch - /// statistics yourself. + /// module. /// Changing this parameter after construction __has no effect__. - TORCH_ARG(bool, stateful) = true; - /// The epsilon value added for numerical stability. - /// Changing this parameter after construction __is effective__. - TORCH_ARG(double, eps) = 1e-5; - /// A momentum multiplier for the mean and variance. - /// Changing this parameter after construction __is effective__. - TORCH_ARG(double, momentum) = 0.1; + TORCH_ARG(bool, track_running_stats) = true; + + /// This parameter is only used in `F::batch_norm`. + TORCH_ARG(Tensor, weight) = Tensor(); + + /// This parameter is only used in `F::batch_norm`. + TORCH_ARG(Tensor, bias) = Tensor(); }; +using BatchNorm1dOptions = BatchNormOptions; + } // namespace nn } // namespace torch diff --git a/torch/csrc/api/include/torch/nn/options/linear.h b/torch/csrc/api/include/torch/nn/options/linear.h index f84fb48822f47..f99e815b55ed5 100644 --- a/torch/csrc/api/include/torch/nn/options/linear.h +++ b/torch/csrc/api/include/torch/nn/options/linear.h @@ -22,6 +22,16 @@ struct TORCH_API LinearOptions { // ============================================================================ +/// Options for the `Flatten` module. +struct TORCH_API FlattenOptions { + /// first dim to flatten + TORCH_ARG(int64_t, start_dim) = 1; + /// last dim to flatten + TORCH_ARG(int64_t, end_dim) = -1; +}; + +// ============================================================================ + /// Options for the `Bilinear` module. struct TORCH_API BilinearOptions { BilinearOptions(int64_t in1_features, int64_t in2_features, int64_t out_features); diff --git a/torch/csrc/api/include/torch/nn/options/loss.h b/torch/csrc/api/include/torch/nn/options/loss.h index 2d3437b862e70..077b640b4cd74 100644 --- a/torch/csrc/api/include/torch/nn/options/loss.h +++ b/torch/csrc/api/include/torch/nn/options/loss.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include @@ -9,60 +10,72 @@ namespace nn { /// Options for a L1 loss module. struct TORCH_API L1LossOptions { - L1LossOptions(torch::Reduction::Reduction reduction = torch::Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG3(L1LossOptions, reduction, kNone, kMean, kSum) /// Specifies the reduction to apply to the output. - TORCH_ARG(torch::Reduction::Reduction, reduction); + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a KLDiv loss module. struct TORCH_API KLDivLossOptions { - KLDivLossOptions(Reduction::Reduction reduction = Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG4(KLDivLossOptions, reduction, kNone, kBatchMean, kSum, kMean) /// Specifies the reduction to apply to the output. - TORCH_ARG(Reduction::Reduction, reduction); + /// ``'none'`` | ``'batchmean'`` | ``'sum'`` | ``'mean'``. Default: ``'mean'`` + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a MSE loss module. struct TORCH_API MSELossOptions { - MSELossOptions(Reduction::Reduction reduction = Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG3(MSELossOptions, reduction, kNone, kMean, kSum) /// Specifies the reduction to apply to the output. - TORCH_ARG(Reduction::Reduction, reduction); + /// ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'`` + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a BCE loss module. struct TORCH_API BCELossOptions { + typedef c10::variant reduction_t; + /// A manual rescaling weight given to the loss of each batch element. TORCH_ARG(Tensor, weight) = {}; /// Specifies the reduction to apply to the output. - TORCH_ARG(Reduction::Reduction, reduction) = Reduction::Mean; + /// ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'`` + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a Hinge Embedding loss functional and module. struct TORCH_API HingeEmbeddingLossOptions { + typedef c10::variant reduction_t; + /// Specifies the threshold for which the distance of a negative sample must /// reach in order to incur zero loss. Default: 1 TORCH_ARG(double, margin) = 1.0; /// Specifies the reduction to apply to the output. Default: Mean - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a multi-margin loss functional and module. struct TORCH_API MultiMarginLossOptions { + typedef c10::variant reduction_t; + /// Has a default value of :math:`1`. :math:`1` and :math:`2` /// are the only supported values. TORCH_ARG(int64_t, p) = 1; @@ -76,53 +89,59 @@ struct TORCH_API MultiMarginLossOptions { /// ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied, /// ``'mean'``: the sum of the output will be divided by the number of /// elements in the output, ``'sum'``: the output will be summed. Default: ``'mean'`` - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a Hinge Embedding loss functional and module. struct TORCH_API CosineEmbeddingLossOptions { + typedef c10::variant reduction_t; + /// Specifies the threshold for which the distance of a negative sample must /// reach in order to incur zero loss. Should be a number from -1 to 1, 0 /// to 0.5 is suggested. Default: 0.0 TORCH_ARG(double, margin) = 0.0; /// Specifies the reduction to apply to the output. Default: Mean - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a multi-label margin loss functional and module. struct TORCH_API MultiLabelMarginLossOptions { - MultiLabelMarginLossOptions(torch::Reduction::Reduction reduction = torch::Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG3(MultiLabelMarginLossOptions, reduction, kNone, kMean, kSum) /// Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. /// 'none': no reduction will be applied, 'mean': the sum of the output will /// be divided by the number of elements in the output, 'sum': the output will /// be summed. Default: 'mean' - TORCH_ARG(torch::Reduction::Reduction, reduction); + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a soft margin loss functional and module. struct TORCH_API SoftMarginLossOptions { - SoftMarginLossOptions(torch::Reduction::Reduction reduction = torch::Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG3(SoftMarginLossOptions, reduction, kNone, kMean, kSum) /// Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. /// 'none': no reduction will be applied, 'mean': the sum of the output will /// be divided by the number of elements in the output, 'sum': the output will /// be summed. Default: 'mean' - TORCH_ARG(torch::Reduction::Reduction, reduction); + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a multi-label soft margin loss functional and module. struct TORCH_API MultiLabelSoftMarginLossOptions { + typedef c10::variant reduction_t; + /// A manual rescaling weight given to each /// class. If given, it has to be a Tensor of size `C`. Otherwise, it is /// treated as if having all ones. @@ -132,13 +151,15 @@ struct TORCH_API MultiLabelSoftMarginLossOptions { /// 'none': no reduction will be applied, 'mean': the sum of the output will /// be divided by the number of elements in the output, 'sum': the output will /// be summed. Default: 'mean' - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a triplet-margin-Loss functional and module. struct TORCH_API TripletMarginLossOptions { + typedef c10::variant reduction_t; + /// Specifies the threshold for which the distance of a negative sample must /// reach in order to incur zero loss. Default: 1 TORCH_ARG(double, margin) = 1.0; @@ -150,7 +171,7 @@ struct TORCH_API TripletMarginLossOptions { /// E. Riba et al. Default: False TORCH_ARG(bool, swap) = false; /// Specifies the reduction to apply to the output. Default: Mean - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; } // namespace nn diff --git a/torch/csrc/api/include/torch/nn/options/upsampling.h b/torch/csrc/api/include/torch/nn/options/upsampling.h new file mode 100644 index 0000000000000..67f62fad47470 --- /dev/null +++ b/torch/csrc/api/include/torch/nn/options/upsampling.h @@ -0,0 +1,73 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + +namespace torch { +namespace nn { + +/// Options for a `D`-dimensional interpolate functional. +struct TORCH_API InterpolateOptions { + typedef c10::variant< + enumtype::kNearest, + enumtype::kLinear, + enumtype::kBilinear, + enumtype::kBicubic, + enumtype::kTrilinear, + enumtype::kArea> mode_t; + + /// output spatial sizes. + TORCH_ARG(std::vector, size) = {}; + + /// multiplier for spatial size. + TORCH_ARG(std::vector, scale_factor) = {}; + + /// the upsampling algorithm: one of "nearest", "linear", "bilinear", + /// "bicubic", "trilinear", and "area". Default: "nearest" + TORCH_ARG(mode_t, mode) = torch::kNearest; + + /// Geometrically, we consider the pixels of the input and output as squares + /// rather than points. If set to "True", the input and output tensors are + /// aligned by the center points of their corner pixels, preserving the values + /// at the corner pixels. If set to "False", the input and output tensors + /// are aligned by the corner points of their corner pixels, and the + /// interpolation uses edge value padding for out-of-boundary values, making + /// this operation *independent* of input size when :attr:`scale_factor` is + /// kept the same. This only has an effect when :attr:`mode` is "linear", + /// "bilinear", "bicubic" or "trilinear". Default: "False" + TORCH_ARG(c10::optional, align_corners) = c10::nullopt; +}; + +/// Options for a `D`-dimensional Upsample module. +struct TORCH_API UpsampleOptions { + /// output spatial sizes. + TORCH_ARG(std::vector, size) = {}; + + /// multiplier for spatial size. + TORCH_ARG(std::vector, scale_factor) = {}; + + /// the upsampling algorithm: one of "nearest", "linear", "bilinear", + /// "bicubic" and "trilinear". Default: "nearest" + typedef c10::variant< + enumtype::kNearest, + enumtype::kLinear, + enumtype::kBilinear, + enumtype::kBicubic, + enumtype::kTrilinear> mode_t; + TORCH_ARG(mode_t, mode) = torch::kNearest; + + /// if "True", the corner pixels of the input and output tensors are + /// aligned, and thus preserving the values at those pixels. This only has + /// effect when :attr:`mode` is "linear", "bilinear", or + /// "trilinear". Default: "False" + TORCH_ARG(c10::optional, align_corners) = c10::nullopt; +}; + +} // namespace nn +} // namespace torch diff --git a/torch/csrc/api/include/torch/serialize/input-archive.h b/torch/csrc/api/include/torch/serialize/input-archive.h index 9965a442f6a7d..d4a4318e7537a 100644 --- a/torch/csrc/api/include/torch/serialize/input-archive.h +++ b/torch/csrc/api/include/torch/serialize/input-archive.h @@ -102,6 +102,7 @@ class TORCH_API InputArchive final { private: jit::script::Module module_; + std::string hierarchy_prefix_; }; } // namespace serialize } // namespace torch diff --git a/torch/csrc/api/src/enum.cpp b/torch/csrc/api/src/enum.cpp index b59eb955e9635..9cb76e0303a92 100644 --- a/torch/csrc/api/src/enum.cpp +++ b/torch/csrc/api/src/enum.cpp @@ -17,6 +17,13 @@ TORCH_ENUM_DEFINE(Constant) TORCH_ENUM_DEFINE(Reflect) TORCH_ENUM_DEFINE(Replicate) TORCH_ENUM_DEFINE(Circular) +TORCH_ENUM_DEFINE(Nearest) +TORCH_ENUM_DEFINE(Bilinear) +TORCH_ENUM_DEFINE(Bicubic) +TORCH_ENUM_DEFINE(Trilinear) +TORCH_ENUM_DEFINE(Area) TORCH_ENUM_DEFINE(Sum) TORCH_ENUM_DEFINE(Mean) TORCH_ENUM_DEFINE(Max) +TORCH_ENUM_DEFINE(None) +TORCH_ENUM_DEFINE(BatchMean) diff --git a/torch/csrc/api/src/nn/modules/batchnorm.cpp b/torch/csrc/api/src/nn/modules/batchnorm.cpp index 806d77bdb2c92..816d07dc31323 100644 --- a/torch/csrc/api/src/nn/modules/batchnorm.cpp +++ b/torch/csrc/api/src/nn/modules/batchnorm.cpp @@ -1,7 +1,9 @@ +#include #include #include #include +#include #include @@ -10,41 +12,45 @@ #include #include +namespace F = torch::nn::functional; + namespace torch { namespace nn { BatchNormImpl::BatchNormImpl(const BatchNormOptions& options_) : options(options_) { + TORCH_WARN("torch::nn::BatchNorm module is deprecated." + "Use BatchNorm{1,2,3}d instead."); reset(); } void BatchNormImpl::reset() { if (options.affine()) { weight = register_parameter( - "weight", torch::empty({options.features()}).uniform_()); - bias = register_parameter("bias", torch::zeros({options.features()})); + "weight", torch::empty({options.num_features()}).uniform_()); + bias = register_parameter("bias", torch::zeros({options.num_features()})); } - if (options.stateful()) { + if (options.track_running_stats()) { running_mean = - register_buffer("running_mean", torch::zeros({options.features()})); + register_buffer("running_mean", torch::zeros({options.num_features()})); running_var = - register_buffer("running_var", torch::ones({options.features()})); + register_buffer("running_var", torch::ones({options.num_features()})); } } void BatchNormImpl::pretty_print(std::ostream& stream) const { stream << std::boolalpha - << "torch::nn::BatchNorm(features=" << options.features() - << ", eps=" << options.eps() << ", momentum=" << options.momentum() - << ", affine=" << options.affine() << ", stateful=" << options.stateful() + << "torch::nn::BatchNorm(num_features=" << options.num_features() + << ", eps=" << options.eps() << ", momentum=" << options.momentum().value() + << ", affine=" << options.affine() << ", track_running_stats=" << options.track_running_stats() << ")"; } Tensor BatchNormImpl::forward(const Tensor& input) { TORCH_CHECK( - options.stateful(), + options.track_running_stats(), "Calling BatchNorm::forward is only permitted when " - "the 'stateful' option is true (was false). " + "the 'track_running_stats' option is true (was false). " "Use BatchNorm::pure_forward instead."); return pure_forward(input, running_mean, running_var); } @@ -67,10 +73,100 @@ Tensor BatchNormImpl::pure_forward( mean, variance, is_training(), - options.momentum(), + options.momentum().value(), options.eps(), torch::cuda::cudnn_is_available()); } +template +BatchNormImplBase::BatchNormImplBase(const BatchNormOptions& options_) + : options(options_) { + reset(); +} + +template +void BatchNormImplBase::reset_running_stats() { + if (options.track_running_stats()) { + running_mean.zero_(); + running_var.fill_(1); + num_batches_tracked.zero_(); + } +} + +template +void BatchNormImplBase::reset() { + if (options.affine()) { + weight = this->register_parameter("weight", torch::empty({options.num_features()})); + bias = this->register_parameter("bias", torch::empty({options.num_features()})); + } else { + weight = this->register_parameter("weight", Tensor()); + bias = this->register_parameter("bias", Tensor()); + } + if (options.track_running_stats()) { + running_mean = this->register_buffer("running_mean", torch::zeros({options.num_features()})); + running_var = this->register_buffer("running_var", torch::ones({options.num_features()})); + num_batches_tracked = this->register_buffer("num_batches_tracked", torch::tensor(0, torch::dtype(torch::kLong))); + } else { + running_mean = this->register_buffer("running_mean", Tensor()); + running_var = this->register_buffer("running_var", Tensor()); + num_batches_tracked = this->register_buffer("num_batches_tracked", Tensor()); + } + + reset_running_stats(); + if (options.affine()) { + torch::nn::init::ones_(weight); + torch::nn::init::zeros_(bias); + } +} + +template +void BatchNormImplBase::pretty_print(std::ostream& stream) const { + stream << std::boolalpha + << "torch::nn::BatchNorm" << D << "d(" + << options.num_features() << ", " + << "eps=" << options.eps() << ", " + << "momentum=" << options.momentum().value() << ", " + << "affine=" << options.affine() << ", " + << "track_running_stats=" << options.track_running_stats() << ")"; +} + +template +Tensor BatchNormImplBase::forward(const Tensor& input) { + _check_input_dim(input); + + double exponential_average_factor; + if (options.momentum() == c10::nullopt) { + exponential_average_factor = 0.0; + } else { + exponential_average_factor = options.momentum().value(); + } + + if (this->is_training() && options.track_running_stats()) { + if (num_batches_tracked.defined()) { + num_batches_tracked += 1; + if (options.momentum() == c10::nullopt) { // use cumulative moving average + exponential_average_factor = 1.0 / num_batches_tracked.item(); + } else { // use exponential moving average + exponential_average_factor = options.momentum().value(); + } + } + } + + return F::batch_norm( + input, + running_mean, + running_var, + BatchNormOptions().weight(weight).bias(bias).momentum(exponential_average_factor).eps(options.eps()), + this->is_training() || !options.track_running_stats()); +} + +void BatchNorm1dImpl::_check_input_dim(const Tensor& input) { + TORCH_CHECK( + input.dim() == 2 || input.dim() == 3, + "expected 2D or 3D input (got ", input.dim(), "D input)"); +} + +template class BatchNormImplBase<1, BatchNorm1dImpl>; + } // namespace nn } // namespace torch diff --git a/torch/csrc/api/src/nn/modules/embedding.cpp b/torch/csrc/api/src/nn/modules/embedding.cpp index a4a583f6ae6f6..81d2b34205b2a 100644 --- a/torch/csrc/api/src/nn/modules/embedding.cpp +++ b/torch/csrc/api/src/nn/modules/embedding.cpp @@ -150,7 +150,7 @@ torch::Tensor EmbeddingBagImpl::forward( !per_sample_weights_.defined() || c10::get_if(&options.mode()), "embedding_bag: per_sample_weights was not null. ", "per_sample_weights is only supported for mode='kSum' (got mode='", - c10::visit(torch::enumtype::enum_name{}, options.mode()), "').Please open a feature request on GitHub."); + torch::enumtype::get_enum_name(options.mode()), "').Please open a feature request on GitHub."); return std::get<0>( torch::embedding_bag( @@ -179,7 +179,7 @@ void EmbeddingBagImpl::pretty_print(std::ostream& stream) const { stream << ", sparse=" << std::boolalpha << options.sparse(); } if (!c10::get_if(&options.mode())) { - stream << ", mode=" << c10::visit(torch::enumtype::enum_name{}, options.mode()); + stream << ", mode=" << torch::enumtype::get_enum_name(options.mode()); } stream << ")"; } diff --git a/torch/csrc/api/src/nn/modules/linear.cpp b/torch/csrc/api/src/nn/modules/linear.cpp index c457cffa52dea..0fc134c1af67a 100644 --- a/torch/csrc/api/src/nn/modules/linear.cpp +++ b/torch/csrc/api/src/nn/modules/linear.cpp @@ -61,6 +61,20 @@ Tensor LinearImpl::forward(const Tensor& input) { // ============================================================================ +FlattenImpl::FlattenImpl(const FlattenOptions& options_) : options(options_) {} + +void FlattenImpl::reset() {} + +void FlattenImpl::pretty_print(std::ostream& stream) const { + stream << "torch::nn::Flatten()"; +} + +Tensor FlattenImpl::forward(const Tensor& input) { + return input.flatten(options.start_dim(), options.end_dim()); +} + +// ============================================================================ + BilinearImpl::BilinearImpl(const BilinearOptions& options_) : options(options_) { reset(); } @@ -90,5 +104,6 @@ void BilinearImpl::pretty_print(std::ostream& stream) const { Tensor BilinearImpl::forward(const Tensor& input1, const Tensor& input2) { return F::bilinear(input1, input2, weight, bias); } + } // namespace nn } // namespace torch diff --git a/torch/csrc/api/src/nn/modules/loss.cpp b/torch/csrc/api/src/nn/modules/loss.cpp index 3eaf2a1ba0b5e..9ca998c616079 100644 --- a/torch/csrc/api/src/nn/modules/loss.cpp +++ b/torch/csrc/api/src/nn/modules/loss.cpp @@ -102,7 +102,7 @@ void MultiMarginLossImpl::reset() { void MultiMarginLossImpl::pretty_print(std::ostream& stream) const { stream << "torch::nn::MultiMarginLoss(p=" << options.p() << ", margin=" << options.margin() << ", weight=" << options.weight() - << ", reduction=" << options.reduction() << ")"; + << ", reduction=" << enumtype::get_enum_name(options.reduction()) << ")"; } Tensor MultiMarginLossImpl::forward(const Tensor& input, const Tensor& target) { diff --git a/torch/csrc/api/src/nn/modules/upsampling.cpp b/torch/csrc/api/src/nn/modules/upsampling.cpp new file mode 100644 index 0000000000000..401cb16ce78b3 --- /dev/null +++ b/torch/csrc/api/src/nn/modules/upsampling.cpp @@ -0,0 +1,49 @@ +#include + +#include + +namespace F = torch::nn::functional; + +namespace torch { +namespace nn { + +UpsampleImpl::UpsampleImpl(const UpsampleOptions& options_) // NOLINT(modernize-pass-by-value) + : options(options_) {} + +void UpsampleImpl::reset() {} + +void UpsampleImpl::pretty_print(std::ostream& stream) const { + stream << "torch::nn::Upsample("; + if (!options.scale_factor().empty()) { + stream << "scale_factor=" << at::ArrayRef(options.scale_factor()); + } else { + stream << "size=" << at::ArrayRef(options.size()); + } + stream << ", mode=" << enumtype::get_enum_name(options.mode()) << ")"; +} + +Tensor UpsampleImpl::forward(const Tensor& input) { + InterpolateOptions::mode_t mode; + if (c10::get_if(&options.mode())) { + mode = torch::kNearest; + } else if (c10::get_if(&options.mode())) { + mode = torch::kLinear; + } else if (c10::get_if(&options.mode())) { + mode = torch::kBilinear; + } else if (c10::get_if(&options.mode())) { + mode = torch::kBicubic; + } else if (c10::get_if(&options.mode())) { + mode = torch::kTrilinear; + } + + return F::interpolate( + input, + InterpolateOptions() + .size(options.size()) + .scale_factor(options.scale_factor()) + .mode(mode) + .align_corners(options.align_corners())); +} + +} // namespace nn +} // namespace torch diff --git a/torch/csrc/api/src/nn/options/batchnorm.cpp b/torch/csrc/api/src/nn/options/batchnorm.cpp index 60c363b8d9e5c..2144443913749 100644 --- a/torch/csrc/api/src/nn/options/batchnorm.cpp +++ b/torch/csrc/api/src/nn/options/batchnorm.cpp @@ -3,7 +3,7 @@ namespace torch { namespace nn { -BatchNormOptions::BatchNormOptions(int64_t features) : features_(features) {} +BatchNormOptions::BatchNormOptions(int64_t num_features) : num_features_(num_features) {} } // namespace nn } // namespace torch diff --git a/torch/csrc/api/src/serialize/input-archive.cpp b/torch/csrc/api/src/serialize/input-archive.cpp index 856ddcb7a4c28..5ebbe1fe7a249 100644 --- a/torch/csrc/api/src/serialize/input-archive.cpp +++ b/torch/csrc/api/src/serialize/input-archive.cpp @@ -63,15 +63,17 @@ void InputArchive::read( Tensor& tensor, bool is_buffer) { TORCH_CHECK( - try_read(key, tensor, is_buffer), - "No such serialized tensor '", - key, - "'"); + try_read(key, tensor, is_buffer), + "No such serialized tensor '", + hierarchy_prefix_, + key, + "'"); } bool InputArchive::try_read(const std::string& key, InputArchive& archive) { if (auto named_module = module_.find_module(key)) { archive.module_ = std::move(*named_module); + archive.hierarchy_prefix_ = hierarchy_prefix_ + key + "."; return true; } else { return false; @@ -80,8 +82,11 @@ bool InputArchive::try_read(const std::string& key, InputArchive& archive) { void InputArchive::read(const std::string& key, InputArchive& archive) { TORCH_CHECK( - try_read(key, archive), - "No such serialized submodule: '", key, "'"); + try_read(key, archive), + "No such serialized submodule: '", + hierarchy_prefix_, + key, + "'"); } void InputArchive::load_from(const std::string& filename, diff --git a/torch/csrc/autograd/VariableTypeManual.cpp b/torch/csrc/autograd/VariableTypeManual.cpp index b141003256807..165c0a06898bd 100644 --- a/torch/csrc/autograd/VariableTypeManual.cpp +++ b/torch/csrc/autograd/VariableTypeManual.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -115,7 +116,7 @@ int64_t _version(const Tensor & self) { Tensor& requires_grad_(Tensor& self, bool _requires_grad) { if (!self.is_leaf() && !_requires_grad) { throw std::runtime_error( - autograd::utils::requires_grad_leaf_error(requires_grad) + autograd::utils::requires_grad_leaf_error(_requires_grad) ); } return self.set_requires_grad(_requires_grad); @@ -149,7 +150,9 @@ Tensor & copy_(Tensor & self, const Tensor & src, bool non_blocking) { check_inplace(self); std::shared_ptr grad_fn; auto requires_grad = compute_requires_grad(self, src); - requires_grad &= isFloatingPoint(self.scalar_type()); + // currently, isFloatingType will return false for (floating) complex types, + // so this might have to be amended when they should be differentiable + requires_grad &= isFloatingType(self.scalar_type()); if (requires_grad) { grad_fn = std::make_shared(); grad_fn->set_next_edges(collect_next_edges(self, src)); diff --git a/torch/csrc/autograd/VariableTypeUtils.h b/torch/csrc/autograd/VariableTypeUtils.h index e42d80aedad13..fa503d47eaa87 100644 --- a/torch/csrc/autograd/VariableTypeUtils.h +++ b/torch/csrc/autograd/VariableTypeUtils.h @@ -80,10 +80,6 @@ inline void increment_version(Tensor & t) { as_variable_ref(t).bump_version(); } -inline bool isFloatingPoint(ScalarType s) { - return s == kFloat || s == kDouble || s == kHalf; -} - struct Flatten : IterArgs { Flatten(variable_list& out) : out(out) {} variable_list& out; diff --git a/torch/csrc/distributed/autograd/utils.cpp b/torch/csrc/distributed/autograd/utils.cpp index 8c0b9602cbc2d..da85418aa0e05 100644 --- a/torch/csrc/distributed/autograd/utils.cpp +++ b/torch/csrc/distributed/autograd/utils.cpp @@ -42,37 +42,45 @@ DistAutogradContext* addRecvRpcBackward( const AutogradMetadata& autogradMetadata, std::vector& tensors, rpc::worker_id_t fromWorkerId) { - TORCH_INTERNAL_ASSERT( - torch::autograd::compute_requires_grad(tensors), - "Received tensors do not require grad, addRecvRpcBackward should not be called"); // Initialize autograd context if necessary. auto& autogradContainer = DistAutogradContainer::getInstance(); DistAutogradContext& autogradContext = autogradContainer.getOrCreateContext(autogradMetadata.autogradContextId); - // Attach the tensors as inputs to the autograd function. - auto grad_fn = std::make_shared( - autogradMetadata, autogradContext, fromWorkerId); - for (auto& tensor : tensors) { - torch::autograd::set_history(tensor, grad_fn); + if (!tensors.empty()) { + TORCH_INTERNAL_ASSERT( + torch::autograd::compute_requires_grad(tensors), + "Received tensors do not require grad, addRecvRpcBackward should not be called"); + + // Attach the tensors as inputs to the autograd function. + auto grad_fn = std::make_shared( + autogradMetadata, autogradContext, fromWorkerId); + for (auto& tensor : tensors) { + torch::autograd::set_history(tensor, grad_fn); + } + + // Now update the autograd context with the necessary information. + autogradContext.addRecvFunction( + grad_fn, autogradMetadata.autogradMessageId); } - // Now update the autograd context with the necessary information. - autogradContext.addRecvFunction(grad_fn, autogradMetadata.autogradMessageId); return &autogradContext; } Message getMessageWithAutograd( const rpc::worker_id_t dstId, torch::distributed::rpc::Message&& wrappedRpcMsg, - MessageType msgType) { + MessageType msgType, + bool forceGradRecording) { auto& autogradContainer = DistAutogradContainer::getInstance(); // If there is no valid context and no tensor requires grads, send original // rpc message. otherwise, attach grad info and grad functions and send // rpcWithAutograd message. + auto tensorsRequireGrad = + torch::autograd::compute_requires_grad(wrappedRpcMsg.tensors()); if (!autogradContainer.hasValidContext() || - !torch::autograd::compute_requires_grad(wrappedRpcMsg.tensors())) { + (!forceGradRecording && !tensorsRequireGrad)) { return std::move(wrappedRpcMsg); } @@ -88,9 +96,11 @@ Message getMessageWithAutograd( autogradMetadata, std::move(wrappedRpcMsg)); - // Record autograd information for 'send'. - addSendRpcBackward( - autogradContext, autogradMetadata, rpcWithAutograd->tensors(), dstId); + if (tensorsRequireGrad) { + // Record autograd information for 'send'. + addSendRpcBackward( + autogradContext, autogradMetadata, rpcWithAutograd->tensors(), dstId); + } return std::move(*rpcWithAutograd).toMessage(); } @@ -98,9 +108,13 @@ Message getMessageWithAutograd( std::shared_ptr sendMessageWithAutograd( RpcAgent& agent, const WorkerInfo& dst, - torch::distributed::rpc::Message&& wrappedRpcMsg) { + torch::distributed::rpc::Message&& wrappedRpcMsg, + bool forceGradRecording) { auto msg = getMessageWithAutograd( - dst.id_, std::move(wrappedRpcMsg), MessageType::FORWARD_AUTOGRAD_REQ); + dst.id_, + std::move(wrappedRpcMsg), + MessageType::FORWARD_AUTOGRAD_REQ, + forceGradRecording); return agent.send(dst, std::move(msg)); } diff --git a/torch/csrc/distributed/autograd/utils.h b/torch/csrc/distributed/autograd/utils.h index 2f368cdf17fa2..cc147c059f1d7 100644 --- a/torch/csrc/distributed/autograd/utils.h +++ b/torch/csrc/distributed/autograd/utils.h @@ -34,19 +34,23 @@ TORCH_API DistAutogradContext* addRecvRpcBackward( // This method is a wrapper utility used internally to wrap autograd info // and attach autograd function for each type of rpc call if it has valid -// context and tensors require grads, in this case, return RpcWithAutograd -// message; otherwise return original rpc message. +// context and tensors require grads or forceGradRecording is true, in this +// case, return RpcWithAutograd message; otherwise return original rpc message. +// NB: forceGradRecording is useful when the request does not contain any tensor +// but the corresponding response does. TORCH_API rpc::Message getMessageWithAutograd( const rpc::worker_id_t dstId, rpc::Message&& wrappedRpcMsg, - rpc::MessageType msgType); + rpc::MessageType msgType, + bool forceGradRecording = false); // Send message after autograd checking TORCH_API std::shared_ptr sendMessageWithAutograd( rpc::RpcAgent& agent, const rpc::WorkerInfo& dst, - rpc::Message&& wrappedRpcMsg); + rpc::Message&& wrappedRpcMsg, + bool forceGradRecording = false); } // namespace autograd } // namespace distributed diff --git a/torch/csrc/distributed/rpc/python_functions.cpp b/torch/csrc/distributed/rpc/python_functions.cpp index 6a13eedf25ad5..50d81fe5c4473 100644 --- a/torch/csrc/distributed/rpc/python_functions.cpp +++ b/torch/csrc/distributed/rpc/python_functions.cpp @@ -37,7 +37,6 @@ std::shared_ptr matchBuiltinOp( // ``createStackForSchema`` to avoid throwing an error. stack = torch::jit::createStackForSchema( op->schema(), args, kwargs, c10::nullopt); - } catch (std::runtime_error& e) { VLOG(1) << "Couldn't match schema: " << op->schema() << " to args: " << args << " and kwargs: " << kwargs @@ -143,11 +142,12 @@ PyRRef pyRemoteBuiltin( ctx.getWorkerId() != dst.id_, "Does not support creating RRef on self yet."); auto userRRef = ctx.createUserRRef(dst.id_); - auto fm = agent.send( - dst, - ScriptRemoteCall( - op, std::move(stack), userRRef->rrefId(), userRRef->forkId()) - .toMessage()); + + auto scriptRemoteCall = c10::guts::make_unique( + op, std::move(stack), userRRef->rrefId(), userRRef->forkId()); + + auto fm = sendMessageWithAutograd( + agent, dst, std::move(*scriptRemoteCall).toMessage()); ctx.addPendingUser(userRRef->forkId(), userRRef); fm->addCallback(finishAcceptUserRRef); @@ -177,13 +177,20 @@ PyRRef pyRemotePythonUdf( ctx.getWorkerId() != dst.id_, "Does not support creating RRef on self yet."); auto userRRef = ctx.createUserRRef(dst.id_); - auto fm = agent.send( + + auto pythonRemoteCall = c10::guts::make_unique( + SerializedPyObj(std::move(pickledPythonUDF), std::move(tensors)), + userRRef->rrefId().toIValue(), + userRRef->forkId().toIValue()); + + // set forceGradRecording to true as even if the args does not contain any + // tensor, the return value might still contain tensors. + auto fm = sendMessageWithAutograd( + agent, dst, - PythonRemoteCall( - SerializedPyObj(std::move(pickledPythonUDF), std::move(tensors)), - userRRef->rrefId().toIValue(), - userRRef->forkId().toIValue()) - .toMessage()); + std::move(*pythonRemoteCall).toMessage(), + true /*forceGradRecording*/ + ); ctx.addPendingUser(userRRef->forkId(), userRRef); fm->addCallback(finishAcceptUserRRef); diff --git a/torch/csrc/distributed/rpc/request_callback_impl.cpp b/torch/csrc/distributed/rpc/request_callback_impl.cpp index 2cc7a8cc81df6..a56d5b4463c74 100644 --- a/torch/csrc/distributed/rpc/request_callback_impl.cpp +++ b/torch/csrc/distributed/rpc/request_callback_impl.cpp @@ -82,7 +82,7 @@ Message RequestCallbackImpl::processRpc( ownerRRef->setValue(std::move(stack.front())); ctx.addForkOfOwner(src.retRRefId(), src.retForkId()); - return std::move(RemoteRet(src.retRRefId(), src.retForkId())).toMessage(); + return RemoteRet(src.retRRefId(), src.retForkId()).toMessage(); } case MessageType::PYTHON_REMOTE_CALL: { auto& prc = static_cast(rpc); @@ -96,7 +96,7 @@ Message RequestCallbackImpl::processRpc( ownerRRef->setValue( PythonRpcHandler::getInstance().runPythonUDF(prc.serializedPyObj())); ctx.addForkOfOwner(rrefId, forkId); - return std::move(RemoteRet(rrefId, forkId)).toMessage(); + return RemoteRet(rrefId, forkId).toMessage(); } case MessageType::SCRIPT_RREF_FETCH_CALL: { auto& srf = static_cast(rpc); @@ -104,7 +104,7 @@ Message RequestCallbackImpl::processRpc( // TODO: make this asynchronous std::shared_ptr> rref = ctx.getOrCreateOwnerRRef(srf.rrefId()); - return std::move(ScriptRRefFetchRet({rref->getValue()})).toMessage(); + return ScriptRRefFetchRet({rref->getValue()}).toMessage(); } case MessageType::PYTHON_RREF_FETCH_CALL: { auto& prf = static_cast(rpc); @@ -114,7 +114,7 @@ Message RequestCallbackImpl::processRpc( ctx.getOrCreateOwnerRRef(prf.rrefId()); SerializedPyObj result = PythonRpcHandler::getInstance().serialize(rref->getValue()); - return std::move(PythonRRefFetchRet(result.toIValues())).toMessage(); + return PythonRRefFetchRet(result.toIValues()).toMessage(); } case MessageType::RREF_USER_DELETE: { auto& rud = static_cast(rpc); @@ -132,7 +132,7 @@ Message RequestCallbackImpl::processRpc( auto& rfr = static_cast(rpc); auto& ctx = RRefContext::getInstance(); ctx.addForkOfOwner(rfr.rrefId(), rfr.forkId()); - return std::move(RRefAck()).toMessage(); + return RRefAck().toMessage(); } case MessageType::FORWARD_AUTOGRAD_REQ: { auto& rpcWithAutograd = static_cast(rpc); diff --git a/torch/csrc/distributed/rpc/rref.cpp b/torch/csrc/distributed/rpc/rref.cpp index 1d08f97ca7b0b..78e884e572397 100644 --- a/torch/csrc/distributed/rpc/rref.cpp +++ b/torch/csrc/distributed/rpc/rref.cpp @@ -1,7 +1,10 @@ #include +#include +#include #include #include +#include namespace torch { namespace distributed { @@ -19,6 +22,26 @@ constexpr int PARENT_IDX = 5; // index of parent in the tuple // NB: if more fields are added, make sure this field is also bumped constexpr int RFD_TUPLE_SIZE = 6; // number of RRefForkData fields in py::tuple +template +T& unwrapAutogradMessage( + const Message& message, + std::unique_ptr& response) { + if (message.type() == MessageType::FORWARD_AUTOGRAD_RESP) { + auto& rpcWithAutograd = static_cast(*response); + + // Attach 'recv' autograd function. + addRecvRpcBackward( + rpcWithAutograd.autogradMetadata(), + rpcWithAutograd.tensors(), + rpcWithAutograd.fromWorkerId()); + + auto& wrappedRpc = rpcWithAutograd.wrappedRpc(); + return static_cast(wrappedRpc); + } else { + return static_cast(*response); + } +} + } // namespace std::atomic RRefContext::nextLocalId_{0}; @@ -133,14 +156,21 @@ template <> std::shared_ptr UserRRef::toHere() { auto future = std::make_shared(nullptr); auto agent = RpcAgent::getDefaultRpcAgent(); - auto futureResponse = agent->send( + + // ScriptRRefFetchCall message always carries autograd context id even if + // the message itself does not contain any tensor, because the response would + // potentially contain tensors. + auto futureResponse = autograd::sendMessageWithAutograd( + *agent, agent->getWorkerInfo(ownerId_), - ScriptRRefFetchCall(rrefId()).toMessage()); + ScriptRRefFetchCall(ownerId_, rrefId()).toMessage(), + true /* forceGradRecording */); futureResponse->addCallback([future](const Message& message) { RRefContext::handleException(message); - auto rfr = ScriptRRefFetchRet::fromMessage(message); - future->markCompleted(rfr->values().front()); + auto response = deserializeResponse(message); + auto& rfr = unwrapAutogradMessage(message, response); + future->markCompleted(rfr.values().front()); }); return future; } @@ -149,14 +179,21 @@ template <> std::shared_ptr UserRRef::toHere() { auto future = std::make_shared(nullptr); auto agent = RpcAgent::getDefaultRpcAgent(); - auto futureResponse = agent->send( + + // PythonRRefFetchCall message always carries autograd context id even if + // the message itself does not contain any tensor, because the response would + // potentially contain tensors. + auto futureResponse = autograd::sendMessageWithAutograd( + *agent, agent->getWorkerInfo(ownerId_), - PythonRRefFetchCall(rrefId()).toMessage()); + PythonRRefFetchCall(ownerId_, rrefId()).toMessage(), + true /* forceGradRecording */); futureResponse->addCallback([future](const Message& message) { RRefContext::handleException(message); - auto rfr = PythonRRefFetchRet::fromMessage(message); - future->markCompleted(c10::ivalue::Tuple::create(rfr->values())); + auto response = deserializeResponse(message); + auto& rfr = unwrapAutogradMessage(message, response); + future->markCompleted(c10::ivalue::Tuple::create(rfr.values())); }); return future; } diff --git a/torch/csrc/distributed/rpc/rref_proto.cpp b/torch/csrc/distributed/rpc/rref_proto.cpp index 740a5b470647f..4c503338d6e56 100644 --- a/torch/csrc/distributed/rpc/rref_proto.cpp +++ b/torch/csrc/distributed/rpc/rref_proto.cpp @@ -77,18 +77,48 @@ std::pair ForkMessageBase::fromMessage( /////////////////////////// RRef Protocol ////////////////////////////////// +Message ScriptRRefFetchCall::toMessage() && { + std::vector ivalues; + ivalues.reserve(2); + ivalues.emplace_back(rrefId_.toIValue()); + ivalues.emplace_back(fromWorkerId_); + return fromIValues(std::move(ivalues), MessageType::SCRIPT_RREF_FETCH_CALL); +} + std::unique_ptr ScriptRRefFetchCall::fromMessage( const Message& message) { + auto values = toIValues(message, MessageType::SCRIPT_RREF_FETCH_CALL); + TORCH_INTERNAL_ASSERT( + values.size() == 2, "ScriptRRefFetchCall expects 2 IValues from message"); + auto id = values[1].toInt(); + TORCH_INTERNAL_ASSERT( + id >= std::numeric_limits::min() && + id <= std::numeric_limits::max(), + "ScriptRRefFetchCall fromWorkerId exceeds worker_id_t limit.") return c10::guts::make_unique( - RRefId::fromIValue(RRefMessageBase::fromMessage( - message, MessageType::SCRIPT_RREF_FETCH_CALL))); + worker_id_t(id), RRefId::fromIValue(values[0])); +} + +Message PythonRRefFetchCall::toMessage() && { + std::vector ivalues; + ivalues.reserve(2); + ivalues.emplace_back(rrefId_.toIValue()); + ivalues.emplace_back(fromWorkerId_); + return fromIValues(std::move(ivalues), MessageType::PYTHON_RREF_FETCH_CALL); } std::unique_ptr PythonRRefFetchCall::fromMessage( const Message& message) { + auto values = toIValues(message, MessageType::PYTHON_RREF_FETCH_CALL); + TORCH_INTERNAL_ASSERT( + values.size() == 2, "PythonRRefFetchCall expects 2 IValues from message"); + auto id = values[1].toInt(); + TORCH_INTERNAL_ASSERT( + id >= std::numeric_limits::min() && + id <= std::numeric_limits::max(), + "PythonRRefFetchCall fromWorkerId exceeds worker_id_t limit.") return c10::guts::make_unique( - RRefId::fromIValue(RRefMessageBase::fromMessage( - message, MessageType::PYTHON_RREF_FETCH_CALL))); + worker_id_t(id), RRefId::fromIValue(values[0])); } const std::vector& RRefFetchRet::values() { diff --git a/torch/csrc/distributed/rpc/rref_proto.h b/torch/csrc/distributed/rpc/rref_proto.h index 55eeb97b183b0..d49daa2300f51 100644 --- a/torch/csrc/distributed/rpc/rref_proto.h +++ b/torch/csrc/distributed/rpc/rref_proto.h @@ -51,20 +51,34 @@ class TORCH_API ForkMessageBase : public RRefMessageBase { // UserRRef uses this message to fetch the remote RRef value from the owner. class TORCH_API ScriptRRefFetchCall final : public RRefMessageBase { public: - explicit ScriptRRefFetchCall(const RRefId& rrefId) - : RRefMessageBase(rrefId, MessageType::SCRIPT_RREF_FETCH_CALL) {} + ScriptRRefFetchCall(worker_id_t fromWorkerId, const RRefId& rrefId) + : RRefMessageBase(rrefId, MessageType::SCRIPT_RREF_FETCH_CALL), + fromWorkerId_(fromWorkerId) {} + inline worker_id_t fromWorkerId() const { + return fromWorkerId_; + } + + Message toMessage() && override; static std::unique_ptr fromMessage( const Message& message); + + private: + const worker_id_t fromWorkerId_; }; class TORCH_API PythonRRefFetchCall final : public RRefMessageBase { public: - explicit PythonRRefFetchCall(const RRefId& rrefId) - : RRefMessageBase(rrefId, MessageType::PYTHON_RREF_FETCH_CALL) {} + PythonRRefFetchCall(worker_id_t fromWorkerId, const RRefId& rrefId) + : RRefMessageBase(rrefId, MessageType::PYTHON_RREF_FETCH_CALL), + fromWorkerId_(fromWorkerId) {} + Message toMessage() && override; static std::unique_ptr fromMessage( const Message& message); + + private: + const worker_id_t fromWorkerId_; }; // OwnerRRef uses this message to send the RRef value to a remote UserRRef diff --git a/torch/csrc/distributed/rpc/utils.cpp b/torch/csrc/distributed/rpc/utils.cpp index df619dcc15829..e40a2ebe51e74 100644 --- a/torch/csrc/distributed/rpc/utils.cpp +++ b/torch/csrc/distributed/rpc/utils.cpp @@ -71,6 +71,12 @@ std::unique_ptr deserializeResponse(const Message& response) { case MessageType::REMOTE_RET: { return RemoteRet::fromMessage(response); } + case MessageType::SCRIPT_RREF_FETCH_RET: { + return ScriptRRefFetchRet::fromMessage(response); + } + case MessageType::PYTHON_RREF_FETCH_RET: { + return PythonRRefFetchRet::fromMessage(response); + } case MessageType::RREF_ACK: { return RRefAck::fromMessage(response); } diff --git a/torch/csrc/jit/autodiff.cpp b/torch/csrc/jit/autodiff.cpp index 6e26375ad660b..af13353414308 100644 --- a/torch/csrc/jit/autodiff.cpp +++ b/torch/csrc/jit/autodiff.cpp @@ -788,12 +788,31 @@ static void lambdaLiftReverse(Gradient& grad_desc, ReverseDetails& rev_info) { for (auto& offset : grad_desc.df_input_captured_outputs) add_capture(graph.outputs()[offset]); - GRAPH_DUMP(" forward graph: ", &graph); - GRAPH_DEBUG(" backward graph: ", *(reverse_block->owningNode())); grad_desc.df = std::make_shared(); grad_desc.df->block()->cloneFrom(reverse_block, [&](Value* v) { return grad_desc.df->inputs()[capture_to_formal_index.at(v)]; }); + + // if we actually profile we can rely on profiling information + // so we don't have to mark every gradient as possibly undefined + if (!getProfilingMode() && getExecutorMode()) { + for (size_t i = 0; i < grad_desc.df_input_vjps.size(); i++) { + auto tt = grad_desc.df->block()->inputs().at(i); + if (auto ttt = tt->type()->cast()) { + tt->setType(ttt->withPossiblyUndefined()); + } else if (auto lt = tt->type()->cast()) { + auto undef_type = + lt->getElementType()->expect()->withPossiblyUndefined(); + tt->setType(ListType::create(undef_type)); + } else { + // unexpected type + TORCH_INTERNAL_ASSERT(false); + } + } + } + + GRAPH_DUMP(" forward graph: ", &graph); + GRAPH_DEBUG(" backward graph: ", *(reverse_block->owningNode())); // reverse_node was just to hold onto reverse_block in a debuggable way // we can remove it now. reverse_block->owningNode()->destroy(); @@ -829,9 +848,6 @@ Gradient differentiate(std::shared_ptr& graph) { // Fills in f, df, f_real_outputs, df_input_captures, // modifies df_input_vjps (new vjps are added for temporaries) lambdaLiftReverse(grad_desc, rev_info); - // It's possible the we've cloned the same constants many times, so - // de-duplicate them - ConstantPooling(grad_desc.df); packReturnValuesIntoTuple(grad_desc.df); return grad_desc; } diff --git a/torch/csrc/jit/docs/OVERVIEW.md b/torch/csrc/jit/docs/OVERVIEW.md index daa5999b3f6ec..8212468343afb 100644 --- a/torch/csrc/jit/docs/OVERVIEW.md +++ b/torch/csrc/jit/docs/OVERVIEW.md @@ -352,7 +352,7 @@ JIT programs are created using either the tracing frontend (`torch.jit.trace`) o The tracer produces graphs by recording what actual operations are done on tensors. The entry point from Python into C++ for tracing using `torch.jit.trace` is `_create_method_from_trace`. -A thread local instance of the TracingState object maintains a mapping between actual data being computing during the trace (e.g. Tensors) stored in IValues, and the abstract `Value*` in the Graph that would compute that value. The functions `void setValueTrace(const IValue&, Value*)` and `Value* getValueTrace(const IValue&)` are used by the tracer to maintain this mapping. +A thread local instance of the TracingState object maintains a mapping between actual data being computed during the trace (e.g. Tensors) stored in IValues, and the abstract `Value*` in the Graph that would compute that value. The functions `void setValueTrace(const IValue&, Value*)` and `Value* getValueTrace(const IValue&)` are used by the tracer to maintain this mapping. An initial IValue to Value mapping is setup up between the inputs to the function being traced and symbolic Value inputs to the Graph being constructed. If we are tracing a `torch.nn.Module`, the tracer also adds Parameters and sub-Modules to the Module being constructed that correspond to the Python `torch.nn.Module` being traced. These values are also added as mapping so that uses of the Parameters in the trace will create uses of the Parameters in the Graph. @@ -384,7 +384,7 @@ The functions `addInputs` and `addOutput` are overloaded to handle the different Currently set/getValueTrace only works on Tensors and Futures. Other types are not natively traced. Instead aggregates like tuples or lists are often flattened into tensors at the end of a trace and explicitly constructed from individual tensors at the beginning of this trace. -The tracer has special behavior when tracing calls to other TorchScript functions. This behavior is implemented in the GraphExecutor right before a Graph is about to be run. If tracing is enabled while running the graph, the GraphExecutor will disable tracing, run the graph as normal, and then inline the Graph into the trace. It then hooks up the IValues computed by running the Graph to inlined Graph's out Values in the inlined graph. +The tracer has special behavior when tracing calls to other TorchScript functions. This behavior is implemented in the GraphExecutor right before a Graph is about to be run. If tracing is enabled while running the graph, the GraphExecutor will disable tracing, run the graph as normal, and then inline the Graph into the trace. It then hooks up the IValues computed by running the Graph to out Values in the inlined graph. > *When a trace calls a TorchScript function, that function is preserved as is, meaning that control-flow is preserved.* This makes it possible to "fix" tracing issues by writing the subset of the program that cannot be traced in script and having the trace invoke it. @@ -392,7 +392,7 @@ The resulting Graph created by tracing is installed as the 'forward' method of t ## Script ## -The script frontend directly converts Python syntax into Modules. Like many compilers this happens in two phases. First, we generate an abstract syntax tree (AST), which is constructed out of Tree objects. The compiler (misnamed, but that is the name of the file) then does semantic analysis on the Tree and lowers it into a Module. We can generate Trees in two ways: (1) using frontend.py, which takes the Python AST and transliterates it into Tree objects, or (2) via the Lexer and Parser which parse python syntax directly. The Lexer/Parser path may seem redundant but it is crucially important. We need to define builtin functions ([script/builtin_functions.cpp](../script/builtin_functions.cpp)) when Python is not linked. We allow users to load TorchScript programs directly from strings without Python ([api/include/torch/jit.h](../../../api/include/torch/jit.h)). We also use this Python syntax as the serialization format for TorchScript, since it allows us to make changes to our IR without breaking backward compatibility. Furthermore, the Lexer is reused to implement the FunctionSchema parser, which turns FunctionSchema declarations from strings into FunctionSchema objects. +The script frontend directly converts Python syntax into Modules. Like many compilers this happens in two phases. First, we generate an abstract syntax tree (AST), which is constructed out of Tree objects. The compiler (misnamed, but that is the name of the file) then does semantic analysis on the Tree and lowers it into a Module. We can generate Trees in two ways: (1) using frontend.py, which takes the Python AST and transliterates it into Tree objects, or (2) via the Lexer and Parser which parse python syntax directly. The Lexer/Parser path may seem redundant but it is crucially important. We need to define builtin functions ([script/builtin_functions.cpp](../script/builtin_functions.cpp)) when Python is not linked. We allow users to load TorchScript programs directly from strings without Python ([api/include/torch/jit.h](../../api/include/torch/jit.h)). We also use this Python syntax as the serialization format for TorchScript, since it allows us to make changes to our IR without breaking backward compatibility. Furthermore, the Lexer is reused to implement the FunctionSchema parser, which turns FunctionSchema declarations from strings into FunctionSchema objects. The following sections look into each the stages in the script frontend in detail. @@ -531,7 +531,7 @@ SugaredValues are how the compiler represents non-first class values during Grap SugaredValues are also how we interact with Python runtime during the compilation process. For instance, `math.pi` is resolved to 3.1415... by first resolving `math` to a SugaredValue representing accesses to Python modules (PythonModuleValue) whose `attr` function turns python numbers into `prim::Constant` Nodes in the graph. -Finally, normal Values are also represented by the SimpleValue SugaredValue in places where it is valid either a SugaredValue or a normal Value to appear. +Finally, normal Values are also represented by the SimpleValue SugaredValue in places where it is valid that either a SugaredValue or a normal Value will appear. ## Resolver ## @@ -601,7 +601,7 @@ In addition to being mutable, tensors also have a set of dynamically determined * size - the precise size of the tensor * requires_grad - whether the tensor is recording its gradient with autograd -Changes in these properties change how operators on tensor will evaluate and would make certain optimization invalid. For instance, if we have fuser capable of generating new cuda kernels but not cpu kernels, it is only valid to fuse operations where the inputs are known to run only on CUDA devices. The GraphExecutor's job is to still enable optimization even when certains combinations of properties prevent optimizations for occurring. +Changes in these properties change how operators on tensor will evaluate and would make certain optimization invalid. For instance, if we have fuser capable of generating new cuda kernels but not cpu kernels, it is only valid to fuse operations where the inputs are known to run only on CUDA devices. The GraphExecutor's job is to still enable optimization even when certains combinations of properties prevent optimizations from occurring. Nodes in a graph are executed *serially* in the order they appear in a block. Nodes may be reordered either during optimization or by the interpreter itself if it can be proven that it is not distinguishable from the original serial execution order. These semantics are necessary since the combination of mutable tensors and potential alias between tensors makes it unsafe to perform arbitrary reordering otherwise. However, the AliasInfo object can accurately track how alias propagate through builtin operators so optimization passes can query when certain reorders or optimizations are safe. @@ -797,7 +797,7 @@ graph(%x : Tensor, return (%30) ``` -Execution starts in `GraphExecutor::run`, which takes takes a Stack of inputs. +Execution starts in `GraphExecutor::run`, which takes a Stack of inputs. *Specialization* The executor *specializes* the Graph for the particular set of inputs. Specialization is handled by the `ArgumentSpec` object which extracts a "signature" composed of all the properties being specialized. We only specialize to the properties of Tensors. The ArgumentSpec only records properties for Tensors that either appear directly in the inputs to the graph or inside Tuples that are inputs to the Graph. The properties recorded are currently: @@ -1050,7 +1050,7 @@ one specifies a file(s) in `PYTORCH_JIT_LOG_LEVEL`. `GRAPH_DEBUG` can be enabled by prefixing a file name with an `>` as in `>alias_analysis`. `>>` and `>>>` are also valid and **currently** are equivalent to `GRAPH_DEBUG` as there is no logging level that is -higher than `GRAPH_DEBUG`. +higher than `GRAPH_DEBUG`. ## DifferentiableGraphOp ## diff --git a/torch/csrc/jit/fuser/executor.cpp b/torch/csrc/jit/fuser/executor.cpp index 69ddad91f016a..d485a6c09a0e2 100644 --- a/torch/csrc/jit/fuser/executor.cpp +++ b/torch/csrc/jit/fuser/executor.cpp @@ -338,6 +338,10 @@ bool runFusion(const int64_t key, Stack& stack, std::string* code_out) { inputs.emplace_back(all_inputs[i].toTensor()); } + if (!inputs.at(0).defined()) { + return false; + } + // Determines device to dispatch to. at::Device device = inputs.at(0).device(); // If there's a device mismatch in the inputs or if one of the input is a diff --git a/torch/csrc/jit/graph_executor.cpp b/torch/csrc/jit/graph_executor.cpp index 2129aa0f897cd..a047a65ebe89e 100644 --- a/torch/csrc/jit/graph_executor.cpp +++ b/torch/csrc/jit/graph_executor.cpp @@ -49,6 +49,8 @@ #include #include +static std::unique_ptr ge; + namespace torch { namespace jit { @@ -206,21 +208,30 @@ static void unpackReturnTuple(Stack &stack) { stack.insert(stack.end(), tuple->elements().begin(), tuple->elements().end()); } +struct DifferentiableGraphOp; + struct DifferentiableGraphBackward : public autograd::Node { DifferentiableGraphBackward( - GraphExecutor executor, + const std::shared_ptr& unspec_graph, size_t input_size, - size_t capture_size) - : executor(std::move(executor)), - captures_(capture_size), - input_instructions_(input_size) {} + size_t capture_size + //c10::optional& grad_executor + ) + : captures_(capture_size), + input_instructions_(input_size), + unspecialized_graph_(unspec_graph->copy()) + // grad_executor_(grad_executor) + {} variable_list apply(variable_list&& inputs) override { Stack stack; - stack.reserve(captures_.size() + inputs.size()); + size_t num_args = captures_.size() + inputs.size(); + stack.reserve(num_args); input_instructions_.unpack(std::move(inputs), stack); captures_.unpack(stack, shared_from_this()); + GraphExecutor& executor = getExecutor(stack); + GRAPH_DEBUG("Running DifferentiableGraphBackward for ", &executor); executor.run(stack); unpackReturnTuple(stack); @@ -258,6 +269,71 @@ struct DifferentiableGraphBackward : public autograd::Node { captures_.capture(val, is_output); } + + static c10::TensorTypePtr getTensorType(bool defined) { + auto tensor_type = TensorType::get(); + + if (defined) { + return tensor_type; + } + + return tensor_type->withUndefined(); + } + + GraphExecutor& getExecutor(Stack& stack) { + + // tensor lists are hashed as a single boolean value + // since all tensors will be either defined or undefined + + std::vector hash; + + for (IValue& v : stack) { + if (v.isTensorList()) { + auto list = v.toTensorListRef(); + hash.push_back(list.size() > 0 ? list[0].defined() : true); + } else if (v.isTensor()) { + hash.push_back(v.toTensor().defined()); + } else { + // assume that every other type is defined + hash.push_back(true); + } + } + + + TORCH_INTERNAL_ASSERT(unspecialized_graph_->inputs().size() == hash.size()); + if (grad_executors_.count(hash) == 0) { + + GRAPH_DEBUG("creating a specialized copy for ", this); + std::shared_ptr spec_copy = unspecialized_graph_->copy(); + for (auto i = 0; i < hash.size(); i++) { + auto input_type = spec_copy->inputs().at(i); + bool defined = hash[i]; + if (input_type->type()->kind() == TensorType::Kind) { + input_type->setType(getTensorType(defined)); + } else if ( + input_type->type()->kind() == ListType::Kind && + input_type->type()->expect()->getElementType()->kind() == + TensorType::Kind) { + input_type->setType(ListType::create(getTensorType(defined))); + } + } + + grad_executors_[hash] = GraphExecutor(spec_copy); + + //grad_executors_.insert({hash, GraphExecutor(spec_copy)}); + } + + + // set last optimized graph + // make a copy because DifferentiableBackward might disappear + // by the time we get to use diff_op_.grad_executor + + //grad_executor_ = GraphExecutor(grad_executors_[hash].graph()) ; + ge.reset(new GraphExecutor(grad_executors_[hash].graph())); + return grad_executors_[hash]; + + } + void addOutputForTensor(const at::Tensor& tensor) { auto v = Variable(tensor); add_next_edge(v.defined() ? v.gradient_edge() : autograd::Edge{}); @@ -318,6 +394,9 @@ struct DifferentiableGraphBackward : public autograd::Node { GraphExecutor executor; CaptureList captures_; UnpackInstructions input_instructions_; + std::unordered_map, GraphExecutor> grad_executors_; + std::shared_ptr unspecialized_graph_; + //c10::optional& grad_executor_; }; // an optimized way of executing the subgraph computed directly on @@ -336,10 +415,11 @@ struct DifferentiableGraphOp { // XXX: keep in mind that stack can be larger than the inputs we need! int operator()(Stack& stack) const { auto grad_fn = std::make_shared( - grad_executor, + this->grad.df, grad.df_input_vjps.size(), grad.df_input_captured_inputs.size() + - grad.df_input_captured_outputs.size()); + grad.df_input_captured_outputs.size()//, + /*grad_executor*/); { auto inputs = last(stack, num_inputs); @@ -377,6 +457,7 @@ struct DifferentiableGraphOp { private: friend GraphExecutor* detail::getGradExecutor(Operation& op); + friend struct DifferentiableGraphBackward; at::Tensor detach(at::Tensor t) const { if (!t.defined()) { @@ -425,6 +506,7 @@ struct DifferentiableGraphOp { Code f; Gradient grad; + //mutable c10::optional grad_executor; GraphExecutor grad_executor; const size_t num_inputs; @@ -456,9 +538,16 @@ RegisterOperators reg_graph_executor_ops({Operator( namespace detail { + + GraphExecutor* getGradExecutor(Operation& op) { if (auto diff_op = op.target()) { + + //TORCH_INTERNAL_ASSERT(diff_op->grad_executor.has_value()) + //return &(*diff_op->grad_executor); + //@#$ to do try this next return &diff_op->grad_executor; + return ge.get(); } return nullptr; } @@ -552,6 +641,7 @@ struct GraphExecutorImpl : public GraphExecutorImplBase { // Phase 0. Inline functions, then clean up any artifacts that the inliner // left in that may inhibit optimization Inline(*opt_graph); + LowerGradOf(*opt_graph); specializeAutogradZero(*opt_graph); LowerSimpleTuples(opt_graph); ConstantPooling(opt_graph); @@ -621,10 +711,10 @@ struct GraphExecutorImpl : public GraphExecutorImplBase { GraphExecutor::GraphExecutor(std::shared_ptr graph) : pImpl( - getProfilingMode() ? dynamic_cast( - new ProfilingGraphExecutorImpl(graph)) - : dynamic_cast( - new GraphExecutorImpl(graph))) {} + getExecutorMode() ? dynamic_cast( + new ProfilingGraphExecutorImpl(graph)) + : dynamic_cast( + new GraphExecutorImpl(graph))) {} void GraphExecutor::run(Stack& inputs) { return pImpl->run(inputs); @@ -643,7 +733,6 @@ GraphExecutorState GraphExecutor::getDebugState() { } void runRequiredPasses(const std::shared_ptr& g) { - LowerGradOf(*g); // implicit inserted expand nodes are not necessarily always valid // when used inside script methods that might have unstable shapes // we remove the implicitly created ones, and have shape analysis @@ -681,14 +770,33 @@ static bool mayIntroduceGradient(const Block* b) { } bool needsGradient(const std::shared_ptr& graph) { - if (!autograd::GradMode::is_enabled()) + if (!autograd::GradMode::is_enabled()) { return false; - if (mayIntroduceGradient(graph->block())) + } + + if (mayIntroduceGradient(graph->block())) { return true; - for (const Value* input : graph->inputs()) { - if (input->type()->requires_grad()) - return true; } + + if (getProfilingMode()) { + for (const Value* input : graph->inputs()) { + for (const auto& use : input->uses()) { + if (use.user->kind() == prim::BailOut) { + auto ptt = use.user->output()->type()->expect(); + if (ptt->requiresGrad() && *ptt->requiresGrad()) { + return true; + } + } + } + } + } else { + for (const Value* input : graph->inputs()) { + if (input->type()->requires_grad()) { + return true; + } + } + } + return false; } diff --git a/torch/csrc/jit/graph_executor.h b/torch/csrc/jit/graph_executor.h index 14cc8376b3f16..8c8dd3697253e 100644 --- a/torch/csrc/jit/graph_executor.h +++ b/torch/csrc/jit/graph_executor.h @@ -61,6 +61,7 @@ TORCH_API void debugSetAutodiffSubgraphInlining(bool state); TORCH_API std::shared_ptr lastExecutedOptimizedGraph(); TORCH_API std::atomic &getProfilingMode(); +TORCH_API std::atomic& getExecutorMode(); struct TORCH_API GraphOptimizerEnabledGuard { GraphOptimizerEnabledGuard(bool state) diff --git a/torch/csrc/jit/init.cpp b/torch/csrc/jit/init.cpp index 8523822923a62..1192b7a616d56 100644 --- a/torch/csrc/jit/init.cpp +++ b/torch/csrc/jit/init.cpp @@ -325,7 +325,18 @@ void initJITBindings(PyObject* module) { }) .def( "_jit_set_profiling_mode", - [](bool profiling_flag) { getProfilingMode() = profiling_flag; }) + [](bool profiling_flag) { + bool oldState = getProfilingMode(); + getProfilingMode() = profiling_flag; + return oldState; + }) + .def( + "_jit_set_profiling_executor", + [](bool profiling_flag) { + bool oldState = getExecutorMode(); + getExecutorMode() = profiling_flag; + return oldState; + }) .def( "_jit_set_inline_everything_mode", [](bool enabled) { script::getInlineEverythingMode() = enabled; }) diff --git a/torch/csrc/jit/ir.cpp b/torch/csrc/jit/ir.cpp index 2241e31e6837e..dee4d1e94e929 100644 --- a/torch/csrc/jit/ir.cpp +++ b/torch/csrc/jit/ir.cpp @@ -244,22 +244,19 @@ std::ostream &Node::print(std::ostream &out, size_t level, auto* pyOp = static_cast(this); out << "^" << pyOp->name(); pyOp->writeScalars(out); - } else if (print_attributes) { - if (hasAttribute(attr::Subgraph) && groups) { + } else if (hasAttribute(attr::Subgraph) && groups) { out << kind().toQualString() << "_" << groups->size(); - if (numAttributes() > 1 && kind() != prim::DifferentiableGraph) { + if (print_attributes && numAttributes() > 1 && kind() != prim::DifferentiableGraph) { printAttributes(out, /*ignore_subgraph=*/true); } groups->push_back(this); } else { out << kind().toQualString(); - if (hasAttributes()) { + if (print_attributes && hasAttributes()) { printAttributes(out); } } - } - out << "(" << inputs() << ")"; if (print_scopes) { @@ -940,6 +937,8 @@ bool Node::hasSideEffects() const { case prim::CallMethod: case prim::BailoutTemplate: case prim::profile: + case prim::BailOut: + case prim::Guard: return true; } diff --git a/torch/csrc/jit/operator.cpp b/torch/csrc/jit/operator.cpp index 02b6a19ae77fd..cc8aaf2dbf8ed 100644 --- a/torch/csrc/jit/operator.cpp +++ b/torch/csrc/jit/operator.cpp @@ -220,8 +220,9 @@ bool Operator::matches(const Node* node) const { const auto& formals = schema().arguments(); // not enough inputs - if (actuals.size() < formals.size()) + if (actuals.size() < formals.size()) { return false; + } TypeEnv type_env; for (size_t i = 0; i < formals.size(); ++i) { @@ -231,6 +232,7 @@ bool Operator::matches(const Node* node) const { if (!matched_type.success()) { return false; } + TypePtr resolved = tryEvalTypeVariables(formal, type_env); if (resolved) { formal = resolved; @@ -239,6 +241,7 @@ bool Operator::matches(const Node* node) const { // not resolved all type variables, e.g. if None was matched to Optional[T] // we will not succeed at matching T. However None <: Optional[T] so this // check can still succeed. + if (!actuals[i]->type()->isSubtypeOf(formal)) { return false; } diff --git a/torch/csrc/jit/passes/alias_analysis.cpp b/torch/csrc/jit/passes/alias_analysis.cpp index 75c3ef2e8bd10..2290a57796107 100644 --- a/torch/csrc/jit/passes/alias_analysis.cpp +++ b/torch/csrc/jit/passes/alias_analysis.cpp @@ -1264,7 +1264,6 @@ bool aliasAnalysisHasSpecialCaseFor(Symbol symbol) { prim::Drop, at::onnx::Reshape, at::onnx::Shape, - prim::AutogradAnyNonZero, prim::AutogradAdd, }; diff --git a/torch/csrc/jit/passes/bailout_graph.cpp b/torch/csrc/jit/passes/bailout_graph.cpp index 95eb6d5e66fa6..e02b4a7765134 100644 --- a/torch/csrc/jit/passes/bailout_graph.cpp +++ b/torch/csrc/jit/passes/bailout_graph.cpp @@ -1,7 +1,8 @@ -#include #include #include +#include #include +#include #include #include #include @@ -23,7 +24,7 @@ static std::unordered_set collectLoopCounts(Node *n) { it = outerNode->owningBlock(); } - return loopCounts; + return std::move(loopCounts); } struct BailOutGraphBuilderForNode { @@ -142,10 +143,14 @@ struct BailOutGraphBuilderForNode { } std::shared_ptr buildBailOutGraphFrom(Node* n) { - + // add graph inputs for guard's input + // and loop counts for loops `n` is contained in + // to make sure we can line bailout grap's inputs up properly + // with arguments to this BailOut node. for (auto bi : n->inputs()) { getOrAddInputForValue(bi); } + buildBailOutBlockFrom(n); // add graph outputs for (auto ov : graph_->outputs()) { @@ -239,11 +244,9 @@ struct BailOutInserter { // currently, there's always one guaded input bailout_node->addInput(it->input()); - // we need to collect loop counts to - // record the number of iterations already run - // however, liveness doesn't capture loop counts - // if they aren't used explicitly in a loop - // so we collect them manually here + // collect loop counts since liveness won't collect them + // if they aren't used explicitly, but they are used + // by BailOut graphs if we trigger a bailout inside a loop auto loopCounts = collectLoopCounts(*it); for (auto lc : loopCounts) { bailout_node->addInput(lc); @@ -296,8 +299,8 @@ void InsertBailOuts(std::shared_ptr graph) { // index matches the given `index` static Node* locateBailOutNodeInUnoptimizedGraph(Block* b, int64_t index) { for (auto n : b->nodes()) { - if (n->kind() == prim::BailOut && n->hasAttribute(attr::index) && - n->i(attr::index) == index) { + if ((n->kind() == prim::BailOut || n->kind() == prim::Guard) && + n->hasAttribute(attr::index) && n->i(attr::index) == index) { return n; } for (auto ib : n->blocks()) { @@ -313,7 +316,7 @@ static Node* locateBailOutNodeInUnoptimizedGraph(Block* b, int64_t index) { // to its users static void removeBailouts(Block* b) { for (auto it = b->nodes().begin(); it != b->nodes().end(); it++) { - if (it->kind() == prim::BailOut) { + if (it->kind() == prim::BailOut || it->kind() == prim::Guard) { // clear profiling information it->inputs().at(0)->setType(TensorType::get()); it->output()->replaceAllUsesWith(it->inputs().at(0)); @@ -333,16 +336,20 @@ TORCH_API std::shared_ptr BuildBailOutGraphFrom( const std::shared_ptr& target) { auto orig_bailout_node = locateBailOutNodeInUnoptimizedGraph(orig->block(), bailout_index); + + GRAPH_DEBUG("bailout triggered for ", *orig_bailout_node); + GRAPH_DUMP("original bailout graph ", orig); TORCH_INTERNAL_ASSERT( orig_bailout_node->inputs().at(0)->type()->cast() == nullptr); TORCH_INTERNAL_ASSERT( - orig_bailout_node && orig_bailout_node->kind() == prim::BailOut && + orig_bailout_node && + (orig_bailout_node->kind() == prim::BailOut || + orig_bailout_node->kind() == prim::Guard) && bailout_index == orig_bailout_node->i(attr::index)); BailOutGraphBuilderForNode bg(orig, target); auto bailout_graph = bg.buildBailOutGraphFrom(orig_bailout_node); removeBailouts(bailout_graph->block()); - ConstantPooling(bailout_graph); return bailout_graph; } diff --git a/torch/csrc/jit/passes/clear_undefinedness.cpp b/torch/csrc/jit/passes/clear_undefinedness.cpp new file mode 100644 index 0000000000000..972f7e62aa516 --- /dev/null +++ b/torch/csrc/jit/passes/clear_undefinedness.cpp @@ -0,0 +1,38 @@ +#include +#include + +namespace torch { +namespace jit { + +void clearUndefinedness(Value* o) { + if (o->type()->kind() == TensorType::Kind) { + o->setType(TensorType::get()); + } else if ( + o->type()->kind() == ListType::Kind && + o->type()->expect()->getElementType()->kind() == + TensorType::Kind) { + o->setType(ListType::create(TensorType::get())); + } +} + +void clearUndefinedness(Block* block) { + for (auto n : block->nodes()) { + for (auto o : n->outputs()) { + clearUndefinedness(o); + } + for (auto ib : n->blocks()) { + clearUndefinedness(ib); + } + } +} + +void ClearUndefinedness(const std::shared_ptr& graph) { + for (auto i : graph->inputs()) { + clearUndefinedness(i); + } + clearUndefinedness(graph->block()); + GRAPH_DUMP("After removeUndefinedness: ", graph); +} + +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/passes/clear_undefinedness.h b/torch/csrc/jit/passes/clear_undefinedness.h new file mode 100644 index 0000000000000..17f78bfda3cba --- /dev/null +++ b/torch/csrc/jit/passes/clear_undefinedness.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace torch { +namespace jit { + +// Undefinedness makes argument matching fail for regular tensor operations +// if 1+ arguments are undefined or possibly undefined tensors. +// Technically, undefined tensors are **not** tensors as the regular tensor +// operations do not know how to handle them. +// However, in practice, there are guards and conversion operators that +// **always** gate regular operations if undefined tensors may be present +// Eventually, we would love to move to the world where we use optionals +// in lieu of undefined tensors. +// When this happens, this pass will be removed +TORCH_API void ClearUndefinedness(const std::shared_ptr& graph); + +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/passes/graph_fuser.cpp b/torch/csrc/jit/passes/graph_fuser.cpp index b07464e4bc8f9..3c1010bf58c0c 100644 --- a/torch/csrc/jit/passes/graph_fuser.cpp +++ b/torch/csrc/jit/passes/graph_fuser.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -199,7 +200,10 @@ struct GraphFuser { fusableDevice &= isFusableDevice(output); } } - return fusableDevice && isFusableMap(node); + + bool is_fusable_map = isFusableMap(node); + GRAPH_DEBUG("isFusableDefault for ", getHeader(node), " fusableDevice = ", fusableDevice, " is_fusable_map = ", is_fusable_map); + return fusableDevice && is_fusable_map; } bool isFusableMap(Node* node) { @@ -211,25 +215,33 @@ struct GraphFuser { } bool isFusableCatNode(Node* node) { - if (node->kind() != aten::cat) + if (node->kind() != aten::cat) { return false; - if (!node->is_constant(attr::dim)) + } + if (!node->is_constant(attr::dim)) { + GRAPH_DEBUG("attr::dim of ", getHeader(node), " isn't a prim::Constant"); return false; + } auto tensors_node = node->namedInput(attr::tensors)->node(); if ((tensors_node->inputs().size() + node->outputs().size()) > subgraph_arg_limit_) { + + GRAPH_DEBUG("too many arguments for ", getHeader(node)); return false; } if (tensors_node->kind() != prim::ListConstruct) + GRAPH_DEBUG("inputs tensors don't come from prim::ListConstruct for ", getHeader(node)); return false; // NB: Note that technically other uses of the list aren't a big problem for // us. It would be enough to place the prim::FusedConcat before the // prim::ListConstruct, and allUsersAreThisConsumerOrOccurAfterIt would // still be satisfied. However, I don't expect this to be necessary any time // soon, and so we're simply assuming that we don't have to deal with it. - if (tensors_node->output()->uses().size() > 1) + if (tensors_node->output()->uses().size() > 1) { + GRAPH_DEBUG("there's more than one user of the input tensors of ", getHeader(node)); return false; + } return true; } @@ -254,6 +266,8 @@ struct GraphFuser { } void mergeFusionGroups(Node* consumer_group, Node* producer_group) { + + GRAPH_UPDATE("Merging a producer group ", getHeader(producer_group), " into ",getHeader(consumer_group)); // Now we have two fusion groups! // Revert the fusion - place all inner nodes of producer back in the outer // graph. @@ -387,6 +401,8 @@ struct GraphFuser { subgraph.eraseInput(p); } } + + GRAPH_UPDATE("Merging ", getHeader(n), " into ",getHeader(group)); return subgraph.insertNode(in_graph); } @@ -402,6 +418,7 @@ struct GraphFuser { auto sel = group->addOutput(); sel->copyMetadata(n->output()); n->replaceAllUsesWith(group); + GRAPH_UPDATE("Replacing ", getHeader(n), " with ",getHeader(group)); n->destroy(); return group; } @@ -1207,6 +1224,7 @@ void PeepholeOptimizeShapeExpressions(Block* block) { void FuseGraph(std::shared_ptr& graph) { GraphFuser(graph->block(), graph).run(); + GRAPH_DUMP("After GraphFuser: ", graph); // After FuseGraph some common subexpressions may come back EliminateCommonSubexpression(graph); // We might have emitted a fair amount of useless shape propagating code, so diff --git a/torch/csrc/jit/passes/guard_elimination.cpp b/torch/csrc/jit/passes/guard_elimination.cpp index 0fefdf5a0a50e..4f160acea99fb 100644 --- a/torch/csrc/jit/passes/guard_elimination.cpp +++ b/torch/csrc/jit/passes/guard_elimination.cpp @@ -16,7 +16,10 @@ struct GuardElimination { aliasDb_(c10::guts::make_unique(graph_)) {} void run() { - moveGuardsToDefs(graph_->block()); + const size_t MAX_ATTEMPTS = 5; + size_t attempts = MAX_ATTEMPTS; + while (attempts-- && moveGuardsToDefs(graph_->block())) { + } GRAPH_DUMP("After moveGuardsToDefs", graph_); coalesceGuards(graph_->block()); GRAPH_DUMP("After coalesceGuards", graph_); @@ -24,7 +27,16 @@ struct GuardElimination { GRAPH_DUMP("After eliminateRedundantGuards", graph_); } - void moveGuardsToDefs(Block* b) { + static bool isLoweredGradOf(Node* n) { + if (n->kind() != prim::If) { + return false; + } + + return n->input(0)->node()->kind() == prim::AutogradAnyNonZero; + } + + bool moveGuardsToDefs(Block* b) { + bool changed = false; for (auto it = b->nodes().begin(); it != b->nodes().end();) { auto n = *it; if (n->kind() == prim::Guard) { @@ -39,6 +51,7 @@ struct GuardElimination { guardee = *n->owningBlock()->nodes().begin(); } bool moved = aliasDb_->moveAfterTopologicallyValid(n, guardee); + changed |= moved; if (moved) { GRAPH_UPDATE( "Moved ", @@ -53,6 +66,21 @@ struct GuardElimination { } } } + + if (b->owningNode() && + isLoweredGradOf( + b->owningNode()) /*b->owningNode()->kind() == prim::If*/) { + for (auto it = b->nodes().begin(); it != b->nodes().end();) { + auto block_node = *it++; + if (block_node->kind() != prim::Guard) { + break; + } + block_node->moveBefore(b->owningNode()); + changed = true; + } + } + + return changed; } void coalesceGuards(Block* b) { @@ -183,7 +211,6 @@ struct GuardElimination { // Guards can be removed if all inputs are guarded and `isSummarized()` // returns // false or inputs are `prim::Constant` - // bool removableGuard(Node *n) { const static auto no_exceptions = std::unordered_set{}; @@ -208,6 +235,17 @@ struct GuardElimination { case aten::neg: case prim::ConstantChunk: case aten::size: + case aten::abs: + case aten::sign: + case aten::pow: + case aten::relu: + case aten::threshold: + case aten::avg_pool2d: + case prim::AutogradAdd: + case prim::AutogradZero: + case aten::rand_like: + case aten::erf: + case aten::erfc: return checkInputs(n, no_exceptions); case aten::cat: // check that the dimension argument is constant @@ -240,7 +278,27 @@ struct GuardElimination { } } return false; + + // this is checked by one of the tests in test_jit_fuser.py + case prim::ListUnpack: { + // check if the input is a constant chunk + // used for LSTM fusions + auto chunk = n->input(0)->node(); + if (chunk->kind() != aten::chunk) { + return false; + } + return checkInputs(chunk, no_exceptions); + } + // this is checked by one of the tests in test_jit_fuser.py + case aten::broadcast_tensors: { + auto list_construct = n->input(0)->node(); + if (list_construct->kind() != prim::ListConstruct) { + return false; + } + return checkInputs(list_construct, no_exceptions); + } case prim::Guard: + case prim::GradOf: return true; default: GRAPH_DEBUG("cannot remove ", n->kind().toQualString()); diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp index 3fdd3d8519e82..bd6f7543b937b 100644 --- a/torch/csrc/jit/passes/quantization.cpp +++ b/torch/csrc/jit/passes/quantization.cpp @@ -145,6 +145,11 @@ class InsertObserversHelper { // Values that are the output of GetAttr[name="bias"] and they // will be propagated through the function call hierarchy std::unordered_set bias_values_; + // Unique id generator for observer module, used for generating + // unique observer names when we insert observer module, we + // record the current unique id used to avoid incrementing from 0 + // every time to find a unique id. + int uid_ = 0; }; // Clone observer module and add it to the original module, @@ -165,12 +170,15 @@ Node* InsertObserversHelper::insertObserverFor( } else { observer_module = std::get<0>(qconfig); } - std::string observer_name = "observer_for_" + v->debugName(); script::Module observer = observer_module.clone(); + std::string observer_name = "_observer_" + std::to_string(uid_++); + while (module.find_module(observer_name)) { + observer_name = "_observer_" + std::to_string(uid_++); + } module.register_module(observer_name, observer); // Get handle of observer module Node* observer_instance = g->create(c10::prim::GetAttr); - // self.observer_for_v + // self._observer_v observer_instance->addInput(g->inputs()[0]); observer_instance->s_(c10::attr::name, observer_name); observer_instance->output()->setDebugName(observer_name); @@ -321,36 +329,35 @@ void InsertObserversHelper::insertObservers( bias_values_.emplace(v); } } - if (v->node()->kind() == prim::CallMethod) { - // If we find a call to a method of a child module, - // we'll recursively insert observers for the forward function to - // the child module. - auto module_instance = v->node()->inputs()[0]; - auto module_method_name = v->node()->s(attr::name); - // TODO: looks like this block is not related to v? maybe we should - // move this outside - script::Module callee_module; - if (module_instance->node()->kind() == prim::GetAttr) { - auto child_module_name = module_instance->node()->s(attr::name); - auto child_module = module.find_module(child_module_name); - TORCH_INTERNAL_ASSERT( - child_module, - "Child module " + child_module_name + " does not exist"); - callee_module = child_module.value(); - } else { - TORCH_INTERNAL_ASSERT( - module_instance == graph->inputs()[0], - "We only support call method either on %self" - "or child instance in insert_observers_pass right now"); - callee_module = module; - } - auto method_graph = - callee_module.get_method(module_method_name).graph(); - propagateValues(v->node(), method_graph); - // Recursively insert observer for the forward function of child - // module - insertObservers(callee_module, module_method_name); + } + + if (n->kind() == prim::CallMethod) { + // If we find a call to a method of a child module, + // we'll recursively insert observers for the forward function to + // the child module. + auto module_instance = n->inputs()[0]; + auto module_method_name = n->s(attr::name); + script::Module callee_module; + if (module_instance->node()->kind() == prim::GetAttr) { + auto child_module_name = module_instance->node()->s(attr::name); + auto child_module = module.find_module(child_module_name); + TORCH_INTERNAL_ASSERT( + child_module, + "Child module " + child_module_name + " does not exist"); + callee_module = child_module.value(); + } else { + TORCH_INTERNAL_ASSERT( + module_instance == graph->inputs()[0], + "We only support call method either on %self" + "or child instance in insert_observers_pass right now"); + callee_module = module; } + auto method_graph = + callee_module.get_method(module_method_name).graph(); + propagateValues(n, method_graph); + // Recursively insert observer for the forward function of child + // module + insertObservers(callee_module, module_method_name); } for (Block* subblock : n->blocks()) { @@ -426,7 +433,7 @@ c10::optional findObserverName(Value* v) { u.user->s(attr::name) == "forward") { auto module_instance = u.user->inputs().at(0); if (module_instance->node()->kind() == prim::GetAttr && - module_instance->node()->s(attr::name).find("observer_for_") != + module_instance->node()->s(attr::name).find("_observer_") != std::string::npos) { return module_instance->node()->s(attr::name); } @@ -544,7 +551,7 @@ c10::optional QuantizeHelper::findChildModuleToQuantize( child_instance->node()->kind() == prim::GetAttr, "Child instance should come from GetAttr."); auto child_module_name = child_instance->node()->s(attr::name); - if (child_module_name.find("observer_for_") == std::string::npos) { + if (child_module_name.find("_observer_") == std::string::npos) { auto child_module = module_.find_module(child_module_name); TORCH_INTERNAL_ASSERT( child_module, @@ -618,6 +625,42 @@ void InsertQuantDeQuantImpl( qh.destroyNodes(); } +void insertPrepackUnpackForLinear(std::shared_ptr& graph) { + std::string linear_with_quant = R"( +graph(%linear, %a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype): + %w_quant = aten::quantize_per_tensor(%w, %w_scale, %w_zero_point, %w_dtype) + %w_dequant = aten::dequantize(%w_quant) + %r = prim::CallFunction(%linear, %a_dequant, %w_dequant, %b) + return (%r) )"; + + std::string linear_with_quant_prepack = R"( +graph(%linear, %a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype): + %w_quant = aten::quantize_per_tensor(%w, %w_scale, %w_zero_point, %w_dtype) + %packed_params = quantized::linear_prepack(%w_quant, %b) + %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::linear_unpack(%packed_params) + %w_dequant = aten::dequantize(%w_quant_unpacked) + %r = prim::CallFunction(%linear, %a_dequant, %w_dequant, %b) + return (%r) )"; + + // Filter to match linear CallFunction + auto filter = [](const Match& match, + const std::unordered_map& vmap) { + const auto& match_vmap = match.values_map; + auto linear_node = match_vmap.at(vmap.at("linear"))->node(); + auto func = + linear_node->output()->type()->expect()->function(); + auto func_name = getFuncName(func->qualname()); + if (func_name == "linear") { + return true; + } + return false; + }; + + SubgraphRewriter rewriter; + rewriter.RegisterRewritePattern(linear_with_quant, linear_with_quant_prepack); + rewriter.runOnGraph(graph, filter); +} + void insertPrepackUnpackForConv2d(std::shared_ptr& graph) { std::string conv_with_quant = R"( graph(%a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype, %stride, %padding, %dilation, %groups): @@ -632,7 +675,7 @@ graph(%a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype, %stride, %padding, %packed_params = quantized::conv_prepack(%w_quant, %b, %stride, %padding, %dilation, %groups) %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::conv_unpack(%packed_params) %w_dequant = aten::dequantize(%w_quant_unpacked) - %r = aten::conv2d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups) + %r = aten::conv2d(%a_dequant, %w_dequant, %b_unpacked, %stride, %padding, %dilation, %groups) return (%r) )"; SubgraphRewriter rewriter; @@ -910,40 +953,7 @@ graph(%self, %scale, %zero_point, %dtype): } void InsertPrepackUnpack(std::shared_ptr& graph) { - std::string linear_with_quant = R"( -graph(%linear, %a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype): - %w_quant = aten::quantize_per_tensor(%w, %w_scale, %w_zero_point, %w_dtype) - %w_dequant = aten::dequantize(%w_quant) - %r = prim::CallFunction(%linear, %a_dequant, %w_dequant, %b) - return (%r) )"; - - std::string linear_with_quant_prepack = R"( -graph(%linear, %a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype): - %w_quant = aten::quantize_per_tensor(%w, %w_scale, %w_zero_point, %w_dtype) - %packed_params = quantized::linear_prepack(%w_quant, %b) - %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::linear_unpack(%packed_params) - %w_dequant = aten::dequantize(%w_quant_unpacked) - %r = prim::CallFunction(%linear, %a_dequant, %w_dequant, %b) - return (%r) )"; - - // Filter to match linear CallFunction - auto filter = [](const Match& match, - const std::unordered_map& vmap) { - const auto& match_vmap = match.values_map; - auto linear_node = match_vmap.at(vmap.at("linear"))->node(); - auto func = - linear_node->output()->type()->expect()->function(); - auto func_name = getFuncName(func->qualname()); - if (func_name == "linear") { - return true; - } - return false; - }; - - SubgraphRewriter rewriter; - rewriter.RegisterRewritePattern(linear_with_quant, linear_with_quant_prepack); - rewriter.runOnGraph(graph, filter); - + insertPrepackUnpackForLinear(graph); insertPrepackUnpackForConv2d(graph); } @@ -1029,8 +1039,11 @@ graph(%a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype, %stride, %padding, } auto w_quant_val = match_vmap.at(vmap.at("w_quant")); // unique name for the module based on %w_quant - auto module_name = - module_name_prefix + std::to_string(w_quant_val->unique()); + int uid = 0; + auto module_name = module_name_prefix + std::to_string(uid++); + while (module.find_module(module_name)) { + module_name_prefix + std::to_string(uid++); + } module.register_module(module_name, wrapper_module); // Add GetAttr of the packed module diff --git a/torch/csrc/jit/passes/specialize_autogradzero.cpp b/torch/csrc/jit/passes/specialize_autogradzero.cpp index f90fde68dd94c..bc0bd59ad59a3 100644 --- a/torch/csrc/jit/passes/specialize_autogradzero.cpp +++ b/torch/csrc/jit/passes/specialize_autogradzero.cpp @@ -1,5 +1,4 @@ #include -#include #include namespace torch { @@ -18,10 +17,14 @@ void specializeAutogradZero(Graph &g) { for (Value* input : g.inputs()) { const auto& tp = input->type(); if (auto tt = tp->cast()) { - if (tt->undefined() && *tt->undefined()) { - state[input] = State::Zero; + if (tt->undefined()) { + if (*tt->undefined()) { + state[input] = State::Zero; + } else { + state[input] = State::Nonzero; + } } else { - state[input] = State::Nonzero; + state[input] = State::Unknown; } } else if ( tp->isSubtypeOf(TensorType::get()) || @@ -34,75 +37,23 @@ void specializeAutogradZero(Graph &g) { for (auto it = g.nodes().begin(); it != g.nodes().end(); ++it) { auto n = *it; - switch (n->kind()) { - case prim::GradOf: { - auto all_zeros = - std::all_of(n->inputs().begin(), n->inputs().end(), [&](Value* v) { - return state[v] == State::Zero; - }); - // Property 1: if all the gradInputs to the GradOf are Zero - // then the gradOutputs are also zero and will be represented as - // AutogradZero nodes - if (all_zeros) { - auto zero = g.createAutogradZero()->insertAfter(n)->output(); - for (auto o : n->outputs()) { - GRAPH_UPDATE("Replacing output %", o->debugName(), - " with AutogradZero %", zero->debugName()); - o->replaceAllUsesWith(zero); - } - } else { - // Property 2: GradOfs are required to correctly handle combinations - // of Nonzero and zero inputs. They are expected to produce - // Nonzero output tensors in this case. - // Remove the GradOf, splicing its body back into the surrounding - // block - auto body = n->blocks().at(0); - for (auto input : n->inputs()) { - // we should never get into a situation when specializing a GradOf - // where we do not know if a value is Nonzero since at the top level - // a gradient graph is composed of Linear nodes and AutogradAdds - // and LinearNodes only appear in these graphs - AT_ASSERT(state[input] != State::Unknown); - } - // hoist the nodes in the GradOf body to be before the linear block - GRAPH_UPDATE("Hoisting out ", getHeader(*it)); - for (auto it = body->nodes().begin(); it != body->nodes().end();) { - auto block_node = *it++; - block_node->moveBefore(n); - } - - for (size_t i = 0; i < n->outputs().size(); ++i) { - GRAPH_UPDATE("Replacing prim::GradOf's use %", - n->outputs().at(i)->debugName(), - " with hoisted value %", - body->outputs().at(i)->debugName()); - n->outputs().at(i)->replaceAllUsesWith(body->outputs().at(i)); - } - } - GRAPH_UPDATE("Destroying ", getHeader(*it)); - it.destroyCurrent(); - } break; + switch (n->kind()) { case prim::AutogradAdd: { auto a = n->input(0); auto b = n->input(1); // if one is Autograd zero, we can just drop the add if (state[a] == State::Zero) { // Zero + b == b - GRAPH_UPDATE("Simplifying ", getHeader(n), " where %", a->debugName(), - " is AutogradZero to %", b->debugName()); n->output()->replaceAllUsesWith(b); it.destroyCurrent(); } else if (state[b] == State::Zero) { // a + Zero == a - GRAPH_UPDATE("Simplifying ", getHeader(n), " where %", b->debugName(), - " is AutogradZero to %", a->debugName()); n->output()->replaceAllUsesWith(a); it.destroyCurrent(); } else if (state[a] == State::Nonzero && state[b] == State::Nonzero) { // when both are Nonzero, we can use a normal, optimizable add // instruction - WithInsertPoint guard(n); auto* g = n->owningGraph(); auto* cOne = g->insertConstant(1); @@ -113,8 +64,6 @@ void specializeAutogradZero(Graph &g) { auto* add_output = add_node->output(); state[add_output] = State::Nonzero; n->output()->replaceAllUsesWith(add_output); - GRAPH_UPDATE("Simplifying ", getHeader(n), " to ", - getHeader(add_node)); it.destroyCurrent(); } else { // otherwise we have conditionally-Nonzero things, and we need @@ -152,6 +101,54 @@ void specializeAutogradZero(Graph &g) { : State::Unknown; } } break; + // Lowered GradOf block + case prim::If: { + auto if_input = n->input(0)->node(); + if (if_input->kind() == prim::AutogradAnyNonZero) { + auto all_zeros = std::all_of( + if_input->inputs().begin(), + if_input->inputs().end(), + [&](Value* v) { return state[v] == State::Zero; }); + + auto all_nonzeros = std::all_of( + if_input->inputs().begin(), + if_input->inputs().end(), + [&](Value* v) { return state[v] == State::Nonzero; }); + // Property 1: if all the gradInputs to the GradOf are Zero + // then the gradOutputs are also zero and will be represented as + // AutogradZero nodes + if (all_zeros) { + auto zero = g.createAutogradZero()->insertAfter(n)->output(); + state[zero] = State::Zero; + for (auto o : n->outputs()) { + o->replaceAllUsesWith(zero); + } + it.destroyCurrent(); + break; + } + + if (all_nonzeros) { + auto body = n->blocks().at(0); + // hoist the nodes in the GradOf body to be before the linear block + for (auto it = body->nodes().begin(); it != body->nodes().end();) { + auto block_node = *it++; + block_node->moveBefore(n); + } + + for (size_t i = 0; i < n->outputs().size(); ++i) { + n->outputs().at(i)->replaceAllUsesWith(body->outputs().at(i)); + state[body->outputs().at(i)] = State::Nonzero; + } + it.destroyCurrent(); + break; + } + } + + for (auto o : n->outputs()) { + state[o] = State::Unknown; + } + break; + } default: for (auto o : n->outputs()) { state[o] = State::Unknown; diff --git a/torch/csrc/jit/profiling_graph_executor_impl.cpp b/torch/csrc/jit/profiling_graph_executor_impl.cpp index 4b3c9779b610c..be7291cbb4865 100644 --- a/torch/csrc/jit/profiling_graph_executor_impl.cpp +++ b/torch/csrc/jit/profiling_graph_executor_impl.cpp @@ -1,9 +1,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -18,34 +20,37 @@ namespace torch { namespace jit { static std::atomic profiling_mode{false}; +static std::atomic executor_mode{false}; + std::atomic& getProfilingMode() { return profiling_mode; } +std::atomic& getExecutorMode() { + return executor_mode; +} + +static bool needsGradientInProfilingMode(Block* b) { + for (auto n : b->nodes()) { + if (n->kind() == prim::BailOut) { + auto ptt = n->output()->type()->expect(); + if (ptt->requiresGrad() && *ptt->requiresGrad()) { + return true; + } + } + + for (auto ib : n->blocks()) { + if (needsGradientInProfilingMode(ib)) { + return true; + } + } + } + return false; +} std::shared_ptr ProfilingGraphExecutorImpl::prepareGraph( const std::shared_ptr& graph, Stack& stack) { auto g = graph->copy(); - ArgumentSpec spec = - arg_spec_creator_.create(autograd::GradMode::is_enabled(), stack); - arg_spec_creator_.specializeTypes(*g, spec); - runRequiredPasses(g); - PropagateRequiresGrad(g); - ConstantPropagation(g); - if (needsGradient(g)) { - auto diff_nodes = CreateAutodiffSubgraphs( - g, getAutodiffSubgraphInlining() ? autodiffSubgraphNodeThreshold : 1); - for (Node* dnode : diff_nodes) { - auto diff_graph = std::move(dnode->g(attr::Subgraph)); - Gradient gradient = differentiate(diff_graph); - // do not optimize DifferentiableGraphs, since - // ideally they will be profiled and then optimized separetely - // when their corresponding DifferentiableGraphOp is called - packGradient(gradient, dnode); - } - InlineAutodiffSubgraphs( - g, getAutodiffSubgraphInlining() ? autodiffSubgraphInlineThreshold : 1); - } return g; } @@ -54,57 +59,66 @@ ProfilingGraphExecutorImpl::ProfilingGraphExecutorImpl( : GraphExecutorImplBase(graph), arg_spec_creator_(*this->graph) {} ExecutionPlan ProfilingGraphExecutorImpl::getPlanFor(Stack& stack) { + GRAPH_DEBUG("Running ProfilingGraphExecutorImpl ", this); if (optimized_plan_) { return *optimized_plan_; } - if (!pr_) { - pr_ = ProfilingRecord::instrumentGraph(prepareGraph(graph, stack)); - auto copy = pr_->graph()->copy(); - LowerGradOf(*copy); - RemoveExpands(copy); - CanonicalizeOps(copy); - EliminateDeadCode(copy); - profiling_plan_ = ExecutionPlan(copy); - // fall-through - } + std::shared_ptr copy; + if (getProfilingMode()) { + if (!pr_) { + pr_ = ProfilingRecord::instrumentGraph(prepareGraph(graph, stack)); + auto copy = pr_->graph()->copy(); + LowerGradOf(*copy); + specializeAutogradZero(*copy); + runRequiredPasses(copy); + GRAPH_DUMP("Profiled Graph: ", copy); + profiling_plan_ = ExecutionPlan(copy); + // fall-through + } - if (!pr_->ready()) { - return *profiling_plan_; + if (!pr_->ready()) { + return *profiling_plan_; + } + copy = pr_->graph()->copy(); + + } else { + copy = graph->copy(); } - // copy already has differentiableGraphs - auto copy = pr_->graph()->copy(); if (!getGraphExecutorOptimize()) { runRequiredPasses(copy); optimized_plan_ = ExecutionPlan(copy); return *optimized_plan_; } - // insert bailouts InsertGuards(copy); - // get rid of autograd specific ops - // we can probably make guard_elimination.cpp - // to handle these ops - specializeAutogradZero(*copy); - // hoist out GradOf blocks - // otherwise we will need to teach - // liveness and buildBailOut graphs - // about them LowerGradOf(*copy); - // constant fold into ConstantChunk - CanonicalizeOps(copy); - EliminateRedundantGuards(copy); - InsertBailOuts(copy); - // TODO: this runs specializeAutogradZero ?? - GRAPH_DUMP("After InsertBailOuts: ", copy); + if (getProfilingMode()) { + EliminateRedundantGuards(copy); + InsertBailOuts(copy); + GRAPH_DUMP("After InsertBailOuts: ", copy); + } + + specializeAutogradZero(*copy); + if (!getProfilingMode()) { + ClearUndefinedness(copy); + } + runRequiredPasses(copy); ConstantPropagation(copy); runOptimization(copy); - if (needsGradient(copy)) { + + // TODO: insert grad propagation + bool needs_gradient = getProfilingMode() + ? needsGradientInProfilingMode(copy->block()) + : needsGradient(copy); + if (needs_gradient) { + GRAPH_DEBUG(this, " needs gradients"); auto diff_nodes = CreateAutodiffSubgraphs( copy, getAutodiffSubgraphInlining() ? autodiffSubgraphNodeThreshold : 1); + GRAPH_DEBUG(" diff_nodes' size is ", diff_nodes.size()); for (Node *dnode : diff_nodes) { auto diff_graph = std::move(dnode->g(attr::Subgraph)); Gradient gradient = differentiate(diff_graph); @@ -120,6 +134,7 @@ ExecutionPlan ProfilingGraphExecutorImpl::getPlanFor(Stack& stack) { runNondiffOptimization(copy); } EliminateDeadCode(copy); + GRAPH_DUMP("Optimized Graph : ", copy); // cache optimized_plan_ = ExecutionPlan(copy); return *optimized_plan_; @@ -127,7 +142,11 @@ ExecutionPlan ProfilingGraphExecutorImpl::getPlanFor(Stack& stack) { GraphExecutorState ProfilingGraphExecutorImpl::getDebugState() { - AT_ERROR("not supported"); + GraphExecutorState state; + TORCH_INTERNAL_ASSERT(optimized_plan_); + auto opt_plan = *optimized_plan_; + state.execution_plans.emplace(ArgumentSpec{0, 0}, opt_plan); + return state; } } // namespace jit diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp index d35c558a4c090..032641c05546b 100644 --- a/torch/csrc/jit/register_prim_ops.cpp +++ b/torch/csrc/jit/register_prim_ops.cpp @@ -931,15 +931,26 @@ RegisterOperators reg( }, aliasAnalysisSpecialCase()), Operator( - prim::AutogradAnyNonZero, + "prim::AutogradAnyNonZero(...) -> int", [](const Node* node) -> Operation { size_t num_inputs = node->inputs().size(); - return [=](Stack& stack) { + return [num_inputs](Stack& stack) { bool result = false; - for (const IValue& t : last(stack, num_inputs)) { - if (t.toTensor().defined()) { - result = true; - break; + for (const IValue& v : last(stack, num_inputs)) { + if (v.isTensor()) { + if (v.toTensor().defined()) { + result = true; + break; + } + } else if (v.isTensorList()) { + for (const at::Tensor& t : v.toTensorListRef()) { + result = true; + } + if (result) { + break; + } + } else { + TORCH_INTERNAL_ASSERT(false); } } drop(stack, num_inputs); @@ -947,18 +958,25 @@ RegisterOperators reg( return 0; }; }, - aliasAnalysisSpecialCase()), + aliasAnalysisFromSchema()), Operator( prim::AutogradAdd, [](Stack& stack) { at::Tensor a, b; pop(stack, a, b); - if (!a.defined()) + if (!a.defined() && !b.defined()) { + // undef + undef == undef + stack.emplace_back(a); + } + else if (!a.defined()) { stack.emplace_back(b); - else if (!b.defined()) + } + else if (!b.defined()) { stack.emplace_back(a); - else + } + else { stack.emplace_back(a + b); + } return 0; }, aliasAnalysisSpecialCase()), diff --git a/torch/custom_class.h b/torch/custom_class.h index c28f6b14d94ea..3b9c73c9da7ef 100644 --- a/torch/custom_class.h +++ b/torch/custom_class.h @@ -60,6 +60,9 @@ detail::types init() { return detail::types{}; } template class class_ { + static_assert(std::is_base_of::value, + "torch::jit::class_ requires T to inherit from CustomClassHolder"); + std::string className; std::string qualClassName; c10::optional> pyClass = c10::nullopt; @@ -70,7 +73,7 @@ class class_ { const std::string topModule = "__torch__.torch"; public: - class_(string className_) : className(std::move(className_)) { + class_(std::string className_) : className(std::move(className_)) { // Currently we register everything as a python class just for convenience. // We'll want to remove this at some point to get rid of the python // dependency. It would require significant changes to class registration, @@ -90,7 +93,7 @@ class class_ { PyObject* rawPyObj = py_object.release().ptr(); return rawPyObj; }; - getClassConverter()[qualClassName] = castToPython; + at::getClassConverter()[qualClassName] = castToPython; // We currently represent custom classes as torchscript classes with a // capsule attribute @@ -100,9 +103,9 @@ class class_ { classTypePtr->addAttribute("capsule", CapsuleType::get()); c10::getCustomClassTypeMap().insert({typeid(c10::intrusive_ptr).name(), - StrongTypePtr(classCu, classTypePtr)}); + c10::StrongTypePtr(classCu, classTypePtr)}); c10::getCustomClassTypeMap().insert({typeid(c10::tagged_capsule).name(), - StrongTypePtr(classCu, classTypePtr)}); + c10::StrongTypePtr(classCu, classTypePtr)}); classCu->register_type(classTypePtr); } @@ -124,7 +127,7 @@ class class_ { return *this; } template - class_& def(string name, Func f) { + class_& def(std::string name, Func f) { auto res = def_(name, f, detail::args_t{}); return *this; } @@ -140,19 +143,19 @@ class class_ { std::vector addInputs_( Func f, std::shared_ptr graph, - guts::index_sequence) { + at::guts::index_sequence) { using argTypes = - typename guts::infer_function_traits_t::parameter_types; + typename at::guts::infer_function_traits_t::parameter_types; std::vector res = { - addInput>::call( + addInput>::call( graph)...}; return res; } template std::vector addInputs(Func f, std::shared_ptr graph) { constexpr auto numArgs = - guts::infer_function_traits_t::number_of_parameters; - return addInputs_(f, graph, guts::make_index_sequence()); + at::guts::infer_function_traits_t::number_of_parameters; + return addInputs_(f, graph, at::guts::make_index_sequence()); } template @@ -192,11 +195,11 @@ class class_ { classTypePtr->addMethod(method); } template - class_& def_(string name, Func f, detail::types funcInfo) { + class_& def_(std::string name, Func f, detail::types funcInfo) { pyClass->def(name.c_str(), f); auto func = [f](c10::intrusive_ptr cur, Types... args) { - return guts::invoke(f, *cur, args...); + return at::guts::invoke(f, *cur, args...); }; defineMethod(name, std::move(func), funcInfo.hasRet); return *this; diff --git a/torch/distributed/rpc/__init__.py b/torch/distributed/rpc/__init__.py index b1e0e17049503..39b76c423a6a2 100644 --- a/torch/distributed/rpc/__init__.py +++ b/torch/distributed/rpc/__init__.py @@ -2,7 +2,7 @@ import sys -from .backend_registry import * # noqa: F401 +from . import backend_registry if sys.version_info >= (3, 0): @@ -13,7 +13,7 @@ def init_model_parallel( self_name, - backend=RpcBackend.PROCESS_GROUP, + backend=backend_registry.BackendType.PROCESS_GROUP, init_method=None, self_rank=-1, worker_name_to_id=None, @@ -43,14 +43,22 @@ def init_model_parallel( init_method(str): backend specific init arguments. num_send_recv_threads(int): Number of threads for send/recv work. """ + # Rendezvous. + world_size = len(worker_name_to_id) + rendezvous_iterator = torch.distributed.rendezvous( + init_method, rank=self_rank, world_size=world_size + ) + store, _, _ = next(rendezvous_iterator) + # Initialize RPC. _init_rpc( backend, - init_method, + store, self_name, self_rank, worker_name_to_id, num_send_recv_threads, ) + # Initialize Autograd. torch.distributed.autograd._init(api._agent.get_worker_info().id) diff --git a/torch/distributed/rpc/api.py b/torch/distributed/rpc/api.py index 4d9534873f6e3..84540654de159 100644 --- a/torch/distributed/rpc/api.py +++ b/torch/distributed/rpc/api.py @@ -2,15 +2,13 @@ from torch.distributed import invoke_remote_builtin, invoke_remote_python_udf from torch.distributed import _start_rpc_agent from torch.distributed import _destroy_rref_context, _cleanup_python_rpc_handler -from torch.distributed import ProcessGroupAgent from torch.distributed import WorkerInfo -from .backend_registry import is_backend_registered, init_backend +from . import backend_registry from .internal import _internal_rpc_pickler, PythonUDF import functools import sys import torch -from enum import Enum _agent = None @@ -58,14 +56,12 @@ def sync_rpc(): _agent.sync() -class RpcBackend(Enum): - PROCESS_GROUP = 1 # TODO: add a context manager to wrap _init_rpc and join_rpc def _init_rpc( - backend=RpcBackend.PROCESS_GROUP, - init_method=None, + backend=backend_registry.BackendType.PROCESS_GROUP, + store=None, self_name=None, self_rank=-1, worker_name_to_id=None, @@ -79,33 +75,15 @@ def _init_rpc( if _agent: raise RuntimeError("RPC is already initialized") - if backend == RpcBackend.PROCESS_GROUP: - from torch.distributed.distributed_c10d import _get_default_group - - group = _get_default_group() - if (self_rank != -1) and (self_rank != group.rank()): - raise RuntimeError("self_rank argument {} doesn't match pg rank {}".format( - self_rank, group.rank())) - if (worker_name_to_id is not None) and (len(worker_name_to_id) != group.size()): - raise RuntimeError("worker_name_to_id argument {} doesn't match pg size {}".format( - worker_name_to_id, group.size())) - # TODO: add try-except and destroy _agent in all processes if any fails. - _agent = ProcessGroupAgent(self_name, group, num_send_recv_threads) - elif is_backend_registered(backend): - # Rendezvous. - world_size = len(worker_name_to_id) - rendezvous_iterator = torch.distributed.rendezvous(init_method, self_rank, world_size) - store, self_rank, world_size = next(rendezvous_iterator) - # Initialize RPC. - _agent = init_backend( - backend, - store=store, - self_name=self_name, - self_rank=self_rank, - worker_name_to_id=worker_name_to_id, - ) - else: - raise RuntimeError("Unrecognized RPC backend ", backend) + # Initialize RPC. + _agent = backend_registry.init_backend( + backend, + store=store, + self_name=self_name, + self_rank=self_rank, + worker_name_to_id=worker_name_to_id, + num_send_recv_threads=num_send_recv_threads, + ) _start_rpc_agent(_agent) diff --git a/torch/distributed/rpc/backend_registry.py b/torch/distributed/rpc/backend_registry.py index e20a9e5d3b681..4fabdf20a934d 100644 --- a/torch/distributed/rpc/backend_registry.py +++ b/torch/distributed/rpc/backend_registry.py @@ -1,15 +1,16 @@ from __future__ import absolute_import, division, print_function, unicode_literals +import collections +import enum -_BACKEND_REGISTRY = {} +import torch.distributed as dist +import torch.distributed.distributed_c10d as dc10d -def _get_backend_registry(): - return _BACKEND_REGISTRY +BackendValue = collections.namedtuple("BackendValue", ["init_backend_handler"]) - -def is_backend_registered(backend_name): - return backend_name in _get_backend_registry() +# Create an enum type, `BackendType`, with empty members. +BackendType = enum.Enum(value="BackendType", names={}) def register_backend(backend_name, init_backend_handler): @@ -21,14 +22,65 @@ def register_backend(backend_name, init_backend_handler): `_init_rpc()` function is called with a backend. This returns the agent. """ - backend_registry = _get_backend_registry() - if backend_name in backend_registry: + global BackendType + if backend_name in BackendType.__members__.keys(): raise RuntimeError("RPC backend {}: already registered".format(backend_name)) - backend_registry[backend_name] = init_backend_handler + # Create a new enum type, `BackendType`, with extended members. + existing_enum_dict = {member.name: member.value for member in BackendType} + extended_enum_dict = dict( + {backend_name: BackendValue(init_backend_handler=init_backend_handler)}, + **existing_enum_dict + ) + BackendType = enum.Enum(value="BackendType", names=extended_enum_dict) + return BackendType[backend_name] + + +def init_backend(backend, *args, **kwargs): + return backend.value.init_backend_handler(*args, **kwargs) + + +def process_group_init_backend_handler( + store, + self_name, + self_rank, + worker_name_to_id, + num_send_recv_threads, + *args, + **kwargs +): + # Initialize ProcessGroup. + if dist.is_initialized(): + raise RuntimeError( + "Default process group must not be initialized before `init_model_parallel`." + ) + + world_size = len(worker_name_to_id) + dist.init_process_group( + backend="gloo", store=store, rank=self_rank, world_size=world_size + ) + + try: + group = dc10d._get_default_group() + assert group is not None, "Failed to initialize default ProcessGroup." + + if (self_rank != -1) and (self_rank != group.rank()): + raise RuntimeError( + "self_rank argument {} doesn't match pg rank {}".format( + self_rank, group.rank() + ) + ) + if (worker_name_to_id is not None) and (len(worker_name_to_id) != group.size()): + raise RuntimeError( + "worker_name_to_id argument {} doesn't match pg size {}".format( + worker_name_to_id, group.size() + ) + ) + # TODO: add try-except and destroy _agent in all processes if any fails. + return dist.ProcessGroupAgent(self_name, group, num_send_recv_threads) + except Exception as ex: + dist.destroy_process_group() + raise ex + -def init_backend(backend_name, *args, **kwargs): - backend_registry = _get_backend_registry() - if backend_name not in backend_registry: - raise RuntimeError("No rpc_init handler for {}.".format(backend_name)) - return backend_registry[backend_name](*args, **kwargs) +register_backend("PROCESS_GROUP", process_group_init_backend_handler) diff --git a/torch/jit/quantized.py b/torch/jit/quantized.py index 0cffc40522d8a..4fe3a738d0388 100644 --- a/torch/jit/quantized.py +++ b/torch/jit/quantized.py @@ -7,6 +7,8 @@ from torch.nn.utils.rnn import PackedSequence +import warnings + class QuantizedLinear(torch.jit.ScriptModule): __constants__ = ['scale', 'zero_point'] @@ -602,6 +604,9 @@ def quantize_rnn_cell_modules(module): def quantize_linear_modules(module, dtype=torch.int8): + warnings.warn("quantize_linear_modules function has been deprecated. " + "Please use torch.quantization.quantize_dynamic API instead.") + reassign = {} for name, mod in module.named_modules(): if mod is module: diff --git a/torch/nn/quantized/dynamic/modules/linear.py b/torch/nn/quantized/dynamic/modules/linear.py index 04b942f6dacd2..7574dd53eb761 100644 --- a/torch/nn/quantized/dynamic/modules/linear.py +++ b/torch/nn/quantized/dynamic/modules/linear.py @@ -17,8 +17,6 @@ class Linear(nnq.Linear): shape :math:`(\text{out\_features}, \text{in\_features})`. bias (Tensor): the non-learnable bias of the module of shape :math:`(\text{out\_features})`. If :attr:`bias` is ``True``, the values are initialized to zero. - scale: `scale` parameter of weight Quantized Tensor, type: double - zero_point: `zero_point` parameter for weight Quantized Tensor, type: long Examples:: @@ -45,6 +43,11 @@ def forward(self, x): def _get_name(self): return 'DynamicQuantizedLinear' + def extra_repr(self): + return 'in_features={}, out_features={}'.format( + self.in_features, self.out_features + ) + @classmethod def from_float(cls, mod): r"""Create a dynamic quantized module from a float module or qparams_dict diff --git a/torch/nn/quantized/modules/activation.py b/torch/nn/quantized/modules/activation.py index d906b0a5c47a8..422e7897bbd91 100644 --- a/torch/nn/quantized/modules/activation.py +++ b/torch/nn/quantized/modules/activation.py @@ -31,10 +31,10 @@ class ReLU(torch.nn.ReLU): """ def __init__(self, inplace=False): super(ReLU, self).__init__(inplace) - assert not inplace, 'torch.nn.quantized.ReLU does not support inplace' + self.inplace = inplace def forward(self, input): - return torch.nn.quantized.functional.relu(input) + return torch.nn.quantized.functional.relu(input, inplace=self.inplace) def _get_name(self): return 'QuantizedReLU' diff --git a/torch/tensor.py b/torch/tensor.py index 618efeb3b875c..4f88d1ba054cf 100644 --- a/torch/tensor.py +++ b/torch/tensor.py @@ -10,6 +10,25 @@ from torch._six import imap from torch._C import _add_docstr from numbers import Number +import functools + + +def _wrap_type_error_to_not_implemented(f): + from torch import _six + import inspect + + # functools.wraps doesn't work well with methods in python 2 + method_assignments = ('__name__', '__doc__') + assigned = (method_assignments if _six.PY2 and inspect.ismethoddescriptor(f) + else functools.WRAPPER_ASSIGNMENTS) + + @functools.wraps(f, assigned=assigned) + def wrapped(*args, **kwargs): + try: + return f(*args, **kwargs) + except TypeError: + return NotImplemented + return wrapped # NB: If you subclass Tensor, and want to share the subclassed class @@ -372,17 +391,20 @@ def __format__(self, format_spec): return object.__format__(self, format_spec) def __ipow__(self, other): - raise NotImplementedError("in-place pow not implemented") + return NotImplemented + @_wrap_type_error_to_not_implemented def __rpow__(self, other): return self.new_tensor(other) ** self + @_wrap_type_error_to_not_implemented def __floordiv__(self, other): result = self / other if result.dtype.is_floating_point: result = result.trunc() return result + @_wrap_type_error_to_not_implemented def __rfloordiv__(self, other): result = other / self if result.dtype.is_floating_point: @@ -391,12 +413,12 @@ def __rfloordiv__(self, other): __neg__ = _C._TensorBase.neg - __eq__ = _C._TensorBase.eq - __ne__ = _C._TensorBase.ne - __lt__ = _C._TensorBase.lt - __le__ = _C._TensorBase.le - __gt__ = _C._TensorBase.gt - __ge__ = _C._TensorBase.ge + __eq__ = _wrap_type_error_to_not_implemented(_C._TensorBase.eq) + __ne__ = _wrap_type_error_to_not_implemented(_C._TensorBase.ne) + __lt__ = _wrap_type_error_to_not_implemented(_C._TensorBase.lt) + __le__ = _wrap_type_error_to_not_implemented(_C._TensorBase.le) + __gt__ = _wrap_type_error_to_not_implemented(_C._TensorBase.gt) + __ge__ = _wrap_type_error_to_not_implemented(_C._TensorBase.ge) __abs__ = _C._TensorBase.abs def __len__(self):