text: // Copyright (c) Microsoft Corporation. All rights reserved.
  // Licensed under the MIT License.
  
  #include "core/optimizer/initializer.h"
  #include "core/optimizer/matmul_transpose_fusion.h"
  #include "core/graph/graph_utils.h"
  #include <deque>
  
  using namespace ONNX_NAMESPACE;
  using namespace ::onnxruntime::common;
  namespace onnxruntime {
  
  static bool GetTransposePerms(const Node& transpose_node, std::vector<int64_t>& perms) {
    ORT_ENFORCE(transpose_node.InputDefs().size() == 1);
  
    // use perms if present
    const auto perm_attr = transpose_node.GetAttributes().find("perm");
    if (perm_attr != transpose_node.GetAttributes().end()) {
      perms = RetrieveValues<int64_t>(perm_attr->second);
      return true;
    }
  
    // otherwise, reverse dimensions
    const NodeArg& input = *transpose_node.InputDefs()[0];
    const TensorShapeProto* shape = input.Shape();
    if (!shape) {
      return false;
    }
  
    perms.resize(shape->dim_size());
    std::iota(perms.rbegin(), perms.rend(), 0);
    return true;
  }
  
  // is_trans is whether to transpose the 2 dims used to MatMul.
  // is_trans_batch is whether to transpose 1st dim and batch dims.
  // Batch here has the same meaning in CUDA's GemmStridedBatched (other than the training batch concept),
  // i.e., all dims except the 2 dims used to MatMul, including from dim[0] to dim[rank-3].
  // For example (take lhs input as example):
  // is_trans=False, is_trans_batch=False:
  //   the input tensor is in shape [b0,...,bn,M,K], no Transpose (no fuse for this case)
  // is_trans=False, is_trans_batch=True:
  //   the input tensor is in shape [M,b0,...,bn,K], Transpose perm=[1,...,rank-2,0,rank-1]
  // is_trans=True , is_trans_batch=False:
  //   the input tensor is in shape [b0,...,bn,K,M], Transpose perm=[0,...,rank-3,rank-1,rank-2]
  // is_trans=True , is_trans_batch=True:
  //   the input tensor is in shape [K,b0,...,bn,M], Transpose perm=[1,...,rank-2,rank-1,0]
  static Node* GetTransposeNodeFromOutput(Graph& graph, NodeArg& node_arg, bool& is_trans, bool& is_trans_batch) {
    is_trans = is_trans_batch = false;
  
    // Skip if not a Transpose node or it has graph output.
    Node* trans_node = graph.GetMutableProducerNode(node_arg.Name());
    if (trans_node == nullptr || trans_node->OpType() != "Transpose" || graph.NodeProducesGraphOutput(*trans_node)) {
      return nullptr;
    }
  
    std::vector<int64_t> perms;
    if (!GetTransposePerms(*trans_node, perms)) {
      return nullptr;
    }
  
    size_t rank = perms.size();
    if (rank < 2) {
      return nullptr;
    }
  
    // We can fuse the Transpose node only when the last axis of original tensor is within the last two dims after transpose.
    int64_t last_axis = static_cast<int64_t>(rank) - 1;
    size_t last_axis_index = perms[rank - 1] == last_axis ? rank - 1 : perms[rank - 2] == last_axis ? rank - 2 : rank;
    if (last_axis_index == rank) {
      return nullptr;
    }
  
    // Transpose node can be fused to MatMul when the batch dims keep same relative orders before and after transpose.
    // But if they are not contiguous, after the fusion, we can only use GemmBatched instead of GemmStridedBatched,
    // which may have perf issue. To keep it simple, we will fuse only when batch dimensions are contiguous.
    if (rank >= 3) {
      if (perms[0] != 0 && perms[0] != 1) {
        return nullptr;
      }
      for (size_t i = 0; i < rank - 3; ++i) {
        if (perms[i] + 1 != perms[i + 1]) {
          return nullptr;
        }
      }
    }
  
    is_trans = last_axis_index == rank - 2;
    is_trans_batch = rank >= 3 && perms[0] == 1;
    return trans_node;
  }
  
  static size_t UpdateConsumerCount(Graph& graph, NodeArg* target, InlinedHashMap<NodeArg*, size_t>& count_map) {
    const auto& node_consumers = graph.GetConsumerNodes(target->Name());
    ORT_ENFORCE(!node_consumers.empty());
    auto it = count_map.find(target);
    if (it == count_map.end()) {
      count_map.insert({target, node_consumers.size() - 1});
      return node_consumers.size() - 1;
    } else {
      count_map[target] -= 1;
      return count_map[target];
    }
  }
  
  /* ReorderCastAndTranspose:
  *  Interchange Cast and Transpose nodes in the graph and return the new Transpose node if possible else nullptr.
  *
  *
  *  Transform the following pattern
  *                              |
  *                         _____|______
  *                         |Transpose |
  *                         |__________|
  *                              |
  *                              |
  *                         _____V______
  *                         |  Cast    |
  *                         |__________|
  *                              |
  *                              V
  *
  *  to
  *                              |
  *                         _____|______
  *                         |  Cast    |
  *                         |__________|
  *                              |
  *                              |
  *                         _____V______
  *                         | Transpose|
  *                         |__________|
  *                              |
  *                              V
  */
  static Node* ReorderCastAndTranspose(Graph& graph, Node* cast,
                                       InlinedHashMap<NodeArg*, size_t>& consumer_count,
                                       std::deque<onnxruntime::NodeIndex>& removed_nodes,
                                       bool& is_trans, bool& is_trans_batch) {
    ORT_ENFORCE(cast != nullptr);
    auto transpose = GetTransposeNodeFromOutput(graph, *cast->MutableInputDefs()[0], is_trans, is_trans_batch);
    if (transpose == nullptr) {
      return nullptr;
    }
    NodeArg* cast_output = cast->MutableOutputDefs()[0];
    NodeArg* transpose_input = transpose->MutableInputDefs()[0];
  
    // Create a new NodeArg to feed the output from the new Cast to the new Transpose.
    // The shape of the new NodeArg is same as the original input to Transport but type
    // should match that of the output from the original Cast.
  
    auto new_cast_output_type_proto = *transpose_input->TypeAsProto();
    const ONNX_NAMESPACE::TensorProto_DataType element_type =
        static_cast<ONNX_NAMESPACE::TensorProto_DataType>(cast_output->TypeAsProto()->tensor_type().elem_type());
    new_cast_output_type_proto.mutable_tensor_type()->set_elem_type(element_type);
    auto& new_cast_output = graph.GetOrCreateNodeArg(cast_output->Name() + "_transformed", &new_cast_output_type_proto);
  
    const std::array new_cast_input_defs{transpose_input};
    const std::array new_cast_output_defs{&new_cast_output};
    const std::array new_transpose_input_defs = {&new_cast_output};
    const std::array new_transpose_output_defs = {cast_output};
  
    Node& new_cast = graph.AddNode(graph.GenerateNodeName(cast->Name() + "_transformed"),
                        cast->OpType(),
                        "Created a new Cast node to interchange Cast and Transpose nodes",
                        new_cast_input_defs,
                        new_cast_output_defs,
                        &cast->GetAttributes(),
                        cast->Domain());
    new_cast.SetExecutionProviderType(cast->GetExecutionProviderType());
  
    Node& new_transpose = graph.AddNode(graph.GenerateNodeName(transpose->Name() + "_transformed"),
                                        transpose->OpType(),
                                        "Created a new Transpose node to interchange Cast and Transpose nodes",
                                        new_transpose_input_defs,
                                        new_transpose_output_defs,
                                        &transpose->GetAttributes(),
                                        transpose->Domain());
    new_transpose.SetExecutionProviderType(transpose->GetExecutionProviderType());
  
    size_t consumers = UpdateConsumerCount(graph, transpose->MutableOutputDefs()[0], consumer_count);
    graph_utils::RemoveNodeOutputEdges(graph, *cast);
    graph.RemoveNode(cast->Index());
    if (consumers == 0) {
      removed_nodes.push_front(transpose->Index());
    }
    return &new_transpose;
  }
  
  // Check whether the element_type is an allowed FusedMatMul data type or not.
  constexpr static bool IsAllowedFusedMatMulDataType(ONNX_NAMESPACE::TensorProto_DataType element_type) {
    return element_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT ||
           element_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ||
           element_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE ||
           element_type == ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16;
  }
  
  /*********************************************************************************************
  
  Case I: The followin is a scenario where Transpose output feeds MatMul. The Transpose input can be either on the left or right.
     The input graph
                           __________                             __________
                           | input0 |                             | input1 |
                           |________|                             |________|
                                |                                      |
                                |                                      |
                                |                                      |
                           _____V______                                |
                           |Transpose |                                |
                           |__________|                                |
                                |                                      |
                                |                                      |
                                |______________           _____________|
                                              |           |
                                              |           |
                                              |           |
                                            __V___________V__
                                            |    MatMul     |
                                            |_______________|
                                                    |
                                                    V
      is transformed to the following
  
                           __________                             __________
                           | input0 |                             | input1 |
                           |________|                             |________|
                                |                                      |
                                |                                      |
                                |                                      |
                                |_____________            _____________|
                                              |           |
                                              |           |
                                              |           |
                                            __V___________V__
                                            |  FusedMatMul  |
                                            |_______________|
                                                    |
                                                    V
  
  Case II: The output of Tanspose feeds Cast and the output from the Cast feeds MatMul
     The input graph
                           __________                             __________
                           | input0 |                             | input1 |
                           |________|                             |________|
                                |                                      |
                                |                                      |
                           _____V______                                |
                           |Transpose |                                |
                           |__________|                                |
                                |                                      |
                                |                                      |
                           _____V______                                |
                           |  Cast    |                                |
                           |__________|                                |
                                |                                      |
                                |______________           _____________|
                                              |           |
                                              |           |
                                              |           |
                                            __V___________V__
                                            |    MatMul     |
                                            |_______________|
                                                    |
                                                    V
      is transformed to the following
  
                           __________                             __________
                           | input0 |                             | input1 |
                           |________|                             |________|
                                |                                      |
                                |                                      |
                                |                                      |
                           _____V______                                |
                           |  Cast    |                                |
                           |__________|                                |
                                |                                      |
                                |______________           _____________|
                                              |           |
                                              |           |
                                              |           |
                                            __V___________V__
                                            |  FusedMatMul  |
                                            |_______________|
                                                    |
                                                    V
  
  ********************************************************************************************************************/
  
  Status MatmulTransposeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const {
    GraphViewer graph_viewer(graph);
    const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder();
    std::deque<onnxruntime::NodeIndex> removed_nodes;
  
    InlinedHashMap<NodeArg*, size_t> consumer_count;
  
    for (auto node_index : node_topology_list) {
      auto& node = *graph.GetNode(node_index);
      ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level, logger));
  
      if ((!graph_utils::IsSupportedOptypeVersionAndDomain(node, "MatMul", {9, 13}) &&
           !graph_utils::IsSupportedOptypeVersionAndDomain(node, "FusedMatMul", {1}, kMSDomain)) ||
          !graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders())) {
        continue;
      }
  
      NodeArg* left_input = node.MutableInputDefs()[0];
      auto left_type = left_input->TypeAsProto()->tensor_type().elem_type();
      if (!IsAllowedFusedMatMulDataType(static_cast<ONNX_NAMESPACE::TensorProto_DataType>(left_type))) {
        continue;
      }
  
      bool is_trans_left = false;
      bool is_trans_batch_left = false;
      Node* left = nullptr;
      // If it's already a FusedMatMul with transBatchA is true, don't fuse it.
      if (node.OpType() != "FusedMatMul" || node.GetAttributes().at("transBatchA").i() == 0) {
        left = GetTransposeNodeFromOutput(graph, *left_input, is_trans_left, is_trans_batch_left);
        if (!left) {
          Node* left_node = graph.GetMutableProducerNode(left_input->Name());
          if (left_node && left_node->OpType() == "Cast") {
            left = ReorderCastAndTranspose(graph, left_node, consumer_count, removed_nodes, is_trans_left,
                                           is_trans_batch_left);
          }
        }
      }
  
      NodeArg* right_input = node.MutableInputDefs()[1];
      auto right_type = right_input->TypeAsProto()->tensor_type().elem_type();
      if (!IsAllowedFusedMatMulDataType(static_cast<ONNX_NAMESPACE::TensorProto_DataType>(right_type))) {
        continue;
      }
  
      bool is_trans_right = false;
      bool is_trans_batch_right = false;
      Node* right = nullptr;
      // If it's already a FusedMatMul with transBatchB is true, don't fuse it.
      if (node.OpType() != "FusedMatMul" || node.GetAttributes().at("transBatchB").i() == 0) {
        right = GetTransposeNodeFromOutput(graph, *right_input, is_trans_right, is_trans_batch_right);
        if (!right) {
          Node* right_node = graph.GetMutableProducerNode(right_input->Name());
          if (right_node && right_node->OpType() == "Cast") {
            right = ReorderCastAndTranspose(graph, right_node, consumer_count, removed_nodes, is_trans_right,
                                            is_trans_batch_right);
          }
        }
      }
  
      // When the rank of two inputs are not equal, we need to pad 1 to one of them for MutMul.
      // For example, if padding 1 is the "M" dim for left input, set is_trans_batch to true is not correct logically.
      // To keep it simple, if any side of is_trans_batch is true, we require both side has same rank.
      if (is_trans_batch_left || is_trans_batch_right) {
        auto shape_left = left_input->Shape();
        auto shape_right = right_input->Shape();
        if (!shape_left || !shape_right || shape_left->dim_size() != shape_right->dim_size()) {
          if (is_trans_batch_left) {
            is_trans_left = is_trans_batch_left = false;
            left = nullptr;
          }
          if (is_trans_batch_right) {
            is_trans_right = is_trans_batch_right = false;
            right= nullptr;
          }
        }
      }
  
      if (!left && !right) {
        continue;
      }
  
      if (left) {
        size_t left_consumers = UpdateConsumerCount(graph, left_input, consumer_count);
        if (left_consumers == 0)
          removed_nodes.push_front(left->Index());
        left_input = left->MutableInputDefs()[0];
      }
  
      if (right) {
        size_t right_consumers = UpdateConsumerCount(graph, right_input, consumer_count);
        if (right_consumers == 0)
          removed_nodes.push_front(right->Index());
        right_input = right->MutableInputDefs()[0];
      }
  
      const std::array input_defs{left_input, right_input};
      const std::array output_defs{node.MutableOutputDefs()[0]};
  
      Node& matmul_node = graph.AddNode(graph.GenerateNodeName("MatMul_With_Transpose"),
                                        "FusedMatMul",
                                        "fused MatMul and Transpose ",
                                        input_defs,
                                        output_defs, {}, kMSDomain);
      float alpha = 1.0f;
      if (node.OpType() == "FusedMatMul") {
        is_trans_left ^= static_cast<bool>(node.GetAttributes().at("transA").i());
        is_trans_right ^= static_cast<bool>(node.GetAttributes().at("transB").i());
        is_trans_batch_left ^= static_cast<bool>(node.GetAttributes().at("transBatchA").i());
        is_trans_batch_right ^= static_cast<bool>(node.GetAttributes().at("transBatchB").i());
        alpha = node.GetAttributes().at("alpha").f();
      }
      matmul_node.AddAttribute("transA", static_cast<int64_t>(is_trans_left));
      matmul_node.AddAttribute("transB", static_cast<int64_t>(is_trans_right));
      matmul_node.AddAttribute("transBatchA", static_cast<int64_t>(is_trans_batch_left));
      matmul_node.AddAttribute("transBatchB", static_cast<int64_t>(is_trans_batch_right));
      matmul_node.AddAttribute("alpha", alpha);
      // Assign provider to this new node. Provider should be same as the provider for old node.
      matmul_node.SetExecutionProviderType(node.GetExecutionProviderType());
  
      graph_utils::FinalizeNodeFusion(graph, matmul_node, node);
  
      modified = true;
    }
  
    // Have to remove node in reversed order for now to walk around the issue in RemoveNode
    for (onnxruntime::NodeIndex removed_node : removed_nodes) {
      graph.RemoveNode(removed_node);
    }
  
    return Status::OK();
  }
  }  // namespace onnxruntime
  