[Phi]Add multi_dot/maxout/multiplex op yaml (#41550) (#41818)

* add multi_dot,maxout,multiplex yaml * add code converage
PaddlePaddle · Apr 15, 2022 · 3449a34 · 3449a34
1 parent 6c067e0
commit 3449a34
Show file tree

Hide file tree

Showing 15 changed files with 273 additions and 28 deletions.
diff --git a/paddle/phi/api/lib/api_custom_impl.cc b/paddle/phi/api/lib/api_custom_impl.cc
@@ -1014,5 +1014,135 @@ std::vector<Tensor> meshgrid_grad_impl(
   return api_output;
 }
 
+std::vector<Tensor> multi_dot_grad_impl(const std::vector<Tensor>& x,
+                                        const Tensor& out_grad) {
+  Backend kernel_backend = Backend::UNDEFINED;
+  DataLayout kernel_layout = DataLayout::UNDEFINED;
+  DataType kernel_data_type = DataType::UNDEFINED;
+
+  if (kernel_backend == Backend::UNDEFINED ||
+      kernel_layout == DataLayout::UNDEFINED ||
+      kernel_data_type == DataType::UNDEFINED) {
+    auto kernel_key_set = ParseKernelKeyByInputArgs(x, out_grad);
+    auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
+    if (kernel_backend == Backend::UNDEFINED) {
+      kernel_backend = kernel_key.backend();
+    }
+    if (kernel_layout == DataLayout::UNDEFINED) {
+      kernel_layout = kernel_key.layout();
+    }
+    if (kernel_data_type == DataType::UNDEFINED) {
+      kernel_data_type = kernel_key.dtype();
+    }
+  }
+
+  VLOG(6) << "multi_dot_grad API kernel key: [" << kernel_backend << ", "
+          << kernel_layout << ", " << kernel_data_type << "]";
+  const auto& kernel = phi::KernelFactory::Instance().SelectKernelOrThrowError(
+      "multi_dot_grad", {kernel_backend, kernel_layout, kernel_data_type});
+  VLOG(6) << "multi_dot_grad API kernel: " << kernel;
+
+  auto* dev_ctx = GetDeviceContextByBackend(kernel_backend);
+
+  auto input_x_vec = PrepareData(x, kernel.InputAt(0), {});
+  std::vector<const phi::DenseTensor*> input_x(input_x_vec->size());
+  for (size_t i = 0; i < input_x.size(); ++i) {
+    input_x[i] = &input_x_vec->at(i);
+  }
+  auto input_out_grad = PrepareData(out_grad, kernel.InputAt(1), {});
+
+  size_t out_number = input_x.size();
+  std::vector<Tensor> api_output;
+  auto kernel_out = SetKernelOutput(out_number, kernel_backend, &api_output);
+
+  auto x_meta_vec = MakeMetaTensor(input_x);
+  std::vector<phi::MetaTensor*> x_metas(x_meta_vec.size());
+  for (size_t i = 0; i < x_meta_vec.size(); ++i) {
+    x_metas[i] = &x_meta_vec[i];
+  }
+
+  std::vector<phi::MetaTensor> meta_outs;
+  meta_outs.reserve(out_number);
+  std::vector<phi::MetaTensor*> meta_out_ptrs;
+  meta_out_ptrs.reserve(out_number);
+  for (size_t i = 0; i < out_number; ++i) {
+    meta_outs.push_back(kernel_out[i]);
+    meta_out_ptrs.push_back(&meta_outs.back());
+  }
+
+  phi::MultiDotGradInferMeta(
+      x_metas, MakeMetaTensor(*input_out_grad), meta_out_ptrs);
+
+  using kernel_signature = void (*)(const platform::DeviceContext&,
+                                    const std::vector<const phi::DenseTensor*>&,
+                                    const phi::DenseTensor&,
+                                    std::vector<phi::DenseTensor*>&);
+  auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
+  (*kernel_fn)(*dev_ctx, input_x, *input_out_grad, kernel_out);
+
+  return api_output;
+}
+
+std::vector<Tensor> multiplex_grad_impl(const std::vector<Tensor>& inputs,
+                                        const Tensor& ids,
+                                        const Tensor& out_grad) {
+  Backend kernel_backend = Backend::UNDEFINED;
+  DataLayout kernel_layout = DataLayout::UNDEFINED;
+  DataType kernel_data_type = DataType::UNDEFINED;
+
+  if (kernel_backend == Backend::UNDEFINED ||
+      kernel_layout == DataLayout::UNDEFINED ||
+      kernel_data_type == DataType::UNDEFINED) {
+    auto kernel_key_set = ParseKernelKeyByInputArgs(out_grad);
+    auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
+    if (kernel_backend == Backend::UNDEFINED) {
+      kernel_backend = kernel_key.backend();
+    }
+    if (kernel_layout == DataLayout::UNDEFINED) {
+      kernel_layout = kernel_key.layout();
+    }
+    if (kernel_data_type == DataType::UNDEFINED) {
+      kernel_data_type = kernel_key.dtype();
+    }
+  }
+
+  VLOG(6) << "multiplex_grad API kernel key: [" << kernel_backend << ", "
+          << kernel_layout << ", " << kernel_data_type << "]";
+  const auto& kernel = phi::KernelFactory::Instance().SelectKernelOrThrowError(
+      "multiplex_grad", {kernel_backend, kernel_layout, kernel_data_type});
+  VLOG(6) << "multiplex_grad API kernel: " << kernel;
+
+  auto* dev_ctx = GetDeviceContextByBackend(kernel_backend);
+
+  auto input_ids = PrepareData(ids, kernel.InputAt(0), {});
+  auto input_out_grad = PrepareData(out_grad, kernel.InputAt(1), {});
+
+  auto out_number = inputs.size();
+  std::vector<Tensor> api_output;
+  auto kernel_out = SetKernelOutput(out_number, kernel_backend, &api_output);
+
+  std::vector<phi::MetaTensor> meta_outs;
+  meta_outs.reserve(out_number);
+  std::vector<phi::MetaTensor*> meta_out_ptrs;
+  meta_out_ptrs.reserve(out_number);
+  for (size_t i = 0; i < out_number; ++i) {
+    meta_outs.push_back(kernel_out[i]);
+    meta_out_ptrs.push_back(&meta_outs.back());
+  }
+
+  phi::MultiplexGradInferMeta(MakeMetaTensor(*input_ids),
+                              MakeMetaTensor(*input_out_grad),
+                              meta_out_ptrs);
+
+  using kernel_signature = void (*)(const platform::DeviceContext&,
+                                    const phi::DenseTensor&,
+                                    const phi::DenseTensor&,
+                                    std::vector<phi::DenseTensor*>&);
+  auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
+  (*kernel_fn)(*dev_ctx, *input_ids, *input_out_grad, kernel_out);
+
+  return api_output;
+}
+
 }  // namespace experimental
 }  // namespace paddle
