From 47ab1762d1c15a7b4c0c068d7294111c5c5f92a2 Mon Sep 17 00:00:00 2001
From: evacha01 <evan.chandler@arm.com>
Date: Mon, 29 Jan 2024 13:23:23 +0000
Subject: Main Compliance testing for CONV3D

Signed-off-by: evacha01 <evan.chandler@arm.com>
Change-Id: Ie05f88db15cd07fd5483ab669329d7048bd3349c
---
 .../src/generate/generate_dot_product.cc           | 143 +++++++++++++++++++++
 1 file changed, 143 insertions(+)

(limited to 'reference_model/src/generate/generate_dot_product.cc')

diff --git a/reference_model/src/generate/generate_dot_product.cc b/reference_model/src/generate/generate_dot_product.cc
index 046007e..7337969 100644
--- a/reference_model/src/generate/generate_dot_product.cc
+++ b/reference_model/src/generate/generate_dot_product.cc
@@ -736,6 +736,147 @@ bool generateTransposeConv2D(const TosaReference::GenerateConfig& cfg,
             return false;
     }
 }
+
+//---------------------------------------------------------------------------//
+//                              Conv3D                                       //
+//---------------------------------------------------------------------------//
+
+template <typename DataType>
+bool generateConv3DInput(const TosaReference::GenerateConfig& cfg,
+                         TosaReference::IDotProductGenerator& generator,
+                         DataType* data,
+                         size_t size)
+{
+    if (cfg.dotProductInfo.kernel.size() != 3 || cfg.dotProductInfo.kernel[0] <= 0 ||
+        cfg.dotProductInfo.kernel[1] <= 0 || cfg.dotProductInfo.kernel[2] <= 0)
+    {
+        WARNING("[Generator][DP][Conv3D][Input] Missing or incorrect kernel size information.");
+        return false;
+    }
+    if (cfg.shape.size() != 5)
+    {
+        WARNING("[Generator][DP][Conv3D][Input] Tensor shape expected 5 dimensions.");
+        return false;
+    }
+
+    const int64_t T   = TosaReference::numElementsFromShape(cfg.shape);
+    const uint32_t ID = cfg.shape[1];
+    const uint32_t IH = cfg.shape[2];
+    const uint32_t IW = cfg.shape[3];
+    const uint32_t IC = cfg.shape[4];
+    const uint32_t KD = cfg.dotProductInfo.kernel[0];
+    const uint32_t KH = cfg.dotProductInfo.kernel[1];
+    const uint32_t KW = cfg.dotProductInfo.kernel[2];
+
+    for (int64_t t = 0; t < T; ++t)
+    {
+        uint32_t ic = t % IC;
+        uint32_t ix = (t / IC) % IW;
+        uint32_t iy = ((t / IC) / IW) % IH;
+        uint32_t id = (((t / IC) / IW) / IH) % ID;
+        uint32_t k  = (((id % KD) * KH + (iy % KH)) * KW + (ix % KW)) * IC + ic;
+
+        data[t] = static_cast<DataType>(generator(k));
+    }
+    return true;
+}
+
+template <typename DataType>
+bool generateConv3DWeight(const TosaReference::GenerateConfig& cfg,
+                          TosaReference::IDotProductGenerator& generator,
+                          DataType* data,
+                          size_t size)
+{
+    if (cfg.shape.size() != 5)
+    {
+        WARNING("[Generator][DP][Conv3D][Weight] Tensor shape expected 5 dimensions.");
+        return false;
+    }
+
+    const int64_t T   = TosaReference::numElementsFromShape(cfg.shape);
+    const uint32_t KD = cfg.shape[0];
+    const uint32_t KH = cfg.shape[1];
+    const uint32_t KW = cfg.shape[2];
+    const uint32_t IC = cfg.shape[3];
+
+    for (int64_t t = 0; t < T; ++t)
+    {
+        uint32_t ic = t % IC;
+        uint32_t kx = (t / IC) % KW;
+        uint32_t ky = ((t / IC) / KW) % KH;
+        uint32_t kd = (((t / IC) / KW) / KH) % KD;
+        uint32_t k  = ((kd * KH + ky) * KW + kx) * IC + ic;
+
+        data[t] = static_cast<DataType>(generator(k));
+    }
+    return true;
+}
+
+template <typename DataType>
+bool generateConv3DBias(const TosaReference::GenerateConfig& cfg,
+                        TosaReference::IDotProductGenerator& generator,
+                        DataType* data,
+                        size_t size)
+{
+    if (cfg.shape.size() != 1)
+    {
+        WARNING("[Generator][DP][Conv3D][Bias] Tensor shape expected 1 dimension.");
+        return false;
+    }
+
+    const uint32_t T = cfg.shape[0];
+
+    for (uint32_t t = 0; t < T; ++t)
+    {
+        data[t] = static_cast<DataType>(generator(2));
+    }
+    return true;
+}
+
+bool generateConv3D(const TosaReference::GenerateConfig& cfg,
+                    TosaReference::IDotProductGenerator& generator,
+                    void* data,
+                    size_t size)
+{
+    switch (cfg.dataType)
+    {
+        case DType::DType_FP32: {
+            float* outData = reinterpret_cast<float*>(data);
+            switch (cfg.inputPos)
+            {
+                case 0:
+                    return generateConv3DInput(cfg, generator, outData, size);
+                case 1:
+                    return generateConv3DWeight(cfg, generator, outData, size);
+                case 2:
+                    return generateConv3DBias(cfg, generator, outData, size);
+                default:
+                    WARNING("[Generator][DP][Conv3D] Invalid input tensor slot position to operator.");
+                    return false;
+            }
+            break;
+        }
+        case DType::DType_FP16: {
+            half_float::half* outData = reinterpret_cast<half_float::half*>(data);
+            switch (cfg.inputPos)
+            {
+                case 0:
+                    return generateConv3DInput(cfg, generator, outData, size);
+                case 1:
+                    return generateConv3DWeight(cfg, generator, outData, size);
+                case 2:
+                    return generateConv3DBias(cfg, generator, outData, size);
+                default:
+                    WARNING("[Generator][DP][Conv3D] Invalid input tensor slot position to operator.");
+                    return false;
+            }
+            break;
+        }
+        default:
+            WARNING("[Generator][DP][Conv3D] Only supports FP32 or FP16.");
+            return false;
+    }
+}
 //---------------------------------------------------------------------------//
 //                              FFT2D                                        //
 //---------------------------------------------------------------------------//
@@ -858,6 +999,8 @@ bool generateDotProduct(const GenerateConfig& cfg, void* data, size_t size)
             return generateDepthwiseConv2D(cfg, *generator, data, size);
         case tosa::Op_TRANSPOSE_CONV2D:
             return generateTransposeConv2D(cfg, *generator, data, size);
+        case tosa::Op_CONV3D:
+            return generateConv3D(cfg, *generator, data, size);
         case tosa::Op_FFT2D:
             return generateFFT2D(cfg, *generator, data, size);
         default:
-- 
cgit v1.2.1