[x86] reduce_sum supports 5-D and 6-D (#10473)

PaddlePaddle · Mar 25, 2024 · 7dbf1ee · 7dbf1ee
1 parent 8795a4d
commit 7dbf1ee
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/lite/kernels/x86/reduce_compute.h b/lite/kernels/x86/reduce_compute.h
@@ -14,6 +14,7 @@
 #pragma once
 
 #include <vector>
+
 #include "lite/backends/x86/fluid/eigen.h"
 #include "lite/core/kernel.h"
 #include "lite/core/op_registry.h"
@@ -92,6 +93,15 @@ class ReduceCompute : public KernelLite<TARGET(kX86), PRECISION(kFloat)> {
     } else {
       int ndim = x_dims.size();
       int rdim = dims.size();
+      HANDLE_DIM(6, 5, Functor);
+      HANDLE_DIM(6, 4, Functor);
+      HANDLE_DIM(6, 3, Functor);
+      HANDLE_DIM(6, 2, Functor);
+      HANDLE_DIM(6, 1, Functor);
+      HANDLE_DIM(5, 4, Functor);
+      HANDLE_DIM(5, 3, Functor);
+      HANDLE_DIM(5, 2, Functor);
+      HANDLE_DIM(5, 1, Functor);
       HANDLE_DIM(4, 3, Functor);
       HANDLE_DIM(4, 2, Functor);
       HANDLE_DIM(4, 1, Functor);