Merge pull request #20618 from vespa-engine/havardpe/sparse-tensor-create

optimize tensor create by making the index up front...

2 files changed, 53 insertions, 54 deletions

diff --git a/eval/src/vespa/eval/eval/fast_value.hpp b/eval/src/vespa/eval/eval/fast_value.hpp
index 8a07b049a69..185529b2f51 100644
--- a/eval/src/vespa/eval/eval/fast_value.hpp
+++ b/eval/src/vespa/eval/eval/fast_value.hpp
@@ -238,7 +238,16 @@ struct FastValue final : Value, ValueBuilder<T> {
     ~FastValue() override;
     const ValueType &type() const override { return my_type; }
     const Value::Index &index() const override { return my_index; }
-    TypedCells cells() const override { return TypedCells(my_cells.memory, get_cell_type<T>(), my_cells.size); }
+    TypedCells cells() const override {
+        if constexpr (std::is_same_v<T, uint32_t>) {
+            // allow use of FastValue templated on types that do not
+            // have a corresponding cell type as long as cells() is
+            // not called
+            abort();
+        } else {
+            return TypedCells(my_cells.memory, get_cell_type<T>(), my_cells.size);
+        }
+    }
     void add_mapping(ConstArrayRef<vespalib::stringref> addr) {
         if constexpr (transient) {
             (void) addr;
@@ -277,6 +286,12 @@ struct FastValue final : Value, ValueBuilder<T> {
         add_mapping(addr);
         return my_cells.add_cells(my_subspace_size);        
     }
+    ArrayRef<T> get_subspace(size_t subspace) {
+        return {my_cells.get(subspace * my_subspace_size), my_subspace_size};
+    }
+    ConstArrayRef<T> get_raw_cells() const {
+        return {my_cells.get(0), my_cells.size};
+    }
     std::unique_ptr<Value> build(std::unique_ptr<ValueBuilder<T>> self) override {
         if (my_index.map.addr_size() == 0) {
             assert(my_index.map.size() == 1);
diff --git a/eval/src/vespa/eval/instruction/generic_create.cpp b/eval/src/vespa/eval/instruction/generic_create.cpp
index c0be28a2475..11fd2437601 100644
--- a/eval/src/vespa/eval/instruction/generic_create.cpp
+++ b/eval/src/vespa/eval/instruction/generic_create.cpp
@@ -2,7 +2,7 @@
 
 #include "generic_create.h"
 #include <vespa/eval/eval/wrap_param.h>
-#include <vespa/eval/eval/array_array_map.h>
+#include <vespa/eval/eval/fast_value.hpp>
 #include <vespa/vespalib/util/stash.h>
 #include <vespa/vespalib/util/typify.h>
 #include <vespa/vespalib/util/shared_string_repo.h>
@@ -19,35 +19,33 @@ using Handle = SharedStringRepo::Handle;
 namespace {
 
 struct CreateParam {
-    const ValueType res_type;
-    size_t num_mapped_dims;
-    size_t dense_subspace_size;
+    static constexpr uint32_t npos = -1;
+    FastValue<uint32_t, false> my_spec;
     size_t num_children;
-    ArrayArrayMap<Handle,size_t> my_spec;
-    const ValueBuilderFactory &factory;
 
-    static constexpr size_t npos = -1;
-
-    ArrayRef<size_t> indexes(ConstArrayRef<Handle> key) {
-        auto [tag, first_time] = my_spec.lookup_or_add_entry(key);
-        auto rv = my_spec.get_values(tag);
-        if (first_time) {
-            for (auto & v : rv) {
-                v = npos;
-            }
+    ArrayRef<uint32_t> indexes(ConstArrayRef<Handle> key) {
+        SmallVector<string_id> my_key;
+        for (const auto &label: key) {
+            my_key.push_back(label.id());
+        }
+        auto old_subspace = my_spec.my_index.map.lookup(ConstArrayRef<string_id>(my_key));
+        if (old_subspace != FastAddrMap::npos()) {
+            return my_spec.get_subspace(old_subspace);
         }
-        return rv;
+        auto new_subspace = my_spec.add_subspace(my_key);
+        for (auto &stack_idx: new_subspace) {
+            stack_idx = npos;
+        }
+        return new_subspace;
     }
 
-    CreateParam(const ValueType &res_type_in,
-                const GenericCreate::SpecMap &spec_in,
-                const ValueBuilderFactory &factory_in)
-        : res_type(res_type_in),
-          num_mapped_dims(res_type.count_mapped_dimensions()),
-          dense_subspace_size(res_type.dense_subspace_size()),
-          num_children(spec_in.size()),
-          my_spec(num_mapped_dims, dense_subspace_size, spec_in.size() / dense_subspace_size),
-          factory(factory_in)
+    CreateParam(const ValueType &res_type,
+                const GenericCreate::SpecMap &spec_in)
+      : my_spec(res_type,
+                res_type.count_mapped_dimensions(),
+                res_type.dense_subspace_size(),
+                spec_in.size() / res_type.dense_subspace_size()),
+        num_children(spec_in.size())
     {
         size_t last_child = num_children - 1;
         for (const auto & entry : spec_in) {
@@ -67,7 +65,6 @@ struct CreateParam {
                 }
             }
             assert(binding == entry.first.end());
-            assert(dense_key < dense_subspace_size);
             // note: reverse order of children on stack
             size_t stack_idx = last_child - entry.second;
             indexes(sparse_key)[dense_key] = stack_idx;
@@ -75,37 +72,24 @@ struct CreateParam {
     }
 };
 
-template <typename T>
+template <typename CT>
 void my_generic_create_op(State &state, uint64_t param_in) {
     const auto &param = unwrap_param<CreateParam>(param_in);
-    auto builder = param.factory.create_transient_value_builder<T>(param.res_type,
-                                                                   param.num_mapped_dims,
-                                                                   param.dense_subspace_size,
-                                                                   param.my_spec.size());
-    SmallVector<string_id> sparse_addr;
-    param.my_spec.each_entry([&](const auto &key, const auto &values)
-        {
-            sparse_addr.clear();
-            for (const auto & label : key) {
-                sparse_addr.push_back(label.id());
-            }
-            T *dst = builder->add_subspace(sparse_addr).begin();
-            for (size_t stack_idx : values) {
-                if (stack_idx == CreateParam::npos) {
-                    *dst++ = T{};
-                } else {
-                    const Value &child = state.peek(stack_idx);
-                    *dst++ = child.as_double();
-                }
-            }
-        });
-    const Value &result = *state.stash.create<Value::UP>(builder->build(std::move(builder)));
+    auto spec = param.my_spec.get_raw_cells();
+    auto cells = state.stash.create_uninitialized_array<CT>(spec.size());
+    CT *dst = cells.begin();
+    for (uint32_t stack_idx: spec) {
+        *dst++ = ((stack_idx != CreateParam::npos)
+                  ? (CT) state.peek(stack_idx).as_double()
+                  : CT{});
+    }
+    const Value &result = state.stash.create<ValueView>(param.my_spec.type(), param.my_spec.my_index, TypedCells(cells));
     state.pop_n_push(param.num_children, result);
 };
 
 struct SelectGenericCreateOp {
-    template <typename T> static auto invoke() {
-        return my_generic_create_op<T>;
+    template <typename CT> static auto invoke() {
+        return my_generic_create_op<CT>;
     }
 };
 
@@ -116,10 +100,10 @@ struct SelectGenericCreateOp {
 Instruction
 GenericCreate::make_instruction(const ValueType &result_type,
                                 const SpecMap &spec,
-                                const ValueBuilderFactory &factory,
+                                const ValueBuilderFactory &,
                                 Stash &stash)
 {
-    const auto &param = stash.create<CreateParam>(result_type, spec, factory);
+    const auto &param = stash.create<CreateParam>(result_type, spec);
     auto fun = typify_invoke<1,TypifyCellType,SelectGenericCreateOp>(result_type.cell_type());
     return Instruction(fun, wrap_param<CreateParam>(param));
 }