#include "LowerParallelTasks.h"

#include <string>

#include "Argument.h"
#include "Closure.h"
#include "DebugArguments.h"
#include "ExprUsesVar.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "Module.h"
#include "Param.h"
#include "Simplify.h"

namespace Halide {
namespace Internal {

namespace {

LoweredArgument make_scalar_arg(const std::string &name, const Type &type) {
    return LoweredArgument(name, Argument::Kind::InputScalar, type, 0, ArgumentEstimates());
}

template<typename T>
LoweredArgument make_scalar_arg(const std::string &name) {
    return make_scalar_arg(name, type_of<T>());
}

std::string task_debug_name(const std::pair<std::string, int> &prefix) {
    if (prefix.second <= 1) {
        return prefix.first;
    } else {
        return prefix.first + "_" + std::to_string(prefix.second - 1);
    }
}

void add_fork(std::pair<std::string, int> &prefix) {
    if (prefix.second == 0) {
        prefix.first += ".fork";
    }
    prefix.second++;
}

void add_suffix(std::pair<std::string, int> &prefix, const std::string &suffix) {
    if (prefix.second > 1) {
        prefix.first += "_" + std::to_string(prefix.second - 1);
        prefix.second = 0;
    }
    prefix.first += suffix;
}

// TODO(zvookin|abadams): This makes multiple passes over the
// IR to cover each node. (One tree walk produces the min
// thread count for all nodes, but we redo each subtree when
// compiling a given node.) Ideally we'd move to a lowering pass
// that converts our parallelism constructs to Call nodes, or
// direct hardware operations in some cases.
// Also, this code has to exactly mirror the logic in get_parallel_tasks.
// It would be better to do one pass on the tree and centralize the task
// deduction logic in one place.
class MinThreads : public IRVisitor {
    using IRVisitor::visit;

    std::pair<Stmt, int> skip_acquires(Stmt first) {
        int count = 0;
        while (first.defined()) {
            const Acquire *acq = first.as<Acquire>();
            if (acq == nullptr) {
                break;
            }
            count++;
            first = acq->body;
        }
        return {first, count};
    }

    void visit(const Fork *op) override {
        int total_threads = 0;
        int direct_acquires = 0;
        // Take the sum of min threads across all
        // cascaded Fork nodes.
        const Fork *node = op;
        while (node != nullptr) {
            result = 0;
            auto after_acquires = skip_acquires(node->first);
            direct_acquires += after_acquires.second;

            after_acquires.first.accept(this);
            total_threads += result;

            const Fork *continued_branches = node->rest.as<Fork>();
            if (continued_branches == nullptr) {
                result = 0;
                after_acquires = skip_acquires(node->rest);
                direct_acquires += after_acquires.second;
                after_acquires.first.accept(this);
                total_threads += result;
            }
            node = continued_branches;
        }
        if (direct_acquires == 0 && total_threads == 0) {
            result = 0;
        } else {
            result = total_threads + 1;
        }
    }

    void visit(const For *op) override {
        result = 0;

        if (op->for_type == ForType::Parallel) {
            IRVisitor::visit(op);
            if (result > 0) {
                result += 1;
            }
        } else if (op->for_type == ForType::Serial) {
            auto after_acquires = skip_acquires(op->body);
            if (after_acquires.second > 0 &&
                !expr_uses_var(op->body.as<Acquire>()->count, op->name)) {
                after_acquires.first.accept(this);
                result++;
            } else {
                IRVisitor::visit(op);
            }
        } else {
            IRVisitor::visit(op);
        }
    }

    // This is a "standalone" Acquire and will result in its own task.
    // Treat it requiring one more thread than its body.
    void visit(const Acquire *op) override {
        result = 0;
        auto after_inner_acquires = skip_acquires(op);
        after_inner_acquires.first.accept(this);
        result = result + 1;
    }

