Skip to content

Commit 469bc9f

Browse files
yuslepukhinyuslepukhin
authored
Address Resize kernel shortcomings (#28402)
This pull request improves the robustness and correctness of the upsampling code in ONNX Runtime, especially for anti-aliased linear and trilinear upsampling on the CPU. The changes focus on safer handling of large tensor dimensions, improved memory safety, and code clarity for interpolation and weight calculation. The most important changes are grouped below. **Dimension and Overflow Handling:** * Added overflow checks for multiplication of large tensor dimensions to prevent integer overflows during output size calculations, using a new `checked_mul_int64` lambda. [[1]](diffhunk://#diff-13eb8371a91e6fab62e63ecc46583049f97e8acc244af6ce8cc1c981d1d72dd3R1106-R1118) [[2]](diffhunk://#diff-13eb8371a91e6fab62e63ecc46583049f97e8acc244af6ce8cc1c981d1d72dd3R1292-R1309) * Ensured all tensor dimensions are validated to fit within the `int32_t` range before narrowing, improving safety for large tensors. [[1]](diffhunk://#diff-13eb8371a91e6fab62e63ecc46583049f97e8acc244af6ce8cc1c981d1d72dd3R1133-R1143) [[2]](diffhunk://#diff-13eb8371a91e6fab62e63ecc46583049f97e8acc244af6ce8cc1c981d1d72dd3L1131-R1234) **Anti-Alias Upsampling Refactor:** * Refactored the anti-alias upsampling filter setup to use a new `InterpolationBound` struct for per-pixel coordinate ranges, replacing the previous flat vector approach. This improves code clarity and reduces indexing errors. [[1]](diffhunk://#diff-051136817a71a65f4763b9f5c6e02c15f9a6aa39189d952717f4f36c6490ee38R25-R35) [[2]](diffhunk://#diff-051136817a71a65f4763b9f5c6e02c15f9a6aa39189d952717f4f36c6490ee38L126-R159) [[3]](diffhunk://#diff-051136817a71a65f4763b9f5c6e02c15f9a6aa39189d952717f4f36c6490ee38L155-L208) * Updated all interpolation and extrapolation routines to use the new `bounds` structure, improving readability and maintainability. [[1]](diffhunk://#diff-051136817a71a65f4763b9f5c6e02c15f9a6aa39189d952717f4f36c6490ee38L272-R290) [[2]](diffhunk://#diff-051136817a71a65f4763b9f5c6e02c15f9a6aa39189d952717f4f36c6490ee38L341-R353) [[3]](diffhunk://#diff-051136817a71a65f4763b9f5c6e02c15f9a6aa39189d952717f4f36c6490ee38L391-R400) * Imlpements CUDA NHWC cubic antialias support **Memory and Type Safety:** * Improved buffer management and type safety in filter weight calculation, including more robust normalization and quantization for int8/uint8 types. * Fixed a minor logic bug in extrapolation handling by ensuring the loop is only entered if there are out-of-bound indices. **General Code Improvements:** * Added missing `<limits>` include and replaced some magic numbers with named variables for clarity. [[1]](diffhunk://#diff-13eb8371a91e6fab62e63ecc46583049f97e8acc244af6ce8cc1c981d1d72dd3R6) [[2]](diffhunk://#diff-13eb8371a91e6fab62e63ecc46583049f97e8acc244af6ce8cc1c981d1d72dd3L1198-R1257) [[3]](diffhunk://#diff-13eb8371a91e6fab62e63ecc46583049f97e8acc244af6ce8cc1c981d1d72dd3L1221-R1272) These changes together make the upsampling code more robust, especially for large or edge-case tensors, and improve maintainability for future development.
1 parent a8260cd commit 469bc9f

8 files changed

Lines changed: 974 additions & 249 deletions

File tree

‎onnxruntime/core/providers/cpu/tensor/upsample.cc‎

Lines changed: 66 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
#include "core/providers/cpu/tensor/upsample.h"
55
#include "core/common/safeint.h"
6+
#include <limits>
67
#include "core/platform/threadpool.h"
78
#include "core/providers/cpu/tensor/upsample_antialias.h"
89

@@ -1102,6 +1103,7 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context,
11021103
}
11031104
AllocatorPtr alloc;
11041105
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc));
1106+
11051107
switch (mode_) {
11061108
case UpsampleMode::NN:
11071109
return UpsampleNearest<T>(X->Data<T>(), Y->MutableData<T>(), X->Shape(), Y->Shape(),
@@ -1116,6 +1118,17 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context,
11161118
bool is_2D = dims.size() == 2;
11171119
bool is_nchw = true;
11181120

1121+
auto is_valid_non_negative_int32 = [](int64_t v) {
1122+
return v >= 0 && v <= std::numeric_limits<int32_t>::max();
1123+
};
1124+
1125+
int64_t batch_size_i64 = 1;
1126+
int64_t num_channels_i64 = 1;
1127+
int64_t input_height_i64 = 0;
1128+
int64_t input_width_i64 = 0;
1129+
int64_t output_height_i64 = 0;
1130+
int64_t output_width_i64 = 0;
1131+
11191132
int32_t batch_size;
11201133
int32_t num_channels;
11211134
int32_t input_height;
@@ -1128,65 +1141,85 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context,
11281141
float width_scale;
11291142

11301143
if (is_2D) {
1131-
batch_size = 1;
1132-
num_channels = 1;
1133-
input_height = static_cast<int32_t>(dims[0]);
1134-
input_width = static_cast<int32_t>(dims[1]);
1135-
1136-
output_height = static_cast<int32_t>(output_dims[0]);
1137-
output_width = static_cast<int32_t>(output_dims[1]);
1144+
input_height_i64 = dims[0];
1145+
input_width_i64 = dims[1];
1146+
output_height_i64 = output_dims[0];
1147+
output_width_i64 = output_dims[1];
11381148

11391149
height_scale = scales[0];
11401150
width_scale = scales[1];
11411151
} else {
11421152
if (scales[1] == 1.0f) {
1143-
batch_size = static_cast<int32_t>(dims[0]);
1144-
num_channels = static_cast<int32_t>(dims[1]);
1145-
input_height = static_cast<int32_t>(dims[2]);
1146-
input_width = static_cast<int32_t>(dims[3]);
1153+
batch_size_i64 = dims[0];
1154+
num_channels_i64 = dims[1];
1155+
input_height_i64 = dims[2];
1156+
input_width_i64 = dims[3];
11471157

1148-
output_height = static_cast<int32_t>(output_dims[2]);
1149-
output_width = static_cast<int32_t>(output_dims[3]);
1158+
output_height_i64 = output_dims[2];
1159+
output_width_i64 = output_dims[3];
11501160

11511161
height_scale = scales[2];
11521162
width_scale = scales[3];
11531163
} else {
11541164
ORT_RETURN_IF_NOT(scales[3] == 1.0f, "4-D input with innermost scale (usually channel of NHWC) as 1.");
11551165
is_nchw = false;
11561166

1157-
batch_size = static_cast<int32_t>(dims[0]);
1158-
num_channels = static_cast<int32_t>(dims[3]);
1159-
input_height = static_cast<int32_t>(dims[1]);
1160-
input_width = static_cast<int32_t>(dims[2]);
1167+
batch_size_i64 = dims[0];
1168+
num_channels_i64 = dims[3];
1169+
input_height_i64 = dims[1];
1170+
input_width_i64 = dims[2];
11611171

1162-
output_height = static_cast<int32_t>(output_dims[1]);
1163-
output_width = static_cast<int32_t>(output_dims[2]);
1172+
output_height_i64 = output_dims[1];
1173+
output_width_i64 = output_dims[2];
11641174

11651175
height_scale = scales[1];
11661176
width_scale = scales[2];
11671177
}
11681178
}
11691179

1180+
ORT_RETURN_IF_NOT(is_valid_non_negative_int32(batch_size_i64) &&
1181+
is_valid_non_negative_int32(num_channels_i64) &&
1182+
is_valid_non_negative_int32(input_height_i64) &&
1183+
is_valid_non_negative_int32(input_width_i64) &&
1184+
is_valid_non_negative_int32(output_height_i64) &&
1185+
is_valid_non_negative_int32(output_width_i64),
1186+
"Resize: dimensions exceed supported int32 range for CPU linear mode.");
1187+
1188+
batch_size = static_cast<int32_t>(batch_size_i64);
1189+
num_channels = static_cast<int32_t>(num_channels_i64);
1190+
input_height = static_cast<int32_t>(input_height_i64);
1191+
input_width = static_cast<int32_t>(input_width_i64);
1192+
output_height = static_cast<int32_t>(output_height_i64);
1193+
output_width = static_cast<int32_t>(output_width_i64);
1194+
1195+
int64_t output_hw = 0;
1196+
ORT_RETURN_IF_NOT(SafeMultiply(output_height_i64, output_width_i64, output_hw),
1197+
"Resize: output height*width overflows int64.");
1198+
1199+
int64_t output_hwc = 0;
1200+
ORT_RETURN_IF_NOT(SafeMultiply(output_hw, num_channels_i64, output_hwc),
1201+
"Resize: output height*width*channels overflows int64.");
1202+
11701203
if (is_nchw) {
11711204
if (antialias_) {
11721205
UpsampleBilinearAntiAlias(batch_size, num_channels, input_height, input_width, output_height, output_width,
11731206
height_scale, width_scale, roi, use_extrapolation_, extrapolation_value_, exclude_outside_,
11741207
X, Y->MutableData<T>(), alloc, get_original_coordinate_,
1175-
output_height * output_width > 64 ? context->GetOperatorThreadPool() : nullptr);
1208+
output_hw > 64 ? context->GetOperatorThreadPool() : nullptr);
11761209
} else {
11771210
UpsampleBilinear(batch_size, num_channels, input_height, input_width, output_height, output_width,
11781211
height_scale, width_scale, roi,
11791212
use_extrapolation_, extrapolation_value_, X->Data<T>(),
11801213
Y->MutableData<T>(), alloc, get_original_coordinate_,
1181-
output_height * output_width > 64 ? context->GetOperatorThreadPool() : nullptr);
1214+
output_hw > 64 ? context->GetOperatorThreadPool() : nullptr);
11821215
}
11831216
} else {
11841217
if (use_extrapolation_) {
11851218
if (antialias_) {
11861219
NhwcUpsampleBilinearAntiAlias(batch_size, num_channels, input_height, input_width, output_height, output_width,
11871220
height_scale, width_scale, roi, use_extrapolation_, extrapolation_value_, exclude_outside_,
11881221
X, Y->MutableData<T>(), alloc, get_original_coordinate_,
1189-
output_height * output_width > 64 ? context->GetOperatorThreadPool() : nullptr);
1222+
output_hw > 64 ? context->GetOperatorThreadPool() : nullptr);
11901223
} else {
11911224
if (!is_2D &&
11921225
(Y->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_UINT8 ||
@@ -1195,21 +1228,21 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context,
11951228
batch_size, num_channels, input_height, input_width, output_height, output_width,
11961229
height_scale, width_scale, roi, extrapolation_value_, X->Data<T>(), Y->MutableData<T>(),
11971230
alloc, get_original_coordinate_,
1198-
output_height * output_width * num_channels > 64 ? context->GetOperatorThreadPool() : nullptr);
1231+
output_hwc > 64 ? context->GetOperatorThreadPool() : nullptr);
11991232
} else {
12001233
NhwcUpsampleBilinear<T, true>(
12011234
batch_size, num_channels, input_height, input_width, output_height, output_width,
12021235
height_scale, width_scale, roi, extrapolation_value_, X->Data<T>(), Y->MutableData<T>(),
12031236
alloc, get_original_coordinate_,
1204-
output_height * output_width * num_channels > 64 ? context->GetOperatorThreadPool() : nullptr);
1237+
output_hwc > 64 ? context->GetOperatorThreadPool() : nullptr);
12051238
}
12061239
}
12071240
} else {
12081241
if (antialias_) {
12091242
NhwcUpsampleBilinearAntiAlias(batch_size, num_channels, input_height, input_width, output_height, output_width,
12101243
height_scale, width_scale, roi, use_extrapolation_, extrapolation_value_, exclude_outside_,
12111244
X, Y->MutableData<T>(), alloc, get_original_coordinate_,
1212-
output_height * output_width > 64 ? context->GetOperatorThreadPool() : nullptr);
1245+
output_hw > 64 ? context->GetOperatorThreadPool() : nullptr);
12131246
} else {
12141247
if (!is_2D &&
12151248
(Y->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_UINT8 ||
@@ -1218,13 +1251,13 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context,
12181251
batch_size, num_channels, input_height, input_width, output_height, output_width,
12191252
height_scale, width_scale, roi, extrapolation_value_, X->Data<T>(), Y->MutableData<T>(),
12201253
alloc, get_original_coordinate_,
1221-
output_height * output_width * num_channels > 64 ? context->GetOperatorThreadPool() : nullptr);
1254+
output_hwc > 64 ? context->GetOperatorThreadPool() : nullptr);
12221255
} else {
12231256
NhwcUpsampleBilinear<T, false>(
12241257
batch_size, num_channels, input_height, input_width, output_height, output_width,
12251258
height_scale, width_scale, roi, extrapolation_value_, X->Data<T>(), Y->MutableData<T>(),
12261259
alloc, get_original_coordinate_,
1227-
output_height * output_width * num_channels > 64 ? context->GetOperatorThreadPool() : nullptr);
1260+
output_hwc > 64 ? context->GetOperatorThreadPool() : nullptr);
12281261
}
12291262
}
12301263
}
@@ -1244,20 +1277,24 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context,
12441277
const int64_t output_height = is_3D ? output_dims[1] : output_dims[3];
12451278
const int64_t output_width = is_3D ? output_dims[2] : output_dims[4];
12461279

