Update on "Add tan_cuda for complex dtypes"

anjali411 · anjali411 · commit d912829babf2 · 2020-05-15T05:28:28.000-07:00
Differential Revision: [D21572209](https://our.internmc.facebook.com/intern/diff/D21572209) [ghstack-poisoned]
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -2883,6 +2883,14 @@ workflows:
           build_environment: "pytorch-linux-xenial-py3.6-gcc5.4-ge_config_profiling-test"
           docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3.6-gcc5.4:9a3986fa-7ce7-4a36-a001-3c9bef9892e2"
           resource_class: large
+      - pytorch_linux_test:
+          name: pytorch_linux_xenial_py3_6_gcc5_4_ge_config_simple_test
+          requires:
+            - setup
+            - pytorch_linux_xenial_py3_6_gcc5_4_build
+          build_environment: "pytorch-linux-xenial-py3.6-gcc5.4-ge_config_simple-test"
+          docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3.6-gcc5.4:9a3986fa-7ce7-4a36-a001-3c9bef9892e2"
+          resource_class: large
       - pytorch_linux_test:
           name: pytorch_linux_xenial_cuda10_2_cudnn7_py3_ge_config_legacy_test
           requires:
diff --git a/.circleci/verbatim-sources/workflows-pytorch-ge-config-tests.yml b/.circleci/verbatim-sources/workflows-pytorch-ge-config-tests.yml
@@ -14,6 +14,14 @@
           build_environment: "pytorch-linux-xenial-py3.6-gcc5.4-ge_config_profiling-test"
           docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3.6-gcc5.4:9a3986fa-7ce7-4a36-a001-3c9bef9892e2"
           resource_class: large
+      - pytorch_linux_test:
+          name: pytorch_linux_xenial_py3_6_gcc5_4_ge_config_simple_test
+          requires:
+            - setup
+            - pytorch_linux_xenial_py3_6_gcc5_4_build
+          build_environment: "pytorch-linux-xenial-py3.6-gcc5.4-ge_config_simple-test"
+          docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3.6-gcc5.4:9a3986fa-7ce7-4a36-a001-3c9bef9892e2"
+          resource_class: large
       - pytorch_linux_test:
           name: pytorch_linux_xenial_cuda10_2_cudnn7_py3_ge_config_legacy_test
           requires:
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -1,7 +1,7 @@
 load("@rules_proto//proto:defs.bzl", "proto_library")
 load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library", "cc_proto_library", "cc_test")
 load("//third_party:substitution.bzl", "template_rule")
-load("//:tools/build_variables.bzl", "torch_cpp_srcs", "libtorch_core_sources", "libtorch_distributed_sources", "libtorch_extra_sources")
+load("//:tools/build_variables.bzl", "torch_cpp_srcs", "libtorch_core_sources", "libtorch_distributed_sources", "libtorch_extra_sources", "jit_core_sources")
 load("//tools/rules:cu.bzl", "cu_library")
 load("//tools/config:defs.bzl", "if_cuda")
 load("//:aten.bzl", "intern_build_aten_ops")
@@ -1979,13 +1979,7 @@ cc_library(
             "torch/csrc/cuda/python_nccl.cpp",
             "torch/csrc/cuda/nccl.cpp",
         ],
-    )) + libtorch_core_sources + libtorch_distributed_sources + torch_cpp_srcs + libtorch_extra_sources + [
-        "torch/csrc/jit/frontend/error_report.cpp",
-        "torch/csrc/jit/frontend/lexer.cpp",
-        "torch/csrc/jit/frontend/function_schema_parser.cpp",
-        "torch/csrc/jit/frontend/strtod.cpp",
-        "torch/csrc/jit/frontend/source_range.cpp",
-        "torch/csrc/jit/frontend/schema_type_parser.cpp",
+    )) + libtorch_core_sources + libtorch_distributed_sources + torch_cpp_srcs + libtorch_extra_sources + jit_core_sources + [
         ":generated_code",
     ],
     copts = TORCH_COPTS + if_cuda(["-DUSE_CUDA=1"]),
diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt
@@ -34,9 +34,13 @@ if(USE_ROCM)
 endif()
 
 # NB: If you edit these globs, you'll have to update setup.py package_data as well
+file(GLOB_RECURSE ATen_CORE_HEADERS  "core/*.h")
+file(GLOB_RECURSE ATen_CORE_SRCS "core/*.cpp")
+file(GLOB_RECURSE ATen_CORE_TEST_SRCS "core/*_test.cpp")
+EXCLUDE(ATen_CORE_SRCS "${ATen_CORE_SRCS}" ${ATen_CORE_TEST_SRCS})
+
 file(GLOB base_h "*.h" "detail/*.h" "cpu/*.h" "cpu/vec256/*.h" "quantized/*.h")
 file(GLOB base_cpp "*.cpp" "detail/*.cpp" "cpu/*.cpp")
-add_subdirectory(core)
 file(GLOB cuda_h "cuda/*.h" "cuda/detail/*.h" "cuda/*.cuh" "cuda/detail/*.cuh")
 file(GLOB cuda_cpp "cuda/*.cpp" "cuda/detail/*.cpp")
 file(GLOB cuda_nvrtc_stub_h "cuda/nvrtc_stub/*.h")
@@ -89,6 +93,26 @@ file(GLOB native_quantized_hip_cpp "native/quantized/hip/*.cpp")
 # XNNPACK
 file(GLOB native_xnnpack "native/xnnpack/*.cpp")
 
+# Add files needed from jit folders
+list(APPEND ATen_CORE_HEADERS
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/source_range.h
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/function_schema_parser.h
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/lexer.h
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/strtod.h
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/parse_string_literal.h
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/schema_type_parser.h
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/error_report.h
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/tree.h
+)
+list(APPEND ATen_CORE_SRCS
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/error_report.cpp
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/function_schema_parser.cpp
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/lexer.cpp
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/strtod.cpp
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/schema_type_parser.cpp
+    ${Caffe2_SOURCE_DIR}/torch/csrc/jit/frontend/source_range.cpp
+)
+
 add_subdirectory(quantized)
 set(all_cpu_cpp ${base_cpp} ${ATen_CORE_SRCS} ${native_cpp} ${native_sparse_cpp} ${native_quantized_cpp} ${native_mkl_cpp} ${native_mkldnn_cpp} ${native_xnnpack} ${generated_cpp} ${core_generated_cpp} ${ATen_CPU_SRCS} ${ATen_QUANTIZED_SRCS} ${cpu_kernel_cpp})
 if(AT_MKL_ENABLED)
diff --git a/aten/src/ATen/core/CMakeLists.txt b/aten/src/ATen/core/CMakeLists.txt
diff --git a/docs/source/torch.rst b/docs/source/torch.rst
@@ -20,6 +20,7 @@ Tensors
     is_storage
     is_complex
     is_floating_point
+    is_nonzero
     set_default_dtype
     get_default_dtype
     set_default_tensor_type
diff --git a/test/distributed/test_distributed.py b/test/distributed/test_distributed.py
@@ -203,7 +203,7 @@ def init(cls):
             os.unlink(os.path.join(barrier_dir, f_name))
 
     @classmethod
-    def sync(cls, wait_for=None, timeout=5):
+    def sync(cls, wait_for=None, timeout=10):
         if wait_for is None:
             wait_for = dist.get_world_size()
         cls.barrier_id += 1
@@ -455,7 +455,6 @@ def _test_group_override_backend(self, initializer):
     @require_backends_available({"gloo", "nccl"})
     @require_world_size(3)
     @skip_if_lt_x_gpu(2)
-    @skip_if_rocm
     def test_backend_group(self):
         self._test_group_override_backend(self._init_group_test)
 
