From ccd17b110f70fb8940a6367301b3876480f6758b Mon Sep 17 00:00:00 2001 From: Kaitlyn Davis Date: Thu, 12 Feb 2026 08:32:51 -0800 Subject: [PATCH] scratch: add reusable aligned scratch buffer reserve/release APIs What: Add a caller-managed scratch buffer type with reserve/release APIs. Why: Enable integrators to reuse aligned scratch storage across calls to reduce malloc/free overhead and latency spikes. Expected impact: Lower allocator pressure and more predictable performance; opt-in API with no default behavior change. Tests: add coverage for reserve/release behavior (tests/testDriver_scratch_buffer_apis.c). Signed-off-by: Kaitlyn Davis Signed-off-by: Kaitlyn Davis --- tests/testDriver_scratch_buffer_apis.c | 101 +++++++++++++++++++++++++ zdnn/scratch_buffer.c | 74 ++++++++++++++++++ zdnn/zdnn.h | 21 +++++ zdnn/zdnn.map | 2 + 4 files changed, 198 insertions(+) create mode 100644 tests/testDriver_scratch_buffer_apis.c create mode 100644 zdnn/scratch_buffer.c diff --git a/tests/testDriver_scratch_buffer_apis.c b/tests/testDriver_scratch_buffer_apis.c new file mode 100644 index 0000000..0b42542 --- /dev/null +++ b/tests/testDriver_scratch_buffer_apis.c @@ -0,0 +1,101 @@ +// 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 + +/* + * Tests for reusable scratch-buffer helper APIs. + */ + +void setUp(void) {} +void tearDown(void) {} + +void test_reserve_scratch_buffer_rejects_null_struct(void) { + TEST_ASSERT_EQUAL(ZDNN_INVALID_BUFFER, + zdnn_reserve_scratch_buffer(NULL, 4096)); +} + +void test_reserve_scratch_buffer_zero_size_is_noop(void) { + zdnn_scratch_buffer scratch = {0}; + + TEST_ASSERT_EQUAL(ZDNN_OK, zdnn_reserve_scratch_buffer(&scratch, 0)); + TEST_ASSERT_NULL(scratch.buffer); + TEST_ASSERT_EQUAL_UINT64(0, scratch.size); +} + +void test_reserve_scratch_buffer_allocates_and_reuses_existing_storage(void) { + zdnn_scratch_buffer scratch = {0}; + + TEST_ASSERT_EQUAL(ZDNN_OK, zdnn_reserve_scratch_buffer(&scratch, 1024)); + TEST_ASSERT_NOT_NULL(scratch.buffer); + TEST_ASSERT_EQUAL_UINT64(1024, scratch.size); + TEST_ASSERT_EQUAL_UINT64(0, + (uintptr_t)scratch.buffer % AIU_PAGESIZE_IN_BYTES); + + void *first_buffer = scratch.buffer; + TEST_ASSERT_EQUAL(ZDNN_OK, zdnn_reserve_scratch_buffer(&scratch, 256)); + TEST_ASSERT_EQUAL_PTR(first_buffer, scratch.buffer); + TEST_ASSERT_EQUAL_UINT64(1024, scratch.size); + + zdnn_release_scratch_buffer(&scratch); +} + +void test_reserve_scratch_buffer_grows_storage_when_needed(void) { + zdnn_scratch_buffer scratch = {0}; + + TEST_ASSERT_EQUAL(ZDNN_OK, zdnn_reserve_scratch_buffer(&scratch, 1024)); + TEST_ASSERT_NOT_NULL(scratch.buffer); + + TEST_ASSERT_EQUAL(ZDNN_OK, zdnn_reserve_scratch_buffer(&scratch, 8192)); + TEST_ASSERT_NOT_NULL(scratch.buffer); + TEST_ASSERT_EQUAL_UINT64(8192, scratch.size); + TEST_ASSERT_EQUAL_UINT64(0, + (uintptr_t)scratch.buffer % AIU_PAGESIZE_IN_BYTES); + + zdnn_release_scratch_buffer(&scratch); +} + +void test_release_scratch_buffer_resets_state_and_is_idempotent(void) { + zdnn_scratch_buffer scratch = {0}; + + TEST_ASSERT_EQUAL(ZDNN_OK, zdnn_reserve_scratch_buffer(&scratch, 4096)); + TEST_ASSERT_NOT_NULL(scratch.buffer); + + zdnn_release_scratch_buffer(&scratch); + TEST_ASSERT_NULL(scratch.buffer); + TEST_ASSERT_EQUAL_UINT64(0, scratch.size); + + zdnn_release_scratch_buffer(&scratch); + TEST_ASSERT_NULL(scratch.buffer); + TEST_ASSERT_EQUAL_UINT64(0, scratch.size); + + zdnn_release_scratch_buffer(NULL); +} + +int main(void) { + UNITY_BEGIN(); + + RUN_TEST(test_reserve_scratch_buffer_rejects_null_struct); + RUN_TEST(test_reserve_scratch_buffer_zero_size_is_noop); + RUN_TEST(test_reserve_scratch_buffer_allocates_and_reuses_existing_storage); + RUN_TEST(test_reserve_scratch_buffer_grows_storage_when_needed); + RUN_TEST(test_release_scratch_buffer_resets_state_and_is_idempotent); + + return UNITY_END(); +} diff --git a/zdnn/scratch_buffer.c b/zdnn/scratch_buffer.c new file mode 100644 index 0000000..9635eaa --- /dev/null +++ b/zdnn/scratch_buffer.c @@ -0,0 +1,74 @@ +// 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 + +#include "zdnn.h" +#include "zdnn_private.h" + +#ifdef __MVS__ +#pragma export(zdnn_reserve_scratch_buffer) +#pragma export(zdnn_release_scratch_buffer) +#endif + +zdnn_status zdnn_reserve_scratch_buffer(zdnn_scratch_buffer *scratch, + uint64_t min_size) { + if (!scratch) { + return ZDNN_STATUS_NO_MSG(ZDNN_INVALID_BUFFER); + } + + // No-op requests are treated as success to simplify caller control flow. + if (min_size == 0) { + return ZDNN_STATUS_OK; + } + + // Existing allocation already satisfies the request. + if (scratch->buffer && scratch->size >= min_size) { + return ZDNN_STATUS_OK; + } + + if (min_size > SIZE_MAX) { + return ZDNN_STATUS_NO_MSG(ZDNN_ALLOCATION_FAILURE); + } + + void *new_buffer = malloc_aligned_4k((size_t)min_size); + if (!new_buffer) { + return ZDNN_STATUS_NO_MSG(ZDNN_ALLOCATION_FAILURE); + } + + if (scratch->buffer) { + free_aligned_4k(scratch->buffer); + } + + scratch->buffer = new_buffer; + scratch->size = min_size; + + return ZDNN_STATUS_OK; +} + +void zdnn_release_scratch_buffer(zdnn_scratch_buffer *scratch) { + if (!scratch) { + return; + } + + if (scratch->buffer) { + free_aligned_4k(scratch->buffer); + } + + scratch->buffer = NULL; + scratch->size = 0; +} diff --git a/zdnn/zdnn.h b/zdnn/zdnn.h index 0c9cacc..fbc41d6 100644 --- a/zdnn/zdnn.h +++ b/zdnn/zdnn.h @@ -179,6 +179,27 @@ typedef enum nnpa_bfp_format { #define ZDNN_SOFTMAX_SAVEAREA_SIZE 8 * 1024 #define ZDNN_8K_SAVEAREA_SIZE 8 * 1024 +// Reusable caller-managed scratch buffer for operations that accept `save_area` +// or `work_area` pointers. +// +// Typical usage: +// 1) call zdnn_reserve_scratch_buffer() with required size helper output +// 2) pass `scratch.buffer` into operation APIs +// 3) call zdnn_release_scratch_buffer() once done +typedef struct zdnn_scratch_buffer { + void *buffer; + uint64_t size; +} zdnn_scratch_buffer; + +// Reserve at least `min_size` bytes in a reusable 4K-aligned scratch buffer. +// +// - If `scratch->size` is already >= `min_size`, this is a no-op. +// - If `min_size` is 0, this is a no-op. +zdnn_status zdnn_reserve_scratch_buffer(zdnn_scratch_buffer *scratch, + uint64_t min_size); + +// Release storage held by a reusable scratch buffer and reset its fields. +void zdnn_release_scratch_buffer(zdnn_scratch_buffer *scratch); // NNPA Hardware defined values for Function Specific Parameters typedef enum nnpa_matmul_operations { NNPA_MATMUL_OP_ADDITION = 0, diff --git a/zdnn/zdnn.map b/zdnn/zdnn.map index d80d3a1..5b6a00c 100644 --- a/zdnn/zdnn.map +++ b/zdnn/zdnn.map @@ -109,6 +109,8 @@ ZDNN_1.0 { zdnn_get_library_version_str; zdnn_get_library_version; zdnn_refresh_nnpa_query_result; + zdnn_reserve_scratch_buffer; + zdnn_release_scratch_buffer; zdnn_add; zdnn_sub; zdnn_mul;