     int remaining_codewords = 0;
     
     // Select Segment
-    int segment_index = blockIdx.x * WARPS_NUM + warpidx;
+    const int block_idx = blockIdx.x + blockIdx.y * gridDim.x;
+    const int segment_index = block_idx * WARPS_NUM + warpidx;
     if ( segment_index >= segment_count )
         return;
     

    uint4 * const s_temp = s_temp_all + threadIdx.x * 2;
    
    // Select Segment
    int segment_index = blockIdx.x * SERIALIZATION_THREADS_PER_TBLOCK + threadIdx.x;
    if ( segment_index >= segment_count )
        return;
    

    
    // Run output compaction kernel (one warp per segment)
    const dim3 compaction_thread(32, WARPS_NUM);
    const dim3 compaction_grid = gpujpeg_huffman_gpu_encoder_grid_size(gpujpeg_div_and_round_up(coder->segment_count, WARPS_NUM));
    gpujpeg_huffman_encoder_compaction_kernel<<<compaction_grid, compaction_thread>>>(
        coder->d_segment,
        coder->segment_count,
