xref: /aosp_15_r20/external/armnn/src/armnn/optimizations/PermuteDepthwiseConv2dWeights.hpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
1 //
2 // Copyright © 2022 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 #pragma once
6 
7 #include "Optimization.hpp"
8 #include "NetworkUtils.hpp"
9 
10 #include <armnnUtils/Permute.hpp>
11 
12 #include <fmt/format.h>
13 
14 namespace armnn
15 {
16 namespace optimizations
17 {
18 
19 class PermuteDepthwiseConv2dWeightsImpl
20 {
21 public:
22 
Run(Graph & graph,Layer & layer) const23     void Run(Graph& graph, Layer& layer) const
24     {
25         if (layer.GetType() == LayerType::DepthwiseConvolution2d)
26         {
27             AddPermuteLayer(graph, PolymorphicDowncast<DepthwiseConvolution2dLayer*>(&layer));
28         }
29     }
30 
31 protected:
32     PermuteDepthwiseConv2dWeightsImpl() = default;
33     ~PermuteDepthwiseConv2dWeightsImpl() = default;
34 
35 private:
36     /// ArmNN format for weights for depthwise is [1, H, W, C] independently of the input/output layout
37     ///
38     /// ACL format for weights for depthwise is:
39     /// - [1, H, W, C] for [N, H, W, C] input/output layout (matches with ArmNN)
40     /// - [1, C, H, W] for [N, C, H, W] input/output layout
41     ///
42     /// Therefore ArmNN weights have to be permuted when input/output layout is [N, C, H, W] to pass them to ACL.
AddPermuteLayer(Graph & graph,DepthwiseConvolution2dLayer * layer)43     static void AddPermuteLayer(Graph& graph, DepthwiseConvolution2dLayer* layer)
44     {
45         TensorInfo inputInfo = layer->GetInputSlot(0).GetConnectedOutputSlot()->GetTensorInfo();
46         TensorInfo weightInfo = layer->GetInputSlot(1).GetConnectedOutputSlot()->GetTensorInfo();
47         if (layer->GetParameters().m_DataLayout == armnn::DataLayout::NHWC)
48         {
49             // No permutation required. Input and weights data layouts are the same.
50             return;
51         }
52         else if (layer->GetParameters().m_DataLayout == armnn::DataLayout::NCHW)
53         {
54             // Weights permutation required. Weights [N,H,W,C] and input [N,C,H,W] data layouts are different.
55             // [ 1, H, W, I*M] --> [ 1, I * M, H, W ]
56             PermutationVector permutationVector = { 0, 2, 3, 1 };
57             TensorInfo weightsPermuted = armnnUtils::Permuted(weightInfo, permutationVector);
58 
59             // Inserts NewLayer so layers don't need to be re-sorted.
60             PermuteLayer* permuteLayer =
61                 graph.InsertNewLayer<PermuteLayer>(layer->GetInputSlot(1),
62                                                    PermuteDescriptor(permutationVector),
63                                                    "permute_layer");
64             permuteLayer->GetOutputSlot().SetTensorInfo(weightsPermuted);
65 
66             // Assign Permute BackendId to be the same as the Depthwise Conv2d BackendId.
67             // Needed as backends have already been assigned at this stage.
68             permuteLayer->SetBackendId(layer->GetBackendId());
69         }
70         else
71         {
72             throw InvalidArgumentException(fmt::format("Unknown data layout for tensor info conversion: {}",
73                                                        GetDataLayoutName(layer->GetParameters().m_DataLayout)));
74         }
75     }
76 };
77 
78 using PermuteDepthwiseConv2dWeights = OptimizeForType<Layer, PermuteDepthwiseConv2dWeightsImpl>;
79 
80 } // namespace optimizations
81 } // namespace armnn
82