diff --git a/test/test_dataloader.py b/test/test_dataloader.py
@@ -1312,9 +1312,23 @@ def test_duplicating_data_with_drop_last(self):
 
         self.assertEqual(scanned_data.size(), scanned_data.unique().size())
 
+    def _test_sampler(self, **kwargs):
+        indices = range(2, 12)  # using a regular iterable
+        dl = DataLoader(self.dataset, sampler=indices, batch_size=2, **kwargs)
+        self.assertEqual(len(dl), 5)
+        for i, (input, _target) in enumerate(dl):
+            self.assertEqual(len(input), 2)
+            self.assertEqual(input, self.data[i * 2 + 2:i * 2 + 4])
+
+    def test_sampler(self):
+        self._test_sampler()
+        self._test_sampler(num_workers=4)
+        if not NO_MULTIPROCESSING_SPAWN and torch.multiprocessing._supports_context:
+            self._test_batch_sampler(num_workers=4, multiprocessing_context='spawn')
+
     def _test_batch_sampler(self, **kwargs):
         # [(0, 1), (2, 3, 4), (5, 6), (7, 8, 9), ...]
-        batches = []
+        batches = []  # using a regular iterable
         for i in range(0, 20, 5):
             batches.append(tuple(range(i, i + 2)))
             batches.append(tuple(range(i + 2, i + 5)))
diff --git a/test/test_jit_simple.py b/test/test_jit_simple.py
@@ -4,6 +4,7 @@
 
 if __name__ == '__main__':
     run_tests()
-    import test_jit_py3
-    suite = unittest.findTestCases(test_jit_py3)
-    unittest.TextTestRunner().run(suite)
+    if not PY2:
+        import test_jit_py3
+        suite = unittest.findTestCases(test_jit_py3)
+        unittest.TextTestRunner().run(suite)
diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl
@@ -39,6 +39,34 @@ libtorch_generated_sources = [
     "torch/csrc/autograd/VariableTypeManual.cpp",
 ]
 