    void visit(const Block *op) override {
        result = 0;
        op->first.accept(this);
        int result_first = result;
        result = 0;
        op->rest.accept(this);
        result = std::max(result, result_first);
    }

public:
    int result = 0;
};

int calculate_min_threads(const Stmt &body) {
    MinThreads min_threads;
    body.accept(&min_threads);
    return min_threads.result;
}

struct LowerParallelTasks : public IRMutator {

    /** Codegen a call to do_parallel_tasks */
    struct ParallelTask {
        Stmt body;
        struct SemAcquire {
            Expr semaphore;
            Expr count;
        };
        std::vector<SemAcquire> semaphores;
        std::string loop_var;
        Expr min, extent;
        Expr serial;
        std::string name;
    };

    using IRMutator::visit;

    Stmt visit(const For *op) override {
        const Acquire *acquire = op->body.as<Acquire>();

        if (op->for_type == ForType::Parallel ||
            (op->for_type == ForType::Serial &&
             acquire &&
             !expr_uses_var(acquire->count, op->name))) {
            return do_as_parallel_task(op);
        }
        return IRMutator::visit(op);
    }

    Stmt visit(const Acquire *op) override {
        return do_as_parallel_task(op);
    }

    Stmt visit(const Fork *op) override {
        return do_as_parallel_task(op);
    }

