// Copyright (c) 2023 CINN Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <queue>

#include "paddle/cinn/hlir/framework/op_lowering.h"

namespace cinn {
namespace hlir {
namespace framework {

std::vector<NodeData*> GetInputNodeData(const Node* node);

ir::Tensor GetTensor(const NodeData* node_data,
                     const absl::flat_hash_map<std::string, Type>& type_dict,
                     const absl::flat_hash_map<std::string, shape_t>& shape_dict);

std::vector<ir::Tensor> CollectInputTensor(const Node* node,
                                           std::vector<ir::Tensor>& func_args,
                                           std::unordered_map<std::string, ir::Tensor>& tensor_map,
                                           const absl::flat_hash_map<std::string, Type>& type_dict,
                                           const absl::flat_hash_map<std::string, shape_t>& shape_dict);

std::unordered_map<Node*, Node*> BuildVirtualConsumer(const GroupPtr& group,
                                                      const absl::flat_hash_map<std::string, shape_t>& shape_dict);

NodeData* GetNodeData(const Node* node);

std::vector<NodeData*> GetAllNodeData(const Node* node);

std::vector<Node*> GetConsumers(const Node* node);

bool IsConstOp(const framework::Node* node);

std::vector<Node*> GetConsumersInSet(const Node* node, const std::unordered_set<Node*>& node_set);

std::vector<Node*> TopologicalOrder(const GroupPtr& group, const std::unordered_map<Node*, Node*>& virtual_consumers);

std::vector<Node*> BFSTopologicalOrderWithPriority(const GroupPtr& group,
                                                   const std::unordered_map<Node*, Node*>& virtual_consumers,
                                                   const absl::flat_hash_map<std::string, shape_t>& shape_dict);

Node* FindGlobalReducer(const std::vector<Node*>& nodes_in_order);

Node* FindNearestReducer(const Node* node, const std::unordered_set<Node*>& nodes_set);

bool CanbeInline(Node* node,
                 const std::vector<Node*> consumers,
                 const Node* reducer,
                 const std::unordered_set<Node*> masters,
                 const GroupPtr& group,
                 const std::unordered_set<Node*>& nodes_set,
                 const absl::flat_hash_map<std::string, shape_t>& shape_dict);

Node* GetMasterToComputeAt(Node* node,
                           const std::vector<Node*>& nodes_in_order,
                           const std::unordered_set<Node*>& nodes_inline,
                           const std::unordered_set<Node*>& nodes_set,
                           const std::unordered_map<Node*, Node*>& virtual_consumers,
                           const absl::flat_hash_map<std::string, shape_t>& shape_dict);

std::unordered_set<Node*> GetMasters(Node* node,
                                     const std::unordered_set<Node*>& nodes_inline,
                                     const std::unordered_set<Node*>& nodes_set);

void LoopAssignReduce(ir::IRSchedule& ir_sch,
                      const Node* node,
                      const Node* reducer,
                      const Target& target,
                      const std::unordered_map<std::string, ir::Tensor>& tensor_map,
                      const absl::flat_hash_map<std::string, shape_t>& shape_dict);

void LoopComputeAt(ir::IRSchedule& ir_sch,
                   Node* node,
                   const Node* master,
                   const GroupPtr& group,
                   const absl::flat_hash_map<std::string, shape_t>& shape_dict,
                   const std::unordered_map<std::string, ir::Tensor>& tensor_map);

void SyncThreadWithShared(ir::IRSchedule& ir_sch,
                          const GroupPtr& group,
                          const std::unordered_set<Node*>& nodes_inline,
                          const std::unordered_set<Node*>& nodes_set,
                          const absl::flat_hash_map<std::string, shape_t>& shape_dict,
                          const std::unordered_map<std::string, ir::Tensor>& tensor_map);

}  // namespace framework
}  // namespace hlir
}  // namespace cinn
