Merge pull request #14293 from vespa-engine/geirst/optimize-imported-direct-tensor-attribute-access

Support accessing tensor refs directly when using an imported DirectT…

19 files changed, 72 insertions, 55 deletions

diff --git a/searchlib/src/vespa/searchlib/features/attributefeature.cpp b/searchlib/src/vespa/searchlib/features/attributefeature.cpp
index 7c0ed1147f7..eac69f87587 100644
--- a/searchlib/src/vespa/searchlib/features/attributefeature.cpp
+++ b/searchlib/src/vespa/searchlib/features/attributefeature.cpp
@@ -469,13 +469,13 @@ createTensorAttributeExecutor(const IAttributeVector *attribute, const vespalib:
                 tensorType.to_spec().c_str());
         return ConstantTensorExecutor::createEmpty(tensorType, stash);
     }
-    if (tensorType.is_dense()) {
-        return stash.create<DenseTensorAttributeExecutor>(tensorAttribute);
+    if (tensorAttribute->supports_extract_dense_view()) {
+        return stash.create<DenseTensorAttributeExecutor>(*tensorAttribute);
     }
-    if (auto direct = dynamic_cast<const DirectTensorAttribute *>(tensorAttribute)) {
-        return stash.create<DirectTensorAttributeExecutor>(*direct);
+    if (tensorAttribute->supports_get_tensor_ref()) {
+        return stash.create<DirectTensorAttributeExecutor>(*tensorAttribute);
     }
-    return stash.create<TensorAttributeExecutor>(tensorAttribute);
+    return stash.create<TensorAttributeExecutor>(*tensorAttribute);
 }
 
 bool
diff --git a/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.cpp b/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.cpp
index 6eff09b65ab..3dfba3092cd 100644
--- a/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.cpp
+++ b/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.cpp
@@ -10,16 +10,16 @@ using vespalib::tensor::MutableDenseTensorView;
 namespace search::features {
 
 DenseTensorAttributeExecutor::
-DenseTensorAttributeExecutor(const ITensorAttribute *attribute)
+DenseTensorAttributeExecutor(const ITensorAttribute& attribute)
     : _attribute(attribute),
-      _tensorView(_attribute->getTensorType())
+      _tensorView(_attribute.getTensorType())
 {
 }
 
 void
 DenseTensorAttributeExecutor::execute(uint32_t docId)
 {
-    _attribute->getTensor(docId, _tensorView);
+    _attribute.extract_dense_view(docId, _tensorView);
     outputs().set_object(0, _tensorView);
 }
 
diff --git a/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.h b/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.h
index 2bcf356689e..a8a84447c88 100644
--- a/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.h
+++ b/searchlib/src/vespa/searchlib/features/dense_tensor_attribute_executor.h
@@ -16,11 +16,11 @@ namespace search::features {
 class DenseTensorAttributeExecutor : public fef::FeatureExecutor
 {
 private:
-    const search::tensor::ITensorAttribute *_attribute;
+    const search::tensor::ITensorAttribute& _attribute;
     vespalib::tensor::MutableDenseTensorView _tensorView;
 
 public:
-    DenseTensorAttributeExecutor(const search::tensor::ITensorAttribute *attribute);
+    DenseTensorAttributeExecutor(const search::tensor::ITensorAttribute& attribute);
     void execute(uint32_t docId) override;
 };
 
diff --git a/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.cpp b/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.cpp
index cd0f4d4d356..56765bd168f 100644
--- a/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.cpp
+++ b/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.cpp
@@ -1,14 +1,13 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "direct_tensor_attribute_executor.h"
-#include <vespa/searchlib/tensor/direct_tensor_attribute.h>
 #include <vespa/eval/tensor/tensor.h>
+#include <vespa/searchlib/tensor/i_tensor_attribute.h>
 
-namespace search {
-namespace features {
+namespace search::features {
 
 DirectTensorAttributeExecutor::
-DirectTensorAttributeExecutor(const DirectTensorAttribute &attribute)
+DirectTensorAttributeExecutor(const ITensorAttribute &attribute)
     : _attribute(attribute)
 {
 }
@@ -19,5 +18,4 @@ DirectTensorAttributeExecutor::execute(uint32_t docId)
     outputs().set_object(0, _attribute.get_tensor_ref(docId));
 }
 
-} // namespace features
-} // namespace search
+}
diff --git a/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.h b/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.h
index e311fba076e..b9a4df739ff 100644
--- a/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.h
+++ b/searchlib/src/vespa/searchlib/features/direct_tensor_attribute_executor.h
@@ -4,17 +4,17 @@
 
 #include <vespa/searchlib/fef/featureexecutor.h>
 
-namespace search::tensor { class DirectTensorAttribute; }
+namespace search::tensor { class ITensorAttribute; }
 namespace search::features {
 
 class DirectTensorAttributeExecutor : public fef::FeatureExecutor
 {
 public:
-    using DirectTensorAttribute = search::tensor::DirectTensorAttribute;
-    DirectTensorAttributeExecutor(const DirectTensorAttribute &attribute);
+    using ITensorAttribute = search::tensor::ITensorAttribute;
+    DirectTensorAttributeExecutor(const ITensorAttribute &attribute);
     void execute(uint32_t docId) override;
 private:
-    const DirectTensorAttribute &_attribute;
+    const ITensorAttribute &_attribute;
 };
 
 }
diff --git a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp
index 51727846f95..802e98dcd66 100644
--- a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp
+++ b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.cpp
@@ -3,13 +3,12 @@
 #include "tensor_attribute_executor.h"
 #include <vespa/searchlib/tensor/i_tensor_attribute.h>
 
-namespace search {
-namespace features {
+namespace search::features {
 
 TensorAttributeExecutor::
-TensorAttributeExecutor(const search::tensor::ITensorAttribute *attribute)
+TensorAttributeExecutor(const search::tensor::ITensorAttribute& attribute)
     : _attribute(attribute),
-      _emptyTensor(attribute->getEmptyTensor()),
+      _emptyTensor(attribute.getEmptyTensor()),
       _tensor()
 {
 }
@@ -17,7 +16,7 @@ TensorAttributeExecutor(const search::tensor::ITensorAttribute *attribute)
 void
 TensorAttributeExecutor::execute(uint32_t docId)
 {
-    _tensor = _attribute->getTensor(docId);
+    _tensor = _attribute.getTensor(docId);
     if (_tensor) {
         outputs().set_object(0, *_tensor);
     } else {
@@ -25,5 +24,4 @@ TensorAttributeExecutor::execute(uint32_t docId)
     }
 }
 
-} // namespace features
-} // namespace search
+}
diff --git a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.h b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.h
index deba5c05a5d..be900e8a3e7 100644
--- a/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.h
+++ b/searchlib/src/vespa/searchlib/features/tensor_attribute_executor.h
@@ -15,12 +15,12 @@ namespace search::features {
 class TensorAttributeExecutor : public fef::FeatureExecutor
 {
 private:
-    const search::tensor::ITensorAttribute *_attribute;
+    const search::tensor::ITensorAttribute& _attribute;
     std::unique_ptr<vespalib::eval::Tensor> _emptyTensor;
     std::unique_ptr<vespalib::eval::Tensor> _tensor;
 
 public:
-    TensorAttributeExecutor(const search::tensor::ITensorAttribute *attribute);
+    TensorAttributeExecutor(const search::tensor::ITensorAttribute& attribute);
     void execute(uint32_t docId) override;
 };
 
diff --git a/searchlib/src/vespa/searchlib/queryeval/nearest_neighbor_iterator.cpp b/searchlib/src/vespa/searchlib/queryeval/nearest_neighbor_iterator.cpp
index 07e10271c55..a64168469f0 100644
--- a/searchlib/src/vespa/searchlib/queryeval/nearest_neighbor_iterator.cpp
+++ b/searchlib/src/vespa/searchlib/queryeval/nearest_neighbor_iterator.cpp
@@ -76,7 +76,7 @@ public:
 
 private:
     double computeDistance(uint32_t docId, double limit) {
-        params().tensorAttribute.getTensor(docId, _fieldTensor);
+        params().tensorAttribute.extract_dense_view(docId, _fieldTensor);
         auto rhs = _fieldTensor.cellsRef();
         return params().distanceFunction->calc_with_limit(_lhs, rhs, limit);
     }
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
index 37a042d4e7f..d63249b816f 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
@@ -178,7 +178,7 @@ DenseTensorAttribute::getTensor(DocId docId) const
 }
 
 void
-DenseTensorAttribute::getTensor(DocId docId, MutableDenseTensorView &tensor) const
+DenseTensorAttribute::extract_dense_view(DocId docId, MutableDenseTensorView &tensor) const
 {
     EntryRef ref;
     if (docId < getCommittedDocIdLimit()) {
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
index 7fd06357114..1de74d4238d 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
@@ -37,7 +37,8 @@ public:
     std::unique_ptr<PrepareResult> prepare_set_tensor(DocId docid, const Tensor& tensor) const override;
     void complete_set_tensor(DocId docid, const Tensor& tensor, std::unique_ptr<PrepareResult> prepare_result) override;
     std::unique_ptr<Tensor> getTensor(DocId docId) const override;
-    void getTensor(DocId docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
+    void extract_dense_view(DocId docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
+    bool supports_extract_dense_view() const override { return true; }
     bool onLoad() override;
     std::unique_ptr<AttributeSaver> onInitSave(vespalib::stringref fileName) override;
     void compactWorst() override;
diff --git a/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.cpp
index 4dc6f5d6917..570be586e88 100644
--- a/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.cpp
@@ -108,12 +108,6 @@ DirectTensorAttribute::get_tensor_ref(DocId docId) const
     return *_emptyTensor;
 }
 
-void
-DirectTensorAttribute::getTensor(DocId, vespalib::tensor::MutableDenseTensorView &) const
-{
-    notImplemented();
-}
-
 std::unique_ptr<AttributeSaver>
 DirectTensorAttribute::onInitSave(vespalib::stringref fileName)
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.h
index a02e2678a48..1429a8b4d5d 100644
--- a/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/direct_tensor_attribute.h
@@ -15,13 +15,13 @@ public:
     virtual ~DirectTensorAttribute();
     virtual void setTensor(DocId docId, const Tensor &tensor) override;
     virtual std::unique_ptr<Tensor> getTensor(DocId docId) const override;
-    virtual void getTensor(DocId docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
     virtual bool onLoad() override;
     virtual std::unique_ptr<AttributeSaver> onInitSave(vespalib::stringref fileName) override;
     virtual void compactWorst() override;
 
     void set_tensor(DocId docId, std::unique_ptr<Tensor> tensor);
-    const Tensor &get_tensor_ref(DocId docId) const;
+    const Tensor &get_tensor_ref(DocId docId) const override;
+    virtual bool supports_get_tensor_ref() const override { return true; }
 };
 
 }  // namespace search::tensor
diff --git a/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h
index 520abb13d06..12b6857d122 100644
--- a/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h
@@ -24,7 +24,11 @@ public:
     virtual ~ITensorAttribute() {}
     virtual std::unique_ptr<Tensor> getTensor(uint32_t docId) const = 0;
     virtual std::unique_ptr<Tensor> getEmptyTensor() const = 0;
-    virtual void getTensor(uint32_t docId, vespalib::tensor::MutableDenseTensorView &tensor) const = 0;
+    virtual void extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const = 0;
+    virtual const Tensor& get_tensor_ref(uint32_t docid) const = 0;
+    virtual bool supports_extract_dense_view() const = 0;
+    virtual bool supports_get_tensor_ref() const = 0;
+
     virtual vespalib::eval::ValueType getTensorType() const = 0;
 
     /**
diff --git a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp
index 593330b7bd4..2e396145138 100644
--- a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp
@@ -51,9 +51,15 @@ ImportedTensorAttributeVectorReadGuard::getEmptyTensor() const
 }
 
 void
-ImportedTensorAttributeVectorReadGuard::getTensor(uint32_t docId, vespalib::tensor::MutableDenseTensorView &tensor) const
+ImportedTensorAttributeVectorReadGuard::extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const
 {
-    _target_tensor_attribute.getTensor(getTargetLid(docId), tensor);
+    _target_tensor_attribute.extract_dense_view(getTargetLid(docid), tensor);
+}
+
+const Tensor&
+ImportedTensorAttributeVectorReadGuard::get_tensor_ref(uint32_t docid) const
+{
+    return _target_tensor_attribute.get_tensor_ref(getTargetLid(docid));
 }
 
 vespalib::eval::ValueType
diff --git a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h
index 74c5cd9a685..e873da510d9 100644
--- a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h
+++ b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h
@@ -30,11 +30,14 @@ public:
 
     const ITensorAttribute *asTensorAttribute() const override;
 
-    virtual std::unique_ptr<Tensor> getTensor(uint32_t docId) const override;
-    virtual std::unique_ptr<Tensor> getEmptyTensor() const override;
-    virtual void getTensor(uint32_t docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
-    virtual vespalib::eval::ValueType getTensorType() const override;
-    virtual void get_state(const vespalib::slime::Inserter& inserter) const override;
+    std::unique_ptr<Tensor> getTensor(uint32_t docId) const override;
+    std::unique_ptr<Tensor> getEmptyTensor() const override;
+    void extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const override;
+    const Tensor& get_tensor_ref(uint32_t docid) const override;
+    bool supports_extract_dense_view() const override { return _target_tensor_attribute.supports_extract_dense_view(); }
+    bool supports_get_tensor_ref() const override { return _target_tensor_attribute.supports_get_tensor_ref(); }
+    vespalib::eval::ValueType getTensorType() const override;
+    void get_state(const vespalib::slime::Inserter& inserter) const override;
 };
 
 }
diff --git a/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.cpp
index d4a20abf2fd..3c0e9b38687 100644
--- a/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.cpp
@@ -55,12 +55,6 @@ SerializedTensorAttribute::getTensor(DocId docId) const
     return _serializedTensorStore.getTensor(ref);
 }
 
-void
-SerializedTensorAttribute::getTensor(DocId, vespalib::tensor::MutableDenseTensorView &) const
-{
-    notImplemented();
-}
-
 bool
 SerializedTensorAttribute::onLoad()
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.h
index 5596341a5b7..7b385359337 100644
--- a/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/serialized_tensor_attribute.h
@@ -17,7 +17,6 @@ public:
     virtual ~SerializedTensorAttribute();
     virtual void setTensor(DocId docId, const Tensor &tensor) override;
     virtual std::unique_ptr<Tensor> getTensor(DocId docId) const override;
-    virtual void getTensor(DocId docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
     virtual bool onLoad() override;
     virtual std::unique_ptr<AttributeSaver> onInitSave(vespalib::stringref fileName) override;
     virtual void compactWorst() override;
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
index 6cf4f6d2689..7abc7bd22b5 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
@@ -202,6 +202,22 @@ TensorAttribute::getEmptyTensor() const
     return _emptyTensor->clone();
 }
 
+void
+TensorAttribute::extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const
+{
+    (void) docid;
+    (void) tensor;
+    notImplemented();
+}
+
+const Tensor&
+TensorAttribute::get_tensor_ref(uint32_t docid) const
+{
+    (void) docid;
+    notImplemented();
+    abort(); // Needed to avoid compile error
+}
+
 vespalib::eval::ValueType
 TensorAttribute::getTensorType() const
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
index 8380e485172..9627ea5b4aa 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
@@ -45,6 +45,10 @@ public:
     void onGenerationChange(generation_t generation) override;
     bool addDoc(DocId &docId) override;
     std::unique_ptr<Tensor> getEmptyTensor() const override;
+    void extract_dense_view(uint32_t docid, vespalib::tensor::MutableDenseTensorView& tensor) const override;
+    const Tensor& get_tensor_ref(uint32_t docid) const override;
+    bool supports_extract_dense_view() const override { return false; }
+    bool supports_get_tensor_ref() const override { return false; }
     vespalib::eval::ValueType getTensorType() const override;
     void get_state(const vespalib::slime::Inserter& inserter) const override;
     void clearDocs(DocId lidLow, DocId lidLimit) override;