Add save-area size helper API

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();
+}

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;
+  }
+}

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,

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;