diff --git a/paddle/phi/api/lib/api_custom_impl.h b/paddle/phi/api/lib/api_custom_impl.h
@@ -62,6 +62,8 @@ std::vector<Tensor> split_impl(const Tensor& x,
                                const IntArray& num_or_sections,
                                const Scalar& axis);
 
+std::vector<Tensor> meshgrid_impl(const std::vector<Tensor>& inputs);
+
 std::tuple<Tensor, Tensor, Tensor> momentum_impl(
     const Tensor& param,
     const Tensor& grad,
@@ -109,9 +111,15 @@ Tensor real_grad_impl(const Tensor& x);
 std::vector<Tensor> stack_grad_impl(const std::vector<Tensor>& x,
                                     const Tensor& out_grad,
                                     int axis);
-std::vector<Tensor> meshgrid_impl(const std::vector<Tensor>& inputs);
 std::vector<Tensor> meshgrid_grad_impl(const std::vector<Tensor>& inputs,
                                        const std::vector<Tensor>& outputs_grad);
 
+std::vector<Tensor> multi_dot_grad_impl(const std::vector<Tensor>& x,
+                                        const Tensor& out_grad);
+
+std::vector<Tensor> multiplex_grad_impl(const std::vector<Tensor>& inputs,
+                                        const Tensor& ids,
+                                        const Tensor& out_grad);
+
 }  // namespace experimental
 }  // namespace paddle
diff --git a/paddle/phi/infermeta/backward.cc b/paddle/phi/infermeta/backward.cc
@@ -329,6 +329,38 @@ void MeshgridGradInferMeta(const std::vector<MetaTensor*>& inputs,
   }
 }
 
