//
//  torchOptimize.cpp
//  MNNConverter
//
//  Created by MNN on 2021/05/12.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#include "torchOptimize.hpp"
#include <torch/csrc/jit/passes/freeze_module.h>
#include <torch/csrc/jit/passes/inliner.h>
#include <torch/csrc/jit/passes/peephole.h>
#include <torch/csrc/jit/passes/dead_code_elimination.h>
#include <torch/csrc/jit/passes/constant_pooling.h>
#include <torch/csrc/jit/passes/constant_propagation.h>
#include <torch/csrc/jit/passes/fold_conv_bn.h>
#include <torch/csrc/jit/passes/remove_dropout.h>
#include <torch/csrc/jit/passes/remove_inplace_ops.h>
#include <torch/csrc/jit/passes/remove_mutation.h>
#include <torch/csrc/jit/passes/normalize_ops.h>
#include <torch/csrc/jit/passes/graph_fuser.h>
#include <torch/csrc/jit/passes/erase_number_types.h>
#include "torchOpConverter.hpp"

namespace torch {
namespace jit {
void removeUselessOps(Block* block) {
    for (auto it = block->nodes().begin(), end = block->nodes().end(); it != end; ++it) {
        for (auto b : it->blocks()) {
            removeUselessOps(b);
        }
        std::set<NodeKind> uselessKind = {
            // prime
            prim::Print,
            prim::RaiseException,
            prim::TimePoint,
            prim::annotate,
            // aten
            aten::warn,
        };
        // useless op
        if (uselessKind.count(it->kind())) {
            for (size_t i = 0; i < it->inputs().size();) {
                auto input = it->inputs().at(i);
                // only handling constants bc of potential side effects
                if (input->uses().size() == 1 &&
                    input->node()->kind() == prim::Constant) {
                    it->removeInput(i);
                    input->node()->destroy();
                } else {
                    ++i;
                }
            }
            it.destroyCurrent();
        }
        if (it->kind() == prim::Loop) {
            if (it->outputs().empty()) {
                it.destroyCurrent();
            }
        }
        if (it->kind().toUnqualString() == std::string("data") ||
            it->kind() == prim::NumToTensor ||
            it->kind() == aten::ScalarImplicit ||
            it->kind() == aten::contiguous ||
            it->kind() == aten::dropout ||
            it->kind() == aten::dropout_ ||
            it->kind() == aten::feature_dropout ||
            it->kind() == aten::clone) {
            it->output()->replaceAllUsesWith(it->input(0));
            for (int i = it->inputs().size()-1; i >= 0; i--) {
                it->removeInput(i);
            }
            it.destroyCurrent();
        }
        if (it->kind() == aten::detach ||
            it->kind() == aten::list ||
            it->kind().toDisplayString() == std::string("aten::cpu")) {
            it->output()->replaceAllUsesWith(it->input(0));
            for (int i = it->inputs().size()-1; i >= 0; i--) {
                it->removeInput(i);
            }
            it.destroyCurrent();
        }
        if (it->kind() == aten::to) {
            auto dst = it->input(1);
            auto ivalue = toIValue(dst);
            if(!ivalue->isInt()) {
                it->output()->replaceAllUsesWith(it->input(0));
                for (int i = it->inputs().size()-1; i >= 0; i--) {
                    it->removeInput(i);
                }
                it.destroyCurrent();
            }
        }
        if (it->kind() == aten::slice) {
            auto start = it->input(2);
            auto end = it->input(3);
            // [0 : 1 : INT_MAX] can remove
            if (toIValue(start) && getValue<int64_t>(start) == 0 &&
                toIValue(end) && getValue<int64_t>(end) == 9223372036854775807) {
                if (it->inputs().size() > 4) {
                    auto stride = it->input(4);
                    if (toIValue(stride) && getValue<int64_t>(stride) != 1) {
                        continue;
                    }
                }
                it->output()->replaceAllUsesWith(it->input(0));
                for (int i = it->inputs().size()-1; i >= 0; i--) {
                    it->removeInput(i);
                }
                it.destroyCurrent();
            }
        }
    }
}

/*
   We rewrite something like:
        x = {v0}
        x.append(v1)
        foo(x)
        x.append(v2)
        bar(x)
   to:
        x1 = {v0, v1}
        foo(x1)
        x2 = {v0, v1, v2}
        bar(x2)
   this is a strengthen version of RemoveListMutation
*/
void removeListAppend(Graph* graph, Block* block) {
    for (auto it = block->nodes().begin(); it != block->nodes().end();) {
        auto* node = *it;
        it++;

        for (Block* sub_block : node->blocks()) {
            removeListAppend(graph, sub_block);
        }

        if (!(node->kind() == aten::append && node->inputs().at(0)->node()->kind() == prim::ListConstruct)) {
            continue;
        }
        Value* mutated_value = node->inputs().at(0);
        Node* list_node = mutated_value->node();
        Node* new_list_node = graph->create(prim::ListConstruct, 1);
        for (Value* input : list_node->inputs()) {
            new_list_node->addInput(input);
        }
        new_list_node->addInput(node->inputs().at(1));
        new_list_node->copyMetadata(list_node);
        new_list_node->insertAfter(node);
        new_list_node->output()->setType(list_node->output()->type());
        mutated_value->replaceAllUsesAfterNodeWith(node, new_list_node->output());
        node->destroy();
    }
}
/*
   We remove all ListConstruct op with only one input and not used by aten::cat, like below:
        %116 : Tensor?[] = prim::ListConstruct(%115)
        %alpha0.1 : Tensor = aten::index_put_(%alpha.1, %116, %x.1, %16)
   ListConstruct used by aten::cat will be reserved like below:
        %features.2 : Tensor[] = prim::ListConstruct(%input3.4)
        %concated_features.380 : Tensor = aten::cat(%features.2, %5)
   Attention: Runing this pass after removeListAppend
 */
void removeListConstructOps(Block* block) {
    for (auto it = block->nodes().begin(), end = block->nodes().end(); it != end; ++it) {
        for (auto b : it->blocks()) {
            removeUselessOps(b);
        }
        if (it->kind() == prim::ListConstruct && it->inputs().size() == 1) {
            bool remove = true;
            for (auto use : it->output()->uses()) {
                if (use.user->kind() == aten::cat) {
                    remove = false;
                    break;
                }
            }
            if (remove) {
                it->output()->replaceAllUsesWith(it->input(0));
                it->removeInput(0);
                it.destroyCurrent();
            }
        }
    }
}

/*
   We rewrite something like:
        y = chunk(x)
        v1, v2, v3 = ListUnpack(y)
   to:
        v1, v2, v3 = chunk(x)
*/
void FuseListUnpack(Block* block) {
    for (auto it = block->nodes().begin(); it != block->nodes().end();) {
        auto* node = *it;
        it++;

        for (Block* sub_block : node->blocks()) {
            FuseListUnpack(sub_block);
        }
        std::set<NodeKind> fusekind = {
            aten::split,
            aten::split_with_sizes,
            aten::split_with_sizes,
            aten::unsafe_split_with_sizes,
            aten::unbind,
            aten::chunk,
            aten::unsafe_chunk,
            aten::where,
        };
        if (fusekind.count(it->kind()) &&
            it->outputs().size() == 1 &&
            it->output()->uses().size() == 1) {
            const auto listunpack = it->output()->uses()[0].user;
            if (listunpack->kind() == prim::ListUnpack) {
                // it->i_(Symbol::fromQualString("attr::_outputs"),
                //         static_cast<int64_t>(listunpack->outputs().size()));
                for (auto i = 0; i < listunpack->outputs().size(); ++i) {
                    auto new_output = it->addOutput();
                    new_output->copyMetadata(listunpack->output(i));
                }
                listunpack->removeAllInputs();
                it->eraseOutput(0);
                listunpack->replaceAllUsesWith(*it);
                listunpack->destroy();
            }
        }
    }
}
/*
   We rewrite something like:
        x = ListConstruct(v1, v2, v3)
        y = stack(y, axis)
   to:
        y = stack(v1, v2, v3, axis)
*/
void FuseListStack(Block* block) {
    for (auto it = block->nodes().begin(); it != block->nodes().end();) {
        auto* node = *it;
        it++;

        for (Block* sub_block : node->blocks()) {
            FuseListUnpack(sub_block);
        }
        std::set<NodeKind> fusekind = {
            aten::stack
        };
        if (it->kind() == aten::stack) {
            auto input = it->input(0)->node();
            if (input->kind() == prim::ListConstruct) {
                auto axis = it->input(1);
                it->removeAllInputs();
                for (int i = 0; i < input->inputs().size(); i++) {
                    it->addInput(input->input(i));
                }
                it->addInput(axis);
                input->destroy();
            }
        }
    }
}
/*
 We rewrite something like:
    %y : int, %z : int = prim::Loop(%6, %2, %y.1, %z.1) # <ipython-input-14-d0a2ead71c2a>:6:4
        block0(%i.1 : int, %y.11 : int, %z.11 : int):
            %y.5 : int = aten::add(%y.11, %i.1) # <ipython-input-14-d0a2ead71c2a>:7:8
            %z.5 : int = aten::mul(%z.11, %5) # <ipython-input-14-d0a2ead71c2a>:8:8
            -> (%2, %y.5, %z.5)
 to:
    %y : int, %z : int = prim::Loop(%6, %2, %y.1, %z.1) # <ipython-input-14-d0a2ead71c2a>:6:4
        block0(%i.1 : int, %y.11 : int, %z.11 : int):
            %y.5 : int = aten::add(%y.1, %i.1) # <ipython-input-14-d0a2ead71c2a>:7:8
            %z.5 : int = aten::mul(%z.1, %5) # <ipython-input-14-d0a2ead71c2a>:8:8
            -> (%2, %y.5, %z.5)
 */
void LoopBodyLegal(Graph* graph, Block* block) {
    for (auto it = block->nodes().begin(); it != block->nodes().end();) {
        auto* node = *it;
        it++;

        for (Block* sub_block : node->blocks()) {
            LoopBodyLegal(graph, sub_block);
        }
        if (node->kind() == prim::Loop) {
            auto body = node->blocks()[0];
            for (int i = body->inputs().size() - 1; i > 0; i--) {
                body->inputs().at(i)->replaceAllUsesWith(node->inputs().at(i + 1));
            }
        }
    }
}
/*
 inference input type, such as below:
    x: Tensor;
    y = aten::embedding(_, x);
 then x's scalar type is int
*/
void InputTypeInfer(Graph* graph) {
    // TODO: add more typeOps and propagateOps
    static std::map<NodeKind, std::vector<ScalarType>> opInputTypes {
        // aten::embedding(Tensor weight, Tensor indices, int padding_idx, bool scale_grad_by_freq, bool sparse) -> Tensor
        { aten::embedding, { ScalarType::Float, ScalarType::Int } },
        // aten::matmul(Tensor self, Tensor other) -> Tensor
        { aten::matmul, { ScalarType::Float, ScalarType::Float } },
        // aten::linear(Tensor input, Tensor weight, Tensor bias) -> Tensor
        { aten::linear, { ScalarType::Float, ScalarType::Float, ScalarType::Float } },
        // aten::conv2d(Tensor input, Tensor weight, Tensor bias, int[] stride, int[] padding, int[] dilation, int groups) -> Tensor
        { aten::conv2d, { ScalarType::Float, ScalarType::Float, ScalarType::Float } },
    };
    static std::set<NodeKind> typePropagateOps {
        // shape change
        aten::slice, aten::view, aten::transpose, aten::permute,
        // compute
        aten::add, aten::sub, aten::mul, aten::div,
    };
    auto mergeType = [](ScalarType type, ScalarType newType) {
        if (type == newType || newType == c10::ScalarType::Undefined) {
            return type;
        }
        if (type == c10::ScalarType::Undefined) {
            return newType;
        }
        MNN_ASSERT(false);
        return c10::ScalarType::Undefined;
    };
    std::function<ScalarType(Value*)> getScalarType = [&](Value* input) -> ScalarType {
        auto inputType = ScalarType::Undefined;
        for (auto use : input->uses()) {
            int idx = -1;
            for (int i = 0; i < use.user->inputs().size(); i++) {
                if (use.user->input(i) == input) {
                    idx = i;
                }
            }
            auto newType = ScalarType::Undefined;
            if (typePropagateOps.find(use.user->kind()) != typePropagateOps.end()) {
                newType = getScalarType(use.user->output());
            } else {
                const auto iter = opInputTypes.find(use.user->kind());
                if (iter != opInputTypes.end() && idx >= 0 && idx < iter->second.size()) {
                    newType = iter->second[idx];
                }
            }
            inputType = mergeType(inputType, newType);
        }
        return inputType;
    };

    for (auto input : graph->inputs()) {
        auto type = input->type()->cast<TensorType>();
        if (!type) {
            continue;
        }
        auto scalarType = getScalarType(input);
        input->setType(type->withScalarType(scalarType));
    }
}

/*
Unpack outputs, such as below:
    return List(x, y); -> return x, y;
    return Dict('x', x); -> return x;
    return Tuple(Tuple(x, y), z); return x, y, z;
*/
void OutputsUnpack(Graph* graph) {
    std::function<void(Node* tuple, std::vector<Node*>&, std::vector<Value*>&)> flattenTuple =
    [&flattenTuple](Node* tuple, std::vector<Node*>& tuples, std::vector<Value*>& values) -> void
    {
        tuples.push_back(tuple);
        for (auto input : tuple->inputs()) {
            auto node = input->node();
            if (node->kind() == prim::TupleConstruct) {
                flattenTuple(node, tuples, values);
            } else {
                values.push_back(input);
            }
        }
    };
    for (int i = 0; i < graph->outputs().size(); i++) {
        auto node = graph->outputs()[i]->node();
        // unpack output
        switch (node->kind()) {
            case prim::TupleConstruct: {
                std::vector<Node*> tuples;
                std::vector<Value*> values;
                flattenTuple(node, tuples, values);
                for (auto realOutput : values) {
                    graph->registerOutput(realOutput);
                }
                graph->eraseOutput(i);
                for (auto tuple : tuples) {
                    if (!tuple->hasUses()) {
                        tuple->destroy();
                    }
                }
                break;
            }
            case prim::DictConstruct: {
                graph->registerOutput(node->input(1));
                graph->eraseOutput(i);
                node->destroy();
                break;
            }
            case prim::ListConstruct: {
                for (int i = 0; i < node->inputs().size(); i++) {
                    graph->registerOutput(node->input(i));
                }
                graph->eraseOutput(i);
                node->destroy();
                break;
            }
        }
    }
}

/*
distinguish overloaded function, such as below:
    torch.max(Tensor, Tensor) is compare
    torch.max(Tensor, int) is reduce
*/
void overloadDistinguish(Block* block) {
    auto symb = c10::Symbol::fromQualString("attr::mnn_tag");
    for (auto it = block->nodes().begin(); it != block->nodes().end();) {
        auto* node = *it;
        it++;

        for (Block* sub_block : node->blocks()) {
            overloadDistinguish(sub_block);
        }
        switch (node->kind()) {
            // min/max(Tensor, Tensor) is compare
            // min/max(Tensor, int) is reduce
            case aten::min:
            case aten::max:
            case aten::sum:
                if (node->inputs().size() > 1 &&
                    (node->input(1)->type()->kind() == c10::TypeKind::IntType ||
                     node->input(1)->type()->kind() == c10::TypeKind::ListType)) {
                    node->s_(symb, "reduce");
                } else {
                    node->s_(symb, "binary");
                }
                break;
            case aten::index:
                if (node->input(1)->node()->kind() == prim::ListConstruct) {
                    node->s_(symb, "stridedslice");
                }
                break;
            default:
                // do nothing
                break;
        }
    }
}
/*
fuse as_tensor, such as below:
    d = prim::dtype(b);
    c = aten::as_tensor(a, d);
 -> c = aten::type_as(a, b)
*/
void FuseAsTensor(Graph* graph, Block* block) {
    for (auto it = block->nodes().begin(); it != block->nodes().end();) {
        auto* node = *it;
        it++;
        for (Block* sub_block : node->blocks()) {
            FuseAsTensor(graph, sub_block);
        }
        if (node->kind() == prim::dtype) {
            for (auto use : node->output(0)->uses()) {
                auto as_tensor = use.user;
                Node* typeAs = graph->create(aten::type_as, 1);
                typeAs->addInput(as_tensor->input(0));
                typeAs->addInput(node->input(0));
                typeAs->output(0)->copyMetadata(as_tensor->output(0));
                as_tensor->replaceAllUsesWith(typeAs);
                as_tensor->removeAllInputs();
                as_tensor->destroy();
            }
        }
    }
}

/*
fuse uniform, such as below:
    d = aten::empty(shape);
    c = aten::uniform_(a, low, hight);
 -> c = aten::uniform_(shape, low, hight)
*/
void FuseUniform(Graph* graph, Block* block) {
    for (auto it = block->nodes().begin(); it != block->nodes().end();) {
        auto* node = *it;
        it++;
        for (Block* sub_block : node->blocks()) {
            FuseUniform(graph, sub_block);
        }
        if (it->kind().toUnqualString() == std::string("uniform_")) {
            auto input = it->input(0)->node();
            if (input->kind() == aten::empty) {
                it->replaceInput(0, input->input(0));
                input->destroy();
            }
        }
    }
}

std::shared_ptr<Graph> torchOptPass(Module& module) {
    module.eval();
    module = torch::jit::freeze_module(module);
    auto graph = module.get_methods()[0].graph();
    Inline(*(graph.get()));
    // normalize, Example: aten::absolute -> aten::abs
    NormalizeOps(graph);
    // remove some ops, Example: prim::RaiseException
    removeUselessOps(graph->block());
    removeDropout(module);
    // Example: x = x + 1; -> x_1 = x + 1;
    RemoveInplaceOps(graph);
    // Example: x = {v0}; x.append(v1); -> x = {v0, v1};
    // RemoveListMutation(graph);
    removeListAppend(graph.get(), graph->block());
    removeListConstructOps(graph->block());
    //RemoveTensorMutation(graph);
    // elimate dead code
    EliminateDeadCode(graph, DCESideEffectPolicy::ALLOW_DELETING_NODES_WITH_SIDE_EFFECTS);
    // constant propagation
    ConstantPooling(graph);
    ConstantPropagation(graph);
    // fuse
    FuseGraph(graph, true);
    PeepholeOptimize(graph);
    FuseAddMM(graph);
    // FoldConvBatchNorm(module);
    FuseListUnpack(graph->block());
    FuseListStack(graph->block());
    // distinguish overload function
    overloadDistinguish(graph->block());
    // legal loop body's var name
    LoopBodyLegal(graph.get(), graph->block());
    // infer input tensor's scalar type by op
    InputTypeInfer(graph.get());
    // split output tensor if wrap with list/tuple
    OutputsUnpack(graph.get());
    // dtype + as_tensor -> type_as
    FuseAsTensor(graph.get(), graph->block());
    // empty + uniform -> uniform
    FuseUniform(graph.get(), graph->block());
#ifdef MNN_DUMP_TORCHSCRIPT
    graph->dump();
#endif
    return graph;
}
}
}