compile tensor function

7 files changed, 266 insertions, 1 deletions

diff --git a/eval/src/vespa/eval/eval/CMakeLists.txt b/eval/src/vespa/eval/eval/CMakeLists.txt
index 3816780d4d9..0eabb4f4219 100644
--- a/eval/src/vespa/eval/eval/CMakeLists.txt
+++ b/eval/src/vespa/eval/eval/CMakeLists.txt
@@ -4,6 +4,7 @@ vespa_add_library(eval_eval OBJECT
     aggr.cpp
     basic_nodes.cpp
     call_nodes.cpp
+    compile_tensor_function.cpp
     delete_node.cpp
     function.cpp
     gbdt.cpp
diff --git a/eval/src/vespa/eval/eval/compile_tensor_function.cpp b/eval/src/vespa/eval/eval/compile_tensor_function.cpp
new file mode 100644
index 00000000000..bfd3696ed98
--- /dev/null
+++ b/eval/src/vespa/eval/eval/compile_tensor_function.cpp
@@ -0,0 +1,83 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "compile_tensor_function.h"
+#include "tensor_function.h"
+
+namespace vespalib::eval {
+
+namespace {
+
+using State = InterpretedFunction::State;
+using Instruction = InterpretedFunction::Instruction;
+
+void op_skip(State &state, uint64_t param) {
+    state.program_offset += param;
+}
+
+void op_skip_if_false(State &state, uint64_t param) {
+    ++state.if_cnt;
+    if (!state.peek(0).as_bool()) {
+        state.program_offset += param;
+    }
+    state.stack.pop_back();
+}
+
+struct Frame {
+    const TensorFunction &node;
+    std::vector<TensorFunction::Child::CREF> children;
+    size_t child_idx;
+    Frame(const TensorFunction &node_in) : node(node_in), children(), child_idx(0) { node.push_children(children); }
+    bool has_next_child() const { return (child_idx < children.size()); }
+    const TensorFunction &next_child() { return children[child_idx++].get().get(); }
+};
+
+struct ProgramCompiler {
+    Stash &stash;
+    std::vector<Frame> stack;
+    std::vector<Instruction> prog;
+    ProgramCompiler(Stash &stash_in) : stash(stash_in), stack(), prog() {}
+
+    void append(const std::vector<Instruction> &other_prog) {
+        prog.insert(prog.end(), other_prog.begin(), other_prog.end());
+    }
+
+    void open(const TensorFunction &node) {
+        if (auto if_node = as<tensor_function::If>(node)) {
+            append(compile_tensor_function(if_node->cond(), stash));
+            auto true_prog = compile_tensor_function(if_node->true_child(), stash);
+            auto false_prog = compile_tensor_function(if_node->false_child(), stash);
+            prog.emplace_back(op_skip_if_false, true_prog.size());
+            append(true_prog);
+            prog.emplace_back(op_skip, false_prog.size());
+            append(false_prog);
+        } else {
+            stack.emplace_back(node);
+        }
+    }
+
+    void close(const TensorFunction &node) {
+        prog.push_back(node.compile_self(stash));
+    }
+
+    std::vector<Instruction> compile(const TensorFunction &function) {
+        open(function);
+        while (!stack.empty()) {
+            if (stack.back().has_next_child()) {
+                open(stack.back().next_child());
+            } else {
+                close(stack.back().node);
+                stack.pop_back();
+            }
+        }
+        return std::move(prog);
+    }
+};
+
+} // namespace vespalib::eval::<unnamed>
+
+std::vector<Instruction> compile_tensor_function(const TensorFunction &function, Stash &stash) {   
+    ProgramCompiler compiler(stash);
+    return compiler.compile(function);
+}
+
+} // namespace vespalib::eval
diff --git a/eval/src/vespa/eval/eval/compile_tensor_function.h b/eval/src/vespa/eval/eval/compile_tensor_function.h
new file mode 100644
index 00000000000..bfac0e0f036
--- /dev/null
+++ b/eval/src/vespa/eval/eval/compile_tensor_function.h
@@ -0,0 +1,16 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "interpreted_function.h"
+#include <vector>
+
+namespace vespalib { class Stash; }
+
+namespace vespalib::eval {
+
+class TensorFunction;
+
+std::vector<InterpretedFunction::Instruction> compile_tensor_function(const TensorFunction &function, Stash &stash);
+
+} // namespace vespalib::eval
diff --git a/eval/src/vespa/eval/eval/interpreted_function.cpp b/eval/src/vespa/eval/eval/interpreted_function.cpp
index 13ab6fe5676..6fcb1829e05 100644
--- a/eval/src/vespa/eval/eval/interpreted_function.cpp
+++ b/eval/src/vespa/eval/eval/interpreted_function.cpp
@@ -6,6 +6,7 @@
 #include "check_type.h"
 #include "tensor_spec.h"
 #include "operation.h"
+#include "tensor_engine.h"
 #include <vespa/vespalib/util/classname.h>
 #include <vespa/eval/eval/llvm/compile_cache.h>
 #include <vespa/vespalib/util/benchmark_timer.h>
diff --git a/eval/src/vespa/eval/eval/interpreted_function.h b/eval/src/vespa/eval/eval/interpreted_function.h
index 2a52a5a8258..1c57b20682f 100644
--- a/eval/src/vespa/eval/eval/interpreted_function.h
+++ b/eval/src/vespa/eval/eval/interpreted_function.h
@@ -3,7 +3,6 @@
 #pragma once
 
 #include "function.h"
-#include "simple_tensor_engine.h"
 #include "node_types.h"
 #include "lazy_params.h"
 #include <vespa/vespalib/util/stash.h>
diff --git a/eval/src/vespa/eval/eval/tensor_function.cpp b/eval/src/vespa/eval/eval/tensor_function.cpp
index 8427cc53a16..62e547cbd7e 100644
--- a/eval/src/vespa/eval/eval/tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/tensor_function.cpp
@@ -11,6 +11,86 @@ namespace vespalib {
 namespace eval {
 namespace tensor_function {
 
+namespace {
+
+using State = InterpretedFunction::State;
+using Instruction = InterpretedFunction::Instruction;
+
+//-----------------------------------------------------------------------------
+
+template <typename T, typename IN>
+uint64_t wrap_param(const IN &value_in) {
+    const T &value = value_in;
+    return (uint64_t)&value;
+}
+
+template <typename T>
+const T &unwrap_param(uint64_t param) { return *((const T *)param); }
+
+//-----------------------------------------------------------------------------
+
+uint64_t to_param(map_fun_t value) { return (uint64_t)value; }
+uint64_t to_param(join_fun_t value) { return (uint64_t)value; }
+map_fun_t to_map_fun(uint64_t param) { return (map_fun_t)param; }
+join_fun_t to_join_fun(uint64_t param) { return (join_fun_t)param; }
+
+//-----------------------------------------------------------------------------
+
+void op_load_const(State &state, uint64_t param) {
+    state.stack.push_back(unwrap_param<Value>(param));
+}
+
+void op_load_param(State &state, uint64_t param) {
+    state.stack.push_back(state.params->resolve(param, state.stash));
+}
+
+//-----------------------------------------------------------------------------
+
+void op_double_map(State &state, uint64_t param) {
+    state.replace(1, state.stash.create<DoubleValue>(to_map_fun(param)(state.peek(0).as_double())));
+}
+
+void op_double_mul(State &state, uint64_t) {
+    state.replace(2, state.stash.create<DoubleValue>(state.peek(1).as_double() * state.peek(0).as_double()));
+}
+
+void op_double_add(State &state, uint64_t) {
+    state.replace(2, state.stash.create<DoubleValue>(state.peek(1).as_double() + state.peek(0).as_double()));
+}
+
+void op_double_join(State &state, uint64_t param) {
+    state.replace(2, state.stash.create<DoubleValue>(to_join_fun(param)(state.peek(1).as_double(), state.peek(0).as_double())));
+}
+
+//-----------------------------------------------------------------------------
+
+void op_tensor_map(State &state, uint64_t param) {
+    state.replace(1, state.engine.map(state.peek(0), to_map_fun(param), state.stash));
+}
+
+void op_tensor_join(State &state, uint64_t param) {
+    state.replace(2, state.engine.join(state.peek(1), state.peek(0), to_join_fun(param), state.stash));
+}
+
+using ReduceParams = std::pair<Aggr,std::vector<vespalib::string>>;
+void op_tensor_reduce(State &state, uint64_t param) {
+    const ReduceParams &params = unwrap_param<ReduceParams>(param);
+    state.replace(1, state.engine.reduce(state.peek(0), params.first, params.second, state.stash));
+}
+
+using RenameParams = std::pair<std::vector<vespalib::string>,std::vector<vespalib::string>>;
+void op_tensor_rename(State &state, uint64_t param) {
+    const RenameParams &params = unwrap_param<RenameParams>(param);
+    state.replace(1, state.engine.rename(state.peek(0), params.first, params.second, state.stash));
+}
+
+void op_tensor_concat(State &state, uint64_t param) {
+    const vespalib::string &dimension = unwrap_param<vespalib::string>(param);
+    state.replace(2, state.engine.concat(state.peek(1), state.peek(0), dimension, state.stash));
+}
+
+} // namespace vespalib::eval::tensor_function
+
 //-----------------------------------------------------------------------------
 
 void
@@ -39,6 +119,12 @@ ConstValue::eval(const TensorEngine &, const LazyParams &, Stash &) const
     return _value;
 }
 
+Instruction
+ConstValue::compile_self(Stash &) const
+{
+    return Instruction(op_load_const, wrap_param<Value>(_value));
+}
+
 //-----------------------------------------------------------------------------
 
 const Value &
@@ -47,6 +133,12 @@ Inject::eval(const TensorEngine &, const LazyParams &params, Stash &stash) const
     return params.resolve(_param_idx, stash);
 }
 
+Instruction
+Inject::compile_self(Stash &) const
+{
+    return Instruction(op_load_param, _param_idx);
+}
+
 //-----------------------------------------------------------------------------
 
 const Value &
@@ -56,6 +148,13 @@ Reduce::eval(const TensorEngine &engine, const LazyParams &params, Stash &stash)
     return engine.reduce(a, _aggr, _dimensions, stash);
 }
 
+Instruction
+Reduce::compile_self(Stash &stash) const
+{
+    ReduceParams &params = stash.create<ReduceParams>(_aggr, _dimensions);
+    return Instruction(op_tensor_reduce, wrap_param<ReduceParams>(params));
+}
+
 //-----------------------------------------------------------------------------
 
 const Value &
@@ -65,6 +164,15 @@ Map::eval(const TensorEngine &engine, const LazyParams &params, Stash &stash) co
     return engine.map(a, _function, stash);
 }
 
