optimize dense tensor create

9 files changed, 225 insertions, 2 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index d91e35f2d63..b3d4585a176 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -25,8 +25,8 @@ vespa_define_module(
     src/tests/eval/value_cache
     src/tests/eval/value_type
     src/tests/gp/ponder_nov2017
-    src/tests/tensor/dense_dimension_combiner
     src/tests/tensor/dense_add_dimension_optimizer
+    src/tests/tensor/dense_dimension_combiner
     src/tests/tensor/dense_dot_product_function
     src/tests/tensor/dense_fast_rename_optimizer
     src/tests/tensor/dense_generic_join
@@ -34,6 +34,7 @@ vespa_define_module(
     src/tests/tensor/dense_inplace_map_function
     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_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_create_function/CMakeLists.txt b/eval/src/tests/tensor/dense_tensor_create_function/CMakeLists.txt
new file mode 100644
index 00000000000..883f331bda8
--- /dev/null
+++ b/eval/src/tests/tensor/dense_tensor_create_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_create_function_test_app TEST
+    SOURCES
+    dense_tensor_create_function_test.cpp
+    DEPENDS
+    vespaeval
+)
+vespa_add_test(NAME eval_dense_tensor_create_function_test_app COMMAND eval_dense_tensor_create_function_test_app)
diff --git a/eval/src/tests/tensor/dense_tensor_create_function/dense_tensor_create_function_test.cpp b/eval/src/tests/tensor/dense_tensor_create_function/dense_tensor_create_function_test.cpp
new file mode 100644
index 00000000000..0fa09e6c46e
--- /dev/null
+++ b/eval/src/tests/tensor/dense_tensor_create_function/dense_tensor_create_function_test.cpp
@@ -0,0 +1,64 @@
+// 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_create_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));
+}
+EvalFixture::ParamRepo param_repo = make_params();
+
+void verify(const vespalib::string &expr, size_t expect_optimized_cnt, size_t expect_not_optimized_cnt) {
+    EvalFixture fixture(prod_engine, expr, param_repo, true);
+    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
+    auto info = fixture.find_all<DenseTensorCreateFunction>();
+    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<Create>().size(), expect_not_optimized_cnt);
+}
+
+//-----------------------------------------------------------------------------
+
+TEST("require that tensor create can be optimized") {
+    TEST_DO(verify("tensor(x[3]):{{x:0}:1,{x:1}:2,{x:2}:3}", 1, 0));
+    TEST_DO(verify("tensor(x[3]):{{x:0}:a,{x:1}:b,{x:2}:c}", 1, 0));
+    TEST_DO(verify("tensor<float>(x[3]):{{x:0}:a,{x:1}:b,{x:2}:c}", 1, 0));
+    TEST_DO(verify("tensor(x[3]):{{x:0}:a+b,{x:1}:b-c,{x:2}:c*a}", 1, 0));
+}
+
+TEST("require that tensor create can be optimized with missing cells (padded with 0.0)") {
+    TEST_DO(verify("tensor(x[3],y[5]):{{x:0,y:1}:a,{x:1,y:3}:b,{x:2,y:4}:c}", 1, 0));
+}
+
+TEST("require that tensor create in not optimized for sparse tensor") {
+    TEST_DO(verify("tensor(x{}):{{x:0}:a,{x:1}:b,{x:2}:c}", 0, 1));
+}
+
+TEST("require that tensor create in not optimized for mixed tensor") {
+    TEST_DO(verify("tensor(x{},y[3]):{{x:a,y:0}:a,{x:a,y:1}:b,{x:a,y:2}:c}", 0, 1));
+}
+
+TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/eval/src/vespa/eval/eval/value_type.cpp b/eval/src/vespa/eval/eval/value_type.cpp
index d6ba8e83855..688112bbb8a 100644
--- a/eval/src/vespa/eval/eval/value_type.cpp
+++ b/eval/src/vespa/eval/eval/value_type.cpp
@@ -174,6 +174,18 @@ ValueType::is_dense() const
 }
 
 size_t
+ValueType::dense_subspace_size() const
+{
+    size_t size = 1;
+    for (const auto &dim : dimensions()) {
+        if (dim.is_indexed()) {
+            size *= dim.size;
+        }
+    }
+    return size;
+}
+
+size_t
 ValueType::dimension_index(const vespalib::string &name) const {
     return my_dimension_index(_dimensions, name);
 }
diff --git a/eval/src/vespa/eval/eval/value_type.h b/eval/src/vespa/eval/eval/value_type.h
index 64003e2636e..df9c582aee8 100644
--- a/eval/src/vespa/eval/eval/value_type.h
+++ b/eval/src/vespa/eval/eval/value_type.h
@@ -58,6 +58,7 @@ public:
     bool is_tensor() const { return (_type == Type::TENSOR); }
     bool is_sparse() const;
     bool is_dense() const;
+    size_t dense_subspace_size() const;
     const std::vector<Dimension> &dimensions() const { return _dimensions; }
     size_t dimension_index(const vespalib::string &name) const;
     std::vector<vespalib::string> dimension_names() const;
