Skip to content

Commit 577eac9

Browse files
quic-ashwshanquic-ashwshan
authored
[QNN EP] MaxPool input rank-3 auto pad bug fix (#24827)
- Previously, padding for rank-3 MaxPool was only computed for auto_pad="NOTSET", using the final output shape. - Identified a broader issue during auto_pad="VALID" implementation: padding must be derived from the recalculated output shape. - Added unit tests to cover all use cases of auto_pad. - Enabled the failing unit test in the cpu pool test ### Description This PR fixes an issue in the padding calculation logic for rank-3 MaxPool operations when using auto_pad. The bug stemmed from using the final output shape (rank-3) to compute padding, rather than the correct intermediate shape (rank-4) that MaxPool actually operates on. The logic has been updated to use the reshaped rank-4 output for accurate padding computation. Unit tests have been added to validate behavior across all auto_pad modes. ### Motivation and Context While implementing support for auto_pad="VALID" in MaxPool, we discovered that the padding for MaxPool rank-3 was being calculated using the final output shape, which is rank-3. However, MaxPool internally operates on a reshaped rank-4 tensor (via pre- and post-processing reshapes). As a result, the padding logic was misaligned with the actual shape used during pooling, leading to test failures.
1 parent 4f208b3 commit 577eac9

File tree

3 files changed

+88
-36
lines changed

3 files changed

+88
-36
lines changed

‎onnxruntime/core/providers/qnn/builder/opbuilder/pool_op_builder.cc‎

Lines changed: 42 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,36 @@ Status PoolOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper,
103103
return Status::OK();
104104
}
105105

106+
static std::vector<uint32_t> AmendOutputShapeForRank3Pool(
107+
gsl::span<const uint32_t> input_shape, // {N, H, W, C}
108+
gsl::span<const uint32_t> kernel_shape, // {k_h, k_w}
109+
gsl::span<const uint32_t> strides, // {s_h, s_w}
110+
gsl::span<const uint32_t> pads) {
111+
assert(input_shape.size() == 4 &&
112+
kernel_shape.size() == 2 &&
113+
strides.size() == 2 &&
114+
pads.size() == 4);
115+
116+
const uint32_t N = input_shape[0];
117+
const uint32_t H = input_shape[1];
118+
const uint32_t W = input_shape[2];
119+
const uint32_t C = input_shape[3];
120+
121+
// pad the spatial dims
122+
uint32_t padded_H = H + pads[0] + pads[2];
123+
uint32_t padded_W = W + pads[1] + pads[3];
124+
125+
// floor-mode on NHWC
126+
uint32_t out_H = (padded_H < kernel_shape[0])
127+
? 0
128+
: (padded_H - kernel_shape[0]) / strides[0] + 1;
129+
uint32_t out_W = (padded_W < kernel_shape[1])
130+
? 0
131+
: (padded_W - kernel_shape[1]) / strides[1] + 1;
132+
133+
return {N, out_H, out_W, C};
134+
}
135+
106136
Status PoolOpBuilder::SetCommonPoolParams(const NodeAttrHelper& node_helper,
107137
std::vector<uint32_t>& filter_size,
108138
std::vector<uint32_t>& pad_amount, std::vector<uint32_t>& strides,
@@ -153,6 +183,14 @@ Status PoolOpBuilder::SetCommonPoolParams(const NodeAttrHelper& node_helper,
153183
dilations = raw_dilations;
154184
}
155185

186+
// Max Pool rank 3 input
187+
if (output_shape.size() == 3) {
188+
// Calculate MaxPool output for rank-4 when input is rank 3
189+
output_shape = AmendOutputShapeForRank3Pool(input_shape,
190+
filter_size,
191+
strides,
192+
pad_amount);
193+
}
156194
auto total_pads_0 = (output_shape[1] - 1) * strides[0] + (filter_size[0] - 1) * dilations[0] + 1 - input_shape[1];
157195
auto total_pads_1 = (output_shape[2] - 1) * strides[1] + (filter_size[1] - 1) * dilations[1] + 1 - input_shape[2];
158196
if (auto_pad.compare("SAME_LOWER") != 0) {
@@ -189,36 +227,6 @@ void SetPoolParam(const NodeUnit& node_unit,
189227
qnn_model_wrapper.AddParamWrapper(std::move(qnn_param));
190228
}
191229

192-
std::vector<uint32_t> ComputePoolOutputShape(
193-
const std::vector<uint32_t>& input_shape, // {N, H, W, C}
194-
const std::vector<uint32_t>& kernel_shape, // {k_h, k_w}
195-
const std::vector<uint32_t>& strides, // {s_h, s_w}
196-
const std::vector<uint32_t>& pads) {
197-
assert(input_shape.size() == 4 &&
198-
kernel_shape.size() == 2 &&
199-
strides.size() == 2 &&
200-
pads.size() == 4);
201-
202-
const uint32_t N = input_shape[0];
203-
const uint32_t H = input_shape[1];
204-
const uint32_t W = input_shape[2];
205-
const uint32_t C = input_shape[3];
206-
207-
// pad the spatial dims
208-
uint32_t padded_H = H + pads[0] + pads[2];
209-
uint32_t padded_W = W + pads[1] + pads[3];
210-
211-
// floor-mode on NHWC
212-
uint32_t out_H = (padded_H < kernel_shape[0])
213-
? 0
214-
: (padded_H - kernel_shape[0]) / strides[0] + 1;
215-
uint32_t out_W = (padded_W < kernel_shape[1])
216-
? 0
217-
: (padded_W - kernel_shape[1]) / strides[1] + 1;
218-
219-
return {N, out_H, out_W, C};
220-
}
221-
222230
Status PoolOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper,
223231
const NodeUnit& node_unit,
224232
std::vector<std::string>&& input_names,
@@ -316,10 +324,10 @@ Status PoolOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wra
316324
}
317325

318326
// Calculate MaxPool output for rank-4 when input is rank 3
319-
auto pooled_shape = ComputePoolOutputShape(onnx_in_shape,
320-
filter_size,
321-
stride,
322-
pad_amount);
327+
auto pooled_shape = AmendOutputShapeForRank3Pool(onnx_in_shape,
328+
filter_size,
329+
stride,
330+
pad_amount);
323331

324332
SetPoolParam(node_unit, QNN_OP_POOL_MAX_2D_PARAM_FILTER_SIZE, std::move(filter_size_dim), std::move(filter_size), param_tensor_names, qnn_model_wrapper);
325333
SetPoolParam(node_unit, QNN_OP_POOL_MAX_2D_PARAM_PAD_AMOUNT, std::move(pad_amount_dim), std::move(pad_amount), param_tensor_names, qnn_model_wrapper);

‎onnxruntime/test/providers/cpu/nn/pool_op_test.cc‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -230,8 +230,7 @@ TEST(PoolTest, MaxPool1D_case2) {
230230
test.AddInput<float>("X", x_dims, x_vals);
231231
test.AddOutput<float>("Y", expected_dims, expected_vals);
232232

233-
// QNN test failed. Caused by a combination of most recent changes, will fix it
234-
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kQnnExecutionProvider});
233+
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
235234
}
236235

