Merge pull request #5191 from vespa-engine/arnej/add-inplace-join

add inplace join

9 files changed, 286 insertions, 0 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index ba42ddecef9..a733b73d9d1 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -26,6 +26,7 @@ vespa_define_module(
     src/tests/gp/ponder_nov2017
     src/tests/tensor/dense_dot_product_function
     src/tests/tensor/dense_fast_rename_function
+    src/tests/tensor/dense_inplace_join_function
     src/tests/tensor/dense_inplace_map_function
     src/tests/tensor/dense_tensor_address_combiner
     src/tests/tensor/dense_tensor_builder
diff --git a/eval/src/tests/tensor/dense_inplace_join_function/CMakeLists.txt b/eval/src/tests/tensor/dense_inplace_join_function/CMakeLists.txt
new file mode 100644
index 00000000000..2808675bc78
--- /dev/null
+++ b/eval/src/tests/tensor/dense_inplace_join_function/CMakeLists.txt
@@ -0,0 +1,8 @@
+# Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(eval_dense_inplace_join_function_test_app TEST
+    SOURCES
+    dense_inplace_join_function_test.cpp
+    DEPENDS
+    vespaeval
+)
+vespa_add_test(NAME eval_dense_inplace_join_function_test_app COMMAND eval_dense_inplace_join_function_test_app)
diff --git a/eval/src/tests/tensor/dense_inplace_join_function/dense_inplace_join_function_test.cpp b/eval/src/tests/tensor/dense_inplace_join_function/dense_inplace_join_function_test.cpp
new file mode 100644
index 00000000000..c794b81f573
--- /dev/null
+++ b/eval/src/tests/tensor/dense_inplace_join_function/dense_inplace_join_function_test.cpp
@@ -0,0 +1,147 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/vespalib/testkit/test_kit.h>
+#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/tensor/default_tensor_engine.h>
+#include <vespa/eval/tensor/dense/dense_inplace_join_function.h>
+#include <vespa/eval/tensor/dense/dense_tensor.h>
+#include <vespa/eval/eval/test/tensor_model.hpp>
+#include <vespa/eval/eval/test/eval_fixture.h>
+
+#include <vespa/vespalib/util/stringfmt.h>
+#include <vespa/vespalib/util/stash.h>
+
+using namespace vespalib;
+using namespace vespalib::eval;
+using namespace vespalib::eval::test;
+using namespace vespalib::tensor;
+using namespace vespalib::eval::tensor_function;
+
+const TensorEngine &prod_engine = DefaultTensorEngine::ref();
+
+double seq_value = 0.0;
+
+struct GlobalSequence : public Sequence {
+    GlobalSequence() {}
+    double operator[](size_t) const override {
+        seq_value += 1.0;
+        return seq_value;
+    }
+    ~GlobalSequence() {}
+};
+GlobalSequence seq;
+
+EvalFixture::ParamRepo make_params() {
+    return EvalFixture::ParamRepo()
+        .add("con_x5_A", spec({x(5)}, seq))
+        .add("con_x5_B", spec({x(5)}, seq))
+        .add("con_x5_C", spec({x(5)}, seq))
+        .add("con_x5y3_A", spec({x(5),y(3)}, seq))
+        .add("con_x5y3_B", spec({x(5),y(3)}, seq))
+        .add_mutable("mut_dbl_A", spec(1.5))
+        .add_mutable("mut_dbl_B", spec(2.5))
+        .add_mutable("mut_x5_A", spec({x(5)}, seq))
+        .add_mutable("mut_x5_B", spec({x(5)}, seq))
+        .add_mutable("mut_x5_C", spec({x(5)}, seq))
+        .add_mutable("mut_x4", spec({x(4)}, seq))
+        .add_mutable("mut_x5y3_A", spec({x(5),y(3)}, seq))
+        .add_mutable("mut_x5y3_B", spec({x(5),y(3)}, seq))
+        .add_mutable("mut_x5_unbound", spec({x(5)}, seq), "tensor(x[])")
+        .add_mutable("mut_x_sparse", spec({x({"a", "b", "c"})}, seq));
+}
+EvalFixture::ParamRepo param_repo = make_params();
+
+void verify_optimized(const vespalib::string &expr, size_t cnt, size_t param_idx) {
+    EvalFixture fixture(prod_engine, expr, param_repo, true, true);
+    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
+    for (size_t i = 0; i < fixture.num_params(); ++i) {
+        TEST_STATE(vespalib::make_string("param %zu", i).c_str());
+        if (i == param_idx) {
+            EXPECT_EQUAL(fixture.get_param(i), fixture.result());
+        } else {
+            EXPECT_NOT_EQUAL(fixture.get_param(i), fixture.result());
+        }
+    }
+    auto info = fixture.find_all<DenseInplaceJoinFunction>();
+    ASSERT_EQUAL(info.size(), cnt);
+    for (size_t i = 0; i < cnt; ++i) {
+        EXPECT_TRUE(info[i]->result_is_mutable());
+    }
+}
+
+void verify_p0_optimized(const vespalib::string &expr, size_t cnt) {
+    verify_optimized(expr, cnt, 0);
+}
+
+void verify_p1_optimized(const vespalib::string &expr, size_t cnt) {
+    verify_optimized(expr, cnt, 1);
+}
+
+void verify_p2_optimized(const vespalib::string &expr, size_t cnt) {
+    verify_optimized(expr, cnt, 2);
+}
+
+void verify_not_optimized(const vespalib::string &expr) {
+    EvalFixture fixture(prod_engine, expr, param_repo, true, true);
+    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
+    for (size_t i = 0; i < fixture.num_params(); ++i) {
+        EXPECT_NOT_EQUAL(fixture.get_param(i), fixture.result());
+    }
+    auto info = fixture.find_all<DenseInplaceJoinFunction>();
+    EXPECT_TRUE(info.empty());
+}
+
+TEST("require that mutable dense concrete tensors are optimized") {
+    TEST_DO(verify_p0_optimized("mut_x5_A-mut_x5_B", 1));
+    TEST_DO(verify_p0_optimized("mut_x5_A-con_x5_B", 1));
+    TEST_DO(verify_p1_optimized("con_x5_A-mut_x5_B", 1));
+    TEST_DO(verify_p0_optimized("mut_x5y3_A-mut_x5y3_B", 1));
+    TEST_DO(verify_p0_optimized("mut_x5y3_A-con_x5y3_B", 1));
+    TEST_DO(verify_p1_optimized("con_x5y3_A-mut_x5y3_B", 1));
+}
+
+TEST("require that self-join operations can be optimized") {
+    TEST_DO(verify_p0_optimized("mut_x5_A+mut_x5_A", 1));
+}
+
+TEST("require that join(tensor,scalar) operations are not optimized") {
+    TEST_DO(verify_not_optimized("mut_x5_A-mut_dbl_B"));
+    TEST_DO(verify_not_optimized("mut_dbl_A-mut_x5_B"));
+}
+
+TEST("require that join with different tensor shapes are not optimized") {
+    TEST_DO(verify_not_optimized("mut_x5_A-mut_x4"));
+    TEST_DO(verify_not_optimized("mut_x4-mut_x5_A"));
+    TEST_DO(verify_not_optimized("mut_x5_A*mut_x5y3_B"));
+}
+
+TEST("require that inplace join operations can be chained") {
+    TEST_DO(verify_p0_optimized("mut_x5_A-(mut_x5_B-mut_x5_C)", 2));
+    TEST_DO(verify_p0_optimized("(mut_x5_A-con_x5_B)-con_x5_C", 2));
+    TEST_DO(verify_p1_optimized("con_x5_A-(mut_x5_B-con_x5_C)", 2));
+    TEST_DO(verify_p2_optimized("con_x5_A-(con_x5_B-mut_x5_C)", 2));
+}
+
+TEST("require that abstract tensors are not optimized") {
+    TEST_DO(verify_not_optimized("mut_x5_unbound+mut_x5_A"));
+    TEST_DO(verify_not_optimized("mut_x5_A+mut_x5_unbound"));
+    TEST_DO(verify_not_optimized("mut_x5_unbound+mut_x5_unbound"));
+}
+
+TEST("require that non-mutable tensors are not optimized") {
+    TEST_DO(verify_not_optimized("con_x5_A+con_x5_B"));
+}
+
+TEST("require that scalar values are not optimized") {
+    TEST_DO(verify_not_optimized("mut_dbl_A+mut_dbl_B"));
+    TEST_DO(verify_not_optimized("mut_dbl_A+5"));
+    TEST_DO(verify_not_optimized("5+mut_dbl_B"));
+}
+
+TEST("require that mapped tensors are not optimized") {
+    TEST_DO(verify_not_optimized("mut_x_sparse+mut_x_sparse"));
+}
+
+TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/eval/src/vespa/eval/eval/test/eval_fixture.cpp b/eval/src/vespa/eval/eval/test/eval_fixture.cpp
index 2593b74e300..bd4dcc37e3c 100644
--- a/eval/src/vespa/eval/eval/test/eval_fixture.cpp
+++ b/eval/src/vespa/eval/eval/test/eval_fixture.cpp
@@ -113,4 +113,10 @@ EvalFixture::get_param(size_t idx) const
     return _engine.to_spec(*(_param_values[idx]));
 }
 
