add logsumexp op spmd rule and its test and register elementwise spmd rule for logical_xxx ops (#67840)

Wennie396 · web-flow · commit c8eb4586eb61 · 2024-09-03T20:00:45.000+08:00
* add logsumexp op spmd rule and its test

* add grad inferspmd for logsumexp

* fix bug

* fix bug

* add spmd rules for logical_and &amp; logical_not &amp; logical_or &amp; logical_xor op
diff --git a/paddle/phi/infermeta/spmd_rules/logsumexp.cc b/paddle/phi/infermeta/spmd_rules/logsumexp.cc
@@ -0,0 +1,53 @@
+// Copyright (c) 2024 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
+//
+//     https://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/logsumexp.h"
+#include "glog/logging.h"
+
+namespace phi {
+namespace distributed {
+
+SpmdInfo LogSumExpInferSpmd(const DistMetaTensor& x,
+                            const std::vector<int>& axis,
+                            bool keepdims,
+                            bool reduce_all) {
+  VLOG(4) << "LogSumExpInferSpmd Call ReductionInferSpmd";
+  std::vector<int64_t> new_axis(axis.begin(), axis.end());
+  return ReductionInferSpmd(x, new_axis, keepdims);
+}
+
+SpmdInfo LogSumExpInferSpmdReverse(const DistMetaTensor& x,
+                                   const DistMetaTensor& out,
+                                   const std::vector<int>& axis,
+                                   bool keepdims,
+                                   bool reduce_all) {
+  VLOG(4) << "LogSumExpInferSpmdReverse Call ReductionInferSpmdReverse";
+  std::vector<int64_t> new_axis(axis.begin(), axis.end());
+  return ReductionInferSpmdReverse(x, out, new_axis, keepdims);
+}
+
+SpmdInfo LogSumExpGradInferSpmd(const DistMetaTensor& x,
+                                const DistMetaTensor& out,
+                                const DistMetaTensor& out_grad,
+                                const std::vector<int>& axis,
+                                bool keepdims,
+                                bool reduce_all) {
+  VLOG(4) << "LogSumExpGradInferSpmd Call ReductionGradInferSpmd";
+  std::vector<int64_t> new_axis(axis.begin(), axis.end());
+  return ReductionGradInferSpmd(
+      x, out, out_grad, new_axis, keepdims, reduce_all);
+}
+
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/logsumexp.h b/paddle/phi/infermeta/spmd_rules/logsumexp.h
@@ -0,0 +1,44 @@
+// Copyright (c) 2024 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
+//
+//     https://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 <vector>
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+#include "paddle/phi/infermeta/spmd_rules/reduction.h"
+
+namespace phi {
+namespace distributed {
+
+SpmdInfo LogSumExpInferSpmd(const DistMetaTensor& x,
+                            const std::vector<int>& axis,
+                            bool keepdims,
+                            bool reduce_all);
+
+SpmdInfo LogSumExpInferSpmdReverse(const DistMetaTensor& x,
+                                   const DistMetaTensor& out,
+                                   const std::vector<int>& axis,
+                                   bool keepdims,
+                                   bool reduce_all);
+
+SpmdInfo LogSumExpGradInferSpmd(const DistMetaTensor& x,
+                                const DistMetaTensor& out,
+                                const DistMetaTensor& out_grad,
+                                const std::vector<int>& axis,
+                                bool keepdims,
+                                bool reduce_all);
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc
@@ -681,5 +681,10 @@ PD_REGISTER_SPMD_RULE(unbind,
                       PD_INFER_SPMD(phi::distributed::UnbindInferSpmd),
                       PD_INFER_SPMD(phi::distributed::UnbindInferSpmdReverse));
 
+PD_REGISTER_SPMD_RULE(
+    logsumexp,
+    PD_INFER_SPMD(phi::distributed::LogSumExpInferSpmd),
+    PD_INFER_SPMD(phi::distributed::LogSumExpInferSpmdReverse));
+
 }  // namespace distributed
 }  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h
@@ -34,6 +34,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/gather.h"
 #include "paddle/phi/infermeta/spmd_rules/gather_nd.h"
 #include "paddle/phi/infermeta/spmd_rules/layer_norm.h"
+#include "paddle/phi/infermeta/spmd_rules/logsumexp.h"
 #include "paddle/phi/infermeta/spmd_rules/matmul.h"
 #include "paddle/phi/infermeta/spmd_rules/numel.h"
 #include "paddle/phi/infermeta/spmd_rules/one_hot.h"
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -1881,6 +1881,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param: [x]
+    spmd_rule : LogSumExpGradInferSpmd
   kernel :
     func : logsumexp_grad
 
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2908,6 +2908,7 @@
   output : Tensor(out)
   infer_meta :
     func : LogicalBinaryInferMeta
+    spmd_rule : ElementwiseBinaryInferSpmd
   kernel :
     func : logical_and
     data_type : x
@@ -2921,6 +2922,7 @@
   output : Tensor(out)
   infer_meta :
     func : LogicalNotInferMeta
+    spmd_rule : ElementwiseUnaryInferSpmd
   kernel :
     func : logical_not
     data_type : x
@@ -2934,6 +2936,7 @@
   output : Tensor(out)
   infer_meta :
     func : LogicalBinaryInferMeta
+    spmd_rule : ElementwiseBinaryInferSpmd
   kernel :
     func : logical_or
     data_type : x
@@ -2947,6 +2950,7 @@
   output : Tensor(out)
   infer_meta :
     func : LogicalBinaryInferMeta
+    spmd_rule : ElementwiseBinaryInferSpmd
   kernel :
     func : logical_xor
     data_type : x
@@ -2995,6 +2999,7 @@
   output : Tensor(out)
   infer_meta :
     func : LogsumexpInferMeta
+    spmd_rule : LogSumExpInferSpmd
   kernel :
     func : logsumexp
   backward : logsumexp_grad
diff --git a/test/auto_parallel/spmd_rules/CMakeLists.txt b/test/auto_parallel/spmd_rules/CMakeLists.txt
@@ -39,6 +39,7 @@ if(WITH_DISTRIBUTE)
   py_test_modules(test_gather_nd_rule MODULES test_gather_nd_rule)
   py_test_modules(test_fused_dropout_add_rule MODULES
                   test_fused_dropout_add_rule)
+  py_test_modules(test_logsumexp_rule MODULES test_logsumexp_rule)
   # End of unittests WITH single card WITHOUT timeout
 
 endif()
diff --git a/test/auto_parallel/spmd_rules/test_logsumexp_rule.py b/test/auto_parallel/spmd_rules/test_logsumexp_rule.py
