custom op spmd rule register

Author: liuzhenhai93

cmake/inference_lib.cmake

@@ -328,6 +328,13 @@ copy(
   inference_lib_dist
   SRCS ${PADDLE_SOURCE_DIR}/paddle/phi/core/visit_type.h
   DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/include/paddle/phi/core/)
+
+copy(
+  inference_lib_dist
+  SRCS ${PADDLE_SOURCE_DIR}/paddle/phi/core/distributed/type_defs.h
+  DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/include/paddle/phi/core/distributed/
+)
+
 copy(
   inference_lib_dist
   SRCS ${PADDLE_SOURCE_DIR}/paddle/phi/core/hostdevice.h

paddle/phi/api/ext/op_meta_info.h

@@ -996,8 +996,11 @@ struct TrtGetOutputDimsFuncImpl<Return (*)(Args...), impl_fn> {
 #endif
 
 ////////////////////// Op Meta Info //////////////////////
-class CustomSpmdInferTensorArg;
-class CustomSpmdInferAttrArg;
+
+using CustomSpmdInferTensorArg =
+    paddle::variant<phi::distributed::DistMetaTensor,
+                    std::vector<phi::distributed::DistMetaTensor>>;
+using CustomSpmdInferAttrArg = paddle::any;
 
 using InferSpmdFunc = phi::distributed::SpmdInfo (*)(
     const std::vector<CustomSpmdInferTensorArg>& inputs,

paddle/phi/core/distributed/type_defs.h

@@ -23,6 +23,7 @@
 namespace phi {
 namespace distributed {
 class TensorDistAttr;
+class DistMetaTensor;
 
 using ArgDistAttr =
     paddle::variant<TensorDistAttr, std::vector<TensorDistAttr>>;

test/cpp/auto_parallel/custom_op_spmd_rule_test.cc

@@ -12,6 +12,7 @@ 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/api/ext/op_meta_info.h"
 #include "paddle/phi/api/ext/spmd_infer.h"
 #include "test/cpp/auto_parallel/spmd_rule_test_util.h"
 
@@ -74,6 +75,15 @@ TEST(CustomOp, Ctor) {
   check_dim_mapping(infered_dist_attrs.second[0], {-1, 1, 0});
   check_partial_dims(infered_dist_attrs.second[0], {});
 }
+
+TEST(CustomOp, Register) {
+  OpMetaInfoBuilder builder("test_custom_op_smpd", 0);
+  auto iter = OpMetaInfoMap::Instance().GetMap().find("test_custom_op_smpd");
+  EXPECT_TRUE(iter != OpMetaInfoMap::Instance().GetMap().end());
+  EXPECT_TRUE(OpMetaInfoHelper::GetInferSpmdFn(iter->second[0]) == nullptr);
+  builder.SetInferSpmdFn(PD_INFER_SPMD_RULE(phi::distributed::ConcatInferSpmd));
+  EXPECT_TRUE(OpMetaInfoHelper::GetInferSpmdFn(iter->second[0]) != nullptr);
+}
 }  // namespace auto_parallel
 }  // namespace distributed
 }  // namespace paddle