dense simple map

11 files changed, 231 insertions, 180 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index 353a9b3e701..9a7bd07534c 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -34,12 +34,12 @@ vespa_define_module(
     src/tests/tensor/dense_fast_rename_optimizer
     src/tests/tensor/dense_generic_join
     src/tests/tensor/dense_inplace_join_function
-    src/tests/tensor/dense_inplace_map_function
     src/tests/tensor/dense_matmul_function
     src/tests/tensor/dense_multi_matmul_function
     src/tests/tensor/dense_remove_dimension_optimizer
     src/tests/tensor/dense_replace_type_function
     src/tests/tensor/dense_simple_join_function
+    src/tests/tensor/dense_simple_map_function
     src/tests/tensor/dense_single_reduce_function
     src/tests/tensor/dense_tensor_create_function
     src/tests/tensor/dense_tensor_peek_function
diff --git a/eval/src/tests/tensor/dense_inplace_map_function/CMakeLists.txt b/eval/src/tests/tensor/dense_inplace_map_function/CMakeLists.txt
deleted file mode 100644
index 3bd1dbaf271..00000000000
--- a/eval/src/tests/tensor/dense_inplace_map_function/CMakeLists.txt
+++ /dev/null
@@ -1,8 +0,0 @@
-# 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
deleted file mode 100644
index f85742b4e0f..00000000000
--- a/eval/src/tests/tensor/dense_inplace_map_function/dense_inplace_map_function_test.cpp
+++ /dev/null
@@ -1,79 +0,0 @@
-// 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("_x5f", spec(float_cells({x(5)}), N()))
-        .add_mutable("_x5y3", spec({x(5),y(3)}, N()))
-        .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 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("require that optimization works for float cells") {
-    TEST_DO(verify_optimized("map(_x5f,f(x)(x+10))", 1));
-}
-
-TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/eval/src/tests/tensor/dense_simple_map_function/CMakeLists.txt b/eval/src/tests/tensor/dense_simple_map_function/CMakeLists.txt
new file mode 100644
index 00000000000..766986bd628
--- /dev/null
+++ b/eval/src/tests/tensor/dense_simple_map_function/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(eval_dense_simple_map_function_test_app TEST
+    SOURCES
+    dense_simple_map_function_test.cpp
+    DEPENDS
+    vespaeval
+    gtest
+)
+vespa_add_test(NAME eval_dense_simple_map_function_test_app COMMAND eval_dense_simple_map_function_test_app)
diff --git a/eval/src/tests/tensor/dense_simple_map_function/dense_simple_map_function_test.cpp b/eval/src/tests/tensor/dense_simple_map_function/dense_simple_map_function_test.cpp
new file mode 100644
index 00000000000..99ec719fa2f
--- /dev/null
+++ b/eval/src/tests/tensor/dense_simple_map_function/dense_simple_map_function_test.cpp
@@ -0,0 +1,78 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#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_simple_map_function.h>
+#include <vespa/eval/eval/test/eval_fixture.h>
+#include <vespa/eval/eval/test/tensor_model.hpp>
+#include <vespa/vespalib/gtest/gtest.h>
+
+using namespace vespalib;
+using namespace vespalib::eval;
+using namespace vespalib::eval::test;
+using namespace vespalib::eval::tensor_function;
+using namespace vespalib::tensor;
+
+const TensorEngine &prod_engine = DefaultTensorEngine::ref();
+
+EvalFixture::ParamRepo make_params() {
+    return EvalFixture::ParamRepo()
+        .add("a", spec(1.5))
+        .add("b", spec(2.5))
+        .add("sparse", spec({x({"a"})}, N()))
+        .add("mixed", spec({x({"a"}),y(5)}, N()))
+        .add_matrix("x", 5, "y", 3);
+}
+EvalFixture::ParamRepo param_repo = make_params();
+
+void verify_optimized(const vespalib::string &expr, bool inplace) {
+    EvalFixture slow_fixture(prod_engine, expr, param_repo, false);
+    EvalFixture fixture(prod_engine, expr, param_repo, true, true);
+    EXPECT_EQ(fixture.result(), EvalFixture::ref(expr, param_repo));
+    EXPECT_EQ(fixture.result(), slow_fixture.result());
+    auto info = fixture.find_all<DenseSimpleMapFunction>();
+    ASSERT_EQ(info.size(), 1u);
+    EXPECT_TRUE(info[0]->result_is_mutable());
+    EXPECT_EQ(info[0]->inplace(), inplace);
+    ASSERT_EQ(fixture.num_params(), 1);
+    if (inplace) {
+        EXPECT_EQ(fixture.get_param(0), fixture.result());
+    } else {
+        EXPECT_TRUE(!(fixture.get_param(0) == fixture.result()));
+    }
+}
+
+void verify_not_optimized(const vespalib::string &expr) {
+    EvalFixture slow_fixture(prod_engine, expr, param_repo, false);
+    EvalFixture fixture(prod_engine, expr, param_repo, true);
+    EXPECT_EQ(fixture.result(), EvalFixture::ref(expr, param_repo));
+    EXPECT_EQ(fixture.result(), slow_fixture.result());
+    auto info = fixture.find_all<DenseSimpleMapFunction>();
+    EXPECT_TRUE(info.empty());
+}
+
+TEST(MapTest, dense_map_is_optimized) {
+    verify_optimized("map(x5y3,f(x)(x+10))", false);
+    verify_optimized("map(x5y3f,f(x)(x+10))", false);
+}
+
+TEST(MapTest, simple_dense_map_can_be_inplace) {
+    verify_optimized("map(@x5y3,f(x)(x+10))", true);
+    verify_optimized("map(@x5y3f,f(x)(x+10))", true);
+}
+
+TEST(MapTest, scalar_map_is_not_optimized) {
+    verify_not_optimized("map(a,f(x)(x+10))");
+}
+
+TEST(MapTest, sparse_map_is_not_optimized) {
+    verify_not_optimized("map(sparse,f(x)(x+10))");
+}
+
+TEST(MapTest, mixed_map_is_not_optimized) {
+    verify_not_optimized("map(mixed,f(x)(x+10))");
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
index 2b13c7652c7..d28e420e6f2 100644
--- a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
+++ b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
@@ -18,7 +18,7 @@
 #include "dense/dense_remove_dimension_optimizer.h"
 #include "dense/dense_lambda_peek_optimizer.h"
 #include "dense/dense_simple_join_function.h"
-#include "dense/dense_inplace_map_function.h"
+#include "dense/dense_simple_map_function.h"
 #include "dense/vector_from_doubles_function.h"
 #include "dense/dense_tensor_create_function.h"
 #include "dense/dense_tensor_peek_function.h"
@@ -290,7 +290,7 @@ DefaultTensorEngine::optimize(const TensorFunction &expr, Stash &stash) const
             child.set(DenseTensorPeekFunction::optimize(child.get(), stash));
             child.set(DenseLambdaPeekOptimizer::optimize(child.get(), stash));
             child.set(DenseFastRenameOptimizer::optimize(child.get(), stash));
-            child.set(DenseInplaceMapFunction::optimize(child.get(), stash));
+            child.set(DenseSimpleMapFunction::optimize(child.get(), stash));
             child.set(DenseSimpleJoinFunction::optimize(child.get(), stash));
             child.set(DenseSingleReduceFunction::optimize(child.get(), stash));
             nodes.pop_back();
diff --git a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
index b955b59cc39..e0e73ecb3b8 100644
--- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
+++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
@@ -6,7 +6,6 @@ vespa_add_library(eval_tensor_dense OBJECT
     dense_dimension_combiner.cpp
     dense_dot_product_function.cpp
     dense_fast_rename_optimizer.cpp
-    dense_inplace_map_function.cpp
     dense_lambda_peek_function.cpp
     dense_lambda_peek_optimizer.cpp
     dense_matmul_function.cpp
@@ -14,6 +13,7 @@ vespa_add_library(eval_tensor_dense OBJECT
     dense_remove_dimension_optimizer.cpp
     dense_replace_type_function.cpp
     dense_simple_join_function.cpp
+    dense_simple_map_function.cpp
     dense_single_reduce_function.cpp
     dense_tensor.cpp
     dense_tensor_address_mapper.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
deleted file mode 100644
index 62434073f8e..00000000000
--- a/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.cpp
+++ /dev/null
@@ -1,64 +0,0 @@
-// 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_view.h"
-#include <vespa/eval/eval/value.h>
-
-namespace vespalib::tensor {
-
-using eval::Value;
-using eval::ValueType;
-using eval::TensorFunction;
-using eval::TensorEngine;
-using eval::as;
-using namespace eval::tensor_function;
-
-namespace {
-
-template <typename CT>
-void my_inplace_map_op(eval::InterpretedFunction::State &state, uint64_t param) {
-    map_fun_t function = (map_fun_t)param;
-    ArrayRef<CT> cells = unconstify(DenseTensorView::typify_cells<CT>(state.peek(0)));
-    for (CT &cell: cells) {
-        cell = function(cell);
-    }
-}
-
-struct MyInplaceMapOp {
-    template <typename CT>
-    static auto get_fun() { return my_inplace_map_op<CT>; }
-};
-
-} // namespace vespalib::tensor::<unnamed>
-
-DenseInplaceMapFunction::DenseInplaceMapFunction(const eval::ValueType &result_type,
-                                                 const eval::TensorFunction &child,
-                                                 map_fun_t function_in)
-    : eval::tensor_function::Map(result_type, child, function_in)
-{
-}
-
-DenseInplaceMapFunction::~DenseInplaceMapFunction()
-{
-}
-
-eval::InterpretedFunction::Instruction
-DenseInplaceMapFunction::compile_self(const TensorEngine &, Stash &) const
-{
-    auto op = select_1<MyInplaceMapOp>(result_type().cell_type());
-    return eval::InterpretedFunction::Instruction(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() && map->result_type().is_dense()) {
-            assert(map->result_type().cell_type() == map->child().result_type().cell_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
deleted file mode 100644
index 52122f4d8dc..00000000000
--- a/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.h
+++ /dev/null
@@ -1,25 +0,0 @@
-// 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::Map
-{
-public:
-    using map_fun_t = ::vespalib::eval::tensor_function::map_fun_t;
-    DenseInplaceMapFunction(const eval::ValueType &result_type,
-                            const eval::TensorFunction &child,
-                            map_fun_t function_in);
-    ~DenseInplaceMapFunction();
-    bool result_is_mutable() const override { return true; }
-    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
-    static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
-};
-
-} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_simple_map_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_simple_map_function.cpp
new file mode 100644
index 00000000000..910e8296afe
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_simple_map_function.cpp
@@ -0,0 +1,115 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "dense_simple_map_function.h"
+#include "dense_tensor_view.h"
+#include <vespa/eval/eval/value.h>
+#include <vespa/eval/eval/operation.h>
+
+namespace vespalib::tensor {
+
+using vespalib::ArrayRef;
+
+using eval::Value;
+using eval::ValueType;
+using eval::TensorFunction;
+using eval::TensorEngine;
+using eval::as;
+
+using namespace eval::operation;
+using namespace eval::tensor_function;
+
+using op_function = eval::InterpretedFunction::op_function;
+using Instruction = eval::InterpretedFunction::Instruction;
+using State = eval::InterpretedFunction::State;
+
+namespace {
+
+struct CallFun {
+    map_fun_t function;
+    CallFun(map_fun_t function_in) : function(function_in) {}
+    double eval(double a) const { return function(a); }
+};
+
+template <typename CT, typename Fun>
+void apply_fun_to_n(CT *dst, const CT *src, size_t n, const Fun &fun) {
+    for (size_t i = 0; i < n; ++i) {
+        dst[i] = fun.eval(src[i]);
+    }
+}
+
+template <typename CT, bool inplace>
+ArrayRef<CT> make_dst_cells(ConstArrayRef<CT> src_cells, Stash &stash) {
+    if (inplace) {
+        return unconstify(src_cells);
+    } else {
+        return stash.create_array<CT>(src_cells.size());
+    }
+}
+
+template <typename CT, typename Fun, bool inplace>
+void my_simple_map_op(State &state, uint64_t param) {
+    Fun my_fun((map_fun_t)param);
+    auto const &child = state.peek(0);
+    auto src_cells = DenseTensorView::typify_cells<CT>(child);
+    auto dst_cells = make_dst_cells<CT, inplace>(src_cells, state.stash);
+    apply_fun_to_n(dst_cells.begin(), src_cells.begin(), dst_cells.size(), my_fun);
+    if (!inplace) {
+        state.pop_push(state.stash.create<DenseTensorView>(child.type(), TypedCells(dst_cells)));
+    }
+}
+
+//-----------------------------------------------------------------------------
+
+template <typename Fun, bool inplace>
+struct MySimpleMapOp {
+    template <typename CT>
+    static auto get_fun() { return my_simple_map_op<CT,Fun,inplace>; }
+};
+
+template <typename Fun>
+op_function my_select_2(ValueType::CellType ct, bool inplace) {
+    if (inplace) {
+        return select_1<MySimpleMapOp<Fun,true>>(ct);
+    } else {
+        return select_1<MySimpleMapOp<Fun,false>>(ct);
+    }
+}
+
+op_function my_select(ValueType::CellType ct, bool inplace, map_fun_t fun_hint) {
+    (void) fun_hint; // ready for function inlining
+    return my_select_2<CallFun>(ct, inplace);
+}
+
+} // namespace vespalib::tensor::<unnamed>
+
+//-----------------------------------------------------------------------------
+
+DenseSimpleMapFunction::DenseSimpleMapFunction(const ValueType &result_type,
+                                               const TensorFunction &child,
+                                               map_fun_t function_in)
+    : Map(result_type, child, function_in)
+{
+}
+
+DenseSimpleMapFunction::~DenseSimpleMapFunction() = default;
+
+Instruction
+DenseSimpleMapFunction::compile_self(const TensorEngine &, Stash &) const
+{
+    auto op = my_select(result_type().cell_type(), inplace(), function());
+    static_assert(sizeof(uint64_t) == sizeof(function()));
+    return Instruction(op, (uint64_t)(function()));
+}
+
+const TensorFunction &
+DenseSimpleMapFunction::optimize(const TensorFunction &expr, Stash &stash)
+{
+    if (auto map = as<Map>(expr)) {
+        if (map->child().result_type().is_dense()) {
+            return stash.create<DenseSimpleMapFunction>(map->result_type(), map->child(), map->function());
+        }
+    }
+    return expr;
+}
+
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_simple_map_function.h b/eval/src/vespa/eval/tensor/dense/dense_simple_map_function.h
new file mode 100644
index 00000000000..daf35a44adc
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_simple_map_function.h
@@ -0,0 +1,25 @@
+// 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/tensor_function.h>
+
+namespace vespalib::tensor {
+
+/**
+ * Tensor function for simple map operations on dense tensors.
+ **/
+class DenseSimpleMapFunction : public eval::tensor_function::Map
+{
+public:
+    using map_fun_t = ::vespalib::eval::tensor_function::map_fun_t;
+    DenseSimpleMapFunction(const eval::ValueType &result_type,
+                           const TensorFunction &child,
+                           map_fun_t function_in);
+    ~DenseSimpleMapFunction() override;
+    bool inplace() const { return child().result_is_mutable(); }
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
+    static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
+};
+
+} // namespace vespalib::tensor