Fix inclusive sum overlfow when applied on int8_t buffer in Compress - #9295
Merged
Conversation
yuslepukhin
commented
Oct 6, 2021
yuslepukhin
commented
Oct 6, 2021
edgchen1
reviewed
Oct 7, 2021
|
|
||
| // 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<int8_t, int32_t> { |
Contributor
There was a problem hiding this comment.
"CastToIn32" -> "CastToInt32" #Resolved
edgchen1
reviewed
Oct 7, 2021
| 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()); |
Contributor
There was a problem hiding this comment.
is it possible for input_iter's result to match the type of condition_cumulative_sum? now it's int32_t and int #Resolved
edgchen1
previously approved these changes
Oct 7, 2021
edgchen1
reviewed
Oct 7, 2021
|
|
||
| #include "core/providers/cuda/tensor/compress_impl.h" | ||
|
|
||
| #include <thrust/scan.h> |
Contributor
There was a problem hiding this comment.
is thrust/scan.h needed? #Resolved
edgchen1
approved these changes
Oct 7, 2021
weixingzhang
approved these changes
Oct 7, 2021
jiafatom
added a commit
that referenced
this pull request
Jun 27, 2026
### Description
Bool initializers supplied via `TensorProto` `raw_data` are copied
verbatim by `UnpackTensor<bool>`, so their bytes are not guaranteed to
be the canonical `{0, 1}` (the `int32_data` path normalizes via
`static_cast<bool>`, but the `raw_data` path did not). Kernels across
the codebase assume bool tensors hold `{0, 1}`.
The CUDA `Compress` kernel is concretely affected: its output-sizing
path sign-extends the condition bytes (`int8_t` -> `int32_t`) through
`cub::DeviceScan::InclusiveSum`, while `_CompressKernel` selects
elements using bool truthiness (`condition_data[div]`). For condition
bytes outside `{0, 1}` the two interpretations disagree and the output
is sized inconsistently with how elements are written. The CPU kernel
uses truthiness for both sizing and selection and is unaffected.
### Changes
- `UnpackTensor<bool>` (`tensorprotoutils.cc`): normalize `raw_data`
bytes to `{0, 1}` after copy. The `UnpackTensorWithExternalData<bool>`
specialization does the same for external data read through that path.
- CUDA `Compress` `CastToInt32` (`compress_impl.cu`): normalize to `{0,
1}` (still returns `int32_t`, preserving the accumulator-widening intent
of #9295) so the sizing path matches the kernel's write predicate,
matching the CPU kernel and the CUDA `NonZero` `bool(x)` convention.
This makes the CUDA `Compress` kernel correct independently of how its
bool condition initializer was materialized.
- Shared helper `utils::NormalizeBoolTensorIfNeeded(Tensor&)`
(`tensorprotoutils.{h,cc}`): single normalization point, reused by
`TensorProtoToTensor()` (external branch, after `MakeCpuTensorCopy`) and
by the session-init external device-copy path.
- `session_state_utils.cc::DeserializeTensorProto`: for **external**
bool initializers loaded onto a non-CPU device, normalize the writable
CPU staging copy before `CopyTensorFromCPUToDevice`. The
`GetExtDataFromTensorProto` buffer may be a read-only mmap, so it is
normalized via a writable copy rather than in place.
- Unit tests in `tensorutils_test.cc` for bool `raw_data` with
non-canonical bytes and for the `NormalizeBoolTensorIfNeeded` helper. A
`Compress` `OpTester` test cannot reproduce the original bug because the
test harness itself normalizes bool during input construction, so
coverage is placed at the deserialization layer. Tests use only Status
returns and gtest assertions, so they build and run in no-exception
builds.
### Coverage / scope of bool normalization
Fully covered:
- In-proto `raw_data` bool initializers (all EPs), via
`UnpackTensor<bool>`.
- External bool initializers reaching `TensorProtoToTensor()`.
- External bool initializers copied to a non-CPU device through the
session-init device-copy path.
Intentionally **not** normalized (by design):
- The CPU zero-copy mmap path for external initializers
(`GetExtDataFromTensorProto` returns a read-only/shared mapping that
cannot be safely modified in place).
- The custom external-data-loader path
(`LoadExtDataToTensorFromTensorProto`), which loads directly into a
device tensor.
These remaining paths are safe for the concrete bug this PR targets
because the CUDA `Compress` kernel is hardened in `compress_impl.cu`
regardless of initializer storage. Other byte-comparing bool consumers
fed by an external mmap/custom-loader initializer with non-canonical
bytes are out of scope here.
### Motivation and Context
`CastToInt32` was introduced in #9295 to widen the `cub::InclusiveSum`
accumulator (an int8 overflow fix); it did not normalize the bool
interpretation. The accumulator-width and bool-normalization concerns
are independent. This change addresses the latter at the deserialization
source and hardens the CUDA `Compress` kernel.
---------
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description:
This change inserts a thrust::transform_iterator that casts int8_t to int32_t when feeding int8_t buffer to
cub::DeviceScan::InclusiveSum() in our CUDA Compress implementation. This makes the accumulator type int32_t and prevents the result overflowing.
Motivation and Context
This addresses #9247