Fix past-the-end reads in cub::DeviceSpmv indexing calculations.

This extends the solution in PR NVIDIA#160.
alliepiper · Jul 30, 2021 · 72e8c79 · 72e8c79
1 parent 57528d4
commit 72e8c79
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/cub/agent/agent_spmv_orig.cuh b/cub/agent/agent_spmv_orig.cuh
@@ -294,7 +294,10 @@ struct AgentSpmv
         // Gather the row end-offsets for the merge tile into shared memory
         for (int item = threadIdx.x; item <= tile_num_rows; item += BLOCK_THREADS)
         {
-            s_tile_row_end_offsets[item] = wd_row_end_offsets[tile_start_coord.x + item];
+            const OffsetT offset =
+              (cub::min)(static_cast<OffsetT>(tile_start_coord.x + item),
+                         static_cast<OffsetT>(spmv_params.num_rows - 1));
+            s_tile_row_end_offsets[item] = wd_row_end_offsets[offset];
         }
 
         CTA_SYNC();
@@ -473,7 +476,10 @@ struct AgentSpmv
         #pragma unroll 1
         for (int item = threadIdx.x; item <= tile_num_rows; item += BLOCK_THREADS)
         {
-            s_tile_row_end_offsets[item] = wd_row_end_offsets[tile_start_coord.x + item];
+            const OffsetT offset =
+              (cub::min)(static_cast<OffsetT>(tile_start_coord.x + item),
+                         static_cast<OffsetT>(spmv_params.num_rows - 1));
+            s_tile_row_end_offsets[item] = wd_row_end_offsets[offset];
         }
 
         CTA_SYNC();