// Copyright (C) 2022 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#pragma once

#include <ngraph/pass/graph_rewrite.hpp>

namespace GNAPluginNS {

/**
 * @brief Replace 2d->4d reshape to pair of 2 reshapes (before Convolution)
 * Before:
 *      [N, HW]
 *         |
 *      Reshape
 *         |
 *   [N, C, H, W]
 *         |
 *    Convolution
 *
 * After (TransposeSinking friendly):
 *      [N, HW]
 *         |
 *      Reshape
 *         |
 *   [N, H, W, C]
 *         |
 *      Reshape
 *         |
 *   [N, C, H, W]
 *         |
 *    Convolution
 */
class Unfuse2dto4dReshapeAndTranspose : public ngraph::pass::MatcherPass {
public:
  NGRAPH_RTTI_DECLARATION;
  Unfuse2dto4dReshapeAndTranspose();
};

/**
 * @brief Replace 2d->4d reshape to pair of 2 reshapes (after Convolution)
 * Before:
 *    Convolution (optionally + bias/pooling/activation)
 *         |
 *    [N, C, H, W]
 *         |
 *      Reshape
 *         |
 *      [N, HW]
 *
 * After (TransposeSinking friendly):
 *    Convolution
 *         |
 *    [N, C, H, W]
 *         |
 *      Reshape
 *         |
 *   [N, H, W, C]
 *         |
 *      Reshape
 *         |
 *      [N, HW]
 *
 */
class Unfuse4dto2dReshapeAndTranspose : public ngraph::pass::MatcherPass {
public:
  NGRAPH_RTTI_DECLARATION;
  Unfuse4dto2dReshapeAndTranspose();
};

} // namespace GNAPluginNS