/**
 * \file src/jit/impl/executor_opr.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#include "megbrain/jit/executor_opr.h"
#include "megbrain/common.h"
#include "megbrain/comp_node_env.h"
#include "megbrain/gopt/framework.h"
#include "megbrain/graph/grad_impl.h"
#include "megbrain/graph/helper.h"
#include "megbrain/jit/compiler.h"
#include "megbrain/jit/param_elem_visitor.h"
#include "megbrain/jit/placeholder_opr.h"
#include "megbrain/opr/basic_arith_wrapper.h"
#include "megbrain/opr/tensor_manip.h"
#include "megbrain/opr/utility.h"
#include "megbrain/utils/hash.h"
#include "megbrain/serialization/opr_shallow_copy.h"

#if MGB_JIT

using namespace mgb;
using namespace jit;

using CPFlag = Compiler::Property::Flag;
/* =================== Fusion ==================== */

MGB_DYN_TYPE_OBJ_FINAL_IMPL(JITExecutor);
JITExecutor::JITExecutor(const InternalGraphPtr& internal_graph,
                         const VarNodeArray& inputs,
                         const OperatorNodeConfig& config)
        : Super(internal_graph->output()->owner_graph(), config,
                ssprintf("JIT-Fusion{%zu}",
                         internal_graph->placeholders().size()),
                inputs),
          m_internal_graph{internal_graph},
          m_compiler{Compiler::get(*inputs[0]->owner_graph(),
                                   inputs[0]->comp_node())} {
    for (auto inp : inputs) {
        add_input({inp});
    }
    m_input_broadcastable.resize(inputs.size());
    auto&& placeholders = m_internal_graph->placeholders();
    mgb_assert(placeholders.size() == inputs.size());
    for (size_t i = 0; i < inputs.size(); ++i) {
        mgb_assert(placeholders[i]->output(0) != internal_graph->output());
        if (placeholders[i]->is_host_value_shape_input() ||
            input()[i]
                    ->owner_opr()
                    ->same_type<opr::MarkNoBroadcastElemwise>()) {
            m_input_broadcastable[i] = false;
        } else {
            m_input_broadcastable[i] = true;
        }
    }
    if (inputs.size() == 1) {
        m_input_broadcastable[0] = false;
    } else {
        Maybe<size_t> non_scalar;
        for (size_t i = 0; i < input().size(); ++i) {
            if (placeholders[i]->is_host_value_shape_input())
                continue;
            if (!(cg::is_const_var_shape(input(i)) &&
                  input(i)->shape().is_scalar())) {
                if (non_scalar.valid()) {
                    non_scalar.invalidate();
                    break;
                }
                non_scalar = i;
            }
        }
        if (non_scalar.valid()) {
            // exactly one input is non-scalar
            m_input_broadcastable[non_scalar.val()] = false;
        }
    }
    add_output(None)->dtype(m_internal_graph->output()->dtype());
    add_equivalence_component<ScalarHash<void*>>(internal_graph->output());
    for (size_t i = 0, it = m_compiler->get_nr_workspace_outputs(this); i < it;
         ++i) {
        cg::add_workspace_output(this);
    }

    // check if there is reduce or dimshuffle opr
    cg::DepOprIter{[this](cg::OperatorNodeBase* opr) {
        if (opr->same_type<opr::Reduce>()) {
            m_feature_bits |= JITFeatureBits::REDUCE;
        }
        if (opr->same_type<opr::Dimshuffle>()) {
            m_feature_bits |= JITFeatureBits::DIMSHUFFLE;
        }
    }}.add(internal_graph->output());
}

void JITExecutor::add_input_layout_constraint() {
    if (m_compiler->property().contain_flag(CPFlag::NEED_INPUT_CONTIG)) {
        for (auto i : input()) {
            i->add_layout_constraint_contiguous();
        }
    } else {
        for (auto i : input()) {
            i->add_layout_constraint_monotone();
        }
    }
}

void JITExecutor::init_output_mem_plan(bool dynamic) {
    Super::init_output_mem_plan(dynamic);
    m_args.need_update = true;
}

SymbolVar JITExecutor::make(const InternalGraphPtr& internal_graph,
                            const VarNodeArray& inputs,
                            const OperatorNodeConfig& config) {
    return internal_graph->output()
            ->owner_graph()
            ->insert_opr(std::make_unique<JITExecutor>(internal_graph, inputs,
                                                       config))
            ->output(0);
}