+size_t
+EvalFixture::num_params() const
+{
+    return _param_values.size();
+}
+
 } // namespace vespalib::eval::test
diff --git a/eval/src/vespa/eval/eval/test/eval_fixture.h b/eval/src/vespa/eval/eval/test/eval_fixture.h
index 7d3b95ad144..e20d435f608 100644
--- a/eval/src/vespa/eval/eval/test/eval_fixture.h
+++ b/eval/src/vespa/eval/eval/test/eval_fixture.h
@@ -86,6 +86,7 @@ public:
     }
     const TensorSpec &result() const { return _result; }
     const TensorSpec get_param(size_t idx) const;
+    size_t num_params() const;
     static TensorSpec ref(const vespalib::string &expr, const ParamRepo &param_repo) {
         return EvalFixture(SimpleTensorEngine::ref(), expr, param_repo, false, false).result();
     }
diff --git a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
index dead8ee4870..457e9310b80 100644
--- a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
+++ b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
@@ -10,6 +10,7 @@
 #include "dense/dense_dot_product_function.h"
 #include "dense/dense_xw_product_function.h"
 #include "dense/dense_fast_rename_function.h"
+#include "dense/dense_inplace_join_function.h"
 #include "dense/dense_inplace_map_function.h"
 #include "dense/vector_from_doubles_function.h"
 #include <vespa/eval/eval/value.h>