+# copied from https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/CMakeLists.txt
+jit_core_headers = [
+    "torch/csrc/utils/memory.h",
+    "torch/csrc/WindowsTorchApiMacro.h",
+    "torch/csrc/jit/frontend/source_range.h",
+    "torch/csrc/jit/serialization/source_range_serialization.h",
+    "torch/csrc/jit/frontend/lexer.h",
+    "torch/csrc/jit/frontend/strtod.h",
+    "torch/csrc/jit/frontend/parser_constants.h",
+    "torch/csrc/jit/frontend/function_schema_parser.h",
+    "torch/csrc/jit/frontend/parse_string_literal.h",
+    "torch/csrc/jit/frontend/schema_type_parser.h",
+    "torch/csrc/jit/frontend/error_report.h",
+    "torch/csrc/jit/frontend/tree.h",
+    "torch/custom_class.h",
+    "torch/custom_class_detail.h",
+    "torch/library.h",
+]
+
+jit_core_sources = [
+    "torch/csrc/jit/frontend/error_report.cpp",
+    "torch/csrc/jit/frontend/function_schema_parser.cpp",
+    "torch/csrc/jit/frontend/lexer.cpp",
+    "torch/csrc/jit/frontend/schema_type_parser.cpp",
+    "torch/csrc/jit/frontend/strtod.cpp",
+    "torch/csrc/jit/frontend/source_range.cpp",
+]
+
 # copied from https://github.com/pytorch/pytorch/blob/master/tools/cpp_build/torch/CMakeLists.txt
 libtorch_core_sources = [
     "torch/csrc/autograd/anomaly_mode.cpp",
@@ -238,7 +266,6 @@ libtorch_core_jit_sources = [
 
 libtorch_cmake_sources = libtorch_core_sources + libtorch_core_jit_sources
 
-
 libtorch_extra_sources = libtorch_core_jit_sources + [
     "torch/csrc/autograd/VariableTypeManual.cpp",
     "torch/csrc/jit/api/module_save.cpp",
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
@@ -2701,6 +2701,39 @@ def merge_dicts(*dicts):
     input (Tensor): the PyTorch tensor to test
 """)
 
+add_docstr(torch.is_nonzero,
+           r"""
+is_nonzero(input) -> (bool)
+
+Returns True if the :attr:`input` is a single element tensor which is not equal to zero
+after type conversions.
+i.e. not equal to ``torch.tensor([0.])`` or ``torch.tensor([0])`` or
+``torch.tensor([False])``.
+Throws a ``RuntimeError`` if ``torch.numel() != 1`` (even in case
+of sparse tensors).
+
+Args:
+    input (Tensor): the PyTorch tensor to test
+Example::
+
+    >>> torch.is_nonzero(torch.tensor([0.]))
+    False
+    >>> torch.is_nonzero(torch.tensor([1.5]))
+    True
+    >>> torch.is_nonzero(torch.tensor([False]))
+    False
+    >>> torch.is_nonzero(torch.tensor([3]))
+    True
+    >>> torch.is_nonzero(torch.tensor([1, 3, 5]))
+    Traceback (most recent call last):
+    ...
+    RuntimeError: bool value of Tensor with more than one value is ambiguous
+    >>> torch.is_nonzero(torch.tensor([]))
+    Traceback (most recent call last):
+    ...
+    RuntimeError: bool value of Tensor with no values is ambiguous
+""")
+
 add_docstr(torch.kthvalue,
            r"""
 kthvalue(input, k, dim=None, keepdim=False, out=None) -> (Tensor, LongTensor)
diff --git a/torch/csrc/distributed/rpc/py_rref.cpp b/torch/csrc/distributed/rpc/py_rref.cpp
@@ -108,8 +108,13 @@ PyRRef::PyRRef(const py::object& value, const py::object& type_hint)
     : PyRRef([&value, &type_hint]() {
         TypePtr elem_type = tryInferTypeWithTypeHint(value, type_hint);
         auto rref = RRefContext::getInstance().createOwnerRRef(elem_type);
-        py::object copy(value); // increases refcount
-        IValue ivalue = jit::toIValue(std::move(copy), elem_type);
+        // jit::toIValue takes a py::handle as the first argument, and it calls
+        // py::handle.cast<py::object>() to incref of provided value. The
+        // returned ivalue will keep the reference alive.
+        // NB: the first argument const py::object& value must be kept alive
+        // until the following jit::toIValue returns (i.e., incref done). That's
+        // why this ctor can only be called while holding GIL.
+        IValue ivalue = jit::toIValue(value, elem_type);
         rref->setValue(std::move(ivalue));
         return rref;
       }()) {}
diff --git a/torch/csrc/distributed/rpc/py_rref.h b/torch/csrc/distributed/rpc/py_rref.h
@@ -14,6 +14,8 @@ enum RRefProxyType { RPC_SYNC, RPC_ASYNC, REMOTE };
 // pickle and unpickle.
 class PyRRef {
  public:
+  // The first ctor can only be called while holding GIL. See its implementation
+  // for more explanations.
   explicit PyRRef(const py::object& value, const py::object& type_hint);
   explicit PyRRef(c10::intrusive_ptr<RRef> rref);
 
diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py
@@ -58,6 +58,20 @@ class ProfilingMode(Enum):
     SIMPLE = 2
     PROFILING = 3
 
+def cppProfilingFlagsToProfilingMode():
+    old_prof_exec_state = torch._C._jit_set_profiling_executor(True)
+    old_prof_mode_state = torch._C._jit_set_profiling_mode(True)
+    torch._C._jit_set_profiling_executor(old_prof_exec_state)
+    torch._C._jit_set_profiling_mode(old_prof_mode_state)
+
+    if old_prof_exec_state:
+        if old_prof_mode_state:
+            return ProfilingMode.PROFILING
+        else:
+            return ProfilingMode.SIMPLE
+    else:
+        return ProfilingMode.LEGACY
+
 @contextmanager
 def enable_profiling_mode_for_profiling_tests():
     if GRAPH_EXECUTOR == ProfilingMode.PROFILING:
@@ -126,14 +140,17 @@ def _get_test_report_path():
 parser.add_argument('--log-suffix', type=str, default="")
 parser.add_argument('--run-parallel', type=int, default=1)
 
-GRAPH_EXECUTOR = ProfilingMode.SIMPLE if IS_SANDCASTLE else ProfilingMode.PROFILING
 args, remaining = parser.parse_known_args()
 if args.ge_config == 'legacy':
     GRAPH_EXECUTOR = ProfilingMode.LEGACY
 elif args.ge_config == 'profiling':
     GRAPH_EXECUTOR = ProfilingMode.PROFILING
-else:
+elif args.ge_config == 'simple':
     GRAPH_EXECUTOR = ProfilingMode.SIMPLE
+else:
+    # infer flags based on the default settings
+    GRAPH_EXECUTOR = cppProfilingFlagsToProfilingMode()
+
 
 LOG_SUFFIX = args.log_suffix
 RUN_PARALLEL = args.run_parallel
diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
@@ -74,11 +74,13 @@ class DataLoader(object):
             (default: ``1``).
         shuffle (bool, optional): set to ``True`` to have the data reshuffled
             at every epoch (default: ``False``).
-        sampler (Sampler, optional): defines the strategy to draw samples from
-            the dataset. If specified, :attr:`shuffle` must be ``False``.
-        batch_sampler (Sampler, optional): like :attr:`sampler`, but returns a batch of
-            indices at a time. Mutually exclusive with :attr:`batch_size`,
-            :attr:`shuffle`, :attr:`sampler`, and :attr:`drop_last`.
+        sampler (Sampler or Iterable, optional): defines the strategy to draw
+            samples from the dataset. Can be any ``Iterable`` with ``__len__``
+            implemented. If specified, :attr:`shuffle` must not be specified.
+        batch_sampler (Sampler or Iterable, optional): like :attr:`sampler`, but
+            returns a batch of indices at a time. Mutually exclusive with
+            :attr:`batch_size`, :attr:`shuffle`, :attr:`sampler`,
+            and :attr:`drop_last`.
         num_workers (int, optional): how many subprocesses to use for data
             loading. ``0`` means that the data will be loaded in the main process.
             (default: ``0``)
diff --git a/torch/utils/data/sampler.py b/torch/utils/data/sampler.py
@@ -167,7 +167,8 @@ class BatchSampler(Sampler):
     r"""Wraps another sampler to yield a mini-batch of indices.
 
     Args:
-        sampler (Sampler): Base sampler.
+        sampler (Sampler or Iterable): Base sampler. Can be any iterable object
+            with ``__len__`` implemented.
         batch_size (int): Size of mini-batch.
         drop_last (bool): If ``True``, the sampler will drop the last batch if
             its size would be less than ``batch_size``
@@ -180,10 +181,9 @@ class BatchSampler(Sampler):
     """
 
     def __init__(self, sampler, batch_size, drop_last):
-        if not isinstance(sampler, Sampler):
-            raise ValueError("sampler should be an instance of "
-                             "torch.utils.data.Sampler, but got sampler={}"
-                             .format(sampler))
+        # Since collections.abc.Iterable does not check for `__getitem__`, which
+        # is one way for an object to be an iterable, we don't do an `isinstance`
+        # check here.
         if not isinstance(batch_size, _int_classes) or isinstance(batch_size, bool) or \
                 batch_size <= 0:
             raise ValueError("batch_size should be a positive integer value, "
