less time (PaddlePaddle#70442)

Author: zhink

paddle/phi/infermeta/multiary.cc

@@ -4734,18 +4734,13 @@ void RmsNormInferMeta(const MetaTensor& x,
                       MetaTensor* out,
                       MetaTensor* residual_out,
                       MetaTensor* inv_var) {
-  std::vector<int64_t> x_dims_vec = common::vectorize(x.dims());
-  auto x_dims_size = x_dims_vec.size();
+  size_t x_dims_size = x.dims().size();
 
   size_t normalized_dims = 1;
   for (size_t i = begin_norm_axis; i < x_dims_size; ++i) {
-    normalized_dims *= x_dims_vec[i];
+    normalized_dims *= x.dims().at(i);
   }
 
-  std::vector<int64_t> inv_var_dims;
-  for (size_t i = size_t(0); i < static_cast<size_t>(begin_norm_axis); i++) {
-    inv_var_dims.push_back(x_dims_vec[i]);
-  }
   PADDLE_ENFORCE_EQ(normalized_dims,
                     norm_weight.dims()[0],
                     common::errors::InvalidArgument(
@@ -4756,9 +4751,7 @@ void RmsNormInferMeta(const MetaTensor& x,
                         normalized_dims,
                         norm_weight.dims()[0]));
 
-  auto out_dims = common::make_ddim(x_dims_vec);
-
-  out->set_dims(out_dims);
+  out->set_dims(x.dims());
 
   if (quant_scale > 0) {
     if (fabs(quant_max_bound - 127.0f) < 0.000001) {
@@ -4774,12 +4767,16 @@ void RmsNormInferMeta(const MetaTensor& x,
 
   if (inv_var != nullptr) {
     inv_var->set_dtype(phi::DataType::FLOAT32);
+    std::vector<int64_t> inv_var_dims;
+    for (size_t i = size_t(0); i < static_cast<size_t>(begin_norm_axis); i++) {
+      inv_var_dims.push_back(x.dims().at(i));
+    }
     inv_var->set_dims(common::make_ddim(inv_var_dims));
     inv_var->set_layout(x.layout());
   }
 
   if (residual != nullptr) {
-    residual_out->set_dims(out_dims);
+    residual_out->set_dims(x.dims());
     residual_out->set_dtype(x.dtype());
     residual_out->set_layout(x.layout());
     residual_out->share_lod(x);