path: root/eval/src/vespa/eval/eval/interpreted_function.cpp

// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include "interpreted_function.h"
#include "node_visitor.h"
#include "node_traverser.h"
#include "tensor_nodes.h"
#include "make_tensor_function.h"
#include "optimize_tensor_function.h"
#include "compile_tensor_function.h"
#include <vespa/vespalib/util/classname.h>
#include <vespa/eval/eval/llvm/compile_cache.h>
#include <vespa/eval/eval/llvm/addr_to_symbol.h>
#include <vespa/vespalib/util/benchmark_timer.h>
#include <vespa/vespalib/util/stringfmt.h>
#include <set>

namespace vespalib::eval {

namespace {

const Function *get_simple_lambda(const nodes::Node &node) {
    if (auto ptr = nodes::as<nodes::TensorMap>(node)) {
        return &ptr->lambda();
    }
    if (auto ptr = nodes::as<nodes::TensorJoin>(node)) {
        return &ptr->lambda();
    }
    if (auto ptr = nodes::as<nodes::TensorMerge>(node)) {
        return &ptr->lambda();
    }
    return nullptr;
}

const Function *get_complex_lambda(const nodes::Node &node) {
    if (auto ptr = nodes::as<nodes::TensorLambda>(node)) {
        return &ptr->lambda();
    }
    if (auto ptr = nodes::as<nodes::TensorMapSubspaces>(node)) {
        return &ptr->lambda();
    }
    return nullptr;
}

void my_nop(InterpretedFunction::State &, uint64_t) {}

} // namespace vespalib::<unnamed>


InterpretedFunction::State::State(const ValueBuilderFactory &factory_in)
    : factory(factory_in),
      params(nullptr),
      stash(),
      stack(),
      program_offset(0),
      if_cnt(0)
{
}

InterpretedFunction::State::~State() = default;

void
InterpretedFunction::State::init(const LazyParams &params_in) {
    params = &params_in;
    stash.clear();
    stack.clear();
    program_offset = 0;
    if_cnt = 0;
}

InterpretedFunction::Context::Context(const InterpretedFunction &ifun)
  : _state(ifun._factory)
{
}

InterpretedFunction::ProfiledContext::ProfiledContext(const InterpretedFunction &ifun)
  : context(ifun),
    cost(ifun.program_size(), std::make_pair(size_t(0), duration::zero()))
{
}

InterpretedFunction::ProfiledContext::~ProfiledContext() = default;

vespalib::string
InterpretedFunction::Instruction::resolve_symbol() const
{
    if (function == nullptr) {
        return "<inject_param>";
    }
    return addr_to_symbol((const void *)function);
}

InterpretedFunction::Instruction
InterpretedFunction::Instruction::nop()
{
    return Instruction(my_nop);
}

InterpretedFunction::Options::Options(const ValueBuilderFactory &factory_in)
  : _factory(factory_in),
    _optimize(optimize_tensor_function),
    _meta(nullptr)
{
}

InterpretedFunction::InterpretedFunction(const ValueBuilderFactory &factory, const TensorFunction &function)
  : _program(),
    _stash(),
    _factory(factory)
{
    _program = compile_tensor_function(function, CTFContext(factory, _stash, nullptr));
}

InterpretedFunction::InterpretedFunction(const Options &opts, const nodes::Node &root, const NodeTypes &types)
    : _program(),
      _stash(),
      _factory(opts.factory())
{
    const TensorFunction &plain_fun = make_tensor_function(opts.factory(), root, types, _stash);
    const TensorFunction &optimized = opts.optimize()(opts.factory(), plain_fun, _stash);
    _program = compile_tensor_function(optimized, CTFContext(opts.factory(), _stash, opts.meta()));
}

InterpretedFunction::~InterpretedFunction() = default;

const Value &
InterpretedFunction::eval(Context &ctx, const LazyParams &params) const
{
    State &state = ctx._state;
    state.init(params);
    while (state.program_offset < _program.size()) {
        _program[state.program_offset++].perform(state);
    }
    assert(state.stack.size() == 1);
    return state.stack.back();
}

const Value &
InterpretedFunction::eval(ProfiledContext &pctx, const LazyParams &params) const
{
    auto &ctx = pctx.context;                            // Profiling
    State &state = ctx._state;
    state.init(params);
    while (state.program_offset < _program.size()) {
        auto pos = state.program_offset;                 // Profiling
        auto before = steady_clock::now();               // Profiling
        _program[state.program_offset++].perform(state);
        auto after = steady_clock::now();                // Profiling
        ++pctx.cost[pos].first;                          // Profiling
        pctx.cost[pos].second += (after - before);       // Profiling
    }
    assert(state.stack.size() == 1);
    return state.stack.back();
}

double
InterpretedFunction::estimate_cost_us(const std::vector<double> &params, double budget) const
{
    Context ctx(*this);
    SimpleParams lazy_params(params);
    auto actual = [&](){eval(ctx, lazy_params);};
    return BenchmarkTimer::benchmark(actual, budget) * 1000.0 * 1000.0;
}

Function::Issues
InterpretedFunction::detect_issues(const Function &function)
{
    struct NotSupported : NodeTraverser {
        Function::Issues issues;
        bool open(const nodes::Node &) override { return true; }
        void close(const nodes::Node &node) override {
            // map/join/merge: simple scalar lambdas must be compilable with llvm
            if (auto lambda = get_simple_lambda(node)) {
                auto inner_issues = CompiledFunction::detect_issues(*lambda);
                if (inner_issues) {
                    auto ctx = make_string("within %s simple lambda", getClassName(node).c_str());
                    issues.add_nested_issues(ctx, inner_issues);
                }
            }
            // tensor lambda/map_subspaces: complex lambdas that may be interpreted and use tensor math
            if (auto lambda = get_complex_lambda(node)) {
                auto inner_issues = InterpretedFunction::detect_issues(*lambda);
                if (inner_issues) {
                    auto ctx = make_string("within %s complex lambda", getClassName(node).c_str());
                    issues.add_nested_issues(ctx, inner_issues);
                }
            }
        }
    } checker;
    function.root().traverse(checker);
    return std::move(checker.issues);
}

InterpretedFunction::EvalSingle::EvalSingle(const ValueBuilderFactory &factory, Instruction op, const LazyParams &params)
    : _state(factory),
      _op(op)
{
    _state.params = &params;
}

const Value &
InterpretedFunction::EvalSingle::eval(const std::vector<Value::CREF> &stack)
{
    _state.stash.clear();
    _state.stack = stack;
    _op.perform(_state);
    assert(_state.stack.size() == 1);
    return _state.stack.back();
}

}