Fix the out-of-bounds access in the C+T+dbias kernel (#28)

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

transformer_engine/common/transpose/cast_transpose_fusion.cu

@@ -121,7 +121,7 @@ cast_transpose_dbias_kernel(const Param param,
   extern __shared__ char scratch[];
 
   const int warp_id = threadIdx.x / THREADS_PER_WARP;
-  const int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;
+  const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;
   const size_t num_tiles_x = row_length / (nvec_in * THREADS_PER_WARP);
   // const size_t num_tiles_y = num_rows / (nvec * THREADS_PER_WARP);
   const size_t tile_id = blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) +
@@ -262,7 +262,7 @@ cast_transpose_dbias_kernel_notaligned(const Param param,
   extern __shared__ char scratch[];
 
   const int warp_id = threadIdx.x / THREADS_PER_WARP;
-  const int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;
+  const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;
   const size_t num_tiles_x = (row_length + nvec_in * THREADS_PER_WARP - 1) /
                              (nvec_in * THREADS_PER_WARP);
   const size_t tile_id = blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) +
@@ -399,8 +399,10 @@ cast_transpose_dbias_kernel_notaligned(const Param param,
       }
     }
 
+    if (my_id_in_warp < tile_length) {
       partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);
     }
+  }
 
   /* warp tile amax reduce*/
   max = reduce_max<cast_transpose_num_threads / THREADS_PER_WARP>(max, warp_id);
@@ -630,7 +632,7 @@ cast_transpose_dbias_dgelu_kernel(const Param param,
   extern __shared__ char scratch[];
 
   const int warp_id = threadIdx.x / THREADS_PER_WARP;
-  const int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;
+  const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;
   const size_t num_tiles_x = row_length / (nvec_in * THREADS_PER_WARP);
   // const size_t num_tiles_y = num_rows / (nvec * THREADS_PER_WARP);
   const size_t tile_id = blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) +
@@ -791,7 +793,7 @@ cast_transpose_dbias_dgelu_kernel_notaligned(const Param param,
   extern __shared__ char scratch[];
 
   const int warp_id = threadIdx.x / THREADS_PER_WARP;
-  const int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;
+  const unsigned int my_id_in_warp = threadIdx.x % THREADS_PER_WARP;
   const size_t num_tiles_x = (row_length + nvec_in * THREADS_PER_WARP - 1) /
                              (nvec_in * THREADS_PER_WARP);
   const size_t tile_id = blockIdx.x * blockDim.x / (THREADS_PER_WARP * n_warps_per_tile) +
@@ -948,8 +950,10 @@ cast_transpose_dbias_dgelu_kernel_notaligned(const Param param,
       }
     }
 
+    if (my_id_in_warp < tile_length) {
       partial_dbias.store_to(my_partial_dbias_tile, my_id_in_warp);
     }
+  }
 
   /* warp tile amax reduce*/
   max = reduce_max<cast_transpose_num_threads / THREADS_PER_WARP>(max, warp_id);