void JITExecutor::init_output_static_infer_desc() {
    using namespace cg::static_infer;
    auto&& mgr = owner_graph()->static_infer_manager();
    mgr.register_shape_infer(
            output(0),
            ShapeInferDesc::make_identity(m_internal_graph->shape_infer()));
    m_compiler->init_workspace_size_infer(this);
    if (m_internal_graph->value_infer()) {
        mgr.register_value_infer(
                output(0),
                ValueInferDesc::make_identity(m_internal_graph->value_infer()));
    }
}

void JITExecutor::scn_do_execute() {
    if (m_executable == nullptr || m_args.need_update) {
        m_executable = m_compiler->compile(this);
    }
    m_executable->execute(this);
}

//! change the inputs which depend on dimshuffle opr, make sure dimshuffles
//! can be ignored
void JITExecutor::do_dimshuffle() {

    auto get_dimshuffled_layout = [](const TensorLayout& ily, int32_t* pattern,
                                 size_t pattern_len) {

        TensorLayout oly{ily.dtype};
        oly.ndim = pattern_len;

        bool input_used[TensorLayout::MAX_NDIM] = {0};
        for (uint32_t idx = 0; idx < pattern_len; ++idx) {
            auto i = pattern[idx];
            if (i < 0) {
                oly.shape[idx] = 1;
                oly.stride[idx] = 1;
            } else {
                input_used[i] = true;
                oly.shape[idx] = ily.shape[i];
                oly.stride[idx] = ily.stride[i];
            }
        }

        for (size_t i = 0; i < ily.ndim; ++i) {
            mgb_assert(input_used[i] || ily.shape[i] == 1,
                       "non-1 dim discarded in Dimshuffle: ishp=%s dim=%zd",
                       static_cast<const TensorShape&>(ily).to_string().c_str(),
                       i);
        }
        return oly;
    };

    // DFS to make sure traverse the dimshuffles in one branch
    std::unordered_set<VarNode*> visited;
    std::vector<OperatorNodeBase*> stack(0);
    std::vector<uint8_t> idx(0);  // input index
    stack.push_back(m_internal_graph->output()->owner_opr());
    idx.push_back(0);

    while (!stack.empty()) {
        if (idx.back() < stack.back()->input().size() &&
            !visited.count(stack.back()->input(idx.back()))) {
            visited.insert(stack.back()->input(idx.back()));
            stack.push_back(stack.back()->input(idx.back())->owner_opr());
            if (stack.back()->same_type<jit::JITPlaceholder>()) {
                auto jitph = gopt::try_cast_as_op<JITPlaceholder>(stack.back());
                size_t input_id = jitph->input_id();
                auto&& input = m_args.inputs[input_id];

                for (int i = stack.size() - 1; i >= 0; --i) {
                    if (stack[i]->same_type<opr::Dimshuffle>()) {
                        auto param =
                                stack[i]->cast_final_safe<opr::Dimshuffle>()
                                        .param();

                        mgb_assert(input.layout.ndim == param.ndim,
                                   "input ndim mismatch for Dimshuffle: "
                                   "expect=%u "
                                   "actual=%zu",
                                   param.ndim, input.layout.ndim);
                        auto dimshuffled_layout = get_dimshuffled_layout(
                                input.layout, param.pattern, param.pattern_len);
                        input.layout = dimshuffled_layout;
                    }
                }

                stack.pop_back();
                ++idx.back();
            } else {
                idx.push_back(0);
            }
        } else {
            stack.pop_back();
            idx.pop_back();
            if (!stack.empty())
                ++idx.back();
        }
    }

}