@@ -225,6 +226,7 @@ DefaultTensorEngine::optimize(const TensorFunction &expr, Stash &stash) const
         child.set(DenseXWProductFunction::optimize(child.get(), stash));
         child.set(DenseFastRenameFunction::optimize(child.get(), stash));
         child.set(DenseInplaceMapFunction::optimize(child.get(), stash));
+        child.set(DenseInplaceJoinFunction::optimize(child.get(), stash));
         nodes.pop_back();
     }
     return root.get();
diff --git a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
index cf6d3a3431c..f78e49dc2f3 100644
--- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
+++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
@@ -3,6 +3,7 @@ vespa_add_library(eval_tensor_dense OBJECT
     SOURCES
     dense_dot_product_function.cpp
     dense_fast_rename_function.cpp
+    dense_inplace_join_function.cpp
     dense_inplace_map_function.cpp
     dense_tensor.cpp
     dense_tensor_address_combiner.cpp
diff --git a/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.cpp
new file mode 100644
index 00000000000..53a5fe9bb27
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.cpp
@@ -0,0 +1,87 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "dense_inplace_join_function.h"
+#include "dense_tensor.h"
+#include "dense_tensor_view.h"
+#include <vespa/eval/eval/value.h>
+#include <vespa/eval/tensor/tensor.h>
+
+namespace vespalib::tensor {
+
+using CellsRef = DenseTensorView::CellsRef;
+using eval::Value;
+using eval::ValueType;
+using eval::TensorFunction;
+using eval::as;
+using namespace eval::tensor_function;
+
+namespace {
+
+CellsRef getCellsRef(const eval::Value &value) {
+    const DenseTensorView &denseTensor = static_cast<const DenseTensorView &>(value);
+    return denseTensor.cellsRef();
+}
+
+template <bool write_left>
+void my_inplace_join_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    join_fun_t function = (join_fun_t)param;
+    CellsRef lhs_cells = getCellsRef(state.peek(1));
+    CellsRef rhs_cells = getCellsRef(state.peek(0));
+    auto dst_cells = unconstify(write_left ? lhs_cells : rhs_cells);
+    for (size_t i = 0; i < dst_cells.size(); ++i) {
+        dst_cells[i] = function(lhs_cells[i], rhs_cells[i]);
+    }
+    if (write_left) {
+        state.stack.pop_back();
+    } else {
+        const Value &result = state.stack.back();
+        state.pop_pop_push(result);
+    }
+}
+
+bool sameShapeConcreteDenseTensors(const ValueType &a, const ValueType &b) {
+    return (a.is_dense() && !a.is_abstract() && (a == b));
+}
+
+} // namespace vespalib::tensor::<unnamed>
+
+
+DenseInplaceJoinFunction::DenseInplaceJoinFunction(const ValueType &result_type,
+                                                   const TensorFunction &lhs,
+                                                   const TensorFunction &rhs,
+                                                   join_fun_t function_in,
+                                                   bool write_left_in)
+    : eval::tensor_function::Op2(result_type, lhs, rhs),
+      _function(function_in),
+      _write_left(write_left_in)
+{
+}
+
+DenseInplaceJoinFunction::~DenseInplaceJoinFunction()
+{
+}
+
+eval::InterpretedFunction::Instruction
+DenseInplaceJoinFunction::compile_self(Stash &) const
+{
+    auto op = _write_left ? my_inplace_join_op<true> : my_inplace_join_op<false>;
+    return eval::InterpretedFunction::Instruction(op, (uint64_t)_function);
+}
+
+const TensorFunction &
+DenseInplaceJoinFunction::optimize(const eval::TensorFunction &expr, Stash &stash)
+{
+    if (auto join = as<Join>(expr)) {
+        const TensorFunction &lhs = join->lhs();
+        const TensorFunction &rhs = join->rhs();
+        if ((lhs.result_is_mutable() || rhs.result_is_mutable()) &&
+            sameShapeConcreteDenseTensors(lhs.result_type(), rhs.result_type()))
+        {
+            return stash.create<DenseInplaceJoinFunction>(join->result_type(), lhs, rhs,
+                    join->function(), lhs.result_is_mutable());
+        }
+    }
+    return expr;
+}
+
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.h b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.h
new file mode 100644
index 00000000000..de2cdae3778
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.h
@@ -0,0 +1,33 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/eval/eval/tensor_function.h>
+
+namespace vespalib::tensor {
+
+/**
+ * Tensor function for inplace join operation on mutable dense tensors.
+ **/
+class DenseInplaceJoinFunction : public eval::tensor_function::Op2
+{
+public:
+    using join_fun_t = ::vespalib::eval::tensor_function::join_fun_t;
+private:
+    join_fun_t _function;
+    bool _write_left;
+public:
+    DenseInplaceJoinFunction(const eval::ValueType &result_type,
+                             const TensorFunction &lhs,
+                             const TensorFunction &rhs,
+                             join_fun_t function_in,
+                             bool write_left_in);
+    ~DenseInplaceJoinFunction();
+    join_fun_t function() const { return _function; }
+    bool write_left() const { return _write_left; }
+    bool result_is_mutable() const override { return true; }
+    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
+};
+
+} // namespace vespalib::tensor