From bfd3fc925d704b826733ad207071aca430fcf9cd Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Mon, 4 Oct 2021 14:48:35 -0700 Subject: [PATCH 1/4] Use transform_iterator --- .../core/providers/cuda/tensor/compress.cc | 9 +++++--- .../providers/cuda/tensor/compress_impl.cu | 20 +++++++++++++--- .../providers/cpu/tensor/compress_op.test.cc | 23 +++++++++++++++++++ 3 files changed, 46 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/compress.cc b/onnxruntime/core/providers/cuda/tensor/compress.cc index 5c9537fdce301..6be6478626427 100644 --- a/onnxruntime/core/providers/cuda/tensor/compress.cc +++ b/onnxruntime/core/providers/cuda/tensor/compress.cc @@ -3,6 +3,7 @@ #include "compress.h" #include "compress_impl.h" +#include namespace onnxruntime { namespace cuda { @@ -48,14 +49,16 @@ Status Compress::ComputeInternal(OpKernelContext* ctx) const { int64_t compress_input_length = has_axis_ ? input_dimensions[axis] : input_size; int64_t valid_condition_length = compress_input_length < condition_length ? compress_input_length : condition_length; - auto condition_cumulative_sum_buffer = GetScratchBuffer(valid_condition_length); + auto condition_cumulative_sum_buffer = GetScratchBuffer(gsl::narrow(valid_condition_length)); auto condition_cumulative_sum = condition_cumulative_sum_buffer.get(); + size_t temp_storage_bytes = 0; CUDA_RETURN_IF_ERROR(CompressCalcPrefixSumTempStorageBytes(Stream(), reinterpret_cast(condition_data), condition_cumulative_sum, - static_cast(valid_condition_length), + gsl::narrow(valid_condition_length), temp_storage_bytes)); + auto temp_buffer = GetScratchBuffer(temp_storage_bytes); auto d_temp_storage = temp_buffer.get(); CUDA_RETURN_IF_ERROR(CompressInclusivePrefixSum(Stream(), @@ -63,7 +66,7 @@ Status Compress::ComputeInternal(OpKernelContext* ctx) const { temp_storage_bytes, reinterpret_cast(condition_data), condition_cumulative_sum, - static_cast(valid_condition_length))); + gsl::narrow(valid_condition_length))); // cudaMemcpyAsync from device memory to pageable host memory will return only once the copy has completed. int32_t positive_condition_count = 0; diff --git a/onnxruntime/core/providers/cuda/tensor/compress_impl.cu b/onnxruntime/core/providers/cuda/tensor/compress_impl.cu index a3212cdf9e300..b3f984596f9f9 100644 --- a/onnxruntime/core/providers/cuda/tensor/compress_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/compress_impl.cu @@ -13,16 +13,30 @@ #include "core/providers/cuda/tensor/compress_impl.h" +#include +#include + namespace onnxruntime { namespace cuda { +// This cast is for transform iterator. This type affects the accumulator type width +// in InclusiveSum() and int8_t overflows. +struct CastToIn32 { + __device__ int32_t operator()(int8_t v) const { + return static_cast(v); + } +}; + cudaError_t CompressCalcPrefixSumTempStorageBytes(cudaStream_t stream, const int8_t* condition_data, int* condition_cumulative_sum, int length, size_t& temp_storage_bytes) { - return cub::DeviceScan::InclusiveSum( - nullptr, temp_storage_bytes, condition_data, condition_cumulative_sum, length, stream); + auto input_iter = thrust::make_transform_iterator(condition_data, CastToIn32()); + return cub::DeviceScan::InclusiveSum( + nullptr, temp_storage_bytes, input_iter, condition_cumulative_sum, length, stream); } + cudaError_t CompressInclusivePrefixSum(cudaStream_t stream, void* d_temp_storage, size_t temp_storage_bytes, const int8_t* condition_data, int* condition_cumulative_sum, int length) { + auto input_iter = thrust::make_transform_iterator(condition_data, CastToIn32()); return cub::DeviceScan::InclusiveSum( - d_temp_storage, temp_storage_bytes, condition_data, condition_cumulative_sum, length, stream); + d_temp_storage, temp_storage_bytes, input_iter, condition_cumulative_sum, length, stream); } template diff --git a/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc b/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc index 7f2782ca06d27..cbc62909d4f72 100644 --- a/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc +++ b/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc @@ -125,6 +125,29 @@ TEST(CompressTest, Compress_default_axis) { test.Run(); } +// Test that we accumulate to a buffer that does not overflow +TEST(CompressTest, Compress_default_axis_issue_9247_cumulative_sum_overflow) { + OpTester test("Compress", 9); + + // Generate input longer than 127 + constexpr size_t elements = 150; + std::vector input; + for (int i = 0; i < elements; ++i) { + input.push_back(static_cast(i)); + } + + // Let's select all of the elements + std::unique_ptr all_true = std::make_unique(elements); + std::fill_n(all_true.get(), elements, true); + std::vector output_shape{static_cast(elements)}; + + test.AddInput("input", {2, 75}, input); + test.AddInput("condition", output_shape, all_true.get(), elements); + // Should get all of the input + test.AddOutput("output", output_shape, input); + test.Run(); +} + TEST(CompressTest, Compress0_string) { OpTester test("Compress", 9); From 45819b94e67ef843183ff1d3e4075fef4077fe04 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 6 Oct 2021 16:51:17 -0700 Subject: [PATCH 2/4] Address CI and some preliminary comments --- onnxruntime/core/providers/cuda/tensor/compress.cc | 1 - onnxruntime/core/providers/cuda/tensor/compress_impl.cu | 5 +++-- onnxruntime/test/providers/cpu/tensor/compress_op.test.cc | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/compress.cc b/onnxruntime/core/providers/cuda/tensor/compress.cc index 6be6478626427..45281fb810851 100644 --- a/onnxruntime/core/providers/cuda/tensor/compress.cc +++ b/onnxruntime/core/providers/cuda/tensor/compress.cc @@ -3,7 +3,6 @@ #include "compress.h" #include "compress_impl.h" -#include namespace onnxruntime { namespace cuda { diff --git a/onnxruntime/core/providers/cuda/tensor/compress_impl.cu b/onnxruntime/core/providers/cuda/tensor/compress_impl.cu index b3f984596f9f9..6500f625f3988 100644 --- a/onnxruntime/core/providers/cuda/tensor/compress_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/compress_impl.cu @@ -15,14 +15,15 @@ #include #include +#include namespace onnxruntime { namespace cuda { // This cast is for transform iterator. This type affects the accumulator type width // in InclusiveSum() and int8_t overflows. -struct CastToIn32 { - __device__ int32_t operator()(int8_t v) const { +struct CastToIn32 : public thrust::unary_function { + __host__ __device__ int32_t operator()(int8_t v) const { return static_cast(v); } }; diff --git a/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc b/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc index cbc62909d4f72..173dbcf8ce854 100644 --- a/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc +++ b/onnxruntime/test/providers/cpu/tensor/compress_op.test.cc @@ -132,7 +132,7 @@ TEST(CompressTest, Compress_default_axis_issue_9247_cumulative_sum_overflow) { // Generate input longer than 127 constexpr size_t elements = 150; std::vector input; - for (int i = 0; i < elements; ++i) { + for (size_t i = 0; i < elements; ++i) { input.push_back(static_cast(i)); } From 62c931b1e29f9f155cc0e8ec995f5264a99ac386 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Thu, 7 Oct 2021 10:36:46 -0700 Subject: [PATCH 3/4] Address review comments --- .../core/providers/cuda/tensor/compress_impl.cu | 14 ++++++++------ .../core/providers/cuda/tensor/compress_impl.h | 6 ++++-- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/compress_impl.cu b/onnxruntime/core/providers/cuda/tensor/compress_impl.cu index 6500f625f3988..989ccef086105 100644 --- a/onnxruntime/core/providers/cuda/tensor/compress_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/compress_impl.cu @@ -21,21 +21,23 @@ namespace onnxruntime { namespace cuda { // This cast is for transform iterator. This type affects the accumulator type width -// in InclusiveSum() and int8_t overflows. -struct CastToIn32 : public thrust::unary_function { +// in InclusiveSum(). By default, the accumulator type matches the input, but for int8_t +// the sum overflows quickly, so we want the source type to match the output (int32_t). +// see https://github.com/NVIDIA/cub/issues/384 +struct CastToInt32 : public thrust::unary_function { __host__ __device__ int32_t operator()(int8_t v) const { return static_cast(v); } }; -cudaError_t CompressCalcPrefixSumTempStorageBytes(cudaStream_t stream, const int8_t* condition_data, int* condition_cumulative_sum, int length, size_t& temp_storage_bytes) { - auto input_iter = thrust::make_transform_iterator(condition_data, CastToIn32()); +cudaError_t CompressCalcPrefixSumTempStorageBytes(cudaStream_t stream, const int8_t* condition_data, int32_t* condition_cumulative_sum, int length, size_t& temp_storage_bytes) { + auto input_iter = thrust::make_transform_iterator(condition_data, CastToInt32()); return cub::DeviceScan::InclusiveSum( nullptr, temp_storage_bytes, input_iter, condition_cumulative_sum, length, stream); } -cudaError_t CompressInclusivePrefixSum(cudaStream_t stream, void* d_temp_storage, size_t temp_storage_bytes, const int8_t* condition_data, int* condition_cumulative_sum, int length) { - auto input_iter = thrust::make_transform_iterator(condition_data, CastToIn32()); +cudaError_t CompressInclusivePrefixSum(cudaStream_t stream, void* d_temp_storage, size_t temp_storage_bytes, const int8_t* condition_data, int32_t* condition_cumulative_sum, int length) { + auto input_iter = thrust::make_transform_iterator(condition_data, CastToInt32()); return cub::DeviceScan::InclusiveSum( d_temp_storage, temp_storage_bytes, input_iter, condition_cumulative_sum, length, stream); } diff --git a/onnxruntime/core/providers/cuda/tensor/compress_impl.h b/onnxruntime/core/providers/cuda/tensor/compress_impl.h index 339784147697d..30b6b00447ab2 100644 --- a/onnxruntime/core/providers/cuda/tensor/compress_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/compress_impl.h @@ -9,8 +9,10 @@ namespace onnxruntime { namespace cuda { -cudaError_t CompressCalcPrefixSumTempStorageBytes(cudaStream_t stream, const int8_t* condition_data, int* condition_cumulative_sum, int length, size_t& temp_storage_bytes); -cudaError_t CompressInclusivePrefixSum(cudaStream_t stream, void* d_temp_storage, size_t temp_storage_bytes, const int8_t* condition_data, int* condition_cumulative_sum, int length); +cudaError_t CompressCalcPrefixSumTempStorageBytes(cudaStream_t stream, const int8_t* condition_data, + int32_t* condition_cumulative_sum, int length, size_t& temp_storage_bytes); +cudaError_t CompressInclusivePrefixSum(cudaStream_t stream, void* d_temp_storage, size_t temp_storage_bytes, + const int8_t* condition_data, int32_t* condition_cumulative_sum, int length); Status CompressImpl(cudaStream_t stream, const size_t element_bytes, From 2efdc58cdb56f8b163efef4377f13845ebc65719 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Thu, 7 Oct 2021 11:20:40 -0700 Subject: [PATCH 4/4] Remove redundant header --- onnxruntime/core/providers/cuda/tensor/compress_impl.cu | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/compress_impl.cu b/onnxruntime/core/providers/cuda/tensor/compress_impl.cu index 989ccef086105..b2c7b60866a77 100644 --- a/onnxruntime/core/providers/cuda/tensor/compress_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/compress_impl.cu @@ -13,9 +13,8 @@ #include "core/providers/cuda/tensor/compress_impl.h" -#include -#include #include +#include namespace onnxruntime { namespace cuda {