void JITExecutor::update_args() {
    m_args.outputs.clear();
    for (auto out : output()) {
        m_args.outputs.push_back({out, out->layout(), -1});
    }
    m_args.inputs.resize(input().size());

    auto is_host_value_shape_input = [this](size_t idx) {
        return m_internal_graph->placeholders()
                .at(idx)
                ->is_host_value_shape_input();
    };

    for (size_t i = 0; i < input().size(); i++) {
        auto&& dst_data = m_args.inputs[i];
        dst_data.from = input(i);
        dst_data.idx = i;
        if (is_host_value_shape_input(i)) {
            auto&& mgr = owner_graph()->static_infer_manager();
            auto&& shpval_inp_val = &mgr.infer_value(input(i));
            cg::copy_tensor_value_to_shape(dst_data.layout, *shpval_inp_val);
            dst_data.layout.dtype = {};
            for (size_t i = 0; i < dst_data.layout.ndim; ++i) {
                dst_data.layout.stride[i] = 0;
            }
        } else {
            dst_data.layout = input(i)->layout();
        }
    }

    //! dimshuffle opr need to change the input.
    do_dimshuffle();

    if (m_compiler->property().contain_flag(CPFlag::NEED_INPUT_COLLAPSE)) {
        // collective collapse datum layout, try to reduce the output ndim
        opr::Elemwise::TensorLayoutPtrArray inp_layouts;
        inp_layouts.reserve(m_args.inputs.size());
        for (size_t i = 0; i < m_args.inputs.size(); i++) {
            if (!is_host_value_shape_input(i)) {
                inp_layouts.push_back(&m_args.inputs[i].layout);
            }
        }
        opr::Elemwise::broadcast_collective_collapse(inp_layouts,
                                                     &m_args.outputs[0].layout);
    }

    // compute and update hash
    XXHash hstate;

    //  update layout info
    auto prop = m_compiler->property();
    if (prop.contain_flag(CPFlag::BIND_NDIM | CPFlag::BIND_SHAPE)) {
        mgb_assert(prop.contain_flag(CPFlag::BIND_NDIM),
                   "BIND_NDIM must be set if bind_shape is set");
        std::vector<size_t> buf;
        buf.reserve(1024);
        buf.push_back(m_args.inputs.size());
        for (auto&& i : m_args.inputs) {
            buf.push_back(i.layout.ndim);
            if (prop.contain_flag(CPFlag::BIND_SHAPE)) {
                for (size_t j = 0; j < i.layout.ndim; ++j) {
                    buf.push_back(i.layout[j]);
                }
            }
        }
        hstate.update(buf.data(), sizeof(buf[0]) * buf.size());
    }
    m_args.hash = hstate.digest();

    // update version number
    static std::atomic_uint_fast64_t global_version;
    m_args.version = global_version.fetch_add(1);

    m_args.need_update = false;
}

const JITExecutor::Args& JITExecutor::args() const {
    if (m_args.need_update) {
        const_cast<JITExecutor*>(this)->update_args();
    }
    return m_args;
}

bool JITExecutor::Args::operator==(const Args& rhs) const {
    auto&& lhs = *this;
    mgb_assert(!lhs.need_update && !rhs.need_update);
    if (lhs.hash != rhs.hash) {
        return false;
    }
    if (lhs.version == rhs.version) {
        return true;
    }
    if (lhs.outputs.size() != rhs.outputs.size())
        return false;
    if (lhs.inputs.size() != rhs.inputs.size())
        return false;

    auto prop = owner->m_compiler->property();

    if (prop.contain_flag(CPFlag::BIND_NDIM | CPFlag::BIND_SHAPE)) {
        bool (*chk_layout)(const TensorLayout&, const TensorLayout&);
        if (prop.contain_flag(CPFlag::BIND_SHAPE)) {
            chk_layout = [](const TensorLayout& lhs, const TensorLayout& rhs) {
                return lhs.eq_shape(rhs);
            };
        } else {
            chk_layout = [](const TensorLayout& lhs, const TensorLayout& rhs) {
                return lhs.ndim == rhs.ndim;
            };
        }
        for (size_t i = 0; i < lhs.inputs.size(); i++) {
            if (!chk_layout(lhs.inputs[i].layout, rhs.inputs[i].layout))
                return false;
        }
        for (size_t i = 0; i < lhs.outputs.size(); i++) {
            if (!chk_layout(lhs.outputs[i].layout, rhs.outputs[i].layout))
                return false;
        }
    }

    // elect a common version so next check can be fast
    lhs.version = rhs.version = std::min(lhs.version, rhs.version);

    return true;
}

JITExecutor::NodeProp* JITExecutor::do_make_node_prop() const {
    auto ret = Super::do_make_node_prop();
    using DepType = NodeProp::DepType;
    SmallVector<DepType> dt(input().size());
    auto&& placeholders = internal_graph().placeholders();
    for (size_t i = 0; i < dt.size(); ++i) {
        dt[i] = placeholders[i]->is_host_value_shape_input()
                        ? DepType::HOST_VALUE
                        : DepType::DEV_VALUE;
    }
    ret->reset_dep_type(input(), dt);
    return ret;
}