+void MultiDotGradInferMeta(const std::vector<MetaTensor*>& x,
+                           const MetaTensor& out_grad,
+                           std::vector<MetaTensor*> x_grad) {
+  PADDLE_ENFORCE_EQ(
+      x.size(),
+      x_grad.size(),
+      errors::InvalidArgument(
+          "Number of Inputs(X) should be equal with Outputs(X@Grad)."
+          "But received Inputs(X)' size = %d , Outputs(X@Grad)' size = %d.",
+          x.size(),
+          x_grad.size()));
+  for (size_t i = 0; i < x.size(); i++) {
+    if (x_grad[i] != nullptr) {
+      x_grad[i]->set_dims(x[i]->dims());
+      x_grad[i]->share_lod(*x[i]);
+    }
+  }
+}
+
+void MultiplexGradInferMeta(const MetaTensor& ids,
+                            const MetaTensor& out_grad,
+                            std::vector<MetaTensor*> ins_grad) {
+  PADDLE_ENFORCE_NE(
+      ins_grad.empty(),
+      true,
+      errors::InvalidArgument("Output(X@Grad) should not be null."));
+  auto dout_dim = out_grad.dims();
+  for (auto in_grad : ins_grad) {
+    in_grad->set_dims(dout_dim);
+  }
+}
+
 void NllLossGradInferMeta(const MetaTensor& x,
                           const MetaTensor& label,
                           paddle::optional<const MetaTensor&> weight,

diff --git a/paddle/phi/infermeta/backward.h b/paddle/phi/infermeta/backward.h
@@ -155,6 +155,14 @@ void MeshgridGradInferMeta(const std::vector<MetaTensor*>& inputs,
                            const std::vector<MetaTensor*>& outputs_grad,
                            std::vector<MetaTensor*> inputs_grad);
 
+void MultiDotGradInferMeta(const std::vector<MetaTensor*>& x,
+                           const MetaTensor& out_grad,
+                           std::vector<MetaTensor*> x_grad);
+
+void MultiplexGradInferMeta(const MetaTensor& ids,
+                            const MetaTensor& out_grad,
+                            std::vector<MetaTensor*> ins_grad);
+
 void NllLossGradInferMeta(const MetaTensor& input,
                           const MetaTensor& label,
                           paddle::optional<const MetaTensor&> weight,

diff --git a/paddle/phi/kernels/impl/multi_dot_kernel_impl.h b/paddle/phi/kernels/impl/multi_dot_kernel_impl.h
@@ -339,8 +339,8 @@ void MultiDotGradMatChainOrder(const Context& ctx,
 
 template <typename T, typename Context>
 void MultiDotGradKernel(const Context& ctx,
-                        const DenseTensor& out_grad,
                         const std::vector<const DenseTensor*>& x,
+                        const DenseTensor& out_grad,
                         std::vector<DenseTensor*> x_grad) {
   auto ins = x;
   auto dout = out_grad;

diff --git a/paddle/phi/kernels/multi_dot_grad_kernel.h b/paddle/phi/kernels/multi_dot_grad_kernel.h
@@ -20,8 +20,8 @@ namespace phi {
 
 template <typename T, typename Context>
 void MultiDotGradKernel(const Context& ctx,
-                        const DenseTensor& out_grad,
                         const std::vector<const DenseTensor*>& x,
+                        const DenseTensor& out_grad,
                         std::vector<DenseTensor*> x_grad);
 
 }  // namespace phi
diff --git a/paddle/phi/ops/compat/multi_dot_sig.cc b/paddle/phi/ops/compat/multi_dot_sig.cc
@@ -19,7 +19,7 @@ namespace phi {
 KernelSignature MultiDotGradOpArgumentMapping(
     const ArgumentMappingContext& ctx) {
   return KernelSignature(
-      "multi_dot_grad", {GradVarName("Out"), "X"}, {}, {GradVarName("X")});
+      "multi_dot_grad", {"X", GradVarName("Out")}, {}, {GradVarName("X")});
 }
 
 }  // namespace phi

diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py
@@ -5971,8 +5971,11 @@ def multiplex(inputs, index, name=None):
             print(res) # [array([[5., 6.], [3., 4.]], dtype=float32)]
 
     """
-    if _non_static_mode():
+
+    if _in_legacy_dygraph():
         return _C_ops.multiplex(index, inputs)
+    if in_dygraph_mode():
+        return _C_ops.final_state_multiplex(inputs, index)
     helper = LayerHelper('multiplex', **locals())
 
     check_type(inputs, 'inputs', (list), 'multiplex')

diff --git a/python/paddle/fluid/tests/unittests/test_maxout_op.py b/python/paddle/fluid/tests/unittests/test_maxout_op.py
@@ -21,6 +21,7 @@
 import paddle.fluid.core as core
 import paddle.nn.functional as F
 from op_test import OpTest
+from paddle.fluid.framework import _test_eager_guard
 
 paddle.enable_static()
 np.random.seed(1)
@@ -38,6 +39,7 @@ def maxout_forward_naive(x, groups, channel_axis):
 class TestMaxOutOp(OpTest):
     def setUp(self):
         self.op_type = "maxout"
+        self.python_api = paddle.nn.functional.maxout
         self.dtype = 'float64'
         self.shape = [3, 6, 2, 4]
         self.groups = 2
@@ -55,10 +57,10 @@ def set_attrs(self):
         pass
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_eager=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out')
+        self.check_grad(['X'], 'Out', check_eager=True)
 
 
 class TestMaxOutOpAxis0(TestMaxOutOp):
@@ -144,6 +146,10 @@ def test_errors(self):
             x_float32 = paddle.fluid.data(name='x_float32', shape=[2, 4, 6, 8])
             self.assertRaises(ValueError, F.maxout, x_float32, 2, 2)
 
+    def test_dygraph_final_state_api(self):
+        with _test_eager_guard():
+            self.test_dygraph_api()
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/python/paddle/fluid/tests/unittests/test_multi_dot_op.py b/python/paddle/fluid/tests/unittests/test_multi_dot_op.py
@@ -18,6 +18,7 @@
 from numpy.linalg import multi_dot
 from op_test import OpTest
 import paddle
+from paddle.fluid.framework import _test_eager_guard
 
 paddle.enable_static()
 
@@ -27,6 +28,7 @@
 class TestMultiDotOp(OpTest):
     def setUp(self):
         self.op_type = "multi_dot"
+        self.python_api = paddle.linalg.multi_dot
         self.dtype = self.get_dtype()
         self.get_inputs_and_outputs()
 
@@ -40,11 +42,11 @@ def get_inputs_and_outputs(self):
         self.outputs = {'Out': multi_dot([self.A, self.B])}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_eager=True)
 
     def test_check_grad(self):
-        self.check_grad(['x0'], 'Out')
-        self.check_grad(['x1'], 'Out')
+        self.check_grad(['x0'], 'Out', check_eager=True)
+        self.check_grad(['x1'], 'Out', check_eager=True)
 
 
 #(A*B)*C
@@ -57,9 +59,9 @@ def get_inputs_and_outputs(self):
         self.outputs = {'Out': multi_dot([self.A, self.B, self.C])}
 
     def test_check_grad(self):
-        self.check_grad(['x0'], 'Out')
-        self.check_grad(['x1'], 'Out')
-        self.check_grad(['x2'], 'Out')
+        self.check_grad(['x0'], 'Out', check_eager=True)
+        self.check_grad(['x1'], 'Out', check_eager=True)
+        self.check_grad(['x2'], 'Out', check_eager=True)
 
 
 #A*(B*C)
@@ -72,9 +74,9 @@ def get_inputs_and_outputs(self):
         self.outputs = {'Out': multi_dot([self.A, self.B, self.C])}
 
     def test_check_grad(self):
-        self.check_grad(['x0'], 'Out')
-        self.check_grad(['x1'], 'Out')
-        self.check_grad(['x2'], 'Out')
+        self.check_grad(['x0'], 'Out', check_eager=True)
+        self.check_grad(['x1'], 'Out', check_eager=True)
+        self.check_grad(['x2'], 'Out', check_eager=True)
 
 
 class TestMultiDotOp4Mat(TestMultiDotOp):
@@ -90,10 +92,10 @@ def get_inputs_and_outputs(self):
         self.outputs = {'Out': multi_dot([self.A, self.B, self.C, self.D])}
 
     def test_check_grad(self):
-        self.check_grad(['x0'], 'Out')
-        self.check_grad(['x1'], 'Out')
-        self.check_grad(['x2'], 'Out')
-        self.check_grad(['x3'], 'Out')
+        self.check_grad(['x0'], 'Out', check_eager=True)
+        self.check_grad(['x1'], 'Out', check_eager=True)
+        self.check_grad(['x2'], 'Out', check_eager=True)
+        self.check_grad(['x3'], 'Out', check_eager=True)
 
 
 class TestMultiDotOpFirst1D(TestMultiDotOp):
@@ -143,9 +145,9 @@ def get_inputs_and_outputs(self):
         self.outputs = {'Out': multi_dot([self.A, self.B, self.C])}
 
     def test_check_grad(self):
-        self.check_grad(['x0'], 'Out')
-        self.check_grad(['x1'], 'Out')
-        self.check_grad(['x2'], 'Out')
+        self.check_grad(['x0'], 'Out', check_eager=True)
+        self.check_grad(['x1'], 'Out', check_eager=True)
+        self.check_grad(['x2'], 'Out', check_eager=True)
 
 
 class TestMultiDotOp4MatLast1D(TestMultiDotOp4Mat):
@@ -260,6 +262,10 @@ def test_dygraph_without_out(self):
         expected_result = np.linalg.multi_dot([input_array1, input_array2])
         self.assertTrue(np.allclose(expected_result, out.numpy()))
 
+    def test_dygraph_final_state_api(self):
+        with _test_eager_guard():
+            self.test_dygraph_without_out()
+
 
 if __name__ == "__main__":
     unittest.main()