Add reusable scratch-buffer reserve/release APIs

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 <stdint.h>
+
+/*
+ * 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();
+}

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 <stddef.h>
+
+#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;
+}

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,

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;