diff --git a/onnxruntime/contrib_ops/cpu/nchwc_ops.cc b/onnxruntime/contrib_ops/cpu/nchwc_ops.cc index 5eb0c24c55384..1a0806d5c513d 100644 --- a/onnxruntime/contrib_ops/cpu/nchwc_ops.cc +++ b/onnxruntime/contrib_ops/cpu/nchwc_ops.cc @@ -129,7 +129,10 @@ Status ReorderOutput::Compute(OpKernelContext* context) const { const auto& X_shape = X->Shape().GetDims(); const auto X_rank = X_shape.size(); ORT_ENFORCE(X_rank == 4); - ORT_ENFORCE(channels_ <= X_shape[1]); + const int64_t nchwc_block_size = static_cast(MlasNchwcGetBlockSize()); + ORT_ENFORCE(X_shape[1] % nchwc_block_size == 0 && + channels_ <= X_shape[1] && X_shape[1] - channels_ < nchwc_block_size, + "Input channels must match the NCHWc block-aligned channel count."); // Build the output shape in NCHW or NHWC order. TensorShapeVector Y_shape(X_rank); diff --git a/onnxruntime/test/contrib_ops/nchwc_ops_test.cc b/onnxruntime/test/contrib_ops/nchwc_ops_test.cc index 2bd8feb750561..a24df9b74ea41 100644 --- a/onnxruntime/test/contrib_ops/nchwc_ops_test.cc +++ b/onnxruntime/test/contrib_ops/nchwc_ops_test.cc @@ -13,6 +13,26 @@ namespace onnxruntime { namespace test { namespace { +TEST(NchwcOpsTest, ReorderOutputRejectsUnalignedInputChannels) { + const int64_t block_size = static_cast(MlasNchwcGetBlockSize()); + if (block_size <= 1) { + GTEST_SKIP() << "NCHWc blocking is not enabled on this platform."; + } + + const int64_t input_channels = block_size - 1; + OpTester test("ReorderOutput", 1, kMSNchwcDomain); + test.AddAttribute("channels", int64_t{1}); + test.AddAttribute("channels_last", int64_t{0}); + test.AddInput("X", {1, input_channels, 2, 2}, + std::vector(static_cast(input_channels) * 4, 0.0f)); + test.AddOutput("Y", {1, 1, 2, 2}, {0.0f, 0.0f, 0.0f, 0.0f}); + + test.Config(OpTester::ExpectResult::kExpectFailure, + "Input channels must match the NCHWc block-aligned channel count.") + .ConfigEp(DefaultCpuExecutionProvider()) + .RunWithConfig(); +} + void RunInvalidNchwcConvTest(const std::vector& input_shape, const std::vector& filter_shape, const std::vector* bias_shape, @@ -37,6 +57,27 @@ void RunInvalidNchwcConvTest(const std::vector& input_shape, .RunWithConfig(); } +TEST(NchwcOpsTest, ReorderOutputRejectsExtraChannelBlock) { + const int64_t block_size = static_cast(MlasNchwcGetBlockSize()); + if (block_size <= 1) { + GTEST_SKIP() << "NCHWc blocking is not enabled on this platform."; + } + + const int64_t input_channels = 2 * block_size; + OpTester test("ReorderOutput", 1, kMSNchwcDomain); + test.AddAttribute("channels", block_size); + test.AddAttribute("channels_last", int64_t{0}); + test.AddInput("X", {1, input_channels, 2, 2}, + std::vector(static_cast(input_channels) * 4, 0.0f)); + test.AddOutput("Y", {1, block_size, 2, 2}, + std::vector(static_cast(block_size) * 4, 0.0f)); + + test.Config(OpTester::ExpectResult::kExpectFailure, + "Input channels must match the NCHWc block-aligned channel count.") + .ConfigEp(DefaultCpuExecutionProvider()) + .RunWithConfig(); +} + } // namespace TEST(NchwcOpsTest, ConvRejectsUnalignedOutputChannels) { @@ -44,7 +85,6 @@ TEST(NchwcOpsTest, ConvRejectsUnalignedOutputChannels) { if (block_size <= 1) { GTEST_SKIP() << "NCHWc blocking is not enabled on this platform."; } - RunInvalidNchwcConvTest({0, 1, 1, 1}, {block_size - 1, 1, 1, 1}, nullptr, 1, "NCHWc Conv input and filter shapes do not match a supported blocked layout."); }