From ee3fd8df0cd55e1dc92dd101b86c64694108db82 Mon Sep 17 00:00:00 2001
From: Kaitlyn Davis <k8ika0s@gmail.com>
Date: Thu, 12 Feb 2026 08:22:41 -0800
Subject: [PATCH] save_area: add helper API to query required size

What: Add a public helper to compute required save_area bytes for supported operations.

Why: Integrators need deterministic preallocation to avoid hidden runtime allocations and reduce latency variance.

Expected impact: Simpler caller memory management and reduced allocator pressure; additive API.

Tests: add coverage for save_area sizing and invalid inputs (tests/testDriver_save_area_size_apis.c).
Signed-off-by: Kaitlyn Davis <k8ika0s@gmail.com>
Signed-off-by: Kaitlyn Davis <kaitlyn.davis@ibm.com>
---
 tests/testDriver_save_area_size_apis.c | 46 ++++++++++++++++++++++++++
 zdnn/save_area.c                       | 33 ++++++++++++++++++
 zdnn/zdnn.h                            |  7 ++++
 zdnn/zdnn.map                          |  1 +
 4 files changed, 87 insertions(+)
 create mode 100644 tests/testDriver_save_area_size_apis.c
 create mode 100644 zdnn/save_area.c

diff --git a/tests/testDriver_save_area_size_apis.c b/tests/testDriver_save_area_size_apis.c
new file mode 100644
index 0000000..cf7ed27
--- /dev/null
+++ b/tests/testDriver_save_area_size_apis.c
@@ -0,0 +1,46 @@
+// 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"
+
+/*
+ * Tests for save-area size helper API.
+ */
+
+void setUp(void) {}
+void tearDown(void) {}
+
+void test_save_area_size_softmax_reduce(void) {
+  TEST_ASSERT_EQUAL_UINT64(ZDNN_8K_SAVEAREA_SIZE,
+                           zdnn_get_save_area_size(NNPA_SOFTMAX));
+  TEST_ASSERT_EQUAL_UINT64(ZDNN_8K_SAVEAREA_SIZE,
+                           zdnn_get_save_area_size(NNPA_REDUCE));
+}
+
+void test_save_area_size_non_savearea_ops_return_zero(void) {
+  TEST_ASSERT_EQUAL_UINT64(0, zdnn_get_save_area_size(NNPA_ADD));
+  TEST_ASSERT_EQUAL_UINT64(0, zdnn_get_save_area_size((nnpa_function_code)999));
+}
+
+int main(void) {
+  UNITY_BEGIN();
+
+  RUN_TEST(test_save_area_size_softmax_reduce);
+  RUN_TEST(test_save_area_size_non_savearea_ops_return_zero);
+
+  return UNITY_END();
+}
diff --git a/zdnn/save_area.c b/zdnn/save_area.c
new file mode 100644
index 0000000..10118d5
--- /dev/null
+++ b/zdnn/save_area.c
@@ -0,0 +1,33 @@
+// 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 "zdnn.h"
+#include "zdnn_private.h"
+
+#ifdef __MVS__
+#pragma export(zdnn_get_save_area_size)
+#endif
+
+uint64_t zdnn_get_save_area_size(nnpa_function_code function_code) {
+  switch (function_code) {
+  case NNPA_SOFTMAX:
+  case NNPA_REDUCE:
+    return ZDNN_8K_SAVEAREA_SIZE;
+  default:
+    return 0;
+  }
+}
diff --git a/zdnn/zdnn.h b/zdnn/zdnn.h
index 0c9cacc..3c4e4ed 100644
--- a/zdnn/zdnn.h
+++ b/zdnn/zdnn.h
@@ -179,6 +179,13 @@ typedef enum nnpa_bfp_format {
 #define ZDNN_SOFTMAX_SAVEAREA_SIZE 8 * 1024
 #define ZDNN_8K_SAVEAREA_SIZE 8 * 1024
 
+// Convenience helper to query the required save-area size (in bytes) for
+// operations that accept a `save_area` pointer (e.g., softmax, reduce).
+//
+// Returns 0 when the operation does not require a save area or when the
+// function code is not recognized.
+uint64_t zdnn_get_save_area_size(nnpa_function_code function_code);
+
 // 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..0125062 100644
--- a/zdnn/zdnn.map
+++ b/zdnn/zdnn.map
@@ -109,6 +109,7 @@ ZDNN_1.0 {
     zdnn_get_library_version_str;
     zdnn_get_library_version;
     zdnn_refresh_nnpa_query_result;
+    zdnn_get_save_area_size;
     zdnn_add;
     zdnn_sub;
     zdnn_mul;