[Auto Parallel] Replace paddle::get usage in spmd_rules dir (#74543)

* Replace paddle::get usage in spmd_rules dir

* Fix bug

Author: ooooo-create

paddle/phi/infermeta/spmd_rules/fused_dropout_add.cc

@@ -36,7 +36,8 @@ SpmdInfo FusedDropoutAddSpmdBase(const DistMetaTensor& x,
   VLOG(4) << "x dist_attr: [" << x.dist_attr().to_string() << "]";
   VLOG(4) << "y dist_attr: [" << y.dist_attr().to_string() << "]";
   VLOG(4) << "out dist_attr: ["
-          << paddle::get<0>(out_info.second[0]).to_string() << "]";
+          << PADDLE_GET_CONST(TensorDistAttr, out_info.second[0]).to_string()
+          << "]";
   VLOG(4) << "seed_offset dist_attr: [" << seed_offset_dist_attr.to_string()
           << "]";
   return {{x.dist_attr(), y.dist_attr()},
@@ -51,9 +52,11 @@ SpmdInfo FusedDropoutAddSpmdReverseBase(const DistMetaTensor& x,
 
   VLOG(4) << "out dist_attr: [" << out.dist_attr().to_string() << "]";
   VLOG(4) << "x dist_attr: ["
-          << paddle::get<0>(reverse_info.first[0]).to_string() << "]";
+          << PADDLE_GET_CONST(TensorDistAttr, reverse_info.first[0]).to_string()
+          << "]";
   VLOG(4) << "y dist_attr: ["
-          << paddle::get<0>(reverse_info.first[1]).to_string() << "]";
+          << PADDLE_GET_CONST(TensorDistAttr, reverse_info.first[1]).to_string()
+          << "]";
   return {reverse_info.first,
           {reverse_info.second[0], seed_offset.dist_attr()}};
 }

paddle/phi/infermeta/spmd_rules/index_select.cc

@@ -98,10 +98,12 @@ SpmdInfo IndexSelectGradInferSpmd(const DistMetaTensor& x,
                         out_grad_ndim));
   // now use forward spmd rule to reduce complexity without actual cost eval.
   SpmdInfo fwd_spmd_info = IndexSelectInferSpmd(x, index, axis);
-  TensorDistAttr x_dist_attr_dst = paddle::get<0>(fwd_spmd_info.first[0]);
-  TensorDistAttr index_dist_attr_dst = paddle::get<0>(fwd_spmd_info.first[1]);
-  TensorDistAttr out_grad_dist_attr_dst =
-      paddle::get<0>(fwd_spmd_info.second[0]);
+  const TensorDistAttr& x_dist_attr_dst =
+      PADDLE_GET_CONST(TensorDistAttr, fwd_spmd_info.first[0]);
+  const TensorDistAttr& index_dist_attr_dst =
+      PADDLE_GET_CONST(TensorDistAttr, fwd_spmd_info.first[1]);
+  const TensorDistAttr& out_grad_dist_attr_dst =
+      PADDLE_GET_CONST(TensorDistAttr, fwd_spmd_info.second[0]);
 
   TensorDistAttr x_grad_dist_attr_dst = x_dist_attr_dst;
   x_grad_dist_attr_dst.clean_partial_status();

paddle/phi/infermeta/spmd_rules/matmul.cc

@@ -291,7 +291,7 @@ SpmdInfo MatmulGradInferSpmd(const DistMetaTensor& x_,
                              bool trans_y) {
   DistMetaTensor x = x_, y = y_;
   auto get_attr = [](const ArgDistAttr& attr) -> const TensorDistAttr& {
-    return paddle::get<TensorDistAttr>(attr);
+    return PADDLE_GET_CONST(TensorDistAttr, attr);
   };
 
   auto confirm_dist_attr_same_fn = [&](const ArgDistAttr& x_dist_attr,

paddle/phi/infermeta/spmd_rules/replicated.cc

@@ -164,15 +164,17 @@ SpmdInfo ReplicatedInferDynamic(
 
   for (int64_t i = 0; i < ninputs; i++) {
     if (paddle::holds_alternative<const DistMetaTensor*>(inputs[i])) {
-      auto dist_meta_tensor_ptr = paddle::get<0>(inputs[i]);
-      auto& dist_meta_tensor = *dist_meta_tensor_ptr;
+      const auto* dist_meta_tensor_ptr =
+          PADDLE_GET_CONST(const DistMetaTensor*, inputs[i]);
+      const auto& dist_meta_tensor = *dist_meta_tensor_ptr;
       auto dist_attr_dst = build_tensor_dist_attr(dist_meta_tensor);
       VLOG(4) << "input " << i << ": dist attr: " << dist_attr_dst.to_string();
       spmd_info.first.emplace_back(dist_attr_dst);
     } else {
       std::vector<phi::distributed::TensorDistAttr> list_dist_attr;
-      auto dist_meta_tensors_ptr = paddle::get<1>(inputs[i]);
-      auto& dist_meta_tensors = *dist_meta_tensors_ptr;
+      const auto* dist_meta_tensors_ptr =
+          PADDLE_GET_CONST(const std::vector<DistMetaTensor>*, inputs[i]);
+      const auto& dist_meta_tensors = *dist_meta_tensors_ptr;
       for (const auto& dist_meta_tensor : dist_meta_tensors) {
         auto dist_attr_dst = build_tensor_dist_attr(dist_meta_tensor);
         VLOG(4) << "input " << i

test/cpp/auto_parallel/spmd_rule_test_util.cc

@@ -22,22 +22,25 @@ const std::vector<int64_t>& get_dims_mapping(
     const phi::distributed::ArgDistAttr& dist_attr) {
   EXPECT_TRUE(
       paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
-  const auto& tensor_attr = paddle::get<0>(dist_attr);
+  const auto& tensor_attr =
+      PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
   return tensor_attr.dims_mapping();
 }
 
 bool is_partial(const phi::distributed::ArgDistAttr& dist_attr) {
   EXPECT_TRUE(
       paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
-  const auto& tensor_attr = paddle::get<0>(dist_attr);
+  const auto& tensor_attr =
+      PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
   return tensor_attr.is_partial();
 }
 
 const std::set<int64_t> get_partial_dims(
     const phi::distributed::ArgDistAttr& dist_attr) {
   EXPECT_TRUE(
       paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
-  const auto& tensor_attr = paddle::get<0>(dist_attr);
+  const auto& tensor_attr =
+      PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
   return tensor_attr.partial_dims();
 }
 
@@ -74,7 +77,8 @@ void check_empty_dist_attr(const phi::distributed::ArgDistAttr& dist_attr,
   EXPECT_TRUE(
       paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr))
       << line;
-  EXPECT_EQ(paddle::get<0>(dist_attr), phi::distributed::TensorDistAttr());
+  EXPECT_EQ(PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr),
+            phi::distributed::TensorDistAttr());
 }
 
 void check_partial_dims(const phi::distributed::ArgDistAttr& dist_attr,
@@ -89,23 +93,23 @@ void check_partial_dims(const phi::distributed::ArgDistAttr& dist_attr,
 void clean_partial_status(phi::distributed::ArgDistAttr* dist_attr) {
   EXPECT_TRUE(
       paddle::holds_alternative<phi::distributed::TensorDistAttr>(*dist_attr));
-  auto& tensor_attr = paddle::get<0>(*dist_attr);
+  auto& tensor_attr = PADDLE_GET(phi::distributed::TensorDistAttr, *dist_attr);
   tensor_attr.clean_partial_status();
 }
 
 void clean_partial_dims(phi::distributed::ArgDistAttr* dist_attr,
                         std::vector<int64_t> dims) {
   EXPECT_TRUE(
       paddle::holds_alternative<phi::distributed::TensorDistAttr>(*dist_attr));
-  auto& tensor_attr = paddle::get<0>(*dist_attr);
+  auto& tensor_attr = PADDLE_GET(phi::distributed::TensorDistAttr, *dist_attr);
   tensor_attr.clean_partial_dims(dims);
 }
 
 void set_partial_status(phi::distributed::ArgDistAttr* dist_attr,
                         std::vector<int64_t> dims) {
   EXPECT_TRUE(
       paddle::holds_alternative<phi::distributed::TensorDistAttr>(*dist_attr));
-  auto& tensor_attr = paddle::get<0>(*dist_attr);
+  auto& tensor_attr = PADDLE_GET(phi::distributed::TensorDistAttr, *dist_attr);
   tensor_attr.set_partial_status(dims);
 }