From 46122ed0c5b724ea9d040ba082f5afc9e501715d Mon Sep 17 00:00:00 2001
From: Alexander Bengtsson <Alexander.Bengtsson@arm.com>
Date: Wed, 12 Feb 2025 16:16:35 +0100
Subject: [PATCH] MLBEDSW-9758: Add leading-dimension decomposition for MaxPool

Change-Id: Id9ee6c1e1a40cf68dbf050c4d94bb37754208468
Signed-off-by: Alexander Bengtsson <Alexander.Bengtsson@arm.com>
---
 ethosu/regor/compiler/scheduler_decompose.cpp | 35 ++++++++++++++++++-
 ethosu/regor/compiler/scheduler_decompose.hpp |  1 +
 ethosu/regor/compiler/scheduler_packing.cpp   |  3 ++
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/ethosu/regor/compiler/scheduler_decompose.cpp b/ethosu/regor/compiler/scheduler_decompose.cpp
index defa6679..60c4023a 100644
--- a/ethosu/regor/compiler/scheduler_decompose.cpp
+++ b/ethosu/regor/compiler/scheduler_decompose.cpp
@@ -219,7 +219,7 @@ bool CanRunOnHardware(Architecture *arch, const SchedulerOperation *schedOp)
             return false;
         }
     }
-    if ( IsConvolution(schedOp->Type()) )
+    if ( IsConvolution(schedOp->Type()) || IsPooling(schedOp->Type()) )
     {
         auto &ofmShape = schedOp->OFM()->SliceShape();
         if ( ofmShape.Size() > 3 && ofmShape.Batch() > 1 ) return false;
@@ -279,6 +279,7 @@ bool CanDecompose(Architecture *, const SchedulerOperation *schedOp)
     if ( schedOp->Type() == OpType::ArgMax ) return true;
     if ( schedOp->Type() == OpType::Reverse ) return true;
     if ( schedOp->Type() == OpType::Transpose ) return true;
+    if ( schedOp->Type() == OpType::MaxPool ) return true;
     return false;
 }
 
@@ -1696,4 +1697,36 @@ std::vector<std::unique_ptr<SchedulerOperation>> DecomposeTranspose(Architecture
     return result;
 }
 
+std::vector<std::unique_ptr<SchedulerOperation>> DecomposeMaxPool(Architecture *arch, std::unique_ptr<SchedulerOperation> op)
+{
+    std::vector<std::unique_ptr<SchedulerOperation>> result;
+    auto ofmConn = op->Output(TensorUsage::OFM);
+    auto &ofmShape = ofmConn->SliceShape();
+    auto &ofmSlice = ofmConn->slice;
+    auto ifmConn = op->Input(TensorUsage::IFM);
+    auto &ifmShape = ifmConn->SliceShape();
+    auto &ifmSlice = ifmConn->slice;
+
+    ofmSlice.Initialize(ofmShape.WithZeros(), ofmShape);
+    ifmSlice.Initialize(ifmShape.WithZeros(), ifmShape);
+
+    if ( auto ifm2Conn = op->TryInput(TensorUsage::IFM1) )
+    {
+        auto &ifm2Shape = ifm2Conn->shape;
+        auto &ifm2Slice = ifm2Conn->slice;
+
+        ifm2Slice.Initialize(ifm2Shape.WithZeros(), ifm2Shape);
+    }
+
+    auto ofmRank = ofmShape.Size();
+    if ( ofmRank > 3 && (ofmShape.Elements() > ofmShape.Height() * ofmShape.Width() * ofmShape.Depth()) )
+    {
+        return DecomposeLeadingDimensions(ofmRank - 3, arch, std::move(op), DecomposeMaxPool);
+    }
+
+    result.emplace_back(std::move(op));
+    return result;
+}
+
+
 }  // namespace regor
diff --git a/ethosu/regor/compiler/scheduler_decompose.hpp b/ethosu/regor/compiler/scheduler_decompose.hpp
index 76e4c99c..20c392a7 100644
--- a/ethosu/regor/compiler/scheduler_decompose.hpp
+++ b/ethosu/regor/compiler/scheduler_decompose.hpp
@@ -44,6 +44,7 @@ std::vector<std::unique_ptr<SchedulerOperation>> DecomposeMatmul(Architecture *a
 std::vector<std::unique_ptr<SchedulerOperation>> DecomposeReduce(Architecture *arch, std::unique_ptr<SchedulerOperation> op);
 std::vector<std::unique_ptr<SchedulerOperation>> DecomposeReverse(Architecture *arch, std::unique_ptr<SchedulerOperation> op);
 std::vector<std::unique_ptr<SchedulerOperation>> DecomposeTranspose(Architecture *arch, std::unique_ptr<SchedulerOperation> op);
+std::vector<std::unique_ptr<SchedulerOperation>> DecomposeMaxPool(Architecture *arch, std::unique_ptr<SchedulerOperation> op);
 
 
 // Operator query helpers
diff --git a/ethosu/regor/compiler/scheduler_packing.cpp b/ethosu/regor/compiler/scheduler_packing.cpp
index 53aa06c0..3b56b1a3 100644
--- a/ethosu/regor/compiler/scheduler_packing.cpp
+++ b/ethosu/regor/compiler/scheduler_packing.cpp
@@ -589,6 +589,9 @@ std::vector<std::unique_ptr<SchedulerOperation>> SchedulerPacking::DecomposeSche
         case OpType::Transpose:
             result = DecomposeTranspose(_arch, std::move(op));
             break;
+        case OpType::MaxPool:
+            result = DecomposeMaxPool(_arch, std::move(op));
+            break;
         default:
             if ( DecomposeAsElementwise(op->Type()) || op->Type() == OpType::MemoryCopy )
             {
-- 
GitLab