Merge pull request #15168 from vespa-engine/arnej/add-generic-lambda

Arnej/add generic lambda 

6 files changed, 211 insertions, 9 deletions

diff --git a/eval/src/tests/eval/tensor_lambda/tensor_lambda_test.cpp b/eval/src/tests/eval/tensor_lambda/tensor_lambda_test.cpp
index c8091fd7c6e..a2150c3d1bb 100644
--- a/eval/src/tests/eval/tensor_lambda/tensor_lambda_test.cpp
+++ b/eval/src/tests/eval/tensor_lambda/tensor_lambda_test.cpp
@@ -4,6 +4,8 @@
 #include <vespa/eval/eval/tensor_function.h>
 #include <vespa/eval/eval/simple_tensor.h>
 #include <vespa/eval/eval/simple_tensor_engine.h>
+#include <vespa/eval/eval/simple_value.h>
+#include <vespa/eval/eval/fast_value.h>
 #include <vespa/eval/tensor/default_tensor_engine.h>
 #include <vespa/eval/tensor/dense/dense_replace_type_function.h>
 #include <vespa/eval/tensor/dense/dense_cell_range_function.h>
@@ -40,6 +42,7 @@ std::ostream &operator<<(std::ostream &os, EvalMode eval_mode)
 }
 
 const TensorEngine &prod_engine = DefaultTensorEngine::ref();
+const ValueBuilderFactory &simple_factory = SimpleValueBuilderFactory::get();
 
 EvalFixture::ParamRepo make_params() {
     return EvalFixture::ParamRepo()
@@ -59,7 +62,9 @@ template <typename T, typename F>
 void verify_impl(const vespalib::string &expr, const vespalib::string &expect, F &&inspect) {
     EvalFixture fixture(prod_engine, expr, param_repo, true);
     EvalFixture slow_fixture(prod_engine, expr, param_repo, false);
+    EvalFixture simple_factory_fixture(simple_factory, expr, param_repo, false);
     EXPECT_EQUAL(fixture.result(), slow_fixture.result());
+    EXPECT_EQUAL(fixture.result(), simple_factory_fixture.result());
     EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
     EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expect, param_repo));
     auto info = fixture.find_all<T>();
@@ -210,4 +215,31 @@ TEST("require that type resolving also include nodes in the inner tensor lambda
     EXPECT_EQUAL(types.get_type(*symbol).to_spec(), "double");
 }
 
+size_t count_nodes(const NodeTypes &types) {
+    size_t cnt = 0;
+    types.each([&](const auto &, const auto &){++cnt;});
+    return cnt;
+}
+
+TEST("require that type exporting also include nodes in the inner tensor lambda function") {
+    auto fun = Function::parse("tensor(x[2])(tensor(y[2])((x+y)+a){y:(x)})");
+    NodeTypes types(*fun, {ValueType::from_spec("double")});
+    const auto &root = fun->root();
+    NodeTypes copy = types.export_types(root);
+    EXPECT_TRUE(copy.errors().empty());
+    EXPECT_EQUAL(count_nodes(types), count_nodes(copy));
+
+    auto lambda = nodes::as<nodes::TensorLambda>(root);
+    ASSERT_TRUE(lambda != nullptr);
+    NodeTypes outer = copy.export_types(lambda->lambda().root());
+    EXPECT_TRUE(outer.errors().empty());
+
+    auto inner_lambda = nodes::as<nodes::TensorLambda>(lambda->lambda().root().get_child(0));
+    ASSERT_TRUE(inner_lambda != nullptr);
+    NodeTypes inner = outer.export_types(inner_lambda->lambda().root());
+    EXPECT_TRUE(inner.errors().empty());
+    // [x, y, (x+y), a, (x+y)+a] are the 5 nodes:
+    EXPECT_EQUAL(count_nodes(inner), 5u);
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/eval/src/vespa/eval/eval/node_types.cpp b/eval/src/vespa/eval/eval/node_types.cpp
index cbc96e719e0..9569518417d 100644
--- a/eval/src/vespa/eval/eval/node_types.cpp
+++ b/eval/src/vespa/eval/eval/node_types.cpp
@@ -307,7 +307,12 @@ struct TypeExporter : public NodeTraverser {
         : parent_type_map(parent_type_map_in),
           exported_type_map(exported_type_map_out),
           missing_cnt(0) {}
-    bool open(const Node &) override { return true; }
+    bool open(const Node &node) override {
+        if (auto lambda = as<TensorLambda>(node)) {
+            lambda->lambda().root().traverse(*this);
+        }
+        return true;
+    }
     void close(const Node &node) override {
         auto pos = parent_type_map.find(&node);
         if (pos != parent_type_map.end()) {
diff --git a/eval/src/vespa/eval/eval/tensor_function.cpp b/eval/src/vespa/eval/eval/tensor_function.cpp
index 8464aa14b59..4a5fe50e27b 100644
--- a/eval/src/vespa/eval/eval/tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/tensor_function.cpp
@@ -11,6 +11,7 @@
 #include <vespa/eval/instruction/generic_concat.h>
 #include <vespa/eval/instruction/generic_create.h>
 #include <vespa/eval/instruction/generic_join.h>
+#include <vespa/eval/instruction/generic_lambda.h>
 #include <vespa/eval/instruction/generic_map.h>
 #include <vespa/eval/instruction/generic_merge.h>
 #include <vespa/eval/instruction/generic_peek.h>
@@ -465,6 +466,9 @@ Lambda::create_spec_impl(const ValueType &type, const LazyParams &params, const
 InterpretedFunction::Instruction
 Lambda::compile_self(EngineOrFactory engine, Stash &stash) const
 {
+    if (engine.is_factory()) {
+        return instruction::GenericLambda::make_instruction(*this, engine.factory(), stash);
+    }
     InterpretedFunction fun(engine, _lambda->root(), _lambda_types);
     LambdaParams &params = stash.create<LambdaParams>(*this, std::move(fun));
     return Instruction(op_tensor_lambda, wrap_param<LambdaParams>(params));
diff --git a/eval/src/vespa/eval/instruction/CMakeLists.txt b/eval/src/vespa/eval/instruction/CMakeLists.txt
index 52f411dc543..926a69bd291 100644
--- a/eval/src/vespa/eval/instruction/CMakeLists.txt
+++ b/eval/src/vespa/eval/instruction/CMakeLists.txt
@@ -2,12 +2,13 @@
 
 vespa_add_library(eval_instruction OBJECT
     SOURCES
-    generic_concat
-    generic_create
-    generic_join
-    generic_peek
-    generic_reduce
-    generic_map
-    generic_merge
-    generic_rename
+    generic_concat.cpp
+    generic_create.cpp
+    generic_join.cpp
+    generic_lambda.cpp
+    generic_map.cpp
+    generic_merge.cpp
+    generic_peek.cpp
+    generic_reduce.cpp
+    generic_rename.cpp
 )
diff --git a/eval/src/vespa/eval/instruction/generic_lambda.cpp b/eval/src/vespa/eval/instruction/generic_lambda.cpp
new file mode 100644
index 00000000000..1ba2f710909
--- /dev/null
+++ b/eval/src/vespa/eval/instruction/generic_lambda.cpp
@@ -0,0 +1,143 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "generic_lambda.h"
+#include <vespa/eval/eval/llvm/compiled_function.h>
+#include <vespa/eval/eval/llvm/compile_cache.h>
+#include <assert.h>
+
+using namespace vespalib::eval::tensor_function;
+
+namespace vespalib::eval::instruction {
+
+using Instruction = InterpretedFunction::Instruction;
+using State = InterpretedFunction::State;
+
+namespace {
+
+//-----------------------------------------------------------------------------
+
+bool step_labels(double *labels, const ValueType &type) {
+    for (size_t idx = type.dimensions().size(); idx-- > 0; ) {
+        if ((labels[idx] += 1.0) < type.dimensions()[idx].size) {
+            return true;
+        } else {
+            labels[idx] = 0.0;
+        }
+    }
+    return false;
+}
+
+struct ParamProxy : public LazyParams {
+    const std::vector<double> &labels;
+    const LazyParams          &params;
+    const std::vector<size_t> &bindings;
+    ParamProxy(const std::vector<double> &labels_in, const LazyParams &params_in, const std::vector<size_t> &bindings_in)
+        : labels(labels_in), params(params_in), bindings(bindings_in) {}
+    const Value &resolve(size_t idx, Stash &stash) const override {
+        if (idx < labels.size()) {
+            return stash.create<DoubleValue>(labels[idx]);
+        }
+        return params.resolve(bindings[idx - labels.size()], stash);
+    }
+};
+
+//-----------------------------------------------------------------------------
+
+struct CompiledParams {
+    const ValueType &result_type;
+    const std::vector<size_t> &bindings;
+    size_t num_cells;
+    CompileCache::Token::UP token;
+    CompiledParams(const Lambda &lambda) 
+        : result_type(lambda.result_type()),
+          bindings(lambda.bindings()),
+          num_cells(result_type.dense_subspace_size()),
+          token(CompileCache::compile(lambda.lambda(), PassParams::ARRAY))
+    {
+        assert(lambda.lambda().num_params() == (result_type.dimensions().size() + bindings.size()));
+    }
+};
+
+template <typename CT>
+void my_compiled_lambda_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    const CompiledParams &params = unwrap_param<CompiledParams>(param);
+    std::vector<double> args(params.result_type.dimensions().size() + params.bindings.size(), 0.0);
+    double *bind_next = &args[params.result_type.dimensions().size()];
+    for (size_t binding: params.bindings) {
+        *bind_next++ = state.params->resolve(binding, state.stash).as_double();
+    }
+    auto fun = params.token->get().get_function();
+    ArrayRef<CT> dst_cells = state.stash.create_uninitialized_array<CT>(params.num_cells);
+    CT *dst = &dst_cells[0];
+    do {
+        *dst++ = fun(&args[0]);
+    } while (step_labels(&args[0], params.result_type));
+    state.stack.push_back(state.stash.create<DenseValueView>(params.result_type, TypedCells(dst_cells)));
+}
+
+struct MyCompiledLambdaOp {
+    template <typename CT>
+    static auto invoke() { return my_compiled_lambda_op<CT>; }
+};
+
+//-----------------------------------------------------------------------------
+
+struct InterpretedParams {
+    const ValueType &result_type;
+    const std::vector<size_t> &bindings;
+    size_t num_cells;
+    InterpretedFunction fun;
+    InterpretedParams(const Lambda &lambda, const ValueBuilderFactory &factory)
+        : result_type(lambda.result_type()),
+          bindings(lambda.bindings()),
+          num_cells(result_type.dense_subspace_size()),
+          fun(factory, lambda.lambda().root(), lambda.types())
+    {
+        assert(lambda.lambda().num_params() == (result_type.dimensions().size() + bindings.size()));
+    }
+};
+
+template <typename CT>
+void my_interpreted_lambda_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    const InterpretedParams &params = unwrap_param<InterpretedParams>(param);
+    std::vector<double> labels(params.result_type.dimensions().size(), 0.0);
+    ParamProxy param_proxy(labels, *state.params, params.bindings);
+    InterpretedFunction::Context ctx(params.fun);
+    ArrayRef<CT> dst_cells = state.stash.create_array<CT>(params.num_cells);
+    CT *dst = &dst_cells[0];
+    do {
+        *dst++ = params.fun.eval(ctx, param_proxy).as_double();
+    } while (step_labels(&labels[0], params.result_type));
+    state.stack.push_back(state.stash.create<DenseValueView>(params.result_type, TypedCells(dst_cells)));
+}
+
+struct MyInterpretedLambdaOp {
+    template <typename CT>
+    static auto invoke() { return my_interpreted_lambda_op<CT>; }
+};
+
+//-----------------------------------------------------------------------------
+
+} // namespace <unnamed>
+
+Instruction
+GenericLambda::make_instruction(const eval::tensor_function::Lambda &lambda_in,
+                                const ValueBuilderFactory &factory, Stash &stash)
+{
+    const ValueType & result_type = lambda_in.result_type();
+    assert(result_type.count_mapped_dimensions() == 0);
+    if (!CompiledFunction::detect_issues(lambda_in.lambda()) &&
+        lambda_in.types().all_types_are_double())
+    {
+        // can do compiled version
+        CompiledParams &params = stash.create<CompiledParams>(lambda_in);
+        auto op = typify_invoke<1,TypifyCellType,MyCompiledLambdaOp>(result_type.cell_type());
+        return Instruction(op, wrap_param<CompiledParams>(params));
+    } else {
+        InterpretedParams &params = stash.create<InterpretedParams>(lambda_in, factory);
+        auto op = typify_invoke<1,TypifyCellType,MyInterpretedLambdaOp>(result_type.cell_type());
+        return Instruction(op, wrap_param<InterpretedParams>(params));
+    }
+}
+
+} // namespace
diff --git a/eval/src/vespa/eval/instruction/generic_lambda.h b/eval/src/vespa/eval/instruction/generic_lambda.h
new file mode 100644
index 00000000000..a9a490f0957
--- /dev/null
+++ b/eval/src/vespa/eval/instruction/generic_lambda.h
@@ -0,0 +1,17 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/eval/eval/interpreted_function.h>
+#include <vespa/eval/eval/tensor_function.h>
+#include <vespa/eval/eval/value_type.h>
+
+namespace vespalib::eval::instruction {
+
+struct GenericLambda {
+    static InterpretedFunction::Instruction
+    make_instruction(const eval::tensor_function::Lambda &lambda_in,
+                     const ValueBuilderFactory &factory, Stash &stash);
+};
+
+} // namespace