+Instruction
+Map::compile_self(Stash &) const
+{
+    if (result_type().is_double()) {
+        return Instruction(op_double_map, to_param(_function));
+    }
+    return Instruction(op_tensor_map, to_param(_function));
+}
+
 //-----------------------------------------------------------------------------
 
 const Value &
@@ -75,6 +183,21 @@ Join::eval(const TensorEngine &engine, const LazyParams &params, Stash &stash) c
     return engine.join(a, b, _function, stash);
 }
 
+Instruction
+Join::compile_self(Stash &) const
+{
+    if (result_type().is_double()) {
+        if (_function == operation::Mul::f) {
+            return Instruction(op_double_mul);
+        }
+        if (_function == operation::Add::f) {
+            return Instruction(op_double_add);
+        }
+        return Instruction(op_double_join, to_param(_function));
+    }
+    return Instruction(op_tensor_join, to_param(_function));
+}
+
 //-----------------------------------------------------------------------------
 
 const Value &
@@ -85,6 +208,12 @@ Concat::eval(const TensorEngine &engine, const LazyParams &params, Stash &stash)
     return engine.concat(a, b, _dimension, stash);
 }
 
+Instruction
+Concat::compile_self(Stash &) const
+{
+    return Instruction(op_tensor_concat, wrap_param<vespalib::string>(_dimension));
+}
+
 //-----------------------------------------------------------------------------
 
 const Value &
