Reduce alignment for small dense tensors.

3 files changed, 24 insertions, 9 deletions

diff --git a/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp b/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp
index 032960c3799..149662cd266 100644
--- a/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp
+++ b/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp
@@ -75,10 +75,19 @@ assertArraySize(const vespalib::string &tensorType, uint32_t expArraySize) {
 
 TEST("require that array size is calculated correctly")
 {
-    TEST_DO(assertArraySize("tensor(x[1])", 32));
+    TEST_DO(assertArraySize("tensor(x[1])", 8));
     TEST_DO(assertArraySize("tensor(x[10])", 96));
     TEST_DO(assertArraySize("tensor(x[3])", 32));
     TEST_DO(assertArraySize("tensor(x[10],y[10])", 800));
+    TEST_DO(assertArraySize("tensor<int8>(x[1])", 8));
+    TEST_DO(assertArraySize("tensor<int8>(x[8])", 8));
+    TEST_DO(assertArraySize("tensor<int8>(x[9])", 16));
+    TEST_DO(assertArraySize("tensor<int8>(x[16])", 16));
+    TEST_DO(assertArraySize("tensor<int8>(x[17])", 32));
+    TEST_DO(assertArraySize("tensor<int8>(x[32])", 32));
+    TEST_DO(assertArraySize("tensor<int8>(x[33])", 64));
+    TEST_DO(assertArraySize("tensor<int8>(x[64])", 64));
+    TEST_DO(assertArraySize("tensor<int8>(x[65])", 96));
 }
 
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
index 86090f2ac92..ed3fb737b7d 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
@@ -17,6 +17,8 @@ namespace {
 
 constexpr size_t MIN_BUFFER_ARRAYS = 1024;
 constexpr size_t DENSE_TENSOR_ALIGNMENT = 32;
+constexpr size_t DENSE_TENSOR_ALIGNMENT_SMALL = 16;
+constexpr size_t DENSE_TENSOR_ALIGNMENT_MIN = 8;
 
 size_t my_align(size_t size, size_t alignment) {
     size += alignment - 1;
@@ -27,17 +29,20 @@ size_t my_align(size_t size, size_t alignment) {
 
 DenseTensorStore::TensorSizeCalc::TensorSizeCalc(const ValueType &type)
     : _numCells(1u),
-      _cell_type(type.cell_type())
+      _cell_type(type.cell_type()),
+      _aligned_size(0u)
 {
     for (const auto &dim: type.dimensions()) {
         _numCells *= dim.size;
     }
-}
-
-size_t
-DenseTensorStore::TensorSizeCalc::alignedSize() const
-{
-    return my_align(bufSize(), DENSE_TENSOR_ALIGNMENT);
+    auto buf_size = bufSize();
+    size_t alignment = DENSE_TENSOR_ALIGNMENT;
+    if (buf_size <= DENSE_TENSOR_ALIGNMENT_MIN) {
+        alignment = DENSE_TENSOR_ALIGNMENT_MIN;
+    } else if (buf_size <= DENSE_TENSOR_ALIGNMENT_SMALL) {
+        alignment = DENSE_TENSOR_ALIGNMENT_SMALL;
+    }
+    _aligned_size = my_align(buf_size, alignment);
 }
 
 DenseTensorStore::BufferType::BufferType(const TensorSizeCalc &tensorSizeCalc, std::unique_ptr<vespalib::alloc::MemoryAllocator> allocator)
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
index 06492596f70..47932fbff7e 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
@@ -25,12 +25,13 @@ public:
     {
         size_t   _numCells; // product of dimension sizes
         vespalib::eval::CellType _cell_type;
+        size_t   _aligned_size;
 
         TensorSizeCalc(const ValueType &type);
         size_t bufSize() const {
             return vespalib::eval::CellTypeUtils::mem_size(_cell_type, _numCells);
         }
-        size_t alignedSize() const;
+        size_t alignedSize() const noexcept { return _aligned_size; }
     };
 
     class BufferType : public vespalib::datastore::BufferType<char>