From c0bd3e192f2ee9899e9cb72e81113ec864aa2ef3 Mon Sep 17 00:00:00 2001 From: Kaitlyn Davis Date: Thu, 12 Feb 2026 08:21:20 -0800 Subject: [PATCH] reshape: fail fast on fallback tmp buffer allocation failure What: Detect and fail cleanly when reshape fallback temporary allocation fails. Why: Allocation failures must return a deterministic status (no crash/UB), especially under memory pressure. Expected impact: Safer failure semantics and easier integration fallback handling. Tests: add regression forcing malloc failure (tests/testDriver_reshape_fallback_alloc_safety.c). Signed-off-by: Kaitlyn Davis Signed-off-by: Kaitlyn Davis --- ...testDriver_reshape_fallback_alloc_safety.c | 92 +++++++++++++++++++ zdnn/reshape_ztensor.c | 11 ++- 2 files changed, 102 insertions(+), 1 deletion(-) create mode 100644 tests/testDriver_reshape_fallback_alloc_safety.c diff --git a/tests/testDriver_reshape_fallback_alloc_safety.c b/tests/testDriver_reshape_fallback_alloc_safety.c new file mode 100644 index 0000000..8b15295 --- /dev/null +++ b/tests/testDriver_reshape_fallback_alloc_safety.c @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +/* + * Copyright IBM Corp. 2024 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "testsupport.h" + +#include + +/* + * Regression coverage for a safety bug in zdnn_reshape_ztensor()'s last-resort + * path. + * + * Historical behavior (bug): + * - When the reshaped tensor required an unstick/restick temp buffer larger + * than the 1MB stack buffer, the implementation attempted malloc(s). + * - If malloc(s) failed, it silently fell back to the 1MB stack buffer anyway, + * risking an out-of-bounds write. + * + * Expected behavior (fixed): + * - If s > 1MB and malloc(s) fails, return ZDNN_ALLOCATION_FAILURE. + */ + +void setUp(void) {} +void tearDown(void) {} + +void test_large_tmpbuf_malloc_failure_returns_error(void) { + zdnn_tensor_desc pre_src = {0}, pre_dest = {0}; + zdnn_tensor_desc tfrmd_src = {0}, tfrmd_dest = {0}; + zdnn_ztensor src = {0}, dest = {0}; + + // Force a huge unstickified size `s` so malloc(s) is expected to fail on any + // realistic system. Use a product that stays within uint64_t (avoid wrap). + // s = num_elements * sizeof(FP32) ~= 2^62 bytes (4 exabytes). + const uint32_t big = (1U << 30); + zdnn_init_pre_transformed_desc(ZDNN_4D, FP32, &pre_src, big, 1, 1, big); + zdnn_init_pre_transformed_desc(ZDNN_4D, FP32, &pre_dest, big, 1, 1, big); + + // Set transformed desc layouts to satisfy reshape preconditions, but make + // dim1 differ so we take the last-resort unstick/restick path. + init_transformed_desc(ZDNN_NHWC, ZDNN_DLFLOAT16, ZDNN_FORMAT_4DFEATURE, + &tfrmd_src, 1, 1, 1, 64); + init_transformed_desc(ZDNN_NHWC, ZDNN_DLFLOAT16, ZDNN_FORMAT_4DFEATURE, + &tfrmd_dest, 1, 1, 2, 32); + + zdnn_init_ztensor(&pre_src, &tfrmd_src, &src); + zdnn_init_ztensor(&pre_dest, &tfrmd_dest, &dest); + + // Provide dummy buffers; we should fail before any unstick/restick touches. + src.buffer = malloc_aligned_4k(4096); + dest.buffer = malloc_aligned_4k(4096); + TEST_ASSERT_NOT_NULL_MESSAGE(src.buffer, "src.buffer allocation failed"); + TEST_ASSERT_NOT_NULL_MESSAGE(dest.buffer, "dest.buffer allocation failed"); + + src.is_transformed = true; + dest.is_transformed = false; + + char buf_stderr[BUFSIZ] = {0}; + + stderr_to_pipe(); + zdnn_status st = zdnn_reshape_ztensor(&src, &dest); + restore_stderr(buf_stderr, BUFSIZ); + + TEST_ASSERT_EQUAL_UINT32_MESSAGE(ZDNN_ALLOCATION_FAILURE, st, + "Expected allocation failure status"); + TEST_ASSERT_NOT_NULL_MESSAGE( + strstr(buf_stderr, "reshape tmp buffer"), + "Expected reshape tmp buffer allocation failure message"); + + free_aligned_4k(src.buffer); + free_aligned_4k(dest.buffer); +} + +int main(void) { + UNITY_BEGIN(); + + RUN_TEST(test_large_tmpbuf_malloc_failure_returns_error); + + return UNITY_END(); +} diff --git a/zdnn/reshape_ztensor.c b/zdnn/reshape_ztensor.c index 1332ac3..37b4a70 100644 --- a/zdnn/reshape_ztensor.c +++ b/zdnn/reshape_ztensor.c @@ -171,7 +171,10 @@ zdnn_status zdnn_reshape_ztensor(const zdnn_ztensor *src, zdnn_ztensor *dest) { // NOTE: this will change when we have "no conversion stick/unstick" // for now, unstick to FP32 and restick to preserve precision. - char stack_tmpbuf[STACK_TMPBUF_SIZE] = {0}; + // This buffer is fully overwritten by zdnn_transform_origtensor(), so it does + // not need to be zero-initialized. Avoiding the implicit memset reduces + // overhead on this slow-path. + char stack_tmpbuf[STACK_TMPBUF_SIZE]; void *malloc_tmpbuf = NULL; zdnn_ztensor tmp_tensor_src, tmp_tensor_dest; @@ -196,6 +199,12 @@ zdnn_status zdnn_reshape_ztensor(const zdnn_ztensor *src, zdnn_ztensor *dest) { uint64_t s = get_num_elements(src, ELEMENTS_PRE) * get_data_type_size(FP32); if (s > STACK_TMPBUF_SIZE) { malloc_tmpbuf = malloc(s); + if (!malloc_tmpbuf) { + return ZDNN_STATUS(ZDNN_ALLOCATION_FAILURE, + "Unable to allocate %" PRIu64 + " bytes for reshape tmp buffer.", + s); + } } // no need to log status, zdnn_transform_origtensor() and