optimize dense tensor peek

8 files changed, 234 insertions, 0 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index b3d4585a176..35731ef5cfb 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -35,6 +35,7 @@ vespa_define_module(
     src/tests/tensor/dense_remove_dimension_optimizer
     src/tests/tensor/dense_replace_type_function
     src/tests/tensor/dense_tensor_create_function
+    src/tests/tensor/dense_tensor_peek_function
     src/tests/tensor/dense_xw_product_function
     src/tests/tensor/direct_dense_tensor_builder
     src/tests/tensor/direct_sparse_tensor_builder
diff --git a/eval/src/tests/tensor/dense_tensor_peek_function/CMakeLists.txt b/eval/src/tests/tensor/dense_tensor_peek_function/CMakeLists.txt
new file mode 100644
index 00000000000..a4d1bbf6fac
--- /dev/null
+++ b/eval/src/tests/tensor/dense_tensor_peek_function/CMakeLists.txt
@@ -0,0 +1,8 @@
+# Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(eval_dense_tensor_peek_function_test_app TEST
+    SOURCES
+    dense_tensor_peek_function_test.cpp
+    DEPENDS
+    vespaeval
+)
+vespa_add_test(NAME eval_dense_tensor_peek_function_test_app COMMAND eval_dense_tensor_peek_function_test_app)
diff --git a/eval/src/tests/tensor/dense_tensor_peek_function/dense_tensor_peek_function_test.cpp b/eval/src/tests/tensor/dense_tensor_peek_function/dense_tensor_peek_function_test.cpp
new file mode 100644
index 00000000000..74435491410
--- /dev/null
+++ b/eval/src/tests/tensor/dense_tensor_peek_function/dense_tensor_peek_function_test.cpp
@@ -0,0 +1,73 @@
+// Copyright 2019 Oath Inc. 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_tensor_peek_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("a", spec(1.0))
+        .add("b", spec(2.0))
+        .add("c", spec(3.0))
+        .add("x3", spec(x(3), N()))
+        .add("x3f", spec(float_cells({x(3)}), N()))
+        .add("x3y2", spec({x(3),y(2)}, N()))
+        .add("x3y2f", spec(float_cells({x(3),y(2)}), N()))
+        .add("xm", spec(x({"1","2","3"}), N()))
+        .add("xmy2", spec({x({"1","2","3"}), y(2)}, N()));
+}
+EvalFixture::ParamRepo param_repo = make_params();
+
+void verify(const vespalib::string &expr, double expect, size_t expect_optimized_cnt, size_t expect_not_optimized_cnt) {
+    EvalFixture fixture(prod_engine, expr, param_repo, true);
+    EXPECT_EQUAL(fixture.result(), TensorSpec("double").add({}, expect));
+    auto info = fixture.find_all<DenseTensorPeekFunction>();
+    EXPECT_EQUAL(info.size(), expect_optimized_cnt);
+    for (size_t i = 0; i < info.size(); ++i) {
+        EXPECT_TRUE(info[i]->result_is_mutable());
+    }
+    EXPECT_EQUAL(fixture.find_all<Peek>().size(), expect_not_optimized_cnt);
+}
+
+//-----------------------------------------------------------------------------
+
+TEST("require that tensor peek can be optimized for dense tensors") {
+    TEST_DO(verify("x3{x:0}", 1.0, 1, 0));
+    TEST_DO(verify("x3{x:(a)}", 2.0, 1, 0));
+    TEST_DO(verify("x3f{x:(c-1)}", 3.0, 1, 0));
+    TEST_DO(verify("x3{x:(c+5)}", 0.0, 1, 0));
+    TEST_DO(verify("x3y2{x:(a),y:(a-1)}", 3.0, 1, 0));
+    TEST_DO(verify("x3y2f{x:1,y:(a)}", 4.0, 1, 0));
+    TEST_DO(verify("x3y2f{x:(a-1),y:(b)}", 0.0, 1, 0));
+}
+
+TEST("require that tensor peek is not optimized for sparse tensor") {
+    TEST_DO(verify("xm{x:1}", 1.0, 0, 1));
+    TEST_DO(verify("xm{x:(c)}", 3.0, 0, 1));
+    TEST_DO(verify("xm{x:(c+1)}", 0.0, 0, 1));
+}
+
+TEST("require that tensor peek is not optimized for mixed tensor") {
+    TEST_DO(verify("xmy2{x:3,y:1}", 6.0, 0, 1));
+    TEST_DO(verify("xmy2{x:(c),y:(a)}", 6.0, 0, 1));
+    TEST_DO(verify("xmy2{x:(a),y:(b)}", 0.0, 0, 1));
+}
+
+TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/eval/src/vespa/eval/eval/tensor_function.h b/eval/src/vespa/eval/eval/tensor_function.h
index 7c7e8318305..c95ffd17bbe 100644
--- a/eval/src/vespa/eval/eval/tensor_function.h
+++ b/eval/src/vespa/eval/eval/tensor_function.h
@@ -84,6 +84,16 @@ struct TensorFunction
      **/
     virtual void push_children(std::vector<Child::CREF> &children) const = 0;
 
