add inplace join

7 files changed, 243 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..c59f4956675
--- /dev/null
+++ b/eval/src/tests/tensor/dense_inplace_join_function/dense_inplace_join_function_test.cpp
@@ -0,0 +1,98 @@
+// 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();
+
+EvalFixture::ParamRepo make_params() {
+    return EvalFixture::ParamRepo()
+        .add("x5", spec({x(5)}, N()))
+        .add_mutable("_d_A", spec(17.0))
+        .add_mutable("_d_B", spec(42.0))
+        .add_mutable("_x5", spec({x(5)}, N()))
+        .add_mutable("_x5_A", spec({x(5)}, Seq({10, 11, 12, 13, 14})))
+        .add_mutable("_x5_B", spec({x(5)}, Seq({16, 18, 20, 22, 24})))
+        .add_mutable("_x5_C", spec({x(5)}, Seq({30, 35, 40, 45, 50})))
+        .add_mutable("_x5y3", spec({x(5),y(3)}, N()))
+        .add_mutable("_x5_u", spec({x(5)}, N()), "tensor(x[])")
+        .add_mutable("_x_m", spec({x({"a", "b", "c"})}, N()));
+}
+EvalFixture::ParamRepo param_repo = make_params();
+
+void verify_left_optimized(const vespalib::string &expr, size_t cnt) {
+    EvalFixture fixture(prod_engine, expr, param_repo, true, true);
+    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
+    EXPECT_EQUAL(fixture.get_param(0), 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_right_optimized(const vespalib::string &expr, size_t cnt) {
+    EvalFixture fixture(prod_engine, expr, param_repo, true, true);
+    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
+    EXPECT_EQUAL(fixture.get_param(1), 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_not_optimized(const vespalib::string &expr) {
+    EvalFixture fixture(prod_engine, expr, param_repo, true, true);
+    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
+    EXPECT_NOT_EQUAL(fixture.get_param(0), fixture.result());
+    auto info = fixture.find_all<DenseInplaceJoinFunction>();
+    EXPECT_TRUE(info.empty());
+}
+
+TEST("require that mutable dense concrete tensors are optimized") {
+    TEST_DO(verify_left_optimized("_x5_A+_x5_B", 1));
+    TEST_DO(verify_left_optimized("_x5_A-_x5_B", 1));
+    TEST_DO(verify_left_optimized("_x5-x5", 1));
+    TEST_DO(verify_right_optimized("x5-_x5", 1));
+}
+
+TEST("require that inplace join operations can be chained") {
+    TEST_DO(verify_left_optimized("_x5_C-(_x5_B-_x5_A)", 2));
+}
+
+TEST("require that abstract tensors are not optimized") {
+    TEST_DO(verify_not_optimized("_x5_u+_x5"));
+    TEST_DO(verify_not_optimized("_x5_u+_x5_u"));
+}
+
+TEST("require that non-mutable tensors are not optimized") {
+    TEST_DO(verify_not_optimized("x5+x5"));
+}
+
+TEST("require that scalar values are not optimized") {
+    TEST_DO(verify_not_optimized("_d_A+_d_B"));
+    TEST_DO(verify_not_optimized("join(_d_A,_d_B,f(x,y)(x+y))"));
+}
+
+TEST("require that mapped tensors are not optimized") {
+    TEST_DO(verify_not_optimized("_x_m+_x_m"));
+}
+
+TEST_MAIN() { TEST_RUN_ALL(); }
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..639246b057c
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.cpp
@@ -0,0 +1,105 @@
+// 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 {
+
+ArrayRef<double> getMutableCells(const eval::Value &value) {
+    const DenseTensorView &denseTensor = static_cast<const DenseTensorView &>(value);
+    return unconstify(denseTensor.cellsRef());
+}
+
+ConstArrayRef<double> getConstCells(const eval::Value &value) {
+    const DenseTensorView &denseTensor = static_cast<const DenseTensorView &>(value);
+    return denseTensor.cellsRef();
+}
+
+void my_inplace_left_join_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    const Value &lhs = state.peek(1);
+    const Value &rhs = state.peek(0);
+    join_fun_t function = (join_fun_t)param;
+    ArrayRef<double> left_cells = getMutableCells(lhs);
+    ConstArrayRef<double> right_cells = getConstCells(rhs);
+    auto rhs_iter = right_cells.cbegin();
+    for (double &cell: left_cells) {
+        cell = function(cell, *rhs_iter);
+        ++rhs_iter;
+    }
+    assert(rhs_iter == right_cells.cend());
+    state.pop_pop_push(lhs);
+}
+
+void my_inplace_right_join_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    const Value &lhs = state.peek(1);
+    const Value &rhs = state.peek(0);
+    join_fun_t function = (join_fun_t)param;
+    ConstArrayRef<double> left_cells = getConstCells(lhs);
+    ArrayRef<double> right_cells = getMutableCells(rhs);
+    auto lhs_iter = left_cells.cbegin();
+    for (double &cell: right_cells) {
+        cell = function(cell, *lhs_iter);
+        ++lhs_iter;
+    }
+    assert(lhs_iter == left_cells.cend());
+    state.pop_pop_push(rhs);
+}
+
+bool isConcreteDenseTensor(const ValueType &type) {
+    return (type.is_dense() && !type.is_abstract());
+}
+
+} // namespace vespalib::tensor::<unnamed>
+
+
+DenseInplaceJoinFunction::DenseInplaceJoinFunction(const eval::tensor_function::Join &orig, bool left_is_mutable)
+    : eval::tensor_function::Op2(orig.result_type(), orig.lhs(), orig.rhs()),
+      _function(orig.function()),
+      _left_is_mutable(left_is_mutable)
+{
+}
+
+DenseInplaceJoinFunction::~DenseInplaceJoinFunction()
+{
+}
+
+eval::InterpretedFunction::Instruction
+DenseInplaceJoinFunction::compile_self(Stash &) const
+{
+    if (_left_is_mutable) {
+        return eval::InterpretedFunction::Instruction(my_inplace_left_join_op, (uint64_t)_function);
+    } else {
+        return eval::InterpretedFunction::Instruction(my_inplace_right_join_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())
+            && join->result_type() == lhs.result_type()
+            && join->result_type() == rhs.result_type()
+            && isConcreteDenseTensor(join->result_type()))
+        {
+            return stash.create<DenseInplaceJoinFunction>(*join, 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..176191c995c
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.h
@@ -0,0 +1,28 @@
+// 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 _left_is_mutable;
+public:
+    DenseInplaceJoinFunction(const eval::tensor_function::Join &orig, bool left_is_mutable);
+    ~DenseInplaceJoinFunction();
+    join_fun_t function() const { return _function; }
+    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