[Prim] Mark reduce_as no need target buffer and use by_pass in both dynamic and static shape branch (#68491)

SigureMo · web-flow · commit 26d2d99b6ba1 · 2024-09-29T11:12:47.000+08:00
diff --git a/paddle/common/ddim.cc b/paddle/common/ddim.cc
@@ -282,6 +282,18 @@ DDim ComputeCompatibleDim(const DDim& dim1, const DDim& dim2) {
   return make_ddim(result);
 }
 
+bool AreDimsWithDynamicShapeCompatible(const DDim& dim1, const DDim& dim2) {
+  if (dim1.size() != dim2.size()) {
+    return false;
+  }
+  for (int i = 0; i < dim1.size(); ++i) {
+    if (dim1[i] >= 0 && dim2[i] >= 0 && dim1[i] != dim2[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
 }  // namespace common
 
 namespace std {
diff --git a/paddle/common/ddim.h b/paddle/common/ddim.h
@@ -234,6 +234,9 @@ TEST_API DDim stride(const DDim& ddim);
 
 TEST_API DDim stride_numel(const DDim& ddim);
 
+TEST_API bool AreDimsWithDynamicShapeCompatible(const DDim& dim1,
+                                                const DDim& dim2);
+
 TEST_API DDim ComputeCompatibleDim(const DDim& dim1, const DDim& dim2);
 
 }  // namespace common
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -20,6 +20,7 @@
 
 #include <math.h>
 #include <vector>
+#include "paddle/common/ddim.h"
 #include "paddle/fluid/prim/api/generated_prim/prim_generated_api.h"
 #include "paddle/fluid/primitive/type/lazy_tensor.h"
 #include "paddle/fluid/primitive/utils/utils.h"
@@ -655,37 +656,36 @@ void add_grad(const Tensor& x,
               Tensor* dx,
               Tensor* dy) {
   if (dy) {
-    if (has_dynamic_shape(y.shape()) || has_dynamic_shape(out_grad.shape())) {
-      auto dy_tmp = reduce_as<T>(out_grad, y);
-      set_output<T>(dy_tmp, dy);
-    } else {
-      if (out_grad.dims() != y.dims()) {
+    if (!common::AreDimsWithDynamicShapeCompatible(out_grad.dims(), y.dims())) {
+      if (has_dynamic_shape(y.shape()) || has_dynamic_shape(out_grad.shape())) {
+        auto dy_tmp = reduce_as<T>(out_grad, y);
+        set_output<T>(dy_tmp, dy);
+      } else {
         phi::DDim reduce_dim =
             get_reduce_dims_from_out(out_grad.dims(), y.dims());
         auto dy_reduce_res =
             out_grad.sum(common::vectorize(reduce_dim), y.dtype(), false);
         auto dy_tmp = reshape<T>(dy_reduce_res, common::vectorize(y.dims()));
         set_output<T>(dy_tmp, dy);
-
-      } else {
-        by_pass<T>(out_grad, dy);
       }
+    } else {
+      by_pass<T>(out_grad, dy);
     }
   }
   if (dx) {
-    if (has_dynamic_shape(x.shape()) || has_dynamic_shape(out_grad.shape())) {
-      auto dx_tmp = reduce_as<T>(out_grad, x);
-      set_output<T>(dx_tmp, dx);
-    } else {
-      if (out_grad.dims() != x.dims()) {
+    if (!common::AreDimsWithDynamicShapeCompatible(out_grad.dims(), x.dims())) {
+      if (has_dynamic_shape(x.shape()) || has_dynamic_shape(out_grad.shape())) {
+        auto dx_tmp = reduce_as<T>(out_grad, x);
+        set_output<T>(dx_tmp, dx);
+      } else {
         auto reduce_dim = get_reduce_dims_from_out(out_grad.dims(), x.dims());
         auto dx_reduce_res =
             out_grad.sum(common::vectorize(reduce_dim), x.dtype(), false);
         auto dx_tmp = reshape<T>(dx_reduce_res, common::vectorize(x.dims()));
         set_output<T>(dx_tmp, dx);
-      } else {
-        by_pass<T>(out_grad, dx);
       }
+    } else {
+      by_pass<T>(out_grad, dx);
     }
   }
 }
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3883,6 +3883,7 @@
     func : reduce_as
     data_type : x
   backward : reduce_as_grad
+  no_need_buffer : target
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : reduce_scatter
