From 8f4ac6b5258c33e8d800b4e122014c1c980000d4 Mon Sep 17 00:00:00 2001
From: jiangcheng <thisjiang@qq.com>
Date: Mon, 1 Mar 2021 16:53:56 +0800
Subject: [PATCH] optimize topk op through limit SortTopK kernel entrance,
 test=develop (#30403)

---
 paddle/fluid/operators/top_k_v2_op.cu | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/operators/top_k_v2_op.cu b/paddle/fluid/operators/top_k_v2_op.cu
index a2c97aee92..0f2da4b8f6 100644
--- 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
-- 
GitLab