[Auto Parallel] Add spmd rule No.19 for roi_align and roi_align_grad ops.

paddle/phi/infermeta/spmd_rules/roi_align.cc

@@ -0,0 +1,128 @@
+/* Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/infermeta/spmd_rules/roi_align.h"
+#include "glog/logging.h"
+#include "paddle/phi/infermeta/spmd_rules/spmd_rule_macro_define.h"
+#include "paddle/phi/infermeta/spmd_rules/utils.h"
+
+namespace phi {
+namespace distributed {
+
+SpmdInfo RoiAlignInferSpmd(const DistMetaTensor& x,
+                           const DistMetaTensor& boxes,
+                           const DistMetaTensor& boxes_num,
+                           int pooled_height,
+                           int pooled_width,
+                           float spatial_scale,
+                           int sampling_ratio,
+                           bool aligned) {
+  EXTRACT_SHAPE_AND_DIST_ATTR(x);
+  EXTRACT_SHAPE_AND_DIST_ATTR(boxes);
+
+  std::vector<int64_t> x_dims_mapping_dst(x_ndim, -1);
+  x_dims_mapping_dst[1] = x_dims_mapping_src[1];
+  std::vector<int64_t> boxes_dims_mapping_dst(boxes_ndim, -1);
+
+  std::vector<int64_t> boxes_num_dims_mapping_dst;
+  TensorDistAttr boxes_num_dist_attr_dst;
+
+  TensorDistAttr x_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
+  TensorDistAttr boxes_dist_attr_dst =
+      CopyTensorDistAttrForOutput(boxes_dist_attr_src);
+
+  TensorDistAttr out_dist_attr_dst =
+      CopyTensorDistAttrForOutput(x_dist_attr_dst);
+
+  x_dist_attr_dst.set_dims_mapping(x_dims_mapping_dst);
+  boxes_dist_attr_dst.set_dims_mapping(boxes_dims_mapping_dst);
+  out_dist_attr_dst.set_dims_mapping(x_dims_mapping_dst);
+  if (boxes_num.initialized()) {
+    EXTRACT_SHAPE_AND_DIST_ATTR(boxes_num);
+    boxes_num_dims_mapping_dst = {-1};
+    boxes_num_dist_attr_dst =
+        CopyTensorDistAttrForOutput(boxes_num_dist_attr_src);
+    boxes_num_dist_attr_dst.set_dims_mapping(boxes_num_dims_mapping_dst);
+    VLOG(4) << "RoiAlignInferSpmd: Done.";
+    LOG_SPMD_INPUT(boxes_num);
+  } else {
+    boxes_num_dist_attr_dst = TensorDistAttr();
+    VLOG(4) << "RoiAlignInferSpmd: Done.";
+  }
+  LOG_SPMD_INPUT(x);
+  LOG_SPMD_INPUT(boxes);
+  LOG_SPMD_OUTPUT(out_dist_attr_dst);
+
+  return {{x_dist_attr_dst, boxes_dist_attr_dst, boxes_num_dist_attr_dst},
+          {out_dist_attr_dst}};
+}
+
+SpmdInfo RoiAlignGradInferSpmd(const DistMetaTensor& x,
+                               const DistMetaTensor& boxes,
+                               const DistMetaTensor& boxes_num,
+                               const DistMetaTensor& out_grad,
+                               int pooled_height,
+                               int pooled_width,
+                               float spatial_scale,
+                               int sampling_ratio,
+                               bool aligned) {
+  EXTRACT_SHAPE_AND_DIST_ATTR(x);
+  EXTRACT_SHAPE_AND_DIST_ATTR(boxes);
+  EXTRACT_SHAPE_AND_DIST_ATTR(out_grad);
+  int64_t c_status = ShardingMergeForAxis(
+      "c", x_dims_mapping_src[1], out_grad_dims_mapping_src[1]);
+  std::vector<int64_t> x_dims_mapping_dst(x_ndim, -1);
+  x_dims_mapping_dst[1] = c_status;
+  std::vector<int64_t> boxes_dims_mapping_dst(boxes_ndim, -1);
+  std::vector<int64_t> out_grad_dims_mapping_dst(out_grad_ndim, -1);
+  out_grad_dims_mapping_dst[1] = c_status;
+  TensorDistAttr x_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
+  TensorDistAttr x_grad_attr_dst = x_dist_attr_dst;
+  TensorDistAttr boxes_dist_attr_dst =
+      CopyTensorDistAttrForOutput(boxes_dist_attr_src);
+  TensorDistAttr out_grad_dist_attr_dst =
+      CopyTensorDistAttrForOutput(out_grad_dist_attr_src);
+  x_dist_attr_dst.set_dims_mapping(x_dims_mapping_dst);
+  x_grad_attr_dst.set_dims_mapping(x_dims_mapping_dst);
+  boxes_dist_attr_dst.set_dims_mapping(boxes_dims_mapping_dst);
+  out_grad_dist_attr_dst.set_dims_mapping(out_grad_dims_mapping_dst);
+
+  TensorDistAttr boxes_num_dist_attr_dst;
+  std::vector<int64_t> boxes_num_dims_mapping_dst;
+  if (boxes_num.initialized()) {
+    EXTRACT_SHAPE_AND_DIST_ATTR(boxes_num);
+    boxes_num_dims_mapping_dst = {-1};
+    boxes_num_dist_attr_dst =
+        CopyTensorDistAttrForOutput(boxes_num_dist_attr_src);
+    boxes_num_dist_attr_dst.set_dims_mapping(boxes_num_dims_mapping_dst);
+    VLOG(4) << "RoiAlignGradInferSpmd: Done.";
+    LOG_SPMD_INPUT(boxes_num);
+  } else {
+    boxes_num_dist_attr_dst = TensorDistAttr();
+    VLOG(4) << "RoiAlignGradInferSpmd: Done.";
+  }
+  LOG_SPMD_INPUT(x);
+  LOG_SPMD_INPUT(boxes);
+  LOG_SPMD_INPUT(out_grad);
+  LOG_SPMD_OUTPUT(x_grad_attr_dst);
+
+  return {{x_dist_attr_dst,
+           boxes_dist_attr_dst,
+           boxes_num_dist_attr_dst,
+           out_grad_dist_attr_dst},
+          {x_grad_attr_dst}};
+}
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/infermeta/spmd_rules/roi_align.h

@@ -0,0 +1,43 @@
+/* Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+
+namespace phi {
+namespace distributed {
+
+SpmdInfo RoiAlignInferSpmd(const DistMetaTensor& x,
+                           const DistMetaTensor& boxes,
+                           const DistMetaTensor& boxes_num,
+                           int pooled_height,
+                           int pooled_width,
+                           float spatial_scale,
+                           int sampling_ratio,
+                           bool aligned);
+
+SpmdInfo RoiAlignGradInferSpmd(const DistMetaTensor& x,
+                               const DistMetaTensor& boxes,
+                               const DistMetaTensor& boxes_num,
+                               const DistMetaTensor& out_grad,
+                               int pooled_height,
+                               int pooled_width,
+                               float spatial_scale,
+                               int sampling_ratio,
+                               bool aligned);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/infermeta/spmd_rules/rules.cc

@@ -785,4 +785,9 @@ PD_REGISTER_SPMD_RULE(
     put_along_axis,
     PD_INFER_SPMD(phi::distributed::PutAlongAxisInferSpmd),
     PD_INFER_SPMD(phi::distributed::PutAlongAxisGradInferSpmd));
+
+// roi_align
+PD_REGISTER_SPMD_RULE(roi_align,
+                      PD_INFER_SPMD(phi::distributed::RoiAlignInferSpmd),
+                      PD_INFER_SPMD(phi::distributed::RoiAlignGradInferSpmd));
 }  // namespace phi::distributed

paddle/phi/infermeta/spmd_rules/rules.h

@@ -64,6 +64,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/replicated.h"
 #include "paddle/phi/infermeta/spmd_rules/reshape.h"
 #include "paddle/phi/infermeta/spmd_rules/rms_norm.h"
+#include "paddle/phi/infermeta/spmd_rules/roi_align.h"
 #include "paddle/phi/infermeta/spmd_rules/roll.h"
 #include "paddle/phi/infermeta/spmd_rules/scale.h"
 #include "paddle/phi/infermeta/spmd_rules/scatter.h"

paddle/phi/ops/yaml/backward.yaml

@@ -2889,6 +2889,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : RoiAlignGradInferSpmd
   kernel :
     func : roi_align_grad
     data_type : boxes

paddle/phi/ops/yaml/ops.yaml

@@ -4468,6 +4468,7 @@
   output : Tensor
   infer_meta :
     func : RoiAlignInferMeta
+    spmd_rule : RoiAlignInferSpmd
   kernel :
     func : roi_align
     data_type : x

test/auto_parallel/spmd_rules/CMakeLists.txt

@@ -48,6 +48,7 @@ if(WITH_DISTRIBUTE)
     py_test_modules(test_unique_rule MODULES test_unique_rule)
     py_test_modules(test_topk_rule MODULES test_topk_rule)
     py_test_modules(test_add_n_rule MODULES test_add_n_rule)
+    py_test_modules(test_roi_align_rule MODULES test_roi_align_rule)
     py_test_modules(test_mean_all_rule MODULES test_mean_all_rule)
     py_test_modules(test_argmin_rule MODULES test_argmin_rule)
     py_test_modules(test_cummax_rule MODULES test_cummax_rule)