add index initialization in the block loop for index_sample kernel wh…

…en dealing with a input tensor whose shape is larger than block_dim * grid_dim (#39736)

* add block and grid loop for index_sample kernel to deal with a large-shape tensor

* fix code format

* limit grid dim

* fix the omissive initialization of index_i in the second cycle for index_sample kernel

* fix conflicts

paddle/fluid/operators/index_sample_op.cu

@@ -44,6 +44,7 @@ __global__ void IndexSampleForward(const IndexT* index, const T* in_data,
   unsigned int index_i = blockDim.x * blockIdx.x + threadIdx.x;
   unsigned int index_j = blockDim.y * blockIdx.y + threadIdx.y;
   for (; index_j < batch_size; index_j += blockDim.y * gridDim.y) {
+    index_i = blockDim.x * blockIdx.x + threadIdx.x;
     for (; index_i < index_length; index_i += blockDim.x * gridDim.x) {
       unsigned int index_idx = index_j * index_length + index_i;
       unsigned int in_idx = index_j * input_length + index_i;
@@ -62,6 +63,7 @@ __global__ void IndexSampleGrad(const IndexT* index, T* in_grad,
   unsigned int index_j = blockDim.y * blockIdx.y + threadIdx.y;
 
   for (; index_j < batch_size; index_j += blockDim.y * gridDim.y) {
+    index_i = blockDim.x * blockIdx.x + threadIdx.x;
     for (; index_i < index_length; index_i += blockDim.x * gridDim.x) {
       unsigned int index_idx = index_j * index_length + index_i;
       unsigned int in_idx = index_j * input_length + index_i;