    Stmt rewrite_parallel_tasks(const std::vector<ParallelTask> &tasks) {
        Stmt body;

        Closure closure;
        for (const auto &t : tasks) {
            Stmt s = t.body;
            if (!t.loop_var.empty()) {
                s = LetStmt::make(t.loop_var, 0, s);
            }
            closure.include(s);
        }

        // The same name can appear as a var and a buffer. Remove the var name in this case.
        for (auto const &b : closure.buffers) {
            closure.vars.erase(b.first);
        }

        int num_tasks = (int)(tasks.size());
        std::vector<Expr> tasks_array_args;
        tasks_array_args.reserve(num_tasks * 9);

        std::string closure_name = unique_name("parallel_closure");
        Expr closure_struct_allocation = closure.pack_into_struct();
        Expr closure_struct = Variable::make(Handle(), closure_name);

        const bool has_task_parent = !task_parents.empty() && task_parents.top_ref().defined();

        Expr result;
        for (int i = 0; i < num_tasks; i++) {
            ParallelTask t = tasks[i];

            const int min_threads = calculate_min_threads(t.body);

            // Decide if we're going to call do_par_for or
            // do_parallel_tasks. halide_do_par_for is simpler, but
            // assumes a bunch of things. Programs that don't use async
            // can also enter the task system via do_par_for.
            const bool use_parallel_for = (num_tasks == 1 &&
                                           min_threads == 0 &&
                                           t.semaphores.empty() &&
                                           !has_task_parent);

            Expr closure_task_parent;

            const std::string closure_arg_name = unique_name("closure_arg");
            auto closure_arg = make_scalar_arg<uint8_t *>(closure_arg_name);

            Type closure_function_type;

            std::vector<LoweredArgument> closure_args(use_parallel_for ? 3 : 5);
            closure_args[0] = make_scalar_arg<void *>("__user_context");
            if (use_parallel_for) {
                // The closure will be a halide_task_t, with arguments like:
                //
                //   typedef int (*halide_task_t)(void *user_context, int task_number, uint8_t *closure);
                //
                closure_function_type = type_of<halide_task_t>();

                closure_args[1] = make_scalar_arg<int32_t>(t.loop_var);
                closure_args[2] = closure_arg;
                // closure_task_parent remains undefined here.
            } else {
                // The closure will be a halide_loop_task_t, with arguments like:
                //
                //   typedef int (*halide_loop_task_t)(void *user_context, int min, int extent, uint8_t *closure, void *task_parent);
                //
                closure_function_type = type_of<halide_loop_task_t>();

                const std::string closure_task_parent_name = unique_name("__task_parent");
                closure_task_parent = Variable::make(type_of<void *>(), closure_task_parent_name);
                // We peeled off a loop. Wrap a new loop around the body
                // that just does the slice given by the arguments.
                std::string loop_min_name = unique_name('t');
                std::string loop_extent_name = unique_name('t');
                if (!t.loop_var.empty()) {
                    t.body = For::make(t.loop_var,
                                       Variable::make(Int(32), loop_min_name),
                                       Variable::make(Int(32), loop_extent_name),
                                       ForType::Serial,
                                       DeviceAPI::None,
                                       t.body);
                } else {
                    internal_assert(is_const_one(t.extent));
                }
                closure_args[1] = make_scalar_arg<int32_t>(loop_min_name);
                closure_args[2] = make_scalar_arg<int32_t>(loop_extent_name);
                closure_args[3] = closure_arg;
                closure_args[4] = make_scalar_arg<void *>(closure_task_parent_name);
            }

            {
                ScopedValue<std::string> save_name(function_name, t.name);

                task_parents.push(closure_task_parent);
                t.body = mutate(t.body);
                task_parents.pop();
            }

            const std::string new_function_name = c_print_name(unique_name(t.name), false);
            {
                Expr closure_arg_var = Variable::make(closure_struct_allocation.type(), closure_arg_name);
                Stmt wrapped_body = closure.unpack_from_struct(closure_arg_var, t.body);

                // TODO(zvookin): Figure out how we want to handle name mangling of closures.
                // For now, the C++ backend makes them extern "C" so they have to be NameMangling::C.
                LoweredFunc closure_func{new_function_name, closure_args, std::move(wrapped_body), LinkageType::Internal, NameMangling::C};
                if (target.has_feature(Target::Debug)) {
                    debug_arguments(&closure_func, target);
                }
                closure_implementations.emplace_back(std::move(closure_func));
            }

            // Codegen will add user_context for us
            // Prefix the function name with "::" as we would in C++ to make
            // it clear we're talking about something in global scope in
            // case some joker names an intermediate Func or Var the same
            // name as the pipeline. This prefix works transparently in the
            // C++ backend.
            Expr new_function_name_arg = Variable::make(closure_function_type, "::" + new_function_name);
            Expr closure_struct_arg = Cast::make(type_of<uint8_t *>(), closure_struct);

            if (use_parallel_for) {
                std::vector<Expr> args = {
                    std::move(new_function_name_arg),
                    t.min,
                    t.extent,
                    std::move(closure_struct_arg)};
                result = Call::make(type_of<halide_error_code_t>(), "halide_do_par_for", args, Call::Extern);
            } else {
                const int semaphores_size = (int)t.semaphores.size();
                std::vector<Expr> semaphore_args(semaphores_size * 2);
                for (int i = 0; i < semaphores_size; i++) {
                    semaphore_args[i * 2] = t.semaphores[i].semaphore;
                    semaphore_args[i * 2 + 1] = t.semaphores[i].count;
                }
                Expr semaphores_array = Call::make(type_of<halide_semaphore_acquire_t *>(), Call::make_struct, semaphore_args, Call::PureIntrinsic);

                tasks_array_args.emplace_back(std::move(new_function_name_arg));
                tasks_array_args.emplace_back(std::move(closure_struct_arg));
                tasks_array_args.emplace_back(StringImm::make(t.name));
                tasks_array_args.emplace_back(std::move(semaphores_array));
                tasks_array_args.emplace_back((int)t.semaphores.size());
                tasks_array_args.emplace_back(t.min);
                tasks_array_args.emplace_back(t.extent);
                tasks_array_args.emplace_back(min_threads);
                tasks_array_args.emplace_back(Cast::make(Bool(), t.serial));
            }
        }

        if (!tasks_array_args.empty()) {
            // Allocate task list array
            Expr tasks_list = Call::make(type_of<halide_parallel_task_t *>(), Call::make_struct, tasks_array_args, Call::PureIntrinsic);
            Expr user_context = Call::make(type_of<void *>(), Call::get_user_context, {}, Call::PureIntrinsic);
            Expr task_parent = has_task_parent ? task_parents.top() : make_zero(Handle());
            result = Call::make(type_of<halide_error_code_t>(), "halide_do_parallel_tasks",
                                {user_context, make_const(Int(32), num_tasks), tasks_list, task_parent},
                                Call::Extern);
        }

        std::string closure_result_name = unique_name("closure_result");
        Expr closure_result = Variable::make(Int(32), closure_result_name);
        Stmt stmt = AssertStmt::make(closure_result == 0, closure_result);
        stmt = LetStmt::make(closure_result_name, result, stmt);
        stmt = LetStmt::make(closure_name, closure_struct_allocation, stmt);
        return stmt;
    }

