[PHI] Fix expand_as and expand_as_grad kernel for big tensor (#72753)

ggggxm · web-flow · commit 6b8e783254d5 · 2025-05-29T14:15:22.000+08:00
* Fix expand_as and expand_as_grad kernel for big tensor

* config op_patches

* add infer sym shape test for expand_as

* fix expected

* fix expand_as sym shape test

* fix expand_as gpu kernel
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -777,8 +777,8 @@ bool SparseWeightEmbeddingOpInferSymbolicShape(
 
 bool ExpandAsOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
-  std::vector<int> target_shape =
-      paddle::dialect::details::GetVectorAttr<int>(op, "target_shape");
+  std::vector<int64_t> target_shape =
+      paddle::dialect::details::GetVectorAttr<int64_t>(op, "target_shape");
   const std::vector<symbol::DimExpr> &output_dims = [&] {
     const auto &input_shape_or_data =
         infer_context->GetShapeOrDataForValue(op->operand_source(1));
@@ -787,7 +787,7 @@ bool ExpandAsOpInferSymbolicShape(
     }
     std::vector<symbol::DimExpr> output_dims;
     output_dims.reserve(target_shape.size());
-    for (int shape : target_shape) {
+    for (int64_t shape : target_shape) {
       output_dims.push_back(shape);
     }
     return output_dims;
diff --git a/paddle/fluid/pir/serialize_deserialize/patch/0.yaml b/paddle/fluid/pir/serialize_deserialize/patch/0.yaml
@@ -50,3 +50,10 @@ op_patches:
           - type: pir::Int64Attribute
           - type: pir::Int64Attribute
           - type: pir::Int64Attribute
+  - op_name : pd_op.expand_as
+    actions:
+      - action : modify_attr
+        object : target_shape
+        type : pir::ArrayAttribute
+        default :
+          - type: pir::Int64Attribute
diff --git a/paddle/phi/api/generator/dist_api_gen.py b/paddle/phi/api/generator/dist_api_gen.py
@@ -1612,7 +1612,7 @@ def generate_prepare_data_code(self) -> str:
     def get_shape_type(self, attr_info):
         shape_type = "int"
         for name, info in attr_info.items():
-            if "IntArray" in info[0]:
+            if "IntArray" in info[0] or "int64_t" in info[0]:
                 shape_type = "int64_t"
         return shape_type
 
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -1817,7 +1817,7 @@ void CEmbeddingInferMeta(const MetaTensor& weight,
 
 void ExpandAsInferMeta(const MetaTensor& x,
                        const MetaTensor& y,
-                       const std::vector<int>& target_shape,
+                       const std::vector<int64_t>& target_shape,
                        MetaTensor* out) {
 #define MAX_RANK_SUPPORTED 8
   auto x_dims = x.dims();
diff --git a/paddle/phi/infermeta/binary.h b/paddle/phi/infermeta/binary.h
@@ -348,7 +348,7 @@ void CEmbeddingInferMeta(const MetaTensor& weight,
 
 void ExpandAsInferMeta(const MetaTensor& x,
                        const MetaTensor& y,
-                       const std::vector<int>& target_shape,
+                       const std::vector<int64_t>& target_shape,
                        MetaTensor* out);
 
 void FakeDequantizeMaxAbsInferMeta(const MetaTensor& x,
diff --git a/paddle/phi/kernels/expand_as_grad_kernel.h b/paddle/phi/kernels/expand_as_grad_kernel.h
@@ -22,7 +22,7 @@ template <typename T, typename Context>
 void ExpandAsGradKernel(const Context& ctx,
                         const DenseTensor& x,
                         const DenseTensor& out_grad,
-                        const std::vector<int>& target_shape,
+                        const std::vector<int64_t>& target_shape,
                         DenseTensor* in_grad);
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/expand_as_kernel.h b/paddle/phi/kernels/expand_as_kernel.h
@@ -22,7 +22,7 @@ template <typename T, typename Context>
 void ExpandAsKernel(const Context& ctx,
                     const DenseTensor& x,
                     const paddle::optional<DenseTensor>& y,
-                    const std::vector<int>& target_shape,
+                    const std::vector<int64_t>& target_shape,
                     DenseTensor* out);
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/gpu/expand_as_grad_kernel.cu b/paddle/phi/kernels/gpu/expand_as_grad_kernel.cu
@@ -26,7 +26,7 @@ template <typename T, typename Context>
 void ExpandAsGradKernel(const Context& context,
                         const DenseTensor& x,
                         const DenseTensor& out_grad,
-                        const std::vector<int>& target_shape,
+                        const std::vector<int64_t>& target_shape,
                         DenseTensor* in_grad) {
   auto in_dims = x.dims();
   auto out_dims = out_grad.dims();
diff --git a/paddle/phi/kernels/gpu/expand_as_kernel.cu b/paddle/phi/kernels/gpu/expand_as_kernel.cu
@@ -27,17 +27,17 @@ template <typename T, typename Context>
 void ExpandAsKernel(const Context& ctx,
                     const DenseTensor& x,
                     const paddle::optional<DenseTensor>& y,
-                    const std::vector<int>& target_shape_t,
+                    const std::vector<int64_t>& target_shape_t,
                     DenseTensor* out) {
-  std::vector<int> target_shape = target_shape_t;
+  std::vector<int64_t> target_shape = target_shape_t;
 
   if (y.get_ptr()) {
-    target_shape = phi::vectorize<int>(y.get_ptr()->dims());
+    target_shape = phi::vectorize<int64_t>(y.get_ptr()->dims());
   }
 
   int rank = x.dims().size();
   int target_rank = static_cast<int>(target_shape.size());
-  auto vec_in_dims = common::vectorize<int>(x.dims());
+  auto vec_in_dims = common::vectorize<int64_t>(x.dims());
 
   unsigned int diff = target_rank - rank;
   vec_in_dims.insert(vec_in_dims.begin(), diff, 1);
diff --git a/paddle/phi/kernels/impl/expand_as_grad_kernel_impl.h b/paddle/phi/kernels/impl/expand_as_grad_kernel_impl.h
@@ -21,7 +21,7 @@ namespace phi {
 template <typename Context, typename T, int Dims>
 void ExpandAsBackward(const Context& ctx,
                       const DenseTensor& out_grad,
-                      const std::vector<int>& reshape_dims_vec,
+                      const std::vector<int64_t>& reshape_dims_vec,
                       const std::vector<int>& reduce_dims_vec,
                       DenseTensor* in_grad) {
   size_t reshape_size = reshape_dims_vec.size();
@@ -46,7 +46,7 @@ template <typename T, typename Context>
 void ExpandAsGradKernel(const Context& context,
                         const DenseTensor& x,
                         const DenseTensor& out_grad,
-                        const std::vector<int>& target_shape,
+                        const std::vector<int64_t>& target_shape,
                         DenseTensor* in_grad) {
   auto x_dims = x.dims();
 
@@ -55,14 +55,14 @@ void ExpandAsGradKernel(const Context& context,
     return;
   }
 
-  auto vec_in_dims = common::vectorize<int>(x_dims);
+  auto vec_in_dims = common::vectorize<int64_t>(x_dims);
   auto diff = target_shape.size() - vec_in_dims.size();
   vec_in_dims.insert(vec_in_dims.begin(), diff, 1);
-  std::vector<int> repeat_times(vec_in_dims.size());
+  std::vector<int64_t> repeat_times(vec_in_dims.size());
   for (size_t i = 0; i < vec_in_dims.size(); ++i) {
     repeat_times[i] = target_shape[i] / vec_in_dims[i];
   }
-  std::vector<int> reshape_dims_vec;
+  std::vector<int64_t> reshape_dims_vec;
   std::vector<int> reduce_dims_vec;
   for (size_t i = 0; i < repeat_times.size(); ++i) {
     reduce_dims_vec.push_back(reshape_dims_vec.size());
diff --git a/paddle/phi/kernels/impl/expand_as_kernel_impl.h b/paddle/phi/kernels/impl/expand_as_kernel_impl.h
@@ -27,13 +27,13 @@ namespace phi {
 template <typename Context, typename T, int Rank>
 void ExpandAs(const Context& context,
               const DenseTensor& x,
-              const std::vector<int>& target_shape,
+              const std::vector<int64_t>& target_shape,
               DenseTensor* out) {
   auto in_dims = x.dims();
   auto vec_in_dims = common::vectorize<int>(in_dims);
   auto diff = target_shape.size() - vec_in_dims.size();
   vec_in_dims.insert(vec_in_dims.begin(), diff, 1);
-  std::vector<int> repeat_times(vec_in_dims.size());
+  std::vector<int64_t> repeat_times(vec_in_dims.size());
   if (Rank == 0) {
     phi::Copy<Context>(context, x, context.GetPlace(), false, out);
     return;
@@ -98,7 +98,7 @@ template <typename T, typename Context>
 void ExpandAsKernel(const Context& ctx,
                     const DenseTensor& x,
                     const paddle::optional<DenseTensor>& y,
-                    const std::vector<int>& target_shape,
+                    const std::vector<int64_t>& target_shape,
                     DenseTensor* out) {
   auto rank = x.dims().size();
   auto target_rank = target_shape.size();
@@ -124,12 +124,12 @@ void ExpandAsKernel(const Context& ctx,
                         target_rank,
                         MAX_RANK_SUPPORTED));
 
-  std::vector<int> real_target_shape = target_shape;
+  std::vector<int64_t> real_target_shape = target_shape;
   for (size_t i = 0; i < target_shape.size(); ++i) {
     if (target_shape[i] == -1) {
       if (y) {
         if (y->IsInitialized()) {
-          real_target_shape = common::vectorize<int>(y->dims());
+          real_target_shape = common::vectorize<int64_t>(y->dims());
         }
       }
       break;
diff --git a/paddle/phi/kernels/impl/solve_kernel_impl.h b/paddle/phi/kernels/impl/solve_kernel_impl.h
@@ -76,15 +76,6 @@ static std::vector<int64_t> get_broadcast_batch_portion(
   return batchPortion;
 }
 
-static inline std::vector<int> convert_to_int_vec(std::vector<int64_t> a) {
-  std::vector<int> ret;
-  for (size_t i = 0; i < a.size(); i++) {
-    ret.emplace_back(static_cast<int>(a[i]));
-  }
-
-  return ret;
-}
-
 // broadcast the batch dimensions of tensor x and tensor y.
 static inline std::tuple<std::vector<int64_t>, std::vector<int64_t>>
 get_broadcast_dims(const Tensor& x, const Tensor& y) {
@@ -150,11 +141,11 @@ static void linalg_solve(const Context& dev_ctx,
   Tensor tmp_x_bc;
 
   phi::ExpandAsKernel<T, Context>(
-      dev_ctx, tmp_x, nullptr, convert_to_int_vec(x_broadcast_dims), &tmp_x_bc);
+      dev_ctx, tmp_x, nullptr, x_broadcast_dims, &tmp_x_bc);
 
   Tensor tmp_y_bc;
   phi::ExpandAsKernel<T, Context>(
-      dev_ctx, tmp_y, nullptr, convert_to_int_vec(y_broadcast_dims), &tmp_y_bc);
+      dev_ctx, tmp_y, nullptr, y_broadcast_dims, &tmp_y_bc);
 
   auto x_dim = x.dims();
   auto y_dim = y.dims();
diff --git a/paddle/phi/kernels/xpu/expand_as_kernel.cc b/paddle/phi/kernels/xpu/expand_as_kernel.cc
@@ -24,7 +24,7 @@ namespace phi {
 template <typename Context, typename T>
 void ExpandAs(const Context& context,
               const DenseTensor& x,
-              const std::vector<int>& target_shape_,
+              const std::vector<int64_t>& target_shape_,
               DenseTensor* out) {
   using XPUType = typename XPUTypeTrait<T>::Type;
   auto vec_in_dims = common::vectorize<int64_t>(x.dims());
@@ -87,7 +87,7 @@ template <typename T, typename Context>
 void ExpandAsKernel(const Context& ctx,
                     const DenseTensor& x,
                     const paddle::optional<DenseTensor>& y,
-                    const std::vector<int>& target_shape,
+                    const std::vector<int64_t>& target_shape,
                     DenseTensor* out) {
   auto rank = x.dims().size();
   auto target_rank = target_shape.size();
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -992,8 +992,8 @@
   composite : exp_grad(out, out_grad, x_grad)
 
 - backward_op : expand_as_grad
-  forward : expand_as (Tensor x, Tensor y, int[] target_shape = {}) -> Tensor(out)
-  args : (Tensor x, Tensor out_grad, int[] target_shape)
+  forward : expand_as (Tensor x, Tensor y, int64_t[] target_shape = {}) -> Tensor(out)
+  args : (Tensor x, Tensor out_grad, int64_t[] target_shape)
   output : Tensor(x_grad)
   infer_meta :
     func : UnchangedInferMeta
diff --git a/paddle/phi/ops/yaml/legacy/backward_exclude.yaml b/paddle/phi/ops/yaml/legacy/backward_exclude.yaml
@@ -13,6 +13,7 @@
 - disable_check_model_nan_inf_grad
 - dropout_grad
 - enable_check_model_nan_inf_grad
+- expand_as_grad
 # FIXME(dev): Remove flatten_grad after deprecate flatten_contiguous_range in old IR system.
 # flatten_grad is alias name of flatten_contiguous_range_grad.
 - flatten_grad
diff --git a/paddle/phi/ops/yaml/legacy/ops_exclude.yaml b/paddle/phi/ops/yaml/legacy/ops_exclude.yaml
@@ -28,6 +28,7 @@
 - empty
 - empty_like
 - enable_check_model_nan_inf
+- expand_as
 - exponential_
 - eye
 # FIXME(dev): Remove flatten after deprecate flatten_contiguous_range in old IR system.
diff --git a/paddle/phi/ops/yaml/legacy/static_backward.yaml b/paddle/phi/ops/yaml/legacy/static_backward.yaml
@@ -112,6 +112,17 @@
     data_type : out_grad
   no_need_buffer : weight
 
+- backward_op : expand_as_grad
+  forward : expand_as (Tensor x, Tensor y, int[] target_shape = {}) -> Tensor(out)
+  args : (Tensor x, Tensor out_grad, int[] target_shape)
+  output : Tensor(x_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param : [x]
+  kernel :
+    func : expand_as_grad
+  no_need_buffer : x
+
 - backward_op : exponential__grad
   forward : exponential_ (Tensor x, float lam=1.0f) -> Tensor(out)
   args : (Tensor out_grad)
diff --git a/paddle/phi/ops/yaml/legacy/static_ops.yaml b/paddle/phi/ops/yaml/legacy/static_ops.yaml
@@ -242,6 +242,19 @@
     force_backend : force_cpu
   traits : paddle::dialect::ForwardOnlyTrait
 
+- op : expand_as
+  args : (Tensor x, Tensor y, int[] target_shape = {})
+  output : Tensor(out)
+  infer_meta :
+    func : ExpandAsInferMeta
+    local_shape: out
+  kernel :
+    func : expand_as
+    data_type : x
+  optional : y
+  backward : expand_as_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
+
 - op : exponential_
   args : (Tensor x, float lam = 1.0f)
   output : Tensor(out)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1781,7 +1781,7 @@
   backward : expand_grad
 
 - op : expand_as
-  args : (Tensor x, Tensor y, int[] target_shape = {})
+  args : (Tensor x, Tensor y, int64_t[] target_shape = {})
   output : Tensor(out)
   infer_meta :
     func : ExpandAsInferMeta
diff --git a/test/ir/pir/cinn/symbolic/test_infer_sym_shape_binary_op.py b/test/ir/pir/cinn/symbolic/test_infer_sym_shape_binary_op.py
@@ -23,6 +23,7 @@
 )
 
 import paddle
+from paddle import _C_ops
 from paddle.static import InputSpec
 
 sys.path.append(dirname(dirname(__file__)))
@@ -67,6 +68,35 @@ def test_eval_symbolic(self):
         return out
 
 
+class ExpandAsNet(paddle.nn.Layer):
+    def __init__(self, target_shape):
+        super().__init__()
+        self.target_shape = target_shape
+
+    def forward(self, x):
+        return _C_ops.expand_as(x, None, self.target_shape)
+
+
+class ExpandAsOpInferSymbolicShapeTest(TestBase):
+    def prepare_data(self):
+        self.target_shape = [10, 3]
+
+        self.expected = ['shape[10, 3], data[NULL]']
+
+    def test_eval_symbolic(self):
+        net = ExpandAsNet(self.target_shape)
+        input_spec = [
+            InputSpec(shape=[None, 1], dtype='float32'),  # x
+        ]
+        net = apply_to_static(net, False, input_spec)
+        net.eval()
+
+        check_infer_results(
+            net, input_spec, 'builtin.shadow_output', self.expected
+        )
+        return True
+
+
 class KronNet(paddle.nn.Layer):
     def __init__(self):
         super().__init__()
