added inplace map operation for mutable concrete dense tensors

also added support for specifying mutable inputs in evaluation fixture

9 files changed, 264 insertions, 20 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index 38abfec3c8c..ba42ddecef9 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -26,10 +26,10 @@ 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_map_function
     src/tests/tensor/dense_tensor_address_combiner
     src/tests/tensor/dense_tensor_builder
     src/tests/tensor/dense_xw_product_function
-    src/tests/tensor/vector_from_doubles_function
     src/tests/tensor/sparse_tensor_builder
     src/tests/tensor/tensor_address
     src/tests/tensor/tensor_conformance
@@ -37,6 +37,7 @@ vespa_define_module(
     src/tests/tensor/tensor_performance
     src/tests/tensor/tensor_serialization
     src/tests/tensor/tensor_slime_serialization
+    src/tests/tensor/vector_from_doubles_function
 
     LIBS
     src/vespa/eval
diff --git a/eval/src/tests/tensor/dense_inplace_map_function/CMakeLists.txt b/eval/src/tests/tensor/dense_inplace_map_function/CMakeLists.txt
new file mode 100644
index 00000000000..3bd1dbaf271
--- /dev/null
+++ b/eval/src/tests/tensor/dense_inplace_map_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_map_function_test_app TEST
+    SOURCES
+    dense_inplace_map_function_test.cpp
+    DEPENDS
+    vespaeval
+)
+vespa_add_test(NAME eval_dense_inplace_map_function_test_app COMMAND eval_dense_inplace_map_function_test_app)
diff --git a/eval/src/tests/tensor/dense_inplace_map_function/dense_inplace_map_function_test.cpp b/eval/src/tests/tensor/dense_inplace_map_function/dense_inplace_map_function_test.cpp
new file mode 100644
index 00000000000..77af747a066
--- /dev/null
+++ b/eval/src/tests/tensor/dense_inplace_map_function/dense_inplace_map_function_test.cpp
@@ -0,0 +1,79 @@
+// 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_map_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", spec(5.0))
+        .add_mutable("_x5", spec({x(5)}, N()))
+        .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_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<DenseInplaceMapFunction>();
+    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<DenseInplaceMapFunction>();
+    EXPECT_TRUE(info.empty());
+}
+
+TEST("require that mutable dense concrete tensors are optimized") {
+    TEST_DO(verify_optimized("map(_x5,f(x)(x+10))", 1));
+    TEST_DO(verify_optimized("map(_x5y3,f(x)(x+10))", 1));
+}
+
+TEST("require that inplace map operations can be chained") {
+    TEST_DO(verify_optimized("map(map(_x5,f(x)(x+10)),f(x)(x-5))", 2));
+}
+
+TEST("require that abstract tensors are not optimized") {
+    TEST_DO(verify_not_optimized("map(_x5_u,f(x)(x+10))"));
+}
+
+TEST("require that non-mutable tensors are not optimized") {
+    TEST_DO(verify_not_optimized("map(x5,f(x)(x+10))"));
+}
+
+TEST("require that scalar values are not optimized") {
+    TEST_DO(verify_not_optimized("map(_d,f(x)(x+10))"));
+}
+
+TEST("require that mapped tensors are not optimized") {
+    TEST_DO(verify_not_optimized("map(_x_m,f(x)(x+10))"));
+}
+
+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 b7e3764f6ce..2593b74e300 100644
--- a/eval/src/vespa/eval/eval/test/eval_fixture.cpp
+++ b/eval/src/vespa/eval/eval/test/eval_fixture.cpp
@@ -20,11 +20,44 @@ NodeTypes get_types(const Function &function, const ParamRepo &param_repo) {
     return NodeTypes(function, param_types);
 }
 
-const TensorFunction &make_tfun(bool optimized, const TensorEngine &engine, const Function &function,
-                                const NodeTypes &node_types, Stash &stash)
-{
-    const TensorFunction &plain_fun = make_tensor_function(engine, function.root(), node_types, stash);
-    return optimized ? engine.optimize(plain_fun, stash) : plain_fun;
+std::set<size_t> get_mutable(const Function &function, const ParamRepo &param_repo) {
+    std::set<size_t> mutable_set;
+    for (size_t i = 0; i < function.num_params(); ++i) {
+        auto pos = param_repo.map.find(function.param_name(i));
+        ASSERT_TRUE(pos != param_repo.map.end());
+        if (pos->second.is_mutable) {
+            mutable_set.insert(i);
+        }
+    }
+    return mutable_set;
+}
+
+struct MyMutableInject : public tensor_function::Inject {
+    MyMutableInject(const ValueType &result_type_in, size_t param_idx_in)
+        : Inject(result_type_in, param_idx_in) {}
+    bool result_is_mutable() const override { return true; }
+};
+
+const TensorFunction &maybe_patch(bool allow_mutable, const TensorFunction &plain_fun, const std::set<size_t> &mutable_set, Stash &stash) {
+    using Child = TensorFunction::Child;
+    if (!allow_mutable) {
+        return plain_fun;
+    }
+    Child root(plain_fun);
+    std::vector<Child::CREF> nodes({root});
+    for (size_t i = 0; i < nodes.size(); ++i) {
+        nodes[i].get().get().push_children(nodes);
+    }
+    while (!nodes.empty()) {
+        const Child &child = nodes.back();
+        if (auto inject = as<tensor_function::Inject>(child.get())) {
+            if (mutable_set.count(inject->param_idx()) > 0) {
+                child.set(stash.create<MyMutableInject>(inject->result_type(), inject->param_idx()));
+            }
+        }
+        nodes.pop_back();
+    }
+    return root.get();
 }
 
 std::vector<Value::UP> make_params(const TensorEngine &engine, const Function &function,
@@ -53,12 +86,16 @@ std::vector<Value::CREF> get_refs(const std::vector<Value::UP> &values) {
 EvalFixture::EvalFixture(const TensorEngine &engine,
                          const vespalib::string &expr,
                          const ParamRepo &param_repo,
-                         bool optimized)
+                         bool optimized,
+                         bool allow_mutable)
     : _engine(engine),
       _stash(),
       _function(Function::parse(expr)),
       _node_types(get_types(_function, param_repo)),
-      _tensor_function(make_tfun(optimized, _engine, _function, _node_types, _stash)),
+      _mutable_set(get_mutable(_function, param_repo)),
+      _plain_tensor_function(make_tensor_function(_engine, _function.root(), _node_types, _stash)),
+      _patched_tensor_function(maybe_patch(allow_mutable, _plain_tensor_function, _mutable_set, _stash)),
+      _tensor_function(optimized ? _engine.optimize(_patched_tensor_function, _stash) : _patched_tensor_function),
       _ifun(_engine, _tensor_function),
       _ictx(_ifun),
       _param_values(make_params(_engine, _function, param_repo)),
@@ -69,4 +106,11 @@ EvalFixture::EvalFixture(const TensorEngine &engine,
     ASSERT_TRUE(!result_type.is_error());
 }
 
+const TensorSpec
+EvalFixture::get_param(size_t idx) const
+{
+    ASSERT_LESS(idx, _param_values.size());
+    return _engine.to_spec(*(_param_values[idx]));
+}
+
 } // 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 1f864e980cc..7d3b95ad144 100644
--- a/eval/src/vespa/eval/eval/test/eval_fixture.h
+++ b/eval/src/vespa/eval/eval/test/eval_fixture.h
@@ -9,6 +9,7 @@
 #include <vespa/eval/eval/simple_tensor_engine.h>
 #include <vespa/eval/tensor/default_tensor_engine.h>
 #include <vespa/vespalib/util/stash.h>
+#include <set>
 
 namespace vespalib::eval::test {
 
@@ -18,23 +19,30 @@ public:
     struct Param {
         TensorSpec value; // actual parameter value
         vespalib::string type;  // pre-defined type (could be abstract)
-        Param(TensorSpec value_in)
-            : value(std::move(value_in)), type(value.type()) {}
-        Param(TensorSpec value_in, const vespalib::string &type_in)
-            : value(std::move(value_in)), type(type_in) {}
+        bool is_mutable; // input will be mutable (if allow_mutable is true)
+        Param(TensorSpec value_in, const vespalib::string &type_in, bool is_mutable_in)
+            : value(std::move(value_in)), type(type_in), is_mutable(is_mutable_in) {}
         ~Param() {}
     };
 
     struct ParamRepo {
         std::map<vespalib::string,Param> map;
         ParamRepo() : map() {}
-        ParamRepo &add(const vespalib::string &name, TensorSpec value_in) {
-            map.insert_or_assign(name, Param(std::move(value_in)));
+        ParamRepo &add(const vespalib::string &name, TensorSpec value_in, const vespalib::string &type_in, bool is_mutable_in) {
+            map.insert_or_assign(name, Param(std::move(value_in), type_in, is_mutable_in));
             return *this;
         }
-        ParamRepo &add(const vespalib::string &name, TensorSpec value_in, const vespalib::string &type_in) {
-            map.insert_or_assign(name, Param(std::move(value_in), type_in));
-            return *this;
+        ParamRepo &add(const vespalib::string &name, TensorSpec value, const vespalib::string &type) {
+            return add(name, value, type, false);
+        }
+        ParamRepo &add_mutable(const vespalib::string &name, TensorSpec value, const vespalib::string &type) {
+            return add(name, value, type, true);
+        }
+        ParamRepo &add(const vespalib::string &name, const TensorSpec &value) {
+            return add(name, value, value.type(), false);
+        }
+        ParamRepo &add_mutable(const vespalib::string &name, const TensorSpec &value) {
+            return add(name, value, value.type(), true);
         }
         ~ParamRepo() {}
     };
@@ -44,6 +52,9 @@ private:
     Stash                         _stash;
     Function                      _function;
     NodeTypes                     _node_types;
+    std::set<size_t>              _mutable_set;
+    const TensorFunction         &_plain_tensor_function;
+    const TensorFunction         &_patched_tensor_function;
     const TensorFunction         &_tensor_function;
     InterpretedFunction           _ifun;
     InterpretedFunction::Context  _ictx;
@@ -64,7 +75,8 @@ private:
     }
 
 public:
-    EvalFixture(const TensorEngine &engine, const vespalib::string &expr, const ParamRepo &param_repo, bool optimized);
+    EvalFixture(const TensorEngine &engine, const vespalib::string &expr, const ParamRepo &param_repo,
+                bool optimized = true, bool allow_mutable = false);
     ~EvalFixture() {}
     template <typename T>
     std::vector<const T *> find_all() {
@@ -73,11 +85,12 @@ public:
         return list;
     }
     const TensorSpec &result() const { return _result; }
+    const TensorSpec get_param(size_t idx) const;
     static TensorSpec ref(const vespalib::string &expr, const ParamRepo &param_repo) {
-        return EvalFixture(SimpleTensorEngine::ref(), expr, param_repo, false).result();
+        return EvalFixture(SimpleTensorEngine::ref(), expr, param_repo, false, false).result();
     }
     static TensorSpec prod(const vespalib::string &expr, const ParamRepo &param_repo) {
-        return EvalFixture(tensor::DefaultTensorEngine::ref(), expr, param_repo, true).result();
+        return EvalFixture(tensor::DefaultTensorEngine::ref(), expr, param_repo, true, 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 7b4a502dd4d..dead8ee4870 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_map_function.h"
 #include "dense/vector_from_doubles_function.h"
 #include <vespa/eval/eval/value.h>
 #include <vespa/eval/eval/tensor_spec.h>
@@ -223,6 +224,7 @@ DefaultTensorEngine::optimize(const TensorFunction &expr, Stash &stash) const
         child.set(DenseDotProductFunction::optimize(child.get(), stash));
         child.set(DenseXWProductFunction::optimize(child.get(), stash));
         child.set(DenseFastRenameFunction::optimize(child.get(), stash));
+        child.set(DenseInplaceMapFunction::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 73315b7a120..cf6d3a3431c 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_map_function.cpp
     dense_tensor.cpp
     dense_tensor_address_combiner.cpp
     dense_tensor_builder.cpp
diff --git a/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.cpp
new file mode 100644
index 00000000000..162bdb2ebfe
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.cpp
@@ -0,0 +1,67 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "dense_inplace_map_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());
+}
+
+void my_inplace_map_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    map_fun_t function = (map_fun_t)param;
+    for (double &cell: getMutableCells(state.peek(0))) {
+        cell = function(cell);
+    }
+}
+
+bool isConcreteDenseTensor(const ValueType &type) {
+    return (type.is_dense() && !type.is_abstract());
+}
+
+} // namespace vespalib::tensor::<unnamed>
+
+DenseInplaceMapFunction::DenseInplaceMapFunction(const eval::ValueType &result_type,
+                                                 const eval::TensorFunction &child,
+                                                 map_fun_t function_in)
+    : eval::tensor_function::Op1(result_type, child),
+      _function(function_in)
+{
+}
+
+DenseInplaceMapFunction::~DenseInplaceMapFunction()
+{
+}
+
+eval::InterpretedFunction::Instruction
+DenseInplaceMapFunction::compile_self(Stash &) const
+{
+    return eval::InterpretedFunction::Instruction(my_inplace_map_op, (uint64_t)_function);
+}
+
+const TensorFunction &
+DenseInplaceMapFunction::optimize(const eval::TensorFunction &expr, Stash &stash)
+{
+    if (auto map = as<Map>(expr)) {
+        if (map->child().result_is_mutable() && isConcreteDenseTensor(map->result_type())) {
+            return stash.create<DenseInplaceMapFunction>(map->result_type(), map->child(), map->function());
+        }
+    }
+    return expr;
+}
+
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.h b/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.h
new file mode 100644
index 00000000000..f02f83edae1
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.h
@@ -0,0 +1,29 @@
+// 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 map operation on mutable dense tensors.
+ **/
+class DenseInplaceMapFunction : public eval::tensor_function::Op1
+{
+public:
+    using map_fun_t = ::vespalib::eval::tensor_function::map_fun_t;
+private:
+    map_fun_t _function;
+public:
+    DenseInplaceMapFunction(const eval::ValueType &result_type,
+                            const eval::TensorFunction &child,
+                            map_fun_t function_in);
+    ~DenseInplaceMapFunction();
+    map_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