Port cat kernel to structured kernels.

Tracking issue: pytorch#55070 Pull Request resolved: pytorch#68640 Approved by: https://github.com/ezyang
jakeharmon8 · Apr 14, 2022 · 22a10ce · 22a10ce
1 parent 4c3ee53
commit 22a10ce
Show file tree

Hide file tree

Showing 19 changed files with 254 additions and 349 deletions.
diff --git a/aten/src/ATen/NamedTensorUtils.cpp b/aten/src/ATen/NamedTensorUtils.cpp
@@ -459,7 +459,7 @@ std::vector<Dimname> broadcast_to_outnames(
   return unify_from_right(reference_names, tensor_names);
 }
 
-std::vector<Dimname> compute_cat_outnames(TensorList tensors) {
+std::vector<Dimname> compute_cat_outnames(ITensorListRef tensors) {
   if (!at::has_names(tensors)) {
     return {};
   }

diff --git a/aten/src/ATen/NamedTensorUtils.h b/aten/src/ATen/NamedTensorUtils.h
@@ -10,7 +10,7 @@ namespace at {
 
 using NameVector = SmallVector<Dimname, kDimVectorStaticSize>;
 
-inline bool has_names(TensorList tensors) {
+inline bool has_names(ITensorListRef tensors) {
   return std::any_of(
       tensors.begin(), tensors.end(), [](const Tensor& t) { return t.has_names(); });
 }
@@ -98,7 +98,7 @@ TORCH_API void propagate_names_for_reduction(const Tensor& result, const Tensor&
 
 TORCH_API void propagate_names_for_expand(const Tensor& result, const Tensor& self);
 
-TORCH_API std::vector<Dimname> compute_cat_outnames(TensorList tensors);
+TORCH_API std::vector<Dimname> compute_cat_outnames(ITensorListRef tensors);
 
 TORCH_API std::vector<Dimname> compute_broadcast_outnames(
     const Tensor& self,

diff --git a/aten/src/ATen/WrapDimUtils.h b/aten/src/ATen/WrapDimUtils.h
@@ -4,6 +4,7 @@
 #include <c10/core/TensorImpl.h>
 #include <c10/util/irange.h>
 #include <ATen/core/Tensor.h>
+#include <ATen/core/ITensorListRef.h>
 
 namespace at {
 
@@ -74,7 +75,7 @@ static inline int64_t legacy_cat_wrap_dim(int64_t dim, const std::vector<std::ve
   return dim;
 }
 
-static inline int64_t legacy_cat_wrap_dim(int64_t dim, TensorList tensors) {
+static inline int64_t legacy_cat_wrap_dim(int64_t dim, ITensorListRef tensors) {
   for (auto& tensor : tensors) {
     if (tensor.dim() == 1 && tensor.sizes()[0] == 0) {
       continue;

diff --git a/aten/src/ATen/native/TensorCompare.cpp b/aten/src/ATen/native/TensorCompare.cpp
@@ -295,7 +295,7 @@ static void isin_sorting(
   // 2. Stable sort all elements, maintaining order indices to reverse the
   //    operation. Stable sort is necessary to keep elements before test
   //    elements within the sorted list.
-  Tensor all_elements = at::_cat({elements_flat, test_elements_flat});
+  Tensor all_elements = at::cat({elements_flat, test_elements_flat});
   Tensor sorted_elements, sorted_order;
   std::tie (sorted_elements, sorted_order) = all_elements.sort(
       /*stable=*/ true, /*dim=*/ 0, /*descending=*/ false);