1280+
int64_t output_hw = 0;
1281+
ORT_RETURN_IF_NOT(SafeMultiply(output_height, output_width, output_hw),
1282+
"Resize: output height*width overflows int64.");
1283+
12471284
if (antialias_) {
12481285
UpsampleTrilinearAntiAlias(batch_size, num_channels, input_depth, input_height, input_width,
12491286
output_depth, output_height, output_width,
12501287
is_3D ? scales[0] : scales[2], is_3D ? scales[1] : scales[3],
12511288
is_3D ? scales[2] : scales[4], roi, use_extrapolation_, extrapolation_value_,
12521289
exclude_outside_, X, Y->MutableData<T>(), alloc, get_original_coordinate_,
1253-
output_height * output_width > 64 ? context->GetOperatorThreadPool() : nullptr);
1290+
output_hw > 64 ? context->GetOperatorThreadPool() : nullptr);
12541291
} else {
12551292
UpsampleTrilinear(batch_size, num_channels, input_depth, input_height, input_width,
12561293
output_depth, output_height, output_width,
12571294
is_3D ? scales[0] : scales[2], is_3D ? scales[1] : scales[3],
12581295
is_3D ? scales[2] : scales[4], roi, use_extrapolation_, extrapolation_value_,
12591296
X->Data<T>(), Y->MutableData<T>(), alloc, get_original_coordinate_,
1260-
output_height * output_width > 64 ? context->GetOperatorThreadPool() : nullptr);
1297+
output_hw > 64 ? context->GetOperatorThreadPool() : nullptr);
12611298
}
12621299
return Status::OK();
12631300
} else {

0 commit comments

Comments
 (0)