add div plugin and add filter (PaddlePaddle#41243)

douch · Apr 10, 2022 · 5082266 · 5082266
1 parent e2b429b
commit 5082266
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 4 deletions.
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -1007,6 +1007,14 @@ bool OpTeller::Tell(const framework::ir::Node* node, bool use_no_calib_int8,
       auto* y_var_desc = block->FindVar(desc.Input("Y")[0]);
       const auto x_shape = x_var_desc->GetShape();
       const auto y_shape = y_var_desc->GetShape();
+      if (op_type == "elementwise_add" && y_var_desc->Persistable()) {
+        if (y_shape.size() != 1) {
+          return false;
+        }
+        if (y_shape[0] != x_shape[1]) {
+          return false;
+        }
+      }
       if (x_shape.size() == 1 && y_shape.size() == 1) {
         VLOG(3) << "Now trt may not support two 1d tensor elementwise op.";
         return false;

diff --git a/paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.cu b/paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.cu
@@ -30,6 +30,11 @@ template <typename T>
 struct Mul {
   __device__ T operator()(const T &a, const T &b) const { return a * b; }
 };
+
+template <typename T>
+struct Div {
+  __device__ T operator()(const T &a, const T &b) const { return a / b; }
+};
 }  // namespace details
 
 template <typename T, typename Operator>
@@ -130,6 +135,10 @@ int ElementWisePlugin::enqueue(int batch_size, const void *const *inputs,
     elementwise_kernel<<<block, thread, 0, stream>>>(
         num, x, y, out, prev_size_, batch_size * midd_size_, post_size_,
         details::Mul<float>());
+  } else if (type_ == "div") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(
+        num, x, y, out, prev_size_, batch_size * midd_size_, post_size_,
+        details::Div<float>());
   } else {
     PADDLE_THROW(platform::errors::Fatal(
         "The %s type elementwise is not implemented in trt plugin.", type_));
@@ -242,11 +251,15 @@ int ElementwisePluginDynamic::enqueue(
   } else if (type_ == "mul") {
     elementwise_kernel<<<block, thread, 0, stream>>>(
         num, x, y, out, prev_size, midd_size, post_size, details::Mul<float>());
+  } else if (type_ == "div") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(
+        num, x, y, out, prev_size, midd_size, post_size, details::Div<float>());
   } else {
-    PADDLE_THROW(platform::errors::Unimplemented(
-        "Paddle-TRT only support elementwise operation: {add, mul} currently, "
-        "but got %s.",
-        type_));
+    PADDLE_THROW(
+        platform::errors::Unimplemented("Paddle-TRT only support elementwise "
+                                        "operation: {add, mul, div} currently, "
+                                        "but got %s.",
+                                        type_));
   }
 
   return cudaGetLastError() != cudaSuccess;