xpu elementwise_add int32 (#5539)

PaddlePaddle · Mar 9, 2021 · 40c3744 · 40c3744
1 parent c20e240
commit 40c3744
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 11 deletions.
diff --git a/lite/kernels/host/range_compute.cc b/lite/kernels/host/range_compute.cc
@@ -101,6 +101,7 @@ REGISTER_LITE_KERNEL(range, kHost, kInt32, kAny, range_int32, def)
                                        DATALAYOUT(kAny))})
     .Finalize();
 
+#ifdef LITE_BUILD_EXTRA
 // float kernel has higher score when picking kernel.
 using range_int32_f =
     paddle::lite::kernels::host::RangeCompute<int, PRECISION(kFloat)>;
@@ -122,3 +123,4 @@ REGISTER_LITE_KERNEL(range, kHost, kFloat, kAny, range_int32_f, int32)
                                        PRECISION(kInt32),
                                        DATALAYOUT(kAny))})
     .Finalize();
+#endif  // LITE_BUILD_EXTRA
diff --git a/lite/kernels/xpu/elementwise_compute.cc b/lite/kernels/xpu/elementwise_compute.cc
@@ -26,9 +26,10 @@ namespace lite {
 namespace kernels {
 namespace xpu {
 
-void ElementwiseAddCompute::Run() {
-  auto& param = this->Param<param_t>();
-  auto& ctx = this->ctx_->As<XPUContext>();
+template <class T>
+void ElementwiseAddCompute<T>::Run() {
+  auto& param = this->template Param<param_t>();
+  auto& ctx = this->ctx_->template As<XPUContext>();
 
   auto& x_dim = param.X->dims();
   auto& y_dim = param.Y->dims();
@@ -47,12 +48,12 @@ void ElementwiseAddCompute::Run() {
   }
 
   int ret =
-      xdnn::broadcast_add<float>(ctx.GetRawContext(),
-                                 param.X->data<float>(),
-                                 param.Y->data<float>(),
-                                 param.Out->mutable_data<float>(TARGET(kXPU)),
-                                 x_shape,
-                                 y_shape);
+      xdnn::broadcast_add<T>(ctx.GetRawContext(),
+                             param.X->template data<T>(),
+                             param.Y->template data<T>(),
+                             param.Out->template mutable_data<T>(TARGET(kXPU)),
+                             x_shape,
+                             y_shape);
 
   CHECK_EQ(ret, 0);
   return;
@@ -163,13 +164,24 @@ REGISTER_LITE_KERNEL(elementwise_add,
                      kXPU,
                      kFloat,
                      kNCHW,
-                     paddle::lite::kernels::xpu::ElementwiseAddCompute,
-                     def)
+                     paddle::lite::kernels::xpu::ElementwiseAddCompute<float>,
+                     float32)
     .BindInput("X", {LiteType::GetTensorTy(TARGET(kXPU))})
     .BindInput("Y", {LiteType::GetTensorTy(TARGET(kXPU))})
     .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kXPU))})
     .Finalize();
 
+REGISTER_LITE_KERNEL(elementwise_add,
+                     kXPU,
+                     kFloat,
+                     kNCHW,
+                     paddle::lite::kernels::xpu::ElementwiseAddCompute<int>,
+                     int32)
+    .BindInput("X", {LiteType::GetTensorTy(TARGET(kXPU), PRECISION(kInt32))})
+    .BindInput("Y", {LiteType::GetTensorTy(TARGET(kXPU), PRECISION(kInt32))})
+    .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kXPU), PRECISION(kInt32))})
+    .Finalize();
+
 REGISTER_LITE_KERNEL(elementwise_mul,
                      kXPU,
                      kFloat,

diff --git a/lite/kernels/xpu/elementwise_compute.h b/lite/kernels/xpu/elementwise_compute.h
@@ -21,6 +21,7 @@ namespace lite {
 namespace kernels {
 namespace xpu {
 
+template <class T>
 class ElementwiseAddCompute
     : public KernelLite<TARGET(kXPU), PRECISION(kFloat)> {
  public: