1 files changed, 27 insertions, 0 deletions
diff --git a/src/armnn/optimizations/PermuteAndBatchToSpaceAsDepthToSpace.hpp b/src/armnn/optimizations/PermuteAndBatchToSpaceAsDepthToSpace.hpp
new file mode 100644
index 0000000000..4a73efca40
--- /dev/null
+++ b/src/armnn/optimizations/PermuteAndBatchToSpaceAsDepthToSpace.hpp
@@ -0,0 +1,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