 #include "paddle/fluid/platform/device/gpu/gpu_launch_config.h"
 #include "paddle/fluid/platform/device/gpu/gpu_primitives.h"
 
+#define PREDEFINED_BLOCK_SIZE_X 512
+#define PREDEFINED_BLOCK_SIZE 1024
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
 namespace paddle {
 namespace operators {
+
+namespace {
+void LimitGridDim(const framework::ExecutionContext& ctx, dim3* grid_dim) {
+  dim3 max_grid_dim = ctx.template device_context<platform::CUDADeviceContext>()
+                          .GetCUDAMaxGridDimSize();
+  grid_dim->x = grid_dim->x < max_grid_dim.x ? grid_dim->x : max_grid_dim.x;
+  grid_dim->y = grid_dim->y < max_grid_dim.y ? grid_dim->y : max_grid_dim.y;
+}
+}
 
 using Tensor = framework::Tensor;
 using LoDTensor = framework::LoDTensor;

    bool same_data_in_index_row = index_length == 1 ? false : true;

    auto block_width = platform::RoundToPowerOfTwo(index_length);
    auto block_height =
        platform::RoundToPowerOfTwo(index_length * batch_size) / block_width;
    dim3 block_dim(block_width, block_height);
    dim3 grid_dim((index_length + block_dim.x - 1) / block_dim.x,
                  (batch_size + block_dim.y - 1) / block_dim.y);

    math::SetConstant<platform::CUDADeviceContext, T> set_zero;
    auto& dev_ctx = ctx.template device_context<platform::CUDADeviceContext>();

    size_t index_length = index_dim[1];

    auto block_width = platform::RoundToPowerOfTwo(index_length);
    block_width = MIN(block_width, PREDEFINED_BLOCK_SIZE_X);
    int block_height =
        platform::RoundToPowerOfTwo(index_length * batch_size) / block_width;
    block_height = MIN(block_height, PREDEFINED_BLOCK_SIZE / block_width);
    dim3 block_dim(block_width, block_height);
    dim3 grid_dim((index_length + block_dim.x - 1) / block_dim.x,
                  (batch_size + block_dim.y - 1) / block_dim.y);
    LimitGridDim(ctx, &grid_dim);

    if (index_type == framework::proto::VarType::INT64) {
      const int64_t* index_data = index->data<int64_t>();
