adapt softmax spmd rule to phi

Author: pkuzyc

paddle/fluid/distributed/auto_parallel/spmd_rules/rules.h

@@ -22,7 +22,6 @@
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/reduction_spmd_rule.h"
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.h"
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/reshape_spmd_rule.h"
-#include "paddle/fluid/distributed/auto_parallel/spmd_rules/softmax_spmd_rule.h"
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.h"
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/transpose_spmd_rule.h"
 
@@ -141,10 +140,6 @@ REGISTER_SPMD_RULE(replicated, ReplicatedSPMDRule);
 REGISTER_SPMD_RULE(embedding, EmbeddingSPMDRule);
 REGISTER_SPMD_RULE(lookup_table_v2, EmbeddingSPMDRule);
 
-// softmax rule
-REGISTER_SPMD_RULE(softmax, SoftmaxSPMDRule);
-REGISTER_SPMD_RULE(log_softmax, SoftmaxSPMDRule);
-
 // cross_entropy_with_softmax
 REGISTER_SPMD_RULE(cross_entropy_with_softmax, CrossEntropyWithSoftmaxSPMDRule);
 REGISTER_SPMD_RULE(softmax_with_cross_entropy, CrossEntropyWithSoftmaxSPMDRule);

paddle/fluid/distributed/auto_parallel/spmd_rules/softmax_spmd_rule.cc

@@ -53,6 +53,8 @@ AttrType InferSpmdContext::AttrAt(size_t idx) const {
   }
 }
 
+template int InferSpmdContext::AttrAt(size_t idx) const;
+
 template <>
 bool InferSpmdContext::AttrAt(size_t idx) const {
   try {
@@ -88,10 +90,6 @@ std::vector<int> InferSpmdContext::AttrAt(size_t idx) const {
   }
 }
 
-// template const std::vector<int64_t>& InferSpmdContext::AttrAt(size_t idx)
-// const; template const std::vector<int>& InferSpmdContext::AttrAt(size_t idx)
-// const;
-
 const Attribute& InferSpmdContext::AttrAt(size_t idx) const {
   return attrs_.at(idx);
 }

paddle/phi/core/distributed/auto_parallel/inferspmd_utils.cc

@@ -153,6 +153,7 @@ struct InferSpmdFnImpl<Return (*)(Args...), infer_spmd_fn> {
 
   // TODO(chenweihang): support other attr type later as needed
   PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(bool);
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(int);
   PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_CONST_ATTRIBUTE_REF(std::vector<int>);
   PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_CONST_ATTRIBUTE_REF(
       std::vector<int64_t>);

paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h

@@ -19,6 +19,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/default_data_parallel.h"
 #include "paddle/phi/infermeta/spmd_rules/matmul.h"
 #include "paddle/phi/infermeta/spmd_rules/replicated.h"
+#include "paddle/phi/infermeta/spmd_rules/softmax.h"
 
 /**
  * Design Notes:
@@ -57,5 +58,14 @@ PD_REGISTER_SPMD_RULE(
     PD_INFER_SPMD(phi::distributed::ReplicatedSpmdInferForward),
     PD_INFER_SPMD(phi::distributed::ReplicatedSpmdInferBackward));
 
+// softmax rule
+PD_REGISTER_SPMD_RULE(softmax,
+                      PD_INFER_SPMD(phi::distributed::SoftmaxInferSpmd),
+                      PD_INFER_SPMD(phi::distributed::SoftmaxInferSpmdReverse));
+
+PD_REGISTER_SPMD_RULE(log_softmax,
+                      PD_INFER_SPMD(phi::distributed::SoftmaxInferSpmd),
+                      PD_INFER_SPMD(phi::distributed::SoftmaxInferSpmdReverse));
+
 }  // namespace distributed
 }  // namespace phi