Main Compliance testing for SCATTER and GATHER

Added indices shuffling and random INT32 support to generate lib
with testing of these new random generator modes

Signed-off-by: Jeremy Johnson <jeremy.johnson@arm.com>
Change-Id: I058d8b092470228075e8fe69c2ededa639163003

3 files changed, 207 insertions, 2 deletions

diff --git a/reference_model/src/generate/generate_pseudo_random.cc b/reference_model/src/generate/generate_pseudo_random.cc
index b62c38f..865483b 100644
--- a/reference_model/src/generate/generate_pseudo_random.cc
+++ b/reference_model/src/generate/generate_pseudo_random.cc
@@ -15,6 +15,7 @@
 #include "generate_utils.h"
 #include "half.hpp"
 
+#include <algorithm>
 #include <array>
 #include <iterator>
 #include <limits>
@@ -27,7 +28,7 @@
 namespace
 {
 
-// Random generator
+// Random FP generator
 template <typename FP>
 class PseudoRandomGeneratorFloat
 {
@@ -127,6 +128,123 @@ bool generateFP(const TosaReference::GenerateConfig& cfg, DataType* data, size_t
     return true;
 }
 
+// Random INT generator
+template <typename INT>
+class PseudoRandomGeneratorInteger
+{
+public:
+    PseudoRandomGeneratorInteger(uint64_t seed)
+        : _gen(seed)
+    {
+        constexpr auto min = std::numeric_limits<INT>::min();
+        constexpr auto max = std::numeric_limits<INT>::max();
+
+        setDistribution(min, max);
+    }
+
+    PseudoRandomGeneratorInteger(uint64_t seed, INT min, INT max)
+        : _gen(seed)
+    {
+        setDistribution(min, max);
+    }
+
+    INT getRandomInteger()
+    {
+        return _unidis(_gen);
+    }
+
+    INT getRandomInteger(INT min, INT max)
+    {
+        typename std::uniform_int_distribution<INT>::param_type range(min, max);
+        return _unidis(_gen, range);
+    }
+
+private:
+    void setDistribution(INT min, INT max)
+    {
+        _unidis = std::uniform_int_distribution<INT>(min, max);
+    }
+
+    std::mt19937 _gen;
+    std::uniform_int_distribution<INT> _unidis;
+};
+
+template <typename DataType>
+bool shuffleINTbyRow(const TosaReference::GenerateConfig& cfg, DataType* data, size_t size)
+{
+    const TosaReference::PseudoRandomInfo& prinfo = cfg.pseudoRandomInfo;
+    PseudoRandomGeneratorInteger<DataType>* generator;
+
+    if (cfg.shape.size() != 2)
+    {
+        WARNING("[Generator][PR][INT] Shuffle only supports 2 dimensional tensors.");
+        return false;
+    }
+    if (prinfo.range.size() != 2)
+    {
+        WARNING("[Generator][PR][INT] Cannot create un-ranged shuffle data.");
+        return false;
+    }
+
+    const int32_t min = std::stoi(prinfo.range[0]);
+    const int32_t max = std::stoi(prinfo.range[1]);
+    generator         = new PseudoRandomGeneratorInteger<DataType>(prinfo.rngSeed, min, max);
+
+    // Work out inclusive range
+    const auto range = std::abs(max - min) + 1;
+    const auto N     = cfg.shape[0];    // Number of rows
+    const auto W     = cfg.shape[1];    // Width of rows
+    if (W > range)
+    {
+        WARNING("[Generator][PR][INT] Cannot fill data size %d with given shuffle range %d.", W, range);
+        return false;
+    }
+
+    std::vector<DataType> numbers(range);
+    for (int n = 0; n < N; ++n)
+    {
+        // Fill in the numbers in range
+        std::iota(numbers.begin(), numbers.end(), min);
+
+        // Perform random shuffling
+        for (auto num = numbers.begin(); num < numbers.end(); ++num)
+        {
+            std::swap(*num, numbers[generator->getRandomInteger()]);
+        }
+        // Copy amount of data required
+        for (auto w = 0; w < W; ++w)
+        {
+            data[(n * W) + w] = numbers[w];
+        }
+    }
+    return true;
+}
+
+template <typename DataType>
+bool generateINT(const TosaReference::GenerateConfig& cfg, DataType* data, size_t size)
+{
+    const TosaReference::PseudoRandomInfo& prinfo = cfg.pseudoRandomInfo;
+    PseudoRandomGeneratorInteger<DataType>* generator;
+
+    const auto T = TosaReference::numElementsFromShape(cfg.shape);
+
+    if (prinfo.range.size() == 2)
+    {
+        const int32_t min = std::stoi(prinfo.range[0]);
+        const int32_t max = std::stoi(prinfo.range[1]);
+        generator         = new PseudoRandomGeneratorInteger<DataType>(prinfo.rngSeed, min, max);
+    }
+    else
+    {
+        generator = new PseudoRandomGeneratorInteger<DataType>(prinfo.rngSeed);
+    }
+
+    for (auto t = 0; t < T; ++t)
+    {
+        data[t] = generator->getRandomInteger();
+    }
+    return true;
+}
 }    // namespace
 
 namespace TosaReference
@@ -155,6 +273,18 @@ bool generatePseudoRandom(const GenerateConfig& cfg, void* data, size_t size)
             half_float::half* outData = reinterpret_cast<half_float::half*>(data);
             return generateFP(cfg, outData, size);
         }
