From b5a4f1355faf34db545231d5bd61429912059f6f Mon Sep 17 00:00:00 2001
From: Alexander Bengtsson <Alexander.Bengtsson@arm.com>
Date: Fri, 4 Apr 2025 16:17:00 +0200
Subject: [PATCH] MLBEDSW-10222: Reject transpose with large axes for Ethos-U55

- Add transpose-specific axis constraints to Ethos-U55 constraints.
- Extend query in MergeTransposes to include shapes and dtypes

Change-Id: Icca7475813121ff9f74432ce01ea606f86d94b9a
Signed-off-by: Alexander Bengtsson <Alexander.Bengtsson@arm.com>
---
 .../ethosu55/ethos_u55_constraints.cpp        | 51 +++++++++++--------
 ethosu/regor/compiler/graphir_optimiser.cpp   |  2 +
 2 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/ethosu/regor/architecture/ethosu55/ethos_u55_constraints.cpp b/ethosu/regor/architecture/ethosu55/ethos_u55_constraints.cpp
index 9d93197e..b318ecb9 100644
--- a/ethosu/regor/architecture/ethosu55/ethos_u55_constraints.cpp
+++ b/ethosu/regor/architecture/ethosu55/ethos_u55_constraints.cpp
@@ -382,34 +382,41 @@ Flags<QueryResult> EthosU55Constraints::OperatorQuery(OpType opType, const ArchO
     // Detailed operator queries
     if ( !IsNone(query->transposeMask) )
     {
-        if ( opType == OpType::Transpose )
+        if ( opType != OpType::Transpose )
         {
-            if ( query->transposeMask == TransposeType::NWHC || query->transposeMask == TransposeType::NHCW ||
-                 query->transposeMask == TransposeType::NCWH )
-            {
-                if ( req )
-                {
-                    req->req.Set(ArchRequirement::OutputFormat, ArchRequirement::InputFormat);
-                    req->ifmFormat = TensorFormat::NHWC;
-                    req->ofmFormat = TensorFormat::NHWC;
-                }
-                result.Set(QueryResult::HasRequirements);
-            }
-            else
+            return QueryResult::Unsupported;
+        }
+        // TODO MLBEDSW-10668: Transpose-implementation does not support large-axis decomposition
+        if ( req && req->decomposeProps.Any(ArchProperty::TensorAxis) )
+        {
+            return QueryResult::Unsupported;
+        }
+        // TODO MLBEDSW-10668: channel-axis for 32-bit NWHC-transpose is constrained to 14-bits
+        static constexpr int TRANSPOSE_32_MAX_CHANNEL = (1 << 14);
+        if ( (shapeInfo && typeInfo) && (query->ifm[0].type == DataType::Int32) &&
+             (query->transposeMask == TransposeType::NWHC) && (query->ifm[0].shape.Depth() > TRANSPOSE_32_MAX_CHANNEL) )
+        {
+            return QueryResult::Unsupported;
+        }
+        // Validate supported transpose-masks
+        if ( query->transposeMask != TransposeType::NWHC && query->transposeMask != TransposeType::NHCW && query->transposeMask != TransposeType::NCWH )
+        {
+            // supported with mask-decomposition requirements
+            if ( req )
             {
-                // supported with decomposition requirements
-                if ( req )
-                {
-                    req->req.Set(ArchRequirement::Decompose);
-                    req->decomposeProps.Set(ArchProperty::TransposeMask);
-                }
-                result.Set(QueryResult::HasRequirements);
+                req->req.Set(ArchRequirement::Decompose);
+                req->decomposeProps.Set(ArchProperty::TransposeMask);
             }
+            result.Set(QueryResult::HasRequirements);
         }
-        else
+        // Always set Input/Output format requirements
+        if ( req )
         {
-            return QueryResult::Unsupported;
+            req->req.Set(ArchRequirement::OutputFormat, ArchRequirement::InputFormat);
+            req->ifmFormat = TensorFormat::NHWC;
+            req->ofmFormat = TensorFormat::NHWC;
         }
+        result.Set(QueryResult::HasRequirements);
     }
 
     // reverseType::W and reverseType::H are supported
diff --git a/ethosu/regor/compiler/graphir_optimiser.cpp b/ethosu/regor/compiler/graphir_optimiser.cpp
index 0390f473..4dd2305c 100644
--- a/ethosu/regor/compiler/graphir_optimiser.cpp
+++ b/ethosu/regor/compiler/graphir_optimiser.cpp
@@ -1646,6 +1646,8 @@ Operation *GraphIrOptimiser::MergeTransposes(Graph *const graph, Operation *cons
             ArchOperatorQuery query;
             ArchRequirements req;
             query.transposeMask = mergedTranspose;
+            Set(query.ifm[0], ifmConn->tensor.get());
+            Set(query.ofm, ofmConn->tensor.get());
             if ( _constraints->OperatorQuery(OpType::Transpose, &query, &req).Any(QueryResult::Native) )
             {
                 // only merge the transpose if the new mask is natively supported
-- 
GitLab