    void get_parallel_tasks(const Stmt &s, std::vector<ParallelTask> &result, std::pair<std::string, int> prefix) {
        const For *loop = s.as<For>();
        const Acquire *acquire = loop ? loop->body.as<Acquire>() : s.as<Acquire>();
        if (const Fork *f = s.as<Fork>()) {
            add_fork(prefix);
            get_parallel_tasks(f->first, result, prefix);
            get_parallel_tasks(f->rest, result, prefix);
        } else if (!loop && acquire) {
            const Variable *v = acquire->semaphore.as<Variable>();
            internal_assert(v);
            add_suffix(prefix, "." + v->name);
            ParallelTask t{s, {}, "", 0, 1, const_false(), task_debug_name(prefix)};
            while (acquire) {
                t.semaphores.push_back({acquire->semaphore, acquire->count});
                t.body = acquire->body;
                acquire = t.body.as<Acquire>();
            }
            result.emplace_back(std::move(t));
        } else if (loop && loop->for_type == ForType::Parallel) {
            add_suffix(prefix, ".par_for." + loop->name);
            ParallelTask t{loop->body, {}, loop->name, loop->min, loop->extent, const_false(), task_debug_name(prefix)};
            result.emplace_back(std::move(t));
        } else if (loop &&
                   loop->for_type == ForType::Serial &&
                   acquire &&
                   !expr_uses_var(acquire->count, loop->name)) {
            const Variable *v = acquire->semaphore.as<Variable>();
            internal_assert(v);
            add_suffix(prefix, ".for." + v->name);
            ParallelTask t{loop->body, {}, loop->name, loop->min, loop->extent, const_true(), task_debug_name(prefix)};
            while (acquire) {
                t.semaphores.push_back({acquire->semaphore, acquire->count});
                t.body = acquire->body;
                acquire = t.body.as<Acquire>();
            }
            result.emplace_back(std::move(t));
        } else {
            add_suffix(prefix, "." + std::to_string(result.size()));
            ParallelTask t{s, {}, "", 0, 1, const_false(), task_debug_name(prefix)};
            result.emplace_back(std::move(t));
        }
    }

    Stmt do_as_parallel_task(const Stmt &s) {
        std::vector<ParallelTask> tasks;
        get_parallel_tasks(s, tasks, {function_name, 0});
        return rewrite_parallel_tasks(tasks);
    }

    LowerParallelTasks(const std::string &name, const Target &t)
        : function_name(name), target(t) {
    }

    std::string function_name;
    const Target &target;
    std::vector<LoweredFunc> closure_implementations;
    SmallStack<Expr> task_parents;
};

}  // namespace

Stmt lower_parallel_tasks(const Stmt &s, std::vector<LoweredFunc> &closure_implementations,
                          const std::string &name, const Target &t) {
    LowerParallelTasks lowering_mutator(name, t);
    Stmt result = lowering_mutator.mutate(s);

    // Main body will be dumped as part of standard lowering debugging, but closures will not be.
    if (debug::debug_level() >= 2) {
        for (const auto &lf : lowering_mutator.closure_implementations) {
            debug(2) << "lower_parallel_tasks generated closure lowered function " << lf.name << ":\n"
                     << lf.body << "\n\n";
        }
    }

    // Append to the end rather than replacing the list entirely.
    closure_implementations.insert(closure_implementations.end(),
                                   lowering_mutator.closure_implementations.begin(),
                                   lowering_mutator.closure_implementations.end());

    return result;
}

}  // namespace Internal
}  // namespace Halide