From 27fea8d5aade449dbeb75da8d572990eb44e79f9 Mon Sep 17 00:00:00 2001 From: Sergey Zimin Date: Thu, 21 Nov 2024 16:35:21 -0800 Subject: [PATCH] Add check to ensure that there is enough room in permuted_indices (#3403) Summary: In scope of the MTIA IG Sprint we encountered a crash (T208229934) that turned out to be an incorrectly provided argument for `permuted_lengths_sum` argument. This argument suppose to speed the operation, but actually, the true value it suppose to substitute is computed every time anyway. One radical solution would be to drop the argument entirely, but it requires more thoughtful analysis. This diff just prevent clearly faulty case of allocated `permuted_indices` being less than required by the function logic. Exactly this case was spotted in case of T208229934. Differential Revision: D66274522 --- fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp b/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp index 550d80277e..44aa495006 100644 --- a/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp +++ b/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp @@ -678,6 +678,11 @@ std::tuple> permute_2D_sparse_data_cpu( int64_t permuted_indices_size = 0; if (permuted_lengths_sum.has_value()) { permuted_indices_size = permuted_lengths_sum.value(); + + // Ensure there is enough space. + TORCH_CHECK( + permuted_indices_size >= + output_offsets_per_thread_cumsum[num_threads * FALSE_SHARING_PAD]); } else { permuted_indices_size = output_offsets_per_thread_cumsum[num_threads * FALSE_SHARING_PAD]; @@ -842,6 +847,11 @@ std::tuple> permute_1D_sparse_data_cpu( int64_t permuted_indices_size = 0; if (permuted_lengths_sum.has_value()) { permuted_indices_size = permuted_lengths_sum.value(); + + // Ensure there is enough space. + TORCH_CHECK( + permuted_indices_size >= + output_offsets[permuted_lengths_size].item()); } else { permuted_indices_size = output_offsets[permuted_lengths_size].item();