diff --git a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
index f1eb9ff1523..18733b3c733 100644
--- a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
+++ b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
@@ -16,6 +16,7 @@
 #include "dense/dense_inplace_join_function.h"
 #include "dense/dense_inplace_map_function.h"
 #include "dense/vector_from_doubles_function.h"
+#include "dense/dense_tensor_create_function.h"
 #include <vespa/eval/eval/value.h>
 #include <vespa/eval/eval/tensor_spec.h>
 #include <vespa/eval/eval/tensor_spec.h>
@@ -258,6 +259,7 @@ DefaultTensorEngine::optimize(const TensorFunction &expr, Stash &stash) const
     while (!nodes.empty()) {
         const Child &child = nodes.back();
         child.set(VectorFromDoublesFunction::optimize(child.get(), stash));
+        child.set(DenseTensorCreateFunction::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 78723fddc17..e071de3e925 100644
--- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
+++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
@@ -9,9 +9,10 @@ vespa_add_library(eval_tensor_dense OBJECT
     dense_inplace_map_function.cpp
     dense_remove_dimension_optimizer.cpp
     dense_replace_type_function.cpp
+    dense_tensor.cpp
     dense_tensor_address_mapper.cpp
     dense_tensor_cells_iterator.cpp
-    dense_tensor.cpp
+    dense_tensor_create_function.cpp
     dense_tensor_modify.cpp
     dense_tensor_reduce.cpp
     dense_tensor_view.cpp
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.cpp
new file mode 100644
index 00000000000..d590e870cfe
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.cpp
@@ -0,0 +1,101 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "dense_tensor_create_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_create_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    const auto *self = (const DenseTensorCreateFunction::Self *)(param);
+    size_t pending_cells = self->result_size;
+    ArrayRef<CT> cells = state.stash.create_array<CT>(pending_cells);
+    while (pending_cells-- > 0) {
+        cells[pending_cells] = (CT) state.peek(0).as_double();
+        state.stack.pop_back();
+    }
+    const Value &result = state.stash.create<DenseTensorView>(self->result_type, TypedCells(cells)); 
+    state.stack.push_back(result);
+}
+
+struct MyTensorCreateOp {
+    template <typename CT>
+    static auto get_fun() { return my_tensor_create_op<CT>; }
+};
+
+size_t get_index(const TensorSpec::Address &addr, const ValueType &type) {
+    size_t cell_idx = 0;
+    for (const auto &binding: addr) {
+        size_t dim_idx = type.dimension_index(binding.first);
+        assert(dim_idx != ValueType::Dimension::npos);
+        assert(binding.second.is_indexed());
+        cell_idx *= type.dimensions()[dim_idx].size;
+        cell_idx += binding.second.index;
+    }
+    return cell_idx;
+}
+
+} // namespace vespalib::tensor::<unnamed>
+
+DenseTensorCreateFunction::DenseTensorCreateFunction(const ValueType &res_type, std::vector<Child> children)
+    : TensorFunction(),
+      _self(res_type, children.size()),
+      _children(std::move(children))
+{
+}
+
+DenseTensorCreateFunction::~DenseTensorCreateFunction()
+{
+}
+
+void
+DenseTensorCreateFunction::push_children(std::vector<Child::CREF> &target) const
+{
+    for (const Child &c : _children) {
+        target.push_back(c);
+    }
+}
+
+eval::InterpretedFunction::Instruction
+DenseTensorCreateFunction::compile_self(Stash &) const
+{
+    static_assert(sizeof(uint64_t) == sizeof(&_self));
+    auto op = select_1<MyTensorCreateOp>(result_type().cell_type());
+    return eval::InterpretedFunction::Instruction(op, (uint64_t)&_self);
+}
+
+const TensorFunction &
+DenseTensorCreateFunction::optimize(const eval::TensorFunction &expr, Stash &stash)
+{
+    if (auto create = as<Create>(expr)) {
+        if (expr.result_type().is_dense()) {
+            size_t num_cells = expr.result_type().dense_subspace_size();
+            const auto &zero_value = stash.create<DoubleValue>(0.0);
+            const auto &zero_node = const_value(zero_value, stash);
+            std::vector<Child> children(num_cells, zero_node);
+            for (const auto &cell: create->spec()) {
+                size_t cell_idx = get_index(cell.first, expr.result_type());
+                children[cell_idx] = cell.second;
+            }
+            return stash.create<DenseTensorCreateFunction>(expr.result_type(), std::move(children));
+        }
+    }
+    return expr;
+}
+
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.h b/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.h
new file mode 100644
index 00000000000..6d262f48aa6
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.h
@@ -0,0 +1,33 @@
+// 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 creating a dense tensor from double values.
+ */
+class DenseTensorCreateFunction : public eval::TensorFunction
+{
+public:
+    struct Self {
+        eval::ValueType result_type;
+        size_t result_size;
+        Self(const eval::ValueType &r, size_t n) : result_type(r), result_size(n) {}
+    };
+private:
+    Self _self;
+    std::vector<Child> _children;
+public:
+    DenseTensorCreateFunction(const eval::ValueType &res_type, std::vector<Child> children);
+    ~DenseTensorCreateFunction();
+    const eval::ValueType &result_type() const override { return _self.result_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