move flatten_dense_tensors and unflatten_dense_tensors to Native (#58006)

eqy · facebook-github-bot · commit 645a5f706ac0 · 2021-05-12T18:18:34.000-07:00
Summary: #55240 CC ngimel Pull Request resolved: #58006 Reviewed By: agolynski Differential Revision: D28386749 Pulled By: ngimel fbshipit-source-id: 4860c35d5ff95bcc38a243d7001180e7bd536314
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -2392,4 +2392,31 @@ Tensor& swapdims_(Tensor& self, int64_t dim0, int64_t dim1) {
   return self.transpose_(dim0, dim1);
 }
 
+Tensor flatten_dense_tensors(TensorList tensors) {
+  static auto flatten = [](const Tensor &t) { return t.contiguous().view({-1}); };
+  if (tensors.size() == 1)
+    return flatten(tensors[0]);
+  return at::cat(fmap(tensors, flatten));
+}
+
+std::vector<Tensor> unflatten_dense_tensors(const Tensor& flat, TensorList tensors) {
+  std::vector<Tensor> outputs;
+  outputs.reserve(tensors.size());
+  size_t offset = 0;
+  for (const auto & tensor : tensors) {
+    auto numel = tensor.numel();
+    // If unflatten an empty tensor, create a new empty tensor using
+    // flat tensor Options.
+    // This can avoid the unflattened empty tensor to share the same storage
+    // with other unflatten tensors.
+    if (numel == 0) {
+      outputs.push_back(at::empty({0}, flat.options()));
+    } else {
+      outputs.push_back(flat.narrow(0, offset, numel).view(tensor.sizes()));
+      offset += numel;
+    }
+  }
+  return outputs;
+}
+
 }} // at::native
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -9857,3 +9857,11 @@
 - func: pad_sequence(Tensor[] sequences, bool batch_first=False, float padding_value=0.0) -> Tensor
   python_module: nn
   variants: function
