share inplace version counter for strided kernels (#59422)

wanghuancoder · web-flow · commit c8b1e8b6ef06 · 2023-11-29T10:58:05.000+08:00
* share inplace version counter for strided kernels
diff --git a/paddle/phi/kernels/stride/as_complex_kernel.cc b/paddle/phi/kernels/stride/as_complex_kernel.cc
@@ -35,6 +35,7 @@ void AsComplexStridedKernel(const Context& dev_ctx,
   }
   out->set_offset(x.offset());
   out->ResetHolder(x.Holder());
+  out->ShareInplaceVersionCounterWith(x);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/stride/as_real_kernel.cc b/paddle/phi/kernels/stride/as_real_kernel.cc
@@ -33,6 +33,7 @@ void AsRealStridedKernel(const Context& dev_ctx,
   }
   out->set_offset(x.offset());
   out->ResetHolder(x.Holder());
+  out->ShareInplaceVersionCounterWith(x);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/stride/as_strided_kernel.cc b/paddle/phi/kernels/stride/as_strided_kernel.cc
@@ -28,6 +28,7 @@ void AsStridedKernel(const Context& dev_ctx,
   out->set_strides(DDim(stride.data(), static_cast<int>(stride.size())));
   out->set_offset(offset);
   out->ResetHolder(input.Holder());
+  out->ShareInplaceVersionCounterWith(input);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/stride/complex_kernel.cc b/paddle/phi/kernels/stride/complex_kernel.cc
@@ -35,6 +35,7 @@ void RealStridedKernel(const Context& dev_ctx,
   out->set_offset(x.offset());
   out->set_strides(stride);
   out->ResetHolder(x.Holder());
+  out->ShareInplaceVersionCounterWith(x);
 }
 
 template <typename T, typename Context>
@@ -54,6 +55,7 @@ void ImagStridedKernel(const Context& dev_ctx,
   out->set_strides(stride);
   out->set_offset(x.offset() + phi::SizeOf(out->dtype()));
   out->ResetHolder(x.Holder());
+  out->ShareInplaceVersionCounterWith(x);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/stride/diagonal_kernel.cc b/paddle/phi/kernels/stride/diagonal_kernel.cc
@@ -78,6 +78,7 @@ void DiagonalStridedKernel(const Context& dev_ctx,
   meta.offset = x_offset;
   out->set_meta(meta);
   out->ResetHolder(x.Holder());
+  out->ShareInplaceVersionCounterWith(x);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/stride/index_select_kernel.cc b/paddle/phi/kernels/stride/index_select_kernel.cc
@@ -53,6 +53,7 @@ void IndexSelectStridedKernel(const Context& ctx,
   meta.strides = DDim(stride.data(), static_cast<int>(stride.size()));
   output->set_meta(meta);
   output->ResetHolder(x.Holder());
+  output->ShareInplaceVersionCounterWith(x);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/stride/reshape_kernel.cc b/paddle/phi/kernels/stride/reshape_kernel.cc
@@ -41,6 +41,7 @@ void ReshapeStridedKernel(const Context& dev_ctx,
     out->set_offset(x_offset);
     out->set_strides(stride);
     out->ResetHolder(x.Holder());
+    out->ShareInplaceVersionCounterWith(x);
   } else {
     DenseTensor tmp;
     DenseTensor tmp_x = x;
diff --git a/paddle/phi/kernels/stride/slice_kernel.cc b/paddle/phi/kernels/stride/slice_kernel.cc
@@ -99,6 +99,7 @@ void SliceStridedKernel(const Context& ctx,
       DDim(output_stride.data(), static_cast<int>(output_stride.size()));
   out->set_meta(meta);
   out->ResetHolder(input.Holder());
+  out->ShareInplaceVersionCounterWith(input);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/stride/squeeze_kernel.cc b/paddle/phi/kernels/stride/squeeze_kernel.cc
@@ -111,6 +111,7 @@ void SqueezeInferStridedKernel(const Context& dev_ctx,
   meta.offset = input.offset();
   out->set_meta(meta);
   out->ResetHolder(input.Holder());
+  out->ShareInplaceVersionCounterWith(input);
 }
 
 template <typename Context>
diff --git a/paddle/phi/kernels/stride/strided_slice_kernel.cc b/paddle/phi/kernels/stride/strided_slice_kernel.cc
@@ -121,6 +121,7 @@ void StridedSliceRawStridedKernel(const Context& dev_ctx,
       DDim(output_stride.data(), static_cast<int>(output_stride.size()));
   out->set_meta(meta);
   out->ResetHolder(input.Holder());
+  out->ShareInplaceVersionCounterWith(input);
 }
 
 template <typename Context>
diff --git a/paddle/phi/kernels/stride/tensor_unfold_kernel.cc b/paddle/phi/kernels/stride/tensor_unfold_kernel.cc
@@ -65,6 +65,7 @@ void TensorUnfoldKernel(const Context& dev_ctx,
   out->set_strides(DDim(stride.data(), static_cast<int>(stride.size())));
   out->set_offset(input.offset());
   out->ResetHolder(input.Holder());
+  out->ShareInplaceVersionCounterWith(input);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/stride/transpose_kernel.cc b/paddle/phi/kernels/stride/transpose_kernel.cc
@@ -41,6 +41,7 @@ void TransposeStridedKernel(const Context& ctx,
 
   out->set_meta(meta);
   out->ResetHolder(x.Holder());
+  out->ShareInplaceVersionCounterWith(x);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/stride/unsqueeze_kernel.cc b/paddle/phi/kernels/stride/unsqueeze_kernel.cc
@@ -71,6 +71,7 @@ void UnsqueezeInferStridedKernel(const Context& dev_ctx,
   meta.offset = input.offset();
   out->set_meta(meta);
   out->ResetHolder(input.Holder());
+  out->ShareInplaceVersionCounterWith(input);
 }
 
 template <typename Context>
diff --git a/paddle/phi/kernels/stride/view_kernel.cc b/paddle/phi/kernels/stride/view_kernel.cc
@@ -32,6 +32,7 @@ void ViewShapeKernel(const Context& dev_ctx,
     meta.offset = input.offset();
     out->set_meta(meta);
     out->ResetHolder(input.Holder());
+    out->ShareInplaceVersionCounterWith(input);
   } else {
     PADDLE_THROW(phi::errors::InvalidArgument(
         "The Tensor can not be viewed, please call reshape."));
@@ -51,6 +52,7 @@ void ViewDtypeKernel(const Context& dev_ctx,
     meta.dtype = dtype;
     out->set_meta(meta);
     out->ResetHolder(input.Holder());
+    out->ShareInplaceVersionCounterWith(input);
   } else if (input_dtype_size == 0) {
     PADDLE_THROW(phi::errors::InvalidArgument(
         "The Tensor's shape is [] can not be viewed."));
@@ -82,6 +84,7 @@ void ViewDtypeKernel(const Context& dev_ctx,
     meta.offset = input.offset() * times;
     out->set_meta(meta);
     out->ResetHolder(input.Holder());
+    out->ShareInplaceVersionCounterWith(input);
   } else {
     PADDLE_ENFORCE_EQ(
         input.strides()[input.strides().size() - 1],
@@ -140,6 +143,7 @@ void ViewDtypeKernel(const Context& dev_ctx,
     meta.offset = input.offset() / times;
     out->set_meta(meta);
     out->ResetHolder(input.Holder());
+    out->ShareInplaceVersionCounterWith(input);
   }
 }
 
diff --git a/test/collective/fleet/hybrid_parallel_pp_recompute.py b/test/collective/fleet/hybrid_parallel_pp_recompute.py
@@ -98,6 +98,7 @@ class TransformerNetPipe(TransformerNet):
     def forward(self, tensors):
         if framework.in_dynamic_mode():
             stable, x = tensors
+            stable = paddle.assign(stable)
             output = super().forward(x)
             return stable, output
         else:

Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,7 @@ void AsComplexStridedKernel(const Context& dev_ctx,`
`35`	`35`	`}`
`36`	`36`	`out->set_offset(x.offset());`
`37`	`37`	`out->ResetHolder(x.Holder());`
	`38`	`+ out->ShareInplaceVersionCounterWith(x);`
`38`	`39`	`}`
`39`	`40`
`40`	`41`	`} // namespace phi`
Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,7 @@ void AsRealStridedKernel(const Context& dev_ctx,`
`33`	`33`	`}`
`34`	`34`	`out->set_offset(x.offset());`
`35`	`35`	`out->ResetHolder(x.Holder());`
	`36`	`+ out->ShareInplaceVersionCounterWith(x);`
`36`	`37`	`}`
`37`	`38`
`38`	`39`	`} // namespace phi`
Original file line number	Diff line number	Diff line change
`@@ -28,6 +28,7 @@ void AsStridedKernel(const Context& dev_ctx,`
`28`	`28`	`out->set_strides(DDim(stride.data(), static_cast<int>(stride.size())));`
`29`	`29`	`out->set_offset(offset);`
`30`	`30`	`out->ResetHolder(input.Holder());`
	`31`	`+ out->ShareInplaceVersionCounterWith(input);`
`31`	`32`	`}`
`32`	`33`
`33`	`34`	`} // namespace phi`
Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,7 @@ void RealStridedKernel(const Context& dev_ctx,`
`35`	`35`	`out->set_offset(x.offset());`
`36`	`36`	`out->set_strides(stride);`
`37`	`37`	`out->ResetHolder(x.Holder());`
	`38`	`+ out->ShareInplaceVersionCounterWith(x);`
`38`	`39`	`}`
`39`	`40`
`40`	`41`	`template <typename T, typename Context>`
`@@ -54,6 +55,7 @@ void ImagStridedKernel(const Context& dev_ctx,`
`54`	`55`	`out->set_strides(stride);`
`55`	`56`	`out->set_offset(x.offset() + phi::SizeOf(out->dtype()));`
`56`	`57`	`out->ResetHolder(x.Holder());`
	`58`	`+ out->ShareInplaceVersionCounterWith(x);`
`57`	`59`	`}`
`58`	`60`
`59`	`61`	`} // namespace phi`
Original file line number	Diff line number	Diff line change
`@@ -78,6 +78,7 @@ void DiagonalStridedKernel(const Context& dev_ctx,`
`78`	`78`	`meta.offset = x_offset;`
`79`	`79`	`out->set_meta(meta);`
`80`	`80`	`out->ResetHolder(x.Holder());`
	`81`	`+ out->ShareInplaceVersionCounterWith(x);`
`81`	`82`	`}`
`82`	`83`
`83`	`84`	`} // namespace phi`
Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,7 @@ void IndexSelectStridedKernel(const Context& ctx,`
`53`	`53`	`meta.strides = DDim(stride.data(), static_cast<int>(stride.size()));`
`54`	`54`	`output->set_meta(meta);`
`55`	`55`	`output->ResetHolder(x.Holder());`
	`56`	`+ output->ShareInplaceVersionCounterWith(x);`
`56`	`57`	`}`
`57`	`58`
`58`	`59`	`} // namespace phi`
Original file line number	Diff line number	Diff line change
`@@ -99,6 +99,7 @@ void SliceStridedKernel(const Context& ctx,`
`99`	`99`	`DDim(output_stride.data(), static_cast<int>(output_stride.size()));`
`100`	`100`	`out->set_meta(meta);`
`101`	`101`	`out->ResetHolder(input.Holder());`
	`102`	`+ out->ShareInplaceVersionCounterWith(input);`
`102`	`103`	`}`
`103`	`104`
`104`	`105`	`} // namespace phi`
Original file line number	Diff line number	Diff line change
`@@ -111,6 +111,7 @@ void SqueezeInferStridedKernel(const Context& dev_ctx,`
`111`	`111`	`meta.offset = input.offset();`
`112`	`112`	`out->set_meta(meta);`
`113`	`113`	`out->ResetHolder(input.Holder());`
	`114`	`+ out->ShareInplaceVersionCounterWith(input);`
`114`	`115`	`}`
`115`	`116`
`116`	`117`	`template <typename Context>`
Original file line number	Diff line number	Diff line change
`@@ -121,6 +121,7 @@ void StridedSliceRawStridedKernel(const Context& dev_ctx,`
`121`	`121`	`DDim(output_stride.data(), static_cast<int>(output_stride.size()));`
`122`	`122`	`out->set_meta(meta);`
`123`	`123`	`out->ResetHolder(input.Holder());`
	`124`	`+ out->ShareInplaceVersionCounterWith(input);`
`124`	`125`	`}`
`125`	`126`
`126`	`127`	`template <typename Context>`