From c93cb4a1ab192f920022f0cd092b9323fd1812e4 Mon Sep 17 00:00:00 2001
From: Jerry Ge <jerry.ge@arm.com>
Date: Thu, 18 May 2023 21:16:22 +0000
Subject: Support int16, uint16, int8, uint8 types numpy reader/writer for
 serialization_lib

Signed-off-by: Jerry Ge <jerry.ge@arm.com>
Change-Id: I287fcda4e895ecfc35083d6cb02140a251c0deb5
---
 include/numpy_utils.h | 24 ++++++++++++++++
 src/numpy_utils.cpp   | 80 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+)

diff --git a/include/numpy_utils.h b/include/numpy_utils.h
index 8c2ed71..29d7e11 100644
--- a/include/numpy_utils.h
+++ b/include/numpy_utils.h
@@ -45,6 +45,14 @@ public:
 
     static NPError readFromNpyFile(const char* filename, const uint32_t elems, half_float::half* databuf);
 
+    static NPError readFromNpyFile(const char* filename, const uint32_t elems, uint8_t* databuf);
+
+    static NPError readFromNpyFile(const char* filename, const uint32_t elems, int8_t* databuf);
+
+    static NPError readFromNpyFile(const char* filename, const uint32_t elems, uint16_t* databuf);
+
+    static NPError readFromNpyFile(const char* filename, const uint32_t elems, int16_t* databuf);
+
     static NPError readFromNpyFile(const char* filename, const uint32_t elems, int32_t* databuf);
 
     static NPError readFromNpyFile(const char* filename, const uint32_t elems, int64_t* databuf);
@@ -58,6 +66,22 @@ public:
     static NPError
         writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const half_float::half* databuf);
 
+    static NPError writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const uint8_t* databuf);
+
+    static NPError writeToNpyFile(const char* filename, const uint32_t elems, const uint8_t* databuf);
+
+    static NPError writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const int8_t* databuf);
+
+    static NPError writeToNpyFile(const char* filename, const uint32_t elems, const int8_t* databuf);
+
+    static NPError writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const uint16_t* databuf);
+
+    static NPError writeToNpyFile(const char* filename, const uint32_t elems, const uint16_t* databuf);
+
+    static NPError writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const int16_t* databuf);
+
+    static NPError writeToNpyFile(const char* filename, const uint32_t elems, const int16_t* databuf);
+
     static NPError writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const int32_t* databuf);
 
     static NPError writeToNpyFile(const char* filename, const uint32_t elems, const int32_t* databuf);
diff --git a/src/numpy_utils.cpp b/src/numpy_utils.cpp
index 123908a..65d76e3 100644
--- a/src/numpy_utils.cpp
+++ b/src/numpy_utils.cpp
@@ -28,6 +28,30 @@ NumpyUtilities::NPError NumpyUtilities::readFromNpyFile(const char* filename, co
     return readFromNpyFileCommon(filename, dtype_str, 1, elems, databuf, true);
 }
 
+NumpyUtilities::NPError NumpyUtilities::readFromNpyFile(const char* filename, const uint32_t elems, uint8_t* databuf)
+{
+    const char dtype_str[] = "'|u1'";
+    return readFromNpyFileCommon(filename, dtype_str, sizeof(uint8_t), elems, databuf, false);
+}
+
+NumpyUtilities::NPError NumpyUtilities::readFromNpyFile(const char* filename, const uint32_t elems, int8_t* databuf)
+{
+    const char dtype_str[] = "'|i1'";
+    return readFromNpyFileCommon(filename, dtype_str, sizeof(int8_t), elems, databuf, false);
+}
+
+NumpyUtilities::NPError NumpyUtilities::readFromNpyFile(const char* filename, const uint32_t elems, uint16_t* databuf)
+{
+    const char dtype_str[] = "'<u2'";
+    return readFromNpyFileCommon(filename, dtype_str, sizeof(uint16_t), elems, databuf, false);
+}
+
+NumpyUtilities::NPError NumpyUtilities::readFromNpyFile(const char* filename, const uint32_t elems, int16_t* databuf)
+{
+    const char dtype_str[] = "'<i2'";
+    return readFromNpyFileCommon(filename, dtype_str, sizeof(int16_t), elems, databuf, false);
+}
+
 NumpyUtilities::NPError NumpyUtilities::readFromNpyFile(const char* filename, const uint32_t elems, int32_t* databuf)
 {
     const char dtype_str[] = "'<i4'";
@@ -280,6 +304,62 @@ NumpyUtilities::NPError
     return writeToNpyFileCommon(filename, dtype_str, 1, shape, databuf, true);    // bools written as size 1
 }
 
+NumpyUtilities::NPError
+    NumpyUtilities::writeToNpyFile(const char* filename, const uint32_t elems, const uint8_t* databuf)
+{
+    std::vector<int32_t> shape = { (int32_t)elems };
+    return writeToNpyFile(filename, shape, databuf);
+}
+
+NumpyUtilities::NPError
+    NumpyUtilities::writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const uint8_t* databuf)
+{
+    const char dtype_str[] = "'|u1'";
+    return writeToNpyFileCommon(filename, dtype_str, sizeof(uint8_t), shape, databuf, false);
+}
+
+NumpyUtilities::NPError
+    NumpyUtilities::writeToNpyFile(const char* filename, const uint32_t elems, const int8_t* databuf)
+{
+    std::vector<int32_t> shape = { (int32_t)elems };
+    return writeToNpyFile(filename, shape, databuf);
+}
+
+NumpyUtilities::NPError
+    NumpyUtilities::writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const int8_t* databuf)
+{
+    const char dtype_str[] = "'|i1'";
+    return writeToNpyFileCommon(filename, dtype_str, sizeof(int8_t), shape, databuf, false);
+}
+
+NumpyUtilities::NPError
+    NumpyUtilities::writeToNpyFile(const char* filename, const uint32_t elems, const uint16_t* databuf)
+{
+    std::vector<int32_t> shape = { (int32_t)elems };
+    return writeToNpyFile(filename, shape, databuf);
+}
+
+NumpyUtilities::NPError
+    NumpyUtilities::writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const uint16_t* databuf)
+{
+    const char dtype_str[] = "'<u2'";
+    return writeToNpyFileCommon(filename, dtype_str, sizeof(uint16_t), shape, databuf, false);
+}
+
+NumpyUtilities::NPError
+    NumpyUtilities::writeToNpyFile(const char* filename, const uint32_t elems, const int16_t* databuf)
+{
+    std::vector<int32_t> shape = { (int32_t)elems };
+    return writeToNpyFile(filename, shape, databuf);
+}
+
+NumpyUtilities::NPError
+    NumpyUtilities::writeToNpyFile(const char* filename, const std::vector<int32_t>& shape, const int16_t* databuf)
+{
+    const char dtype_str[] = "'<i2'";
+    return writeToNpyFileCommon(filename, dtype_str, sizeof(int16_t), shape, databuf, false);
+}
+
 NumpyUtilities::NPError
     NumpyUtilities::writeToNpyFile(const char* filename, const uint32_t elems, const int32_t* databuf)
 {
-- 
cgit v1.2.1