+        case DType::DType_INT32: {
+            int32_t* outData = reinterpret_cast<int32_t*>(data);
+            if (cfg.opType == Op::Op_SCATTER && cfg.inputPos == 1)
+            {
+                // Indices for SCATTER must not repeat - perform data shuffle
+                return shuffleINTbyRow(cfg, outData, size);
+            }
+            else
+            {
+                return generateINT(cfg, outData, size);
+            }
+        }
         default:
             WARNING("[Generator][PR] Unsupported type.");
             return false;
diff --git a/reference_model/src/generate/generate_utils.cc b/reference_model/src/generate/generate_utils.cc
index bdc840d..d31048f 100644
--- a/reference_model/src/generate/generate_utils.cc
+++ b/reference_model/src/generate/generate_utils.cc
@@ -51,6 +51,7 @@ NLOHMANN_JSON_SERIALIZE_ENUM(Op,
                                  { Op::Op_EXP, "EXP" },
                                  { Op::Op_FLOOR, "FLOOR" },
                                  { Op::Op_FULLY_CONNECTED, "FULLY_CONNECTED" },
+                                 { Op::Op_GATHER, "GATHER" },
                                  { Op::Op_GREATER, "GREATER" },
                                  { Op::Op_GREATER_EQUAL, "GREATER_EQUAL" },
                                  { Op::Op_IDENTITY, "IDENTITY" },
@@ -69,6 +70,7 @@ NLOHMANN_JSON_SERIALIZE_ENUM(Op,
                                  { Op::Op_REDUCE_MAX, "REDUCE_MAX" },
                                  { Op::Op_REDUCE_MIN, "REDUCE_MIN" },
                                  { Op::Op_REDUCE_SUM, "REDUCE_SUM" },
+                                 { Op::Op_SCATTER, "SCATTER" },
                                  { Op::Op_SIGMOID, "SIGMOID" },
                                  { Op::Op_SUB, "SUB" },
                                  { Op::Op_TANH, "TANH" },
diff --git a/reference_model/test/generate_tests.cpp b/reference_model/test/generate_tests.cpp
index 2c318e0..e4a6d20 100644
--- a/reference_model/test/generate_tests.cpp
+++ b/reference_model/test/generate_tests.cpp
@@ -448,7 +448,7 @@ TEST_CASE("positive - FP32 conv2d dot product (last 3 values)")
         conv2d_test_FP32(tosaName, tosaElements, templateJsonCfg, "5", 2, lastExpected);
     }
 }
-TEST_CASE("positive - pseudo random")
+TEST_CASE("positive - FP32 pseudo random")
 {
     std::string templateJsonCfg = R"({
         "tensors" : {
@@ -823,4 +823,77 @@ TEST_CASE("positive - FP32 avg_pool2d dot product (first 3 values)")
         avg_pool2d_test_FP32(tosaName, tosaElements, templateJsonCfg, "5", expected);
     }
 }
+
+TEST_CASE("positive - INT32 pseudo random")
+{
+    std::string templateJsonCfg = R"({
+        "tensors" : {
+            "input0" : {
+                "generator": "PSEUDO_RANDOM",
+                "data_type": "INT32",
+                "input_type": "VARIABLE",
+                "shape" : [ 2, 12 ],
+                "input_pos": 0,
+                "op" : "SCATTER",
+                "pseudo_random_info": {
+                    "rng_seed": 13,
+                    "range": [ "-5", "5" ]
+                }
+            },
+            "input1" : {
+                "generator": "PSEUDO_RANDOM",
+                "data_type": "INT32",
+                "input_type": "VARIABLE",
+                "shape" : [ 2, 10 ],
+                "input_pos": 1,
+                "op" : "SCATTER",
+                "pseudo_random_info": {
+                    "rng_seed": 14,
+                    "range": [ "0", "9" ]
+                }
+            }
+
+        }
+    })";
+
+    const std::string tosaNameP0 = "input0";
+    const size_t tosaElementsP0  = 2 * 12;
+    const std::string tosaNameP1 = "input1";
+    const size_t tosaElementsP1  = 2 * 10;
+
+    SUBCASE("scatter - int32 random")
+    {
+        std::string jsonCfg = templateJsonCfg;
+
+        std::vector<int32_t> bufferP0(tosaElementsP0);
+        REQUIRE(tgd_generate_data(jsonCfg.c_str(), tosaNameP0.c_str(), (void*)bufferP0.data(), tosaElementsP0 * 4));
+        for (auto e = bufferP0.begin(); e < bufferP0.end(); ++e)
+        {
+            // Check the values are within range
+            bool withinRange = (*e >= -5 && *e <= 5);
+            REQUIRE(withinRange);
+        }
+    }
+
+    SUBCASE("scatter - int32 row shuffle")
+    {
+        std::string jsonCfg = templateJsonCfg;
+
+        std::vector<int32_t> bufferP1(tosaElementsP1);
+        REQUIRE(tgd_generate_data(jsonCfg.c_str(), tosaNameP1.c_str(), (void*)bufferP1.data(), tosaElementsP1 * 4));
+
+        std::vector<bool> set;
+        for (int32_t n = 0; n < 2; ++n)
+        {
+            set.assign(10, false);
+            for (int32_t i = 0; i < 10; ++i)
+            {
+                auto idx = bufferP1[i];
+                // Check that the values in the buffer only occur once
+                REQUIRE(!set[idx]);
+                set[idx] = true;
+            }
+        }
+    }
+}
 TEST_SUITE_END();    // generate