[3/N] Avoid copy in std::get (pytorch#141843)

Fixes #ISSUE_NUMBER

Pull Request resolved: pytorch#141843
Approved by: https://github.com/Skylion007

Author: cyyever

aten/src/ATen/LegacyBatchedTensorImpl.cpp

@@ -76,7 +76,7 @@ void BatchedTensorImpl::checkInvariants() const {
   }
 }
 
-// The following are publically exposed as methods of Tensor
+// The following are publicly exposed as methods of Tensor
 
 IntArrayRef BatchedTensorImpl::strides_custom() const {
   return strides_default();
@@ -113,7 +113,7 @@ const char* BatchedTensorImpl::tensorimpl_type_name() const {
   return "BatchedTensorImpl";
 }
 
-Tensor makeBatched(const Tensor& tensor, BatchDims bdims) {
+Tensor makeBatched(Tensor tensor, BatchDims bdims) {
   TORCH_INTERNAL_ASSERT(!isBatchedTensor(tensor));
   auto tensor_dim = tensor.dim();
   TORCH_CHECK(
@@ -124,15 +124,15 @@ Tensor makeBatched(const Tensor& tensor, BatchDims bdims) {
       std::all_of(bdims.begin(), bdims.end(),
           [](const BatchDim& bdim) { return bdim.level() < kVmapNumLevels; }),
       "We only support up to ", kVmapNumLevels, " nested vmaps");
-  return at::detail::make_tensor<BatchedTensorImpl>(tensor, std::move(bdims));
+  return at::detail::make_tensor<BatchedTensorImpl>(std::move(tensor), std::move(bdims));
 }
 
-Tensor addBatchDim(const Tensor& tensor, int64_t level, int64_t dim) {
+Tensor addBatchDim(Tensor tensor, int64_t level, int64_t dim) {
   const auto* batched = maybeGetBatchedImpl(tensor);
   if (!batched) {
     BatchDims bdims;
     bdims.emplace_back(level, dim);
-    return at::detail::make_tensor<BatchedTensorImpl>(tensor, std::move(bdims));
+    return at::detail::make_tensor<BatchedTensorImpl>(std::move(tensor), std::move(bdims));
   }
   BatchDims new_bdims(batched->bdims().begin(), batched->bdims().end());
   auto actual_bdim = batched->actualDim(dim, /*wrap_dim=*/true);

aten/src/ATen/LegacyBatchedTensorImpl.h

@@ -148,10 +148,10 @@ inline std::ostream& operator<<(std::ostream& out, const BatchDim& bdim) {
 }
 
 // Use this to construct a BatchedTensor from a regular Tensor
-TORCH_API Tensor makeBatched(const Tensor& tensor, BatchDims bdims);
+TORCH_API Tensor makeBatched(Tensor tensor, BatchDims bdims);
 
 // Adds a batch dim to `tensor`, returning a BatchedTensor
-TORCH_API Tensor addBatchDim(const Tensor& tensor, int64_t level, int64_t dim);
+TORCH_API Tensor addBatchDim(Tensor tensor, int64_t level, int64_t dim);
 
 // Checks if an inplace operation on self and other is "vmap compatible".
 // See NOTE: [vmap-incompatible in-place operations] for the definition of this.

aten/src/ATen/core/boxing/KernelFunction_test.cpp

@@ -353,19 +353,17 @@ void expectOutOfPlaceMultiUnboxedCallingWorks(const KernelFunction& func) {
   auto t1 = at::zeros({1});
   auto t2 = at::zeros({1});
 
-  std::tuple<at::Tensor&, at::Tensor&> tup = func.call<
+  auto [t1_out, t2_out] = func.call<
     std::tuple<at::Tensor&, at::Tensor&>, at::Scalar, at::Scalar, at::Tensor&, at::Tensor&
   >(dummy, CPU_TEST_SET, s1, s2, t1, t2);
 
   // kernel should have updated out args and returned them in a tuple
   EXPECT_EQ(t1.item().toFloat(), 1.0f);
   EXPECT_EQ(t2.item().toFloat(), 2.0f);
 
-  auto t1_out = std::get<0>(tup);
   EXPECT_EQ(t1_out.item().toFloat(), 1.0f);
   EXPECT_TRUE(t1_out.is_same(t1));
 
-  auto t2_out = std::get<1>(tup);
   EXPECT_EQ(t2_out.item().toFloat(), 2.0f);
   EXPECT_TRUE(t2_out.is_same(t2));
 }

aten/src/ATen/functorch/BatchRulesBinaryOps.cpp

@@ -218,47 +218,43 @@ static std::tuple<Tensor, std::optional<int64_t>> masked_select_backward_batch_r
 
 static std::tuple<Tensor, std::optional<int64_t>> cdist_backward_batch_rule(
     const Tensor& grad, std::optional<int64_t> grad_bdim,
-    const Tensor& x1, std::optional<int64_t> x1_bdim,
-    const Tensor& x2, std::optional<int64_t> x2_bdim,
+    Tensor x1, std::optional<int64_t> x1_bdim,
+    Tensor x2, std::optional<int64_t> x2_bdim,
     const double p,
     const Tensor& cdist, std::optional<int64_t> cdist_bdim) {
 
-  auto x1_ = x1;
   if (cdist_bdim && !x1_bdim) {
     // We need to make sure that x1 has batch dim if cdist has one
     // otherwise, we get
     // RuntimeError: Function CdistBackward0 returned an invalid gradient at index 1 - got [5]
     // but expected shape compatible with [4, 5]
     auto bs = cdist.size(*cdist_bdim);
-    x1_ = ensure_has_bdim(x1, false, bs);
-    x1_ = x1_.contiguous();
+    x1 = ensure_has_bdim(x1, false, bs).contiguous();
     x1_bdim = 0;
   }
 
   // We need to apply the same preprocessing on x1 and x2 as in the forward pass
   // _binary_pointwise_batch_rule
-  auto x12 = _binary_pointwise_helper(x1_, x1_bdim, x2, x2_bdim);
-  x1_ = std::move(std::get<0>(x12));
-  auto& x2_ = std::get<1>(x12);
+  std::tie(x1, x2)= _binary_pointwise_helper(x1, x1_bdim, x2, x2_bdim);
 
   auto grad_ = moveBatchDimToFront(grad, grad_bdim);
   if ((x1_bdim || x2_bdim) && !grad_bdim) {
     // We need to make sure that grad has batch dim if x1 or x2 have one
     // Probably, there is an assumption on the strides.
     // Otherwise grad input contains thrash values, e.g. -7.0816e+29, 7.0816e+29
-    auto bs = get_bdim_size2(x1_, 0, x2_, 0);
+    auto bs = get_bdim_size2(x1, 0, x2, 0);
     grad_ = ensure_has_bdim(grad_, grad_bdim.has_value(), bs);
     grad_ = grad_.contiguous();
   }
 
-  auto out = at::_cdist_backward(grad_, x1_, x2_, p, cdist);
+  auto out = at::_cdist_backward(grad_, x1, x2, p, cdist);
 
   std::optional<int64_t> out_bdim = std::nullopt;
   if (x1_bdim || x2_bdim) {
     out_bdim = 0;
   }
 
-  return std::make_tuple(out, out_bdim);
+  return std::make_tuple(std::move(out), out_bdim);
 }
 
 static void fill__Tensor_batch_rule(

aten/src/ATen/functorch/BatchRulesNorm.cpp

@@ -42,6 +42,7 @@ static Tensor padRight(const Tensor& tensor, std::optional<int64_t> has_bdim, in
 }
 
 template<typename F, F Func>
+static
 std::tuple<Tensor, std::optional<int64_t>,Tensor, std::optional<int64_t>,Tensor, std::optional<int64_t>>
 batch_norm_batch_rule(
     const Tensor& input, std::optional<int64_t> input_bdim,
@@ -70,10 +71,10 @@ batch_norm_batch_rule(
   if (!input_bdim && !running_mean_bdim && !running_var_bdim) {
     const auto dummy_weight = at::ones(input.size(1), input.options());  // cudnn and miopen require a weight
     const auto dummy_bias = at::zeros(input.size(1), input.options());   // without this, get "strides() called on undefined Tensor" on cuda
-    const auto result = Func(input, dummy_weight, dummy_bias, running_mean_opt, running_var_opt, training, momentum, eps);
+    auto result = Func(input, dummy_weight, dummy_bias, running_mean_opt, running_var_opt, training, momentum, eps);
     result0 = std::get<0>(result).transpose(0, 1);          // [C, B, *]
-    mean = std::get<1>(result);
-    rstd = std::get<2>(result);
+    mean = std::move(std::get<1>(result));
+    rstd = std::move(std::get<2>(result));
   } else {
     bdim_size = get_bdim_size3(input, input_bdim, running_mean, running_mean_bdim, running_var, running_var_bdim);
     auto input_ = moveBatchDimToFront(input, input_bdim);
@@ -95,12 +96,12 @@ batch_norm_batch_rule(
 
     const auto dummy_weight = at::ones(input_.size(1), input_.options());  // cudnn and miopen require a weight
     const auto dummy_bias = at::zeros(input_.size(1), input_.options());   // without this, get "strides() called on undefined Tensor" on cuda
-    const auto result = Func(input_, dummy_weight, dummy_bias, running_mean_, running_var_, training, momentum, eps);
+    auto result = Func(input_, dummy_weight, dummy_bias, running_mean_, running_var_, training, momentum, eps);
     result0 = std::get<0>(result).transpose(0, 1);                // [(B0, C), B, *]
+    mean = std::move(std::get<1>(result));
+    rstd = std::move(std::get<2>(result));
     result0 = reshape_dim_outof(0, bdim_size.value(), result0);   // [B0, C, B, *]
-    mean = std::get<1>(result);
     mean = reshape_dim_outof(0, bdim_size.value(), mean);         // [B0, C]
-    rstd = std::get<2>(result);
     rstd = reshape_dim_outof(0, bdim_size.value(), rstd);         // [B0, C]
   }
 
@@ -124,6 +125,7 @@ batch_norm_batch_rule(
 }
 
 template<typename F, F Func>
+static
 std::tuple<at::Tensor, std::optional<int64_t>> batch_norm_backward_no_weight_bias_batch_rule(
     const at::Tensor & grad_out, std::optional<int64_t> grad_out_bdim,
     const at::Tensor & input, std::optional<int64_t> input_bdim,
@@ -142,9 +144,9 @@ std::tuple<at::Tensor, std::optional<int64_t>> batch_norm_backward_no_weight_bia
     TORCH_INTERNAL_ASSERT(!mean_bdim);
     TORCH_INTERNAL_ASSERT(!rstd_bdim);
     const auto dummy_weight = at::ones(input.size(1), input.options());
-    const auto result = Func(
+    auto result =Func(
         grad_out, input, dummy_weight, running_mean_opt, running_var_opt, mean, rstd, training, eps, {true, false, false});
-    return std::make_tuple(std::get<0>(result), std::nullopt);
+    return {std::move(std::get<0>(result)), std::nullopt};
   }
 
   auto grad_out_ = moveBatchDimToFront(grad_out, grad_out_bdim);
@@ -196,6 +198,7 @@ std::tuple<at::Tensor, std::optional<int64_t>> batch_norm_backward_no_weight_bia
 }
 
 template<typename F, F Func>
+static
 std::tuple<at::Tensor,at::Tensor,at::Tensor> batch_norm_backward_plumbing(
     const at::Tensor & grad_out,
     const at::Tensor & input,
@@ -270,15 +273,15 @@ std::tuple<at::Tensor,at::Tensor,at::Tensor> batch_norm_backward_plumbing(
         unwrapTensorAtLevel(grad_normalized_input.transpose(0, 1), cur_level);       // [B0, B, C, *]
 
     c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
-    const auto results = batch_norm_backward_no_weight_bias_batch_rule<F, Func>(
+    auto results = batch_norm_backward_no_weight_bias_batch_rule<F, Func>(
         grad_normalized_input_value, grad_normalized_input_bdim,
         input_value, input_bdim,
         running_mean_value, running_mean_bdim,
         running_var_value, running_var_bdim,
         save_mean_value, save_mean_bdim,
         save_rstd_value, save_rstd_bdim,
         training, eps);
-    grad_input = makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
+    grad_input = makeBatched(std::move(std::get<0>(results)), std::get<1>(results), cur_level);
   }
   return std::make_tuple(grad_input, grad_weight, grad_bias);
 }
@@ -312,16 +315,13 @@ static std::tuple<Tensor,Tensor,Tensor> native_group_norm_plumbing(
     const auto bdim_size = input_value.size(*input_bdim);
 
     c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
-    const auto result = at::native_group_norm(input_, std::nullopt, std::nullopt, N * bdim_size, C, HxW, group, eps);
-    result0 = makeBatched(reshape_dim_outof(0, bdim_size, std::get<0>(result)), 0, cur_level);
-    mean = makeBatched(reshape_dim_outof(0, bdim_size, std::get<1>(result)), 0, cur_level);
-    rstd = makeBatched(reshape_dim_outof(0, bdim_size, std::get<2>(result)), 0, cur_level);
+    std::tie(result0, mean, rstd) = at::native_group_norm(input_, std::nullopt, std::nullopt, N * bdim_size, C, HxW, group, eps);
+    result0 = makeBatched(reshape_dim_outof(0, bdim_size, result0), 0, cur_level);
+    mean = makeBatched(reshape_dim_outof(0, bdim_size, mean), 0, cur_level);
+    rstd = makeBatched(reshape_dim_outof(0, bdim_size, rstd), 0, cur_level);
   } else {
     c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
-    const auto result = at::native_group_norm(input_value, std::nullopt, std::nullopt, N, C, HxW, group, eps);
-    result0 = std::get<0>(result);
-    mean = std::get<1>(result);
-    rstd = std::get<2>(result);
+    std::tie(result0, mean, rstd) = at::native_group_norm(input_value, std::nullopt, std::nullopt, N, C, HxW, group, eps);
   }
 
   if (weight.defined()) {
@@ -334,10 +334,10 @@ static std::tuple<Tensor,Tensor,Tensor> native_group_norm_plumbing(
     result0 = result0 + padded_bias;
   }
 
-  return std::make_tuple(result0, mean, rstd);
+  return std::make_tuple(std::move(result0), std::move(mean), std::move(rstd));
 }
 
-static std::tuple<at::Tensor, std::optional<int64_t>> group_norm_backward_no_weight_bias_batch_rule(
+static at::Tensor group_norm_backward_no_weight_bias_batch_rule(
     const at::Tensor & grad_out, std::optional<int64_t> grad_out_bdim,
     const at::Tensor & input, std::optional<int64_t> input_bdim,
     const at::Tensor & mean, std::optional<int64_t> mean_bdim,
@@ -359,15 +359,13 @@ static std::tuple<at::Tensor, std::optional<int64_t>> group_norm_backward_no_wei
   mean_ = reshape_dim_into(0, 0, mean_);         // [B0 * N, G]
   rstd_ = reshape_dim_into(0, 0, rstd_);         // [B0 * N, G]
 
-  const auto result = native_group_norm_backward(
+  auto result0 = std::get<0>(native_group_norm_backward(
       grad_out_.contiguous(),
       input_.contiguous(),
       mean_.contiguous(),
       rstd_.contiguous(),
-      std::nullopt, N * bdim_size, C, HxW, group, {true, false, false});
-  auto result0 = std::get<0>(result);
-  result0 = reshape_dim_outof(0, bdim_size, result0);
-  return std::make_tuple(result0, 0);
+      std::nullopt, N * bdim_size, C, HxW, group, {true, false, false}));
+  return reshape_dim_outof(0, bdim_size, result0);
 }
 
 static std::tuple<Tensor,Tensor,Tensor> native_group_norm_backward_plumbing(
@@ -422,19 +420,19 @@ static std::tuple<Tensor,Tensor,Tensor> native_group_norm_backward_plumbing(
         unwrapTensorAtLevel(grad_normalized_input, cur_level);
 
     c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
-    const auto res = group_norm_backward_no_weight_bias_batch_rule(
+    auto tensor = group_norm_backward_no_weight_bias_batch_rule(
         grad_normalized_input_value, grad_normalized_input_bdim,
         input_value, input_bdim,
         mean_value, mean_bdim,
         rstd_value, rstd_bdim,
         N, C, HxW, group
     );
-    grad_input = makeBatched(std::get<0>(res), std::get<1>(res), cur_level);
+    grad_input = makeBatched(std::move(tensor), 0, cur_level);
   }
   return std::make_tuple(grad_input, grad_weight, grad_bias);
 }
 
-C10_ALWAYS_INLINE bool has_same_shape(
+static bool has_same_shape(
     const Tensor& tensor, std::optional<int64_t> tensor_bdim,
     c10::SymIntArrayRef normalized_shape) {
   if (!tensor.defined()) {
@@ -457,7 +455,7 @@ C10_ALWAYS_INLINE bool has_same_shape(
   return true;
 }
 
-C10_ALWAYS_INLINE void check_same_shape(
+static C10_ALWAYS_INLINE void check_same_shape(
     const Tensor& tensor, std::optional<int64_t> tensor_bdim,
     c10::SymIntArrayRef normalized_shape, const std::string& name) {
   TORCH_CHECK(has_same_shape(tensor, tensor_bdim, normalized_shape),
@@ -469,7 +467,7 @@ C10_ALWAYS_INLINE void check_same_shape(
 }
 
 // Ugh, hard to deduplicate
-C10_ALWAYS_INLINE void _check_layer_norm_inputs(
+static C10_ALWAYS_INLINE void _check_layer_norm_inputs(
     SymIntArrayRef normalized_shape,
     const Tensor& weight, std::optional<int64_t> weight_bdim,
     const Tensor& bias, std::optional<int64_t> bias_bdim) {
@@ -493,11 +491,9 @@ native_layer_norm_batch_rule(
     double eps) {
   auto input_ = moveBatchDimToFront(input, input_bdim);
   if (!weight_bdim && !bias_bdim) {
-    const auto result = at::native_layer_norm_symint(input_, normalized_shape, weight_opt, bias_opt, eps);
-    const auto mean = std::get<1>(result);
-    const auto rstd = std::get<2>(result);
+    auto [result0, mean, rstd] = at::native_layer_norm_symint(input_, normalized_shape, weight_opt, bias_opt, eps);
     const auto stats_bdim = compute_stat_bdim(input_bdim, mean);
-    return std::make_tuple(std::get<0>(result), 0, mean, stats_bdim, rstd, stats_bdim);
+    return std::make_tuple(std::move(result0), 0, std::move(mean), stats_bdim, std::move(rstd), stats_bdim);
   }
 
   // See [Note: hacky wrapper removal for optional tensor]
@@ -509,9 +505,7 @@ native_layer_norm_batch_rule(
 
   const auto input_logical_rank = rankWithoutBatchDim(input, input_bdim);
   const auto result = at::native_layer_norm_symint(input_, normalized_shape, std::nullopt, std::nullopt, eps);
-  auto result0 = std::get<0>(result);
-  const auto mean = std::get<1>(result);
-  const auto rstd = std::get<2>(result);
+  auto [result0, mean, rstd] = result;
   const auto stats_bdim = compute_stat_bdim(input_bdim, mean);
 
   if (weight.defined()) {
@@ -638,7 +632,7 @@ static std::tuple<at::Tensor,at::Tensor,at::Tensor> native_layer_norm_backward_p
         unwrapTensorAtLevel(grad_normalized_input, cur_level);
 
     c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
-    const auto results = native_layer_norm_backward_no_weight_bias_batch_rule(
+    auto results = native_layer_norm_backward_no_weight_bias_batch_rule(
         grad_normalized_input_value, grad_normalized_input_bdim,
         input_value, input_bdim,
         normalized_shape,

aten/src/ATen/functorch/BatchedTensorImpl.cpp

@@ -171,18 +171,18 @@ void BatchedTensorImpl::shallow_copy_from(const c10::intrusive_ptr<TensorImpl>&
   TORCH_CHECK(false, "mutating directly with `.data` under vmap transform is not allowed.");
 }
 
-Tensor makeBatched(const Tensor& tensor, int64_t bdim, int64_t level) {
+Tensor makeBatched(Tensor tensor, int64_t bdim, int64_t level) {
   DispatchKeySet key_set = getKeysToPropagateToWrapper(tensor);
   auto* batched = maybeGetBatchedImpl(tensor);
   if (batched) {
     auto batched_level = batched->level();
     TORCH_INTERNAL_ASSERT(level > batched_level, " batched_level: ", batched_level, " level: ", level);
   }
-  return at::detail::make_tensor<BatchedTensorImpl>(key_set, tensor, bdim, level);
+  return at::detail::make_tensor<BatchedTensorImpl>(key_set, std::move(tensor), bdim, level);
 }
 
-Tensor addBatchDim(const Tensor& tensor, int64_t dim, int64_t level) {
-  return makeBatched(tensor, dim, level);
+Tensor addBatchDim(Tensor tensor, int64_t dim, int64_t level) {
+  return makeBatched(std::move(tensor), dim, level);
 }
 
 } // namespace at::functorch

aten/src/ATen/functorch/BatchedTensorImpl.h

@@ -144,10 +144,10 @@ inline std::bitset<kVmapNumLevels> createVmapLevelsBitset(int64_t level) {
 }
 
 // Use this to construct a BatchedTensor from a regular Tensor
-TORCH_API Tensor makeBatched(const Tensor& tensor, int64_t dim, int64_t level);
+TORCH_API Tensor makeBatched(Tensor tensor, int64_t dim, int64_t level);
 
 // Adds a batch dim to `tensor`, returning a BatchedTensor
-TORCH_API Tensor addBatchDim(const Tensor& tensor, int64_t dim, int64_t level);
+TORCH_API Tensor addBatchDim(Tensor tensor, int64_t dim, int64_t level);
 
 // Certain dispatch keys must be propagated to the BatchedTensor (or, in general,
 // any wrapper Tensor subclasses). This is because there are methods on Tensor