237236
TEST(PoolTest, MaxPool1D_case3) {

‎onnxruntime/test/providers/qnn/pool_op_test.cpp‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,51 @@ TEST_F(QnnHTPBackendTests, MaxPool_Rank3_Ceil_HTP_u8) {
262262
ExpectedEPNodeAssignment::All);
263263
}
264264

265+
// 1-D MaxPool HTP test for rank-3 with ceil_mode=1 and auto_pad='VALID'
266+
TEST_F(QnnHTPBackendTests, MaxPool_Rank3_Ceil_HTP_u8_auto_pad_VALID) {
267+
RunQDQPoolOpTest<uint8_t>(
268+
"MaxPool",
269+
TestInputDef<float>({1, 3, 3}, false, -10.0f, 10.0f),
270+
{utils::MakeAttribute("kernel_shape", std::vector<int64_t>{3}),
271+
utils::MakeAttribute("strides", std::vector<int64_t>{3}),
272+
utils::MakeAttribute("pads", std::vector<int64_t>{0, 0}),
273+
utils::MakeAttribute("dilations", std::vector<int64_t>{1}),
274+
utils::MakeAttribute("ceil_mode", static_cast<int64_t>(1)),
275+
utils::MakeAttribute("storage_order", static_cast<int64_t>(0)),
276+
utils::MakeAttribute("auto_pad", "VALID")},
277+
ExpectedEPNodeAssignment::All);
278+
}
279+
280+
// 1-D MaxPool HTP test for rank-3 with ceil_mode=1 and auto_pad='SAME_UPPER'
281+
TEST_F(QnnHTPBackendTests, MaxPool_Rank3_Ceil_HTP_u8_auto_pad_SAME_UPPER) {
282+
RunQDQPoolOpTest<uint8_t>(
283+
"MaxPool",
284+
TestInputDef<float>({1, 3, 3}, false, -10.0f, 10.0f),
285+
{utils::MakeAttribute("kernel_shape", std::vector<int64_t>{3}),
286+
utils::MakeAttribute("strides", std::vector<int64_t>{3}),
287+
utils::MakeAttribute("pads", std::vector<int64_t>{0, 0}),
288+
utils::MakeAttribute("dilations", std::vector<int64_t>{1}),
289+
utils::MakeAttribute("ceil_mode", static_cast<int64_t>(1)),
290+
utils::MakeAttribute("storage_order", static_cast<int64_t>(0)),
291+
utils::MakeAttribute("auto_pad", "SAME_UPPER")},
292+
ExpectedEPNodeAssignment::All);
293+
}
294+
295+
// 1-D MaxPool HTP test for rank-3 with ceil_mode=1 and auto_pad='SAME_LOWER'
296+
TEST_F(QnnHTPBackendTests, MaxPool_Rank3_Ceil_HTP_u8_auto_pad_SAME_LOWER) {
297+
RunQDQPoolOpTest<uint8_t>(
298+
"MaxPool",
299+
TestInputDef<float>({1, 3, 3}, false, -10.0f, 10.0f),
300+
{utils::MakeAttribute("kernel_shape", std::vector<int64_t>{3}),
301+
utils::MakeAttribute("strides", std::vector<int64_t>{3}),
302+
utils::MakeAttribute("pads", std::vector<int64_t>{0, 0}),
303+
utils::MakeAttribute("dilations", std::vector<int64_t>{1}),
304+
utils::MakeAttribute("ceil_mode", static_cast<int64_t>(1)),
305+
utils::MakeAttribute("storage_order", static_cast<int64_t>(0)),
306+
utils::MakeAttribute("auto_pad", "SAME_LOWER")},
307+
ExpectedEPNodeAssignment::All);
308+
}
309+
265310
TEST_F(QnnHTPBackendTests, MaxPool_Ceil_HTP_u8) {
266311
RunQDQPoolOpTest<uint8_t>("MaxPool",
267312
TestInputDef<float>({1, 2, 3, 3}, false, -10.0f, 10.0f), // Dynamic input with range [-10, 10]

0 commit comments

Comments
 (0)