aboutsummaryrefslogtreecommitdiff
path: root/src/armnn/optimizations/PermuteAndBatchToSpaceAsDepthToSpace.hpp
blob: 4a73efca40d10d7e204dea0cfcff8df4a15fa38f (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
//
// Copyright © 2019 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
#pragma once

#include "Optimization.hpp"

namespace armnn
{
namespace optimizations
{

/// Replaces Permute leading into BatchToSpace with a DepthToSpace
/// in the case where the Permute swaps the batch and channels dimensions
/// such that the replacement is valid.
class PermuteAndBatchToSpaceAsDepthToSpaceImpl
{
public:
    void Run(Graph& graph, InputSlot& connection) const;
};

using PermuteAndBatchToSpaceAsDepthToSpace =
    OptimizeForConnection<PermuteLayer, BatchToSpaceNdLayer, PermuteAndBatchToSpaceAsDepthToSpaceImpl>;

}    // namespace optimizations
}    // namespace armnn