From c1e22955e44c3fe3a9e6eac4183049cb09e2efbb Mon Sep 17 00:00:00 2001 From: Johan Gunnarsson Date: Wed, 12 Mar 2025 15:44:06 +0100 Subject: [PATCH] MLBEDSW-10553: Improve transpose decomposition when batch = 1 We can handle all 2-axis swaps in a 3D shape. We can also handle all 2-axis swaps among the innermost 3 axes as long as all other axes are ones. Signed-off-by: Johan Gunnarsson Change-Id: If96581e3229579a12cdef5c3cd4628e45feb5c96 --- ethosu/regor/compiler/scheduler_decompose.cpp | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/ethosu/regor/compiler/scheduler_decompose.cpp b/ethosu/regor/compiler/scheduler_decompose.cpp index 096bae9a..deb1aada 100644 --- a/ethosu/regor/compiler/scheduler_decompose.cpp +++ b/ethosu/regor/compiler/scheduler_decompose.cpp @@ -1453,8 +1453,8 @@ static std::vector> SwapAxes(Architecture *a assert(shape.IsValid()); assert(tail); - assert(a < shape.Size()); - assert(b < shape.Size()); + assert(a >= 0 && a < shape.Size()); + assert(b >= 0 && b < shape.Size()); assert(a < b); LOG_TRACE2("SwapAxes: Swap ({}), {} <-> {}\n", shape.ToString(), a, b); @@ -1470,8 +1470,12 @@ static std::vector> SwapAxes(Architecture *a // 4. Swap back axis N-1 to position B, like in step 2. // 5. Swap back axis 0 to position A, like in step 1. - // We can handle all swaps in a 3D shape - if ( shape.Size() < 4 ) + // We can swap any two axes of the three innermost axes (H/W/C) of a shape if any of the following is true: + // + // A) The shape is 3D or less. I.e. the shape is (H, W, C) or (W, C). + // B) The shape is 4D or more and all axes other than the three innermost axes (H/W/C) are ones. I.e. the shape + // is (1, H, W, C), (1, 1, H, W, C) or (1, 1, 1, H, W, C). + if ( shape.Size() <= 3 || (a >= shape.Size() - 3 && b >= shape.Size() - 3 && shape.AxisProduct(0, -3) == 1) ) { // Build transpose type for this swap Shape perm(nullptr, shape.Size()); -- GitLab