[NPU] reduce_prod modify (PaddlePaddle#234)

Author: Aganlengzi

backends/npu/kernels/reduce_prod_kernel.cc

@@ -25,13 +25,15 @@ void ProdRawKernel(const Context& dev_ctx,
                    bool reduce_all,
                    phi::DenseTensor* out) {
   auto dims = axes.GetData();
+  auto x_dims = x.dims();
+  auto x_dims_size = x_dims.size();
   dev_ctx.template Alloc<T>(out);
 
   NPUAttributeMap attr_input = {{"axes", dims}, {"keep_dims", keep_dim}};
 
   if (reduce_all) {
     std::vector<int> dim_vec;
-    for (int i = 0; i < x.dims().size(); i++) {
+    for (int i = 0; i < x_dims_size; i++) {
       dim_vec.push_back(i);
     }
 
@@ -56,8 +58,37 @@ void ProdRawKernel(const Context& dev_ctx,
                              {phi::DenseTensorMeta::DataType::INT32},
                              {phi::DenseTensorMeta::DataType::INT32});
   } else {
-    const auto& runner = NpuOpRunner("ReduceProdD", {x}, {*out}, attr_input);
-    runner.Run(dev_ctx.stream());
+    // TODO(Aganlengzi): remove this branch when performance of ReduceProdD
+    // is good enough for big shapes.
+    // Here, we use SplitV and Mul to deal with special cases.
+    if (x_dims[x_dims_size - 1] == 2 && dims.size() == 1 &&
+        (dims[0] == -1 || dims[0] == x_dims_size - 1)) {
+      auto stream = dev_ctx.stream();
+      phi::DenseTensor x1, x2;
+      x1.set_meta(out->meta());
+      x2.set_meta(out->meta());
+      dev_ctx.template Alloc<T>(&x1);
+      dev_ctx.template Alloc<T>(&x2);
+      // split
+      std::vector<phi::DenseTensor> outputs;
+      outputs.push_back(x1);
+      outputs.push_back(x2);
+      std::vector<int> sections = {1, 1};
+      NpuOpRunner runner_split;
+      runner_split.SetType("SplitV")
+          .AddInput(x)
+          .AddInput(dev_ctx, std::move(sections))
+          .AddInput(dev_ctx, std::vector<int32_t>({-1}))
+          .AddOutputs(outputs)
+          .AddAttrs({{"num_split", static_cast<int32_t>(sections.size())}})
+          .Run(stream);
+      // elementwise mul
+      const auto& runner = NpuOpRunner("Mul", {x1, x2}, {*out}, {});
+      runner.Run(stream);
+    } else {
+      const auto& runner = NpuOpRunner("ReduceProdD", {x}, {*out}, attr_input);
+      runner.Run(dev_ctx.stream());
+    }
   }
 }
 

backends/npu/tests/unittests/test_reduce_prod_op_npu.py

@@ -71,6 +71,19 @@ def setUp(self):
         self.outputs = {'Out': self.inputs['X'].prod(axis=tuple([0]))}
 
 
+class TestNPUReduceProd4(TestNPUReduceProd):
+    def setUp(self):
+        self.op_type = "reduce_prod"
+        self.set_npu()
+        self.init_dtype()
+
+        self.inputs = {
+            'X': np.random.random((32, 888, 50, 2)).astype(self.dtype)
+        }
+        self.attrs = {'dim': [-1]}
+        self.outputs = {'Out': self.inputs['X'].prod(axis=tuple([-1]))}
+
+
 class TestNPUReduceProd6D(TestNPUReduceProd):
     def setUp(self):
         self.op_type = "reduce_prod"