optimize topk op through limit SortTopK kernel entrance, test=develop (#30403)

4 years ago · 8f4ac6b525
parent bfb8a64234
commit 8f4ac6b525
1 changed files with 2 additions and 1 deletions
--- a/paddle/fluid/operators/top_k_v2_op.cu
+++ b/paddle/fluid/operators/top_k_v2_op.cu
@ -150,7 +150,8 @@ class TopkV2OpCUDAKernel : public framework::OpKernel<T> {

      if (k > input_width) k = input_width;

-      if ((input_width <= 1024 || k >= 128 || k == input_width)) {
+      if (((input_width <= 1024 && input_height <= 2048) || k >= 128 ||
+           k == input_width)) {
        if (SortTopk<T>(dev_ctx, &trans_input, input_width, input_height, k,
                        &trans_out, &trans_ind, largest)) {
          // last step, tranpose back the indices and output