Prepare for dynamic cluster size in dense tensor store.

4 files changed, 22 insertions, 9 deletions

diff --git a/searchlib/src/vespa/searchlib/datastore/entryref.h b/searchlib/src/vespa/searchlib/datastore/entryref.h
index 457ffac4e26..242522c2bc4 100644
--- a/searchlib/src/vespa/searchlib/datastore/entryref.h
+++ b/searchlib/src/vespa/searchlib/datastore/entryref.h
@@ -36,6 +36,7 @@ public:
     static uint32_t numBuffers() { return 1 << BufferBits; } 
     static uint64_t align(uint64_t val) { return val; }
     static uint64_t pad(uint64_t val) { (void) val; return 0ul; }
+    static constexpr bool isAlignedEntryRefType = false;
 };
 
 /**
@@ -56,6 +57,7 @@ public:
     static uint64_t offsetSize() { return ParentType::offsetSize() << OffsetAlign; }
     static uint64_t align(uint64_t val) { return val + pad(val); }
     static uint64_t pad(uint64_t val) { return (-val & PadConstant); }
+    static constexpr bool isAlignedEntryRefType = true;
 };
 
 }
diff --git a/searchlib/src/vespa/searchlib/datastore/raw_allocator.hpp b/searchlib/src/vespa/searchlib/datastore/raw_allocator.hpp
index 1c72d793ec6..69b3cb300b9 100644
--- a/searchlib/src/vespa/searchlib/datastore/raw_allocator.hpp
+++ b/searchlib/src/vespa/searchlib/datastore/raw_allocator.hpp
@@ -25,7 +25,15 @@ RawAllocator<EntryT, RefT>::alloc(size_t numElems, size_t extraElems)
     size_t oldBufferSize = state.size();
     EntryT *buffer = _store.getBufferEntry<EntryT>(activeBufferId, oldBufferSize);
     state.pushed_back(numElems);
-    return HandleType(RefT(oldBufferSize, activeBufferId), buffer);
+    if (RefT::isAlignedEntryRefType) {
+        // AlignedEntryRef constructor scales down offset by alignment
+        return HandleType(RefT(oldBufferSize, activeBufferId), buffer);
+    } else {
+        // Must perform scaling ourselves, according to cluster size
+        size_t clusterSize = state.getClusterSize();
+        assert((numElems % clusterSize) == 0u);
+        return HandleType(RefT(oldBufferSize / clusterSize, activeBufferId), buffer);
+    }
 }
 
 }
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
index 1c4e07e38ee..3cb31534bb5 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
@@ -18,11 +18,10 @@ using vespalib::eval::ValueType;
 namespace search::tensor {
 
 constexpr size_t MIN_BUFFER_CLUSTERS = 1024;
+constexpr size_t DENSE_TENSOR_ALIGNMENT = 32;
 
 DenseTensorStore::BufferType::BufferType()
-    : datastore::BufferType<char>(RefType::align(1),
-                              MIN_BUFFER_CLUSTERS,
-                              RefType::offsetSize() / RefType::align(1)),
+    : datastore::BufferType<char>(DENSE_TENSOR_ALIGNMENT, MIN_BUFFER_CLUSTERS, RefType::offsetSize()),
       _unboundDimSizesSize(0u)
 {}
 
@@ -40,7 +39,7 @@ size_t
 DenseTensorStore::BufferType::getReservedElements(uint32_t bufferId) const
 {
     return datastore::BufferType<char>::getReservedElements(bufferId) +
-        RefType::align(_unboundDimSizesSize);
+        align(_unboundDimSizesSize);
 }
 
 DenseTensorStore::DenseTensorStore(const ValueType &type)
@@ -75,7 +74,7 @@ const void *
 DenseTensorStore::getRawBuffer(RefType ref) const
 {
     return _store.getBufferEntry<char>(ref.bufferId(),
-                                         ref.offset());
+                                       ref.offset() * _bufferType.getClusterSize());
 }
 
 
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
index 67a2dc7b8c0..9810ff43a19 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
@@ -26,8 +26,7 @@ namespace search::tensor {
 class DenseTensorStore : public TensorStore
 {
 public:
-    // 32 entry alignment, entry type is char => 32 bytes alignment
-    using RefType = datastore::AlignedEntryRefT<22, 5>;
+    using RefType = datastore::EntryRefT<22>;
     using DataStoreType = datastore::DataStoreT<RefType>;
     using ValueType = vespalib::eval::ValueType;
 
@@ -44,6 +43,11 @@ public:
             _unboundDimSizesSize = unboundDimSizesSize_in;
         }
         size_t getReservedElements(uint32_t bufferId) const override;
+        static size_t align(size_t size, size_t alignment) {
+            size += alignment - 1;
+            return (size - (size % alignment));
+        }
+        size_t align(size_t size) const { return align(size, _clusterSize); }
     };
 private:
     DataStoreType _concreteStore;
@@ -61,7 +65,7 @@ private:
     setDenseTensor(const TensorType &tensor);
     datastore::Handle<char> allocRawBuffer(size_t numCells);
     size_t alignedSize(size_t numCells) const {
-        return RefType::align(numCells * _cellSize + unboundDimSizesSize());
+        return _bufferType.align(numCells * _cellSize + unboundDimSizesSize());
     }
 
 public: