Update on "[ONNX] Fix bug in unfold symbolic (#50504)"

Fix bug in unfold symbolic

[ghstack-poisoned]

.circleci/docker/common/install_rocm.sh

@@ -2,6 +2,22 @@
 
 set -ex
 
+install_magma() {
+    # "install" hipMAGMA into /opt/rocm/magma by copying after build
+    git clone https://bitbucket.org/icl/magma.git -b hipMAGMA
+    pushd magma
+    cp make.inc-examples/make.inc.hip-mkl-gcc make.inc
+    echo 'LIBDIR += -L$(MKLROOT)/lib' >> make.inc
+    echo 'LIB += -Wl,--enable-new-dtags -Wl,--rpath,/opt/rocm/lib -Wl,--rpath,$(MKLROOT)/lib -Wl,--rpath,/opt/rocm/magma/lib' >> make.inc
+    echo 'DEVCCFLAGS += --amdgpu-target=gfx803 --amdgpu-target=gfx900 --amdgpu-target=gfx906 --amdgpu-target=gfx908' >> make.inc
+    export PATH="${PATH}:/opt/rocm/bin"
+    make -f make.gen.hipMAGMA -j $(nproc)
+    make lib/libmagma.so -j $(nproc) MKLROOT=/opt/conda
+    make testing/testing_dgemm -j $(nproc) MKLROOT=/opt/conda
+    popd
+    mv magma /opt/rocm
+}
+
 install_ubuntu() {
     apt-get update
     if [[ $UBUNTU_VERSION == 18.04 ]]; then
@@ -10,7 +26,6 @@ install_ubuntu() {
     fi
     apt-get install -y kmod
     apt-get install -y wget
-    apt-get install -y libopenblas-dev
 
     # Need the libc++1 and libc++abi1 libraries to allow torch._C to load at runtime
     apt-get install -y libc++1
@@ -24,13 +39,7 @@ install_ubuntu() {
     DEBIAN_FRONTEND=noninteractive apt-get install -y --allow-unauthenticated \
                    rocm-dev \
                    rocm-utils \
-                   rocfft \
-                   miopen-hip \
-                   rocblas \
-                   hipsparse \
-                   rocrand \
-                   hipcub \
-                   rocthrust \
+                   rocm-libs \
                    rccl \
                    rocprofiler-dev \
                    roctracer-dev
@@ -44,9 +53,11 @@ install_ubuntu() {
       DEBIAN_FRONTEND=noninteractive apt-get install -y --allow-unauthenticated ${MIOPENKERNELS}
     fi
 
-  # Cleanup
-  apt-get autoclean && apt-get clean
-  rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+    install_magma
+
+    # Cleanup
+    apt-get autoclean && apt-get clean
+    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 }
 
 install_centos() {
@@ -70,17 +81,13 @@ install_centos() {
   yum install -y \
                    rocm-dev \
                    rocm-utils \
-                   rocfft \
-                   miopen-hip \
-                   rocblas \
-                   hipsparse \
-                   rocrand \
+                   rocm-libs \
                    rccl \
-                   hipcub \
-                   rocthrust \
                    rocprofiler-dev \
                    roctracer-dev
 
+  install_magma
+
   # Cleanup
   yum clean all
   rm -rf /var/cache/yum

.circleci/docker/ubuntu-rocm/Dockerfile

@@ -58,6 +58,7 @@ ENV PATH /opt/rocm/hcc/bin:$PATH
 ENV PATH /opt/rocm/hip/bin:$PATH
 ENV PATH /opt/rocm/opencl/bin:$PATH
 ENV PATH /opt/rocm/llvm/bin:$PATH
+ENV MAGMA_HOME /opt/rocm/magma
 ENV LANG C.UTF-8
 ENV LC_ALL C.UTF-8
 

.jenkins/pytorch/build.sh

@@ -53,6 +53,11 @@ if [[ "$BUILD_ENVIRONMENT" == *coverage* ]]; then
   export USE_CPP_CODE_COVERAGE=ON
 fi
 
+if [[ "$BUILD_ENVIRONMENT" == *cuda11* ]]; then
+  # enable split torch_cuda build option in CMake
+  export BUILD_SPLIT_CUDA=ON
+fi
+
 # TODO: Don't run this...
 pip_install -r requirements.txt || true
 

.jenkins/pytorch/test.sh

@@ -22,6 +22,10 @@ if [[ "$BUILD_ENVIRONMENT" == *coverage* ]]; then
   export PYTORCH_COLLECT_COVERAGE=1
 fi
 
+if [[ "$BUILD_ENVIRONMENT" == *cuda11.1* ]]; then
+  export BUILD_SPLIT_CUDA=ON
+fi
+
 if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
   # Print GPU info
   rocminfo | grep -E 'Name:.*\sgfx|Marketing'

.jenkins/pytorch/win-test-helpers/build_pytorch.bat

@@ -110,6 +110,10 @@ if "%REBUILD%" == "" (
     aws s3 cp "s3://ossci-windows/Restore PyTorch Environment.lnk" "C:\Users\circleci\Desktop\Restore PyTorch Environment.lnk"
   )
 )
+:: tests if BUILD_ENVIRONMENT contains cuda11 as a substring
+if not x%BUILD_ENVIRONMENT:cuda11=%==x%BUILD_ENVIRONMENT% (
+   set BUILD_SPLIT_CUDA=ON
+)
 
 python setup.py install --cmake && sccache --show-stats && (
   if "%BUILD_ENVIRONMENT%"=="" (
@@ -118,4 +122,3 @@ python setup.py install --cmake && sccache --show-stats && (
     7z a %TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\torch %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\caffe2 && copy /Y "%TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z" "%PYTORCH_FINAL_PACKAGE_DIR%\"
   )
 )
-

.jenkins/pytorch/win-test.sh

@@ -40,6 +40,10 @@ if [ -n "$CIRCLE_PULL_REQUEST" ]; then
   file_diff_from_base "$DETERMINE_FROM"
 fi
 
+if [[ "${CIRCLE_JOB}" == *11.1* ]]; then
+  export BUILD_SPLIT_CUDA=ON
+fi
+
 run_tests() {
     if [ -z "${JOB_BASE_NAME}" ] || [[ "${JOB_BASE_NAME}" == *-test ]]; then
         $SCRIPT_HELPERS_DIR/test_python_nn.bat "$DETERMINE_FROM"

CMakeLists.txt

@@ -161,6 +161,13 @@ option(COLORIZE_OUTPUT "Colorize output during compilation" ON)
 option(USE_ASAN "Use Address Sanitizer" OFF)
 option(USE_TSAN "Use Thread Sanitizer" OFF)
 option(USE_CUDA "Use CUDA" ON)
+# BUILD_SPLIT_CUDA must also be exported as an environment variable before building, with
+# `export BUILD_SPLIT_CUDA=1` because cpp_extension.py can only work properly if this variable
+# also exists in the environment.
+# This option is incompatible with CUDA_SEPARABLE_COMPILATION.
+cmake_dependent_option(
+    BUILD_SPLIT_CUDA "Split torch_cuda library into torch_cuda_cu and torch_cuda_cpp" OFF
+    "USE_CUDA AND NOT CUDA_SEPARABLE_COMPILATION" OFF)
 option(USE_FAST_NVCC "Use parallel NVCC build" OFF)
 option(USE_ROCM "Use ROCm" ON)
 option(CAFFE2_STATIC_LINK_CUDA "Statically link CUDA libraries" OFF)

aten/src/ATen/CMakeLists.txt

@@ -339,9 +339,20 @@ if(USE_CUDA AND NOT USE_ROCM)
   if(CAFFE2_USE_CUDNN)
     list(APPEND ATen_CUDA_DEPENDENCY_LIBS ${CUDNN_LIBRARIES})
   endif()
+  if($ENV{ATEN_STATIC_CUDA})
+    list(APPEND ATen_CUDA_DEPENDENCY_LIBS "${CUDA_TOOLKIT_ROOT_DIR}/lib64/libculibos.a")
+    list(APPEND ATen_CUDA_DEPENDENCY_LIBS "${CUDA_TOOLKIT_ROOT_DIR}/lib64/libcudart_static.a")
+  endif($ENV{ATEN_STATIC_CUDA})
+endif()
+
 
   if(USE_MAGMA)
-    list(APPEND ATen_CUDA_DEPENDENCY_LIBS ${MAGMA_LIBRARIES})
+    if(USE_CUDA)
+      list(APPEND ATen_CUDA_DEPENDENCY_LIBS ${MAGMA_LIBRARIES})
+    endif(USE_CUDA)
+    if(USE_ROCM)
+      list(APPEND ATen_HIP_DEPENDENCY_LIBS ${MAGMA_LIBRARIES})
+    endif(USE_ROCM)
     if(MSVC)
       if($ENV{TH_BINARY_BUILD})
         # Do not do this on Linux: see Note [Extra MKL symbols for MAGMA in torch_cpu]
@@ -351,11 +362,6 @@ if(USE_CUDA AND NOT USE_ROCM)
       endif($ENV{TH_BINARY_BUILD})
     endif(MSVC)
   endif(USE_MAGMA)
-  if($ENV{ATEN_STATIC_CUDA})
-    list(APPEND ATen_CUDA_DEPENDENCY_LIBS "${CUDA_TOOLKIT_ROOT_DIR}/lib64/libculibos.a")
-    list(APPEND ATen_CUDA_DEPENDENCY_LIBS "${CUDA_TOOLKIT_ROOT_DIR}/lib64/libcudart_static.a")
-  endif($ENV{ATEN_STATIC_CUDA})
-endif()
 
 # NB: We're relying on cmake/Dependencies.cmake to appropriately setup HIP dependencies.
 # In principle we could duplicate them, but handling the rocblas

aten/src/ATen/core/Dict.h

@@ -21,6 +21,7 @@ using valid_dict_key_types = guts::typelist::typelist<
   int64_t,
   std::string,
   double,
+  c10::complex<double>,
   bool,
   at::Tensor
 >;

aten/src/ATen/core/Dict_inl.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <ATen/core/ivalue.h>
+#include <c10/util/hash.h>
 
 namespace c10 {
 namespace detail {
@@ -43,6 +44,8 @@ inline size_t DictKeyHash::operator()(const IValue& ivalue) const {
     return std::hash<std::string>()(ivalue.toStringRef());
   } else if (ivalue.isDouble()) {
     return std::hash<double>()(ivalue.toDouble());
+  } else if (ivalue.isComplexDouble()) {
+    return c10::hash<c10::complex<double>>()(ivalue.toComplexDouble());
   } else if (ivalue.isBool()) {
     return std::hash<bool>()(ivalue.toBool());
   } else if (ivalue.isTensor()) {

aten/src/ATen/core/function_schema.h

@@ -85,10 +85,16 @@ struct Argument {
   }
 
   Argument cloneWithType(TypePtr new_type) const {
-    return Argument(name_, new_type, N_, default_value_, kwarg_only_, alias_info_);
+    return Argument(
+        name_,
+        std::move(new_type),
+        N_,
+        default_value_,
+        kwarg_only_,
+        alias_info_);
   }
 
-  // this function check whether this Argument is backward compatible with
+  // this function checks whether this Argument is backward compatible with
   // the old one. we consider the following cases are backward compatible:
   //   1) two arguments are equal
   //   2) this arg's type should be subtype of old
@@ -97,7 +103,7 @@ struct Argument {
       const Argument& old,
       std::ostream* why_not=nullptr) const;
 
-private:
+ private:
   std::string name_;
   TypePtr type_;
   // for list types, an optional statically known length for the list
@@ -223,7 +229,7 @@ struct FunctionSchema {
     }
   }
 
-public:
+ public:
 
   void dump() const;
 
@@ -265,13 +271,13 @@ struct FunctionSchema {
   }
   FunctionSchema cloneWithName(std::string name, std::string overload_name) const {
     return FunctionSchema(
-      std::move(name),
-      std::move(overload_name),
-      arguments(),
-      returns(),
-      is_vararg(),
-      is_varret()
-      );
+        std::move(name),
+        std::move(overload_name),
+        arguments(),
+        returns(),
+        is_vararg(),
+        is_varret()
+        );
   }
   FunctionSchema cloneWithArguments(std::vector<Argument> new_arguments) const {
     return FunctionSchema(
@@ -305,7 +311,8 @@ struct FunctionSchema {
   // values.
   void checkAndNormalizeInputs(
       std::vector<IValue>& inputs,
-      const std::unordered_map<std::string, IValue>& kwargs) const;
+      const std::unordered_map<std::string, IValue>& kwargs =
+          std::unordered_map<std::string, IValue>{}) const;
 
   std::string findErrorInKwargs(const std::vector<std::string>& kwargs) const;
 
@@ -354,11 +361,11 @@ struct FunctionSchema {
 
 inline bool operator==(const FunctionSchema& lhs, const FunctionSchema& rhs) {
   return lhs.name() == rhs.name()
-      && lhs.overload_name() == rhs.overload_name()
-      && lhs.arguments() == rhs.arguments()
-      && lhs.returns() == rhs.returns()
-      && lhs.is_vararg() == rhs.is_vararg()
-      && lhs.is_varret() == rhs.is_varret();
+     && lhs.overload_name() == rhs.overload_name()
+     && lhs.arguments() == rhs.arguments()
+     && lhs.returns() == rhs.returns()
+     && lhs.is_vararg() == rhs.is_vararg()
+     && lhs.is_varret() == rhs.is_varret();
 }
 
 inline bool operator!=(const FunctionSchema& lhs, const FunctionSchema& rhs) {
@@ -375,7 +382,7 @@ inline std::ostream& operator<<(std::ostream& out, const Argument& arg) {
   // so we always use Type(alias)? format
   auto type = arg.type();
   bool is_opt = type->kind() == OptionalType::Kind;
-  auto unopt_type = is_opt ? type->cast<OptionalType>()->getElementType() : type;
+  auto unopt_type = is_opt ? type->castRaw<OptionalType>()->getElementType() : type;
 
   if (unopt_type->kind() == ListType::Kind && arg.N()) {
     // sized lists get size N from arg, not type

aten/src/ATen/core/ivalue.cpp

@@ -463,7 +463,7 @@ std::ostream& printMaybeAnnotatedDict(
     std::ostream& out,
     const IValue& the_dict,
     IValueFormatter formatter) {
-  auto value_type = the_dict.type()->cast<DictType>()->getValueType();
+  auto value_type = the_dict.type()->castRaw<DictType>()->getValueType();
   if (the_dict.toGenericDict().size() == 0 ||
       !elementTypeCanBeInferredFromMembers(value_type)) {
     out << "annotate(" << the_dict.type()->annotation_str() << ",";

aten/src/ATen/core/jit_type.h

@@ -738,14 +738,15 @@ struct TORCH_API DictType : public Type {
       case TypeKind::IntType:
       case TypeKind::BoolType:
       case TypeKind::FloatType:
+      case TypeKind::ComplexDoubleType:
       case TypeKind::StringType:
       case TypeKind::TensorType:
         return DictTypePtr(new DictType(key, value));
       default:
         AT_ERROR(
             "Cannot create dict for key type '",
             key->str(),
-            "', only int, float, Tensor and string keys are supported");
+            "', only int, float, complex, Tensor and string keys are supported");
     }
   }
 
@@ -974,7 +975,7 @@ struct TORCH_API TupleType : public NamedType {
     }
 
     const auto& l_elements = elements();
-    const auto& r_elements = rhs.cast<TupleType>()->elements();
+    const auto& r_elements = rhs.castRaw<TupleType>()->elements();
     if (l_elements.size() != r_elements.size())
       return false;
     for (size_t i = 0; i < l_elements.size(); ++i) {

aten/src/ATen/core/jit_type_base.h

@@ -162,7 +162,7 @@ struct TORCH_API Type : std::enable_shared_from_this<Type> {
   template <typename T>
   const T* castRaw() const {
     if (T::Kind == kind()) {
-      return static_cast<T*>(this);
+      return static_cast<const T*>(this);
     }
     return nullptr;
   }

aten/src/ATen/core/op_registration/infer_schema.cpp

@@ -21,6 +21,10 @@ std::vector<Argument> createArgumentVector(c10::ArrayRef<ArgumentDef> args) {
 C10_EXPORT FunctionSchema make_function_schema(std::string&& name, std::string&& overload_name, c10::ArrayRef<ArgumentDef> arguments, c10::ArrayRef<ArgumentDef> returns) {
   return FunctionSchema(std::move(name), std::move(overload_name), createArgumentVector(arguments), createArgumentVector(returns));
 }
+
+C10_EXPORT FunctionSchema make_function_schema(c10::ArrayRef<ArgumentDef> arguments, c10::ArrayRef<ArgumentDef> returns) {
+  return make_function_schema("", "", arguments, returns);
+}
 }
 }
 

aten/src/ATen/core/op_registration/infer_schema.h

@@ -107,11 +107,12 @@ struct createSingleReturn {
 };
 
 C10_API FunctionSchema make_function_schema(std::string&& name, std::string&& overload_name, c10::ArrayRef<ArgumentDef> arguments, c10::ArrayRef<ArgumentDef> returns);
+C10_API FunctionSchema make_function_schema(c10::ArrayRef<ArgumentDef> arguments, c10::ArrayRef<ArgumentDef> returns);
 
 /// Creates a `FunctionSchema` object from a `FunctionTraits` type for a
 /// function. Flattens std::tuple returns into multiple return types
 template <typename FunctionTraits>
-FunctionSchema createFunctionSchemaFromTraitsFlattenedReturns(std::string&& name, std::string&& overload_name) {
+FunctionSchema createFunctionSchemaFromTraitsFlattenedReturns() {
  using ReturnType = typename FunctionTraits::return_type;
  using ParameterTypes = typename FunctionTraits::parameter_types;
 
@@ -121,7 +122,7 @@ FunctionSchema createFunctionSchemaFromTraitsFlattenedReturns(std::string&& name
  constexpr auto arguments = createArguments<ParameterTypes>::call();
  constexpr auto returns = createReturns<ReturnType>::call();
 
- return make_function_schema(std::move(name), std::move(overload_name), arguments, returns);
+ return make_function_schema(arguments, returns);
 }
 
 /// Creates a `FunctionSchema` object from a `FunctionTraits` type for a
@@ -144,8 +145,8 @@ FunctionSchema createFunctionSchemaFromTraitsSingleReturn(std::string&& name, st
 }
 
 template<class FuncType>
-FunctionSchema inferFunctionSchemaFlattenedReturns(std::string&& name, std::string&& overload_name) {
-  return detail::infer_schema::createFunctionSchemaFromTraitsFlattenedReturns<guts::infer_function_traits_t<FuncType>>(std::move(name), std::move(overload_name));
+FunctionSchema inferFunctionSchemaFlattenedReturns() {
+  return detail::infer_schema::createFunctionSchemaFromTraitsFlattenedReturns<guts::infer_function_traits_t<FuncType>>();
 }
 
 template<class FuncType>