@@ -94,6 +223,13 @@ Rename::eval(const TensorEngine &engine, const LazyParams &params, Stash &stash)
     return engine.rename(a, _from, _to, stash);
 }
 
+Instruction
+Rename::compile_self(Stash &stash) const
+{
+    RenameParams &params = stash.create<RenameParams>(_from, _to);
+    return Instruction(op_tensor_rename, wrap_param<RenameParams>(params));
+}
+
 //-----------------------------------------------------------------------------
 
 void
@@ -112,6 +248,14 @@ If::eval(const TensorEngine &engine, const LazyParams &params, Stash &stash) con
             : false_child().eval(engine, params, stash));
 }
 
+Instruction
+If::compile_self(Stash &) const
+{
+    // 'if' is handled directly by compile_tensor_function to enable
+    // lazy-evaluation of true/false sub-expressions.
+    abort();
+}
+
 //-----------------------------------------------------------------------------
 
 const Node &const_value(const Value &value, Stash &stash) {
diff --git a/eval/src/vespa/eval/eval/tensor_function.h b/eval/src/vespa/eval/eval/tensor_function.h
index d9ee5cc068c..c739ea8cba9 100644
--- a/eval/src/vespa/eval/eval/tensor_function.h
+++ b/eval/src/vespa/eval/eval/tensor_function.h
@@ -12,6 +12,8 @@
 #include "value.h"
 #include "aggr.h"
 
+#include "interpreted_function.h"
+
 namespace vespalib {
 
 class Stash;
@@ -75,6 +77,17 @@ struct TensorFunction
     virtual void push_children(std::vector<Child::CREF> &children) const = 0;
 
     /**
+     * Compile this node into a single instruction that can be run by
+     * an interpreted function. Sub-expressions are compiled as
+     * separate instructions and their results will be available on
+     * the value stack during execution.
+     *
+     * @return instruction representing the operation of this node
+     * @param stash heterogeneous object store
+     **/
+    virtual InterpretedFunction::Instruction compile_self(Stash &stash) const = 0;
+
+    /**
      * Evaluate this tensor function based on the given
      * parameters. The given stash can be used to store temporary
      * objects that need to be kept alive for the return value to be
@@ -157,6 +170,7 @@ private:
 public:
     ConstValue(const Value &value_in) : Leaf(value_in.type()), _value(value_in) {}
     const Value &eval(const TensorEngine &engine, const LazyParams &params, Stash &) const final override;
+    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
 };
 
 //-----------------------------------------------------------------------------
@@ -170,6 +184,7 @@ public:
         : Leaf(result_type_in), _param_idx(param_idx_in) {}
     size_t param_idx() const { return _param_idx; }
     const Value &eval(const TensorEngine &engine, const LazyParams &params, Stash &) const final override;
+    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
 };
 
 //-----------------------------------------------------------------------------
@@ -188,6 +203,7 @@ public:
     Aggr aggr() const { return _aggr; }
     const std::vector<vespalib::string> &dimensions() const { return _dimensions; }
     const Value &eval(const TensorEngine &engine, const LazyParams &params, Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
 };
 
 //-----------------------------------------------------------------------------
@@ -203,6 +219,7 @@ public:
         : Op1(result_type_in, child_in), _function(function_in) {}
     map_fun_t function() const { return _function; }
     const Value &eval(const TensorEngine &engine, const LazyParams &params, Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
 };
 
 //-----------------------------------------------------------------------------
@@ -219,6 +236,7 @@ public:
         : Op2(result_type_in, lhs_in, rhs_in), _function(function_in) {}
     join_fun_t function() const { return _function; }
     const Value &eval(const TensorEngine &engine, const LazyParams &params, Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
 };
 
 //-----------------------------------------------------------------------------
@@ -235,6 +253,7 @@ public:
         : Op2(result_type_in, lhs_in, rhs_in), _dimension(dimension_in) {}
     const vespalib::string &dimension() const { return _dimension; }
     const Value &eval(const TensorEngine &engine, const LazyParams &params, Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
 };
 
 //-----------------------------------------------------------------------------
@@ -253,6 +272,7 @@ public:
     const std::vector<vespalib::string> &from() const { return _from; }
     const std::vector<vespalib::string> &to() const { return _to; }
     const Value &eval(const TensorEngine &engine, const LazyParams &params, Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
 };
 
 //-----------------------------------------------------------------------------
@@ -274,6 +294,7 @@ public:
     const TensorFunction &false_child() const { return _false_child.get(); }
     void push_children(std::vector<Child::CREF> &children) const final override;    
     const Value &eval(const TensorEngine &engine, const LazyParams &params, Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
 };
 
 //-----------------------------------------------------------------------------