megdnn::TensorShape JITExecutor::broadcasted_input_shape() const {
    megdnn::TensorShapeArray inp_shps;
    megdnn::TensorShape brdcast_shp;
    auto placeholders = m_internal_graph->placeholders();
    for (auto ph : placeholders) {
        if (!ph->is_host_value_shape_input()) {
            inp_shps.push_back(input(ph->input_id())->shape());
        }
    }
    megdnn::Elemwise::deduce_shape(inp_shps, brdcast_shp);
    return brdcast_shp;
}


#if MGB_ENABLE_GRAD
MGB_IMPL_OPR_GRAD(JITExecutor) {
    VarNodeArray grad_inputs;
    for (auto input : opr.input())
        grad_inputs.push_back(input);
    mgb_assert(out_grad[0]);
    grad_inputs.push_back(opr.output(0));
    grad_inputs.push_back(out_grad[0]);
    auto fwd_igraph_ptr = opr.internal_graph_ptr();
    auto output_ph = JITPlaceholder::make(
            fwd_igraph_ptr->output(), fwd_igraph_ptr->placeholders().size());
    auto og_ph = JITPlaceholder::make(
            out_grad[0], fwd_igraph_ptr->placeholders().size() + 1);
    auto loss = opr::VirtualLoss::make({fwd_igraph_ptr->output()}, {og_ph});
    auto gx = cg::grad(loss, fwd_igraph_ptr->placeholders()[wrt_idx]->output(0),
                       false, false);
    if (!gx.node()) {
        return nullptr;
    }
    if (gx.node()->owner_opr()->same_type<opr::InvalidGrad>()) {
        return opr::InvalidGrad::make(opr, wrt_idx);
    }
    if (opr.compiler()->property().feature_bits & JITFeatureBits::REDUCE) {
        // expand the gradient graph into the original graph to handle bcast
        // oprs
        ThinHashMap<VarNode*, VarNode*> old2new;
        VarNodeArray new_inp;
        auto on_opr = [&old2new, &grad_inputs,
                       &new_inp](cg::OperatorNodeBase* opr) {
            if (auto ph = gopt::try_cast_as_op<JITPlaceholder>(opr)) {
                old2new[opr->output(0)] = grad_inputs.at(ph->input_id());
                return;
            }
            if (auto imm = gopt::try_cast_as_op<opr::ImmutableTensor>(opr)) {
                HostTensorND hval{grad_inputs[0]->comp_node()};
                hval.copy_from(imm->value()).sync();
                old2new[opr->output(0)] =
                        opr::ImmutableTensor::make(*opr->owner_graph(), hval)
                                .node();
                return;
            }
            new_inp.clear();
            for (auto inp : opr->input()) {
                new_inp.push_back(old2new.at(inp));
            }
            auto new_opr = serialization::copy_opr_shallow(*opr, new_inp);
            old2new[opr->output(0)] = new_opr->output(0);
        };
        cg::DepOprIter{on_opr}.add(gx.node());
        return old2new.at(gx.node());
    } else {
        PlaceholderArray placeholders = fwd_igraph_ptr->placeholders();
        for (SymbolVar i : {output_ph, og_ph}) {
            placeholders.push_back(
                    &i.node()->owner_opr()->cast_final_safe<JITPlaceholder>());
        }
        for (size_t i = 0; i < placeholders.size(); ++i) {
            if (gx.node() == placeholders[i]->output(0)) {
                return grad_inputs[i];
            }
        }
        auto grad_ig = std::make_shared<InternalGraph>(
                gx.node(), fwd_igraph_ptr->shape_infer(), nullptr,
                std::move(placeholders));
        auto grad_jit = JITExecutor::make(grad_ig, grad_inputs);

        if (opr.input_broadcastable()[wrt_idx]) {
            grad_jit = opr::reduce_sum(
                    grad_jit, opr::GetVarShape::make(opr.input(wrt_idx)));
        }
        return grad_jit.node();
    }
}
#endif  // MGB_ENABLE_GRAD

#endif  // MGB_JIT

// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}