Fix possible overflow in WriteCoalescingCallbackWrapper::TearDown (#1…

…7642)

Fixes possible overflow in `WriteCoalescingCallbackWrapper::TearDown` function if the `tile_out_count` is sufficiently large enough. The `out_char += blockDim.x` could overflow when within block-size of the max of `tile_out_count`.

Authors:
  - David Wendt (https://github.com/davidwendt)
  - Yunsong Wang (https://github.com/PointKernel)
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Yunsong Wang (https://github.com/PointKernel)

URL: #17642

cpp/src/io/fst/agent_dfa.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -308,12 +308,14 @@ class WriteCoalescingCallbackWrapper {
   {
     __syncthreads();
     if constexpr (!DiscardTranslatedOutput) {
-      for (uint32_t out_char = threadIdx.x; out_char < tile_out_count; out_char += blockDim.x) {
+      for (thread_index_type out_char = threadIdx.x; out_char < tile_out_count;
+           out_char += blockDim.x) {
         out_it[tile_out_offset + out_char] = temp_storage.compacted_symbols[out_char];
       }
     }
     if constexpr (!DiscardIndexOutput) {
-      for (uint32_t out_char = threadIdx.x; out_char < tile_out_count; out_char += blockDim.x) {
+      for (thread_index_type out_char = threadIdx.x; out_char < tile_out_count;
+           out_char += blockDim.x) {
         out_idx_it[tile_out_offset + out_char] =
           temp_storage.compacted_offset[out_char] + tile_in_offset;
       }