From a4a8ee41d2076e93ba1ce13b6f3f1620a364f0c3 Mon Sep 17 00:00:00 2001
From: Alexander Bengtsson <Alexander.Bengtsson@arm.com>
Date: Wed, 23 Apr 2025 17:07:37 +0200
Subject: [PATCH] MLBEDSW-10733: Add Ethos-U85 supported-op checks for GatherV2

- Add ConstraintGather to Ethos-U85 supported-operator checks
- Don't list GatherND as a supported operator

Change-Id: I94c9d7e9a105992adce643b52f89b966a36273fc
Signed-off-by: Alexander Bengtsson <Alexander.Bengtsson@arm.com>
---
 .../regor/compiler/tflite_graph_optimiser.cpp |  8 ++---
 .../tflite/tflite_supported_operators_u85.cpp | 31 ++++++++++++++++++-
 .../tflite/tflite_supported_operators_u85.hpp |  1 +
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/ethosu/regor/compiler/tflite_graph_optimiser.cpp b/ethosu/regor/compiler/tflite_graph_optimiser.cpp
index 92f10a88..b303615f 100644
--- a/ethosu/regor/compiler/tflite_graph_optimiser.cpp
+++ b/ethosu/regor/compiler/tflite_graph_optimiser.cpp
@@ -857,7 +857,7 @@ Operation *TFLiteGraphOptimiser::ConvertReverse(Graph *const graph, Operation *c
     return returnOp;
 }
 
-// Replace TFLite GatherV2 and GatherNd with GraphIR Gather, if possible.
+// Replace TFLite GatherV2 with GraphIR Gather
 Operation *TFLiteGraphOptimiser::ConvertGather(Graph *const graph, Operation *const operation)
 {
     UNUSED(graph);
@@ -899,12 +899,8 @@ Operation *TFLiteGraphOptimiser::ConvertGather(Graph *const graph, Operation *co
                 assert(batchDimsParam <= axisParam);
             }
         }
-        // TODO: Convert to supported ops check
         // TODO: MLBEDSW-10279 Investigate if constraint can be relaxed
-        if ( axisParam != batchDimsParam )
-        {
-            return returnOp;
-        }
+        assert(axisParam == batchDimsParam);
 
         // Calculate GraphIR Gather N dim
         int N = 1;
diff --git a/ethosu/regor/tflite/tflite_supported_operators_u85.cpp b/ethosu/regor/tflite/tflite_supported_operators_u85.cpp
index 79f9d44d..f4c8c9ad 100644
--- a/ethosu/regor/tflite/tflite_supported_operators_u85.cpp
+++ b/ethosu/regor/tflite/tflite_supported_operators_u85.cpp
@@ -92,7 +92,6 @@ TfLiteSupportedOperatorsU85::TfLiteSupportedOperatorsU85(IArchitectureConstraint
         OpType::Abs,
         OpType::SplitV,
         OpType::ReverseV2,
-        OpType::GatherNd,
         OpType::Quantize,
         OpType::HardSwish,
         OpType::SelectV2,
@@ -120,6 +119,7 @@ TfLiteSupportedOperatorsU85::TfLiteSupportedOperatorsU85(IArchitectureConstraint
     _checks = {
         &TfLiteSupportedOperatorsU85::ConstraintResizeCommon,
         &TfLiteSupportedOperatorsU85::ConstraintResizeBilinear,
+        &TfLiteSupportedOperatorsU85::ConstraintGather,
     };
 }
 
@@ -299,4 +299,33 @@ bool TfLiteSupportedOperatorsU85::ConstraintResizeBilinear(const Operation *op)
     }
     return true;
 }
+
+bool TfLiteSupportedOperatorsU85::ConstraintGather(const Operation *op)
+{
+    OpType opType = op->Type();
+    if ( opType != OpType::GatherV2 )
+    {
+        return true;
+    }
+    const tflite::Operator *const passthrough = static_cast<const tflite::Operator *>(op->Passthrough());
+    const auto options = passthrough->builtin_options_as_GatherOptions();
+    auto *params = op->Input(TensorUsage::IFM0);
+    assert(params);
+    int paramsRank = params->shape.Size();
+    int batchDimsParam = 0;
+    int axisParam = 0;
+    if ( options )
+    {
+        axisParam = options->axis();
+        if ( axisParam < 0 ) axisParam = paramsRank - (-axisParam);
+        batchDimsParam = options->batch_dims();
+    }
+
+    if ( axisParam != batchDimsParam )
+    {
+        Failure(op, fmt::format("axis: {} != batch_dims: {}", axisParam, batchDimsParam), "axis must be equal to batch_dims");
+        return false;
+    }
+    return true;
+}
 }  // namespace regor
diff --git a/ethosu/regor/tflite/tflite_supported_operators_u85.hpp b/ethosu/regor/tflite/tflite_supported_operators_u85.hpp
index 5a76570e..8aab9e66 100644
--- a/ethosu/regor/tflite/tflite_supported_operators_u85.hpp
+++ b/ethosu/regor/tflite/tflite_supported_operators_u85.hpp
@@ -41,5 +41,6 @@ public:
 private:
     bool ConstraintResizeCommon(const Operation *op);
     bool ConstraintResizeBilinear(const Operation *op);
+    bool ConstraintGather(const Operation *op);
 };
 }  // namespace regor
-- 
GitLab