+
+- func: flatten_dense_tensors(Tensor[] tensors) -> Tensor
+  variants: function
+  python_module: nn
+
+- func: unflatten_dense_tensors(Tensor flat, Tensor[] tensors) -> Tensor[]
+  variants: function
+  python_module: nn
diff --git a/torch/_C/_nn.pyi.in b/torch/_C/_nn.pyi.in
@@ -27,3 +27,7 @@ def _parse_to(tensor: Tensor, non_blocking: _bool, copy: _bool, *,
 # Defined in aten/src/ATen/naitve/PadSequence.cpp
 def pad_sequence(sequences: List[Tensor], batch_first: bool = False,
                  padding_value: float = ...) -> Tensor: ...
+
+def flatten_dense_tensors(tensors: List[Tensor]) -> Tensor: ...
+
+def unflatten_dense_tensors(flat: Tensor, tensors: List[Tensor]) -> List[Tensor]: ...
diff --git a/torch/_utils.py b/torch/_utils.py
@@ -259,10 +259,7 @@ def _flatten_dense_tensors(tensors):
     Returns:
         A contiguous 1D buffer containing input tensors.
     """
-    if len(tensors) == 1:
-        return tensors[0].contiguous().view(-1)
-    flat = torch.cat([t.contiguous().view(-1) for t in tensors], dim=0)
-    return flat
+    return torch._C._nn.flatten_dense_tensors(tensors)
 
 
 def _flatten_sparse_tensors(tensors):
@@ -276,8 +273,8 @@ def _flatten_sparse_tensors(tensors):
         A tuple of two contiguous 1D buffers, one containing input tensors'
         indices and the other containing the values.
     """
-    flat_indices = _flatten_dense_tensors([torch.Tensor._indices(t) for t in tensors])
-    flat_values = _flatten_dense_tensors([torch.Tensor._values(t) for t in tensors])
+    flat_indices = torch._C._nn.flatten_dense_tensors([torch.Tensor._indices(t) for t in tensors])
+    flat_values = torch._C._nn.flatten_dense_tensors([torch.Tensor._values(t) for t in tensors])
     return flat_indices, flat_values
 
 
@@ -294,13 +291,7 @@ def _unflatten_dense_tensors(flat, tensors):
         Unflattened dense tensors with sizes same as tensors and values from
         flat.
     """
-    outputs = []
-    offset = 0
-    for tensor in tensors:
-        numel = tensor.numel()
-        outputs.append(flat.narrow(0, offset, numel).view_as(tensor))
-        offset += numel
-    return tuple(outputs)
+    return torch._C._nn.unflatten_dense_tensors(flat, tensors)
 
 
 def _unflatten_sparse_tensors(flat, tensors):
@@ -319,8 +310,8 @@ def _unflatten_sparse_tensors(flat, tensors):
         flat.
     """
     flat_indices, flat_values = flat
-    indices = _unflatten_dense_tensors(flat_indices, [torch.Tensor._indices(t) for t in tensors])
-    values = _unflatten_dense_tensors(flat_values, [torch.Tensor._values(t) for t in tensors])
+    indices = torch._C._nn.unflatten_dense_tensors(flat_indices, [torch.Tensor._indices(t) for t in tensors])
+    values = torch._C._nn.unflatten_dense_tensors(flat_values, [torch.Tensor._values(t) for t in tensors])
     outputs = []
     for t, i, v in zip(tensors, indices, values):
         outputs.append(t.new(i, v, t.size()))
diff --git a/torch/csrc/utils/tensor_flatten.cpp b/torch/csrc/utils/tensor_flatten.cpp
@@ -94,8 +94,8 @@ at::Tensor get_values(const at::Tensor& t) {
 }
 
 std::pair<at::Tensor, at::Tensor> flatten_sparse_tensors(at::TensorList tensors) {
-  auto flat_indices = flatten_dense_tensors(fmap(tensors, &get_indices));
-  auto flat_values = flatten_dense_tensors(fmap(tensors, &get_values));
+  auto flat_indices = utils::flatten_dense_tensors(fmap(tensors, &get_indices));
+  auto flat_values = utils::flatten_dense_tensors(fmap(tensors, &get_values));
   return std::make_pair(flat_indices, flat_values);
 }
 
@@ -104,8 +104,8 @@ std::vector<at::Tensor> unflatten_sparse_tensors(
         at::TensorList tensors) {
   if (tensors.size() == 0) return {};
 
-  auto indices = unflatten_dense_tensors(flat_indices, fmap(tensors, &get_indices));
-  auto values = unflatten_dense_tensors(flat_values, fmap(tensors, &get_values));
+  auto indices = utils::unflatten_dense_tensors(flat_indices, fmap(tensors, &get_indices));
+  auto values = utils::unflatten_dense_tensors(flat_values, fmap(tensors, &get_values));
 
   std::vector<at::Tensor> outputs;
   outputs.reserve(tensors.size());
diff --git a/torch/csrc/utils/tensor_flatten.h b/torch/csrc/utils/tensor_flatten.h
@@ -18,30 +18,11 @@ inline size_t type_id(const at::Tensor& tensor) {
 }
 
 inline at::Tensor flatten_dense_tensors(at::TensorList tensors) {
-  static auto flatten = [](const at::Tensor &t) { return t.contiguous().view({-1}); };
-  if (tensors.size() == 1)
-    return flatten(tensors[0]);
-  return at::cat(fmap(tensors, flatten));
+  return at::flatten_dense_tensors(tensors);
 }
 
 inline std::vector<at::Tensor> unflatten_dense_tensors(const at::Tensor& flat, at::TensorList tensors) {
-  std::vector<at::Tensor> outputs;
-  outputs.reserve(tensors.size());
-  size_t offset = 0;
-  for (const auto & tensor : tensors) {
-    auto numel = tensor.numel();
-    // If unflatten an empty tensor, create a new empty tensor using
-    // flat tensor Options.
-    // This can avoid the unflattened empty tensor to share the same storage
-    // with other unflatten tensors.
-    if (numel == 0) {
-      outputs.push_back(at::empty({0}, flat.options()));
-    } else {
-      outputs.push_back(flat.narrow(0, offset, numel).view(tensor.sizes()));
-      offset += numel;
-    }
-  }
-  return outputs;
+  return at::unflatten_dense_tensors(flat, tensors);
 }
 
 