+    std::vector<Child> copy_children() const {
+        std::vector<Child::CREF> child_refs;
+        std::vector<Child> children_copy;
+        push_children(child_refs);
+        for (const auto &child_ref: child_refs) {
+            children_copy.emplace_back(child_ref.get().get());
+        }
+        return children_copy;
+    }
+
     /**
      * Compile this node into a single instruction that can be run by
      * an interpreted function. Sub-expressions are compiled as
diff --git a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
index 18733b3c733..b310ad72fc1 100644
--- a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
+++ b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
@@ -17,6 +17,7 @@
 #include "dense/dense_inplace_map_function.h"
 #include "dense/vector_from_doubles_function.h"
 #include "dense/dense_tensor_create_function.h"
+#include "dense/dense_tensor_peek_function.h"
 #include <vespa/eval/eval/value.h>
 #include <vespa/eval/eval/tensor_spec.h>
 #include <vespa/eval/eval/tensor_spec.h>
@@ -260,6 +261,7 @@ DefaultTensorEngine::optimize(const TensorFunction &expr, Stash &stash) const
         const Child &child = nodes.back();
         child.set(VectorFromDoublesFunction::optimize(child.get(), stash));
         child.set(DenseTensorCreateFunction::optimize(child.get(), stash));
+        child.set(DenseTensorPeekFunction::optimize(child.get(), stash));
         child.set(DenseDotProductFunction::optimize(child.get(), stash));
         child.set(DenseXWProductFunction::optimize(child.get(), stash));
         child.set(DenseFastRenameOptimizer::optimize(child.get(), stash));
diff --git a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
index e071de3e925..2ad54d48ab3 100644
--- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
+++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
@@ -14,6 +14,7 @@ vespa_add_library(eval_tensor_dense OBJECT
     dense_tensor_cells_iterator.cpp
     dense_tensor_create_function.cpp
     dense_tensor_modify.cpp
+    dense_tensor_peek_function.cpp
     dense_tensor_reduce.cpp
     dense_tensor_view.cpp
     dense_xw_product_function.cpp
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.cpp
new file mode 100644
index 00000000000..f9d15a377e9
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.cpp
@@ -0,0 +1,105 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "dense_tensor_peek_function.h"
+#include "dense_tensor_view.h"
+#include <vespa/eval/eval/operation.h>
+#include <vespa/eval/eval/value.h>
+#include <vespa/eval/tensor/tensor.h>
+
+namespace vespalib::tensor {
+
+using eval::Value;
+using eval::DoubleValue;
+using eval::ValueType;
+using eval::TensorSpec;
+using eval::TensorFunction;
+using Child = eval::TensorFunction::Child;
+using eval::as;
+using namespace eval::tensor_function;
+using namespace eval::operation;
+
+namespace {
+
+template <typename CT>
+void my_tensor_peek_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    const auto *spec = (const std::vector<std::pair<int64_t,size_t>> *)(param);
+    size_t idx = 0;
+    size_t factor = 1;
+    bool valid = true;
+    for (const auto &dim: *spec) {
+        if (dim.first >= 0) {
+            idx += (dim.first * factor);
+        } else {
+            size_t dim_idx(round(state.peek(0).as_double()));
+            state.stack.pop_back();
+            valid &= (dim_idx < dim.second);
+            idx += (dim_idx * factor);
+        }
+        factor *= dim.second;
+    }
+    auto cells = DenseTensorView::typify_cells<CT>(state.peek(0));
+    state.stack.pop_back();
+    const Value &result = state.stash.create<DoubleValue>(valid ? cells[idx] : 0.0);
+    state.stack.push_back(result);
+}
+
+struct MyTensorPeekOp {
+    template <typename CT>
+    static auto get_fun() { return my_tensor_peek_op<CT>; }
+};
+
+} // namespace vespalib::tensor::<unnamed>
+
+DenseTensorPeekFunction::DenseTensorPeekFunction(std::vector<Child> children,
+                                                 const std::vector<std::pair<int64_t,size_t>> &spec)
+    : TensorFunction(),
+      _children(std::move(children)),
+      _spec(spec)
+{
+}
+
+DenseTensorPeekFunction::~DenseTensorPeekFunction() = default;
+
+void
+DenseTensorPeekFunction::push_children(std::vector<Child::CREF> &target) const
+{
+    for (const Child &c: _children) {
+        target.push_back(c);
+    }
+}
+
+eval::InterpretedFunction::Instruction
+DenseTensorPeekFunction::compile_self(Stash &) const
+{
+    static_assert(sizeof(uint64_t) == sizeof(&_spec));
+    auto op = select_1<MyTensorPeekOp>(_children[0].get().result_type().cell_type());
+    return eval::InterpretedFunction::Instruction(op, (uint64_t)&_spec);
+}
+
+const TensorFunction &
+DenseTensorPeekFunction::optimize(const eval::TensorFunction &expr, Stash &stash)
+{
+    if (auto peek = as<Peek>(expr)) {
+        const ValueType &peek_type = peek->param_type();
+        if (expr.result_type().is_double() && peek_type.is_dense()) {
+            std::vector<std::pair<int64_t,size_t>> spec;
+            assert(peek_type.dimensions().size() == peek->spec().size());
+            for (auto dim = peek_type.dimensions().rbegin(); dim != peek_type.dimensions().rend(); ++dim) {
+                auto dim_spec = peek->spec().find(dim->name);
+                assert(dim_spec != peek->spec().end());
+                if (std::holds_alternative<TensorSpec::Label>(dim_spec->second)) {
+                    const auto &label = std::get<TensorSpec::Label>(dim_spec->second);
+                    assert(label.is_indexed());
+                    spec.emplace_back(label.index, dim->size);
+                } else {
+                    assert(std::holds_alternative<TensorFunction::Child>(dim_spec->second));
+                    spec.emplace_back(-1, dim->size);
+                }
+            }
+            return stash.create<DenseTensorPeekFunction>(peek->copy_children(), spec);
+        }
+    }
+    return expr;
+}
+
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.h b/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.h
new file mode 100644
index 00000000000..4924a0de329
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.h
@@ -0,0 +1,34 @@
+// Copyright 2019 Oath Inc. 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 looking at the cell value of a dense tensor.
+ */
+class DenseTensorPeekFunction : public eval::TensorFunction
+{
+private:
+    // first child is the tensor we want to peek
+    // other children are dimension index expressions
+    // (index expressions are sorted by normalized dimension order)
+    std::vector<Child> _children;
+
+    // index and size of all dimensions in reverse order
+    // when index is -1, use result of next child expression
+    // (note that child expression order is inverted by the stack)
+    std::vector<std::pair<int64_t,size_t>> _spec;
+public:
+    DenseTensorPeekFunction(std::vector<Child> children, const std::vector<std::pair<int64_t,size_t>> &spec);
+    ~DenseTensorPeekFunction();
+    const eval::ValueType &result_type() const override { return eval::DoubleValue::double_type(); }
+    void push_children(std::vector<Child::CREF> &children) const override;
+    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    bool result_is_mutable() const override { return true; }
+    static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
+};
+
+} // namespace vespalib::tensor