Merge pull request #4595 from vespa-engine/balder/track-common-cell-index

Balder/track common cell index

3 files changed, 93 insertions, 66 deletions

diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_address_combiner.cpp b/eval/src/vespa/eval/tensor/dense/dense_tensor_address_combiner.cpp
index 5ad4181bbf5..df6ac162d7f 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_address_combiner.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_address_combiner.cpp
@@ -8,13 +8,10 @@ namespace vespalib::tensor {
 
 DenseTensorAddressCombiner::~DenseTensorAddressCombiner() = default;
 
-DenseTensorAddressCombiner::DenseTensorAddressCombiner(const eval::ValueType &lhs, const eval::ValueType &rhs,
-                                                       CellsRef rhsCells)
-    : _rightType(rhs),
-      _combinedAddress(),
-      _rightCells(rhsCells),
-      _rightAddress(rhs.dimensions().size(), 0),
-      _accumulatedSize(_rightAddress.size()),
+DenseTensorAddressCombiner::DenseTensorAddressCombiner(const eval::ValueType &combined, const eval::ValueType &lhs,
+                                                       const eval::ValueType &rhs)
+    : _rightAddress(rhs),
+      _combinedAddress(combined),
       _left(),
       _commonRight(),
       _right()
@@ -40,14 +37,23 @@ DenseTensorAddressCombiner::DenseTensorAddressCombiner(const eval::ValueType &lh
         _right.emplace_back(numDimensions++, rhsItr-rhs.dimensions().cbegin());
         ++rhsItr;
     }
-    _combinedAddress.resize(numDimensions);
+}
+
+AddressContext::AddressContext(const eval::ValueType &type)
+    : _type(type),
+      _accumulatedSize(_type.dimensions().size()),
+      _address(type.dimensions().size(), 0)
+
+{
     size_t multiplier = 1;
-    for (int32_t i(_rightAddress.size() - 1); i >= 0; i--) {
+    for (int32_t i(_address.size() - 1); i >= 0; i--) {
         _accumulatedSize[i] = multiplier;
-        multiplier *= _rightType.dimensions()[i].size;
+        multiplier *= type.dimensions()[i].size;
     }
 }
 
+AddressContext::~AddressContext() = default;
+
 eval::ValueType
 DenseTensorAddressCombiner::combineDimensions(const eval::ValueType &lhs, const eval::ValueType &rhs)
 {
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_address_combiner.h b/eval/src/vespa/eval/tensor/dense/dense_tensor_address_combiner.h
index fdd8da467cc..325242bbd9b 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_address_combiner.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_address_combiner.h
@@ -9,6 +9,42 @@
 
 namespace vespalib::tensor {
 
+class AddressContext {
+public:
+    using Address = DenseTensorCellsIterator::Address;
+    using size_type = eval::ValueType::Dimension::size_type;
+    using Mapping = std::vector<std::pair<uint32_t, uint32_t>>;
+    AddressContext(const eval::ValueType &type);
+    ~AddressContext();
+    size_type dimSize(uint32_t dim) const { return _type.dimensions()[dim].size; }
+    size_type wholeDimStep(uint32_t dim) const { return _accumulatedSize[dim] * dimSize(dim); }
+    size_t index() const {
+        size_t cellIdx(0);
+        for (uint32_t i(0); i < _address.size(); i++) {
+            cellIdx += _address[i]*_accumulatedSize[i];
+        }
+        return cellIdx;
+    }
+    void update(const Address & addr, const Mapping & mapping) {
+        for (const auto & m : mapping) {
+            _address[m.first] = addr[m.second];
+        }
+    }
+    bool updateCommon(const Address & addr, const Mapping & mapping) {
+        for (const auto & m : mapping) {
+            if (addr[m.first] >= dimSize(m.second)) {
+                return false;
+            }
+            _address[m.second] = addr[m.first];
+        }
+        return true;
+    }
+
+    const eval::ValueType &_type;
+    std::vector<size_t>    _accumulatedSize;
+    Address                _address;
+
+};
 
 /**
  * Combines two dense tensor addresses to a new tensor address.
@@ -18,78 +54,61 @@ namespace vespalib::tensor {
 class DenseTensorAddressCombiner
 {
 public:
-    using Mapping = std::vector<std::pair<uint32_t, uint32_t>>;
+    using Mapping = AddressContext::Mapping;
 
 private:
     using Address = DenseTensorCellsIterator::Address;
     using CellsRef = vespalib::ConstArrayRef<double>;
     using size_type = eval::ValueType::Dimension::size_type;
 
-    const eval::ValueType &_rightType;
-    Address                _combinedAddress;
-    CellsRef               _rightCells;
-    Address                _rightAddress;
-    std::vector<size_t>    _accumulatedSize;
+    AddressContext         _rightAddress;
+    AddressContext         _combinedAddress;
+
     Mapping                _left;
     Mapping                _commonRight;
     Mapping                _right;
-    void update(const Address & addr, const Mapping & mapping) {
-        for (const auto & m : mapping) {
-            _combinedAddress[m.first] = addr[m.second];
-        }
-    }
-    double rightCell(size_t cellIdx) const { return _rightCells[cellIdx]; }
-    size_t rightIndex(const Address &address) const {
-        size_t cellIdx(0);
-        for (uint32_t i(0); i < address.size(); i++) {
-            cellIdx += address[i]*_accumulatedSize[i];
-        }
-        return cellIdx;
-    }
+
 public:
-    DenseTensorAddressCombiner(const eval::ValueType &lhs, const eval::ValueType &rhs, CellsRef rhsCells);
+    DenseTensorAddressCombiner(const eval::ValueType &combined, const eval::ValueType &lhs, const eval::ValueType &rhs);
     ~DenseTensorAddressCombiner();
-    void updateLeftAndCommon(const Address & addr) { update(addr, _left); }
-
+    void updateLeftAndCommon(const Address & addr) { _combinedAddress.update(addr, _left); }
+    bool updateCommon() { return _rightAddress.updateCommon(_combinedAddress._address, _commonRight); }
     bool hasAnyRightOnlyDimensions() const { return ! _right.empty(); }
 
-    const Address &address() const { return _combinedAddress; }
-
-    bool updateCommonRight() {
-        for (const auto & m : _commonRight) {
-            if (_combinedAddress[m.first] >= _rightType.dimensions()[m.second].size) {
-                return false;
-            }
-            _rightAddress[m.second] = _combinedAddress[m.first];
-        }
-        return true;
-    }
-    double rightCell() { return rightCell(rightIndex(_rightAddress)); }
+    const Address & address() const { return _combinedAddress._address; }
+    size_t rightCellIndex() const { return _rightAddress.index(); }
 
     template <typename Func>
-    void for_each(Func && func) {
+    void for_each_right(const CellsRef & rhsCells, Func && func) {
+        // The rightAddress oly holds the starting point for iteration and what is need to efficiently maintain
+        // an index for addressing th ecells.
         const int32_t lastDimension = _right.size() - 1;
         int32_t curDimension = lastDimension;
-        size_t cellIdx = rightIndex(_rightAddress);
+        size_t rightCellIdx = _rightAddress.index();
+        size_t combinedCellIdx = _combinedAddress.index();
         while (curDimension >= 0) {
             const uint32_t rdim = _right[curDimension].second;
             const uint32_t cdim = _right[curDimension].first;
-            size_type & cindex = _combinedAddress[cdim];
+            size_type & cindex = _combinedAddress._address[cdim];
             if (curDimension == lastDimension) {
-                for (cindex = 0; cindex < _rightType.dimensions()[rdim].size; cindex++) {
-                    func(_combinedAddress, rightCell(cellIdx));
-                    cellIdx += _accumulatedSize[rdim];
+                for (cindex = 0; cindex < _rightAddress.dimSize(rdim); cindex++) {
+                    func(combinedCellIdx, rhsCells[rightCellIdx]);
+                    rightCellIdx += _rightAddress._accumulatedSize[rdim];
+                    combinedCellIdx += _combinedAddress._accumulatedSize[cdim];
                 }
                 cindex = 0;
-                cellIdx -= _accumulatedSize[rdim] * _rightType.dimensions()[rdim].size;
+                rightCellIdx -= _rightAddress.wholeDimStep(rdim);
+                combinedCellIdx -= _combinedAddress.wholeDimStep(cdim);
                 curDimension--;
             } else {
-                if (cindex < _rightType.dimensions()[rdim].size) {
+                if (cindex < _rightAddress.dimSize(rdim)) {
                     cindex++;
-                    cellIdx += _accumulatedSize[rdim];
+                    rightCellIdx += _rightAddress._accumulatedSize[rdim];
+                    combinedCellIdx += _combinedAddress._accumulatedSize[cdim];
                     curDimension++;
                 } else {
-                    cellIdx -= _accumulatedSize[rdim] * _rightType.dimensions()[rdim].size;
+                    rightCellIdx -= _rightAddress.wholeDimStep(rdim);
+                    combinedCellIdx -= _combinedAddress.wholeDimStep(cdim);
                     cindex = 0;
                     curDimension--;
                 }
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_apply.hpp b/eval/src/vespa/eval/tensor/dense/dense_tensor_apply.hpp
index 80720bd10e1..25478510587 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_apply.hpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_apply.hpp
@@ -11,17 +11,17 @@ namespace vespalib::tensor::dense {
 template <typename Function>
 std::unique_ptr<Tensor>
 apply(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
-      const DenseTensorView &lhs, Function &&func) __attribute__((noinline));
+      const DenseTensorView &lhs, const DenseTensorView::CellsRef & rhsCells, Function &&func) __attribute__((noinline));
 
 template <typename Function>
 std::unique_ptr<Tensor>
 apply(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
-      const DenseTensorView &lhs, Function &&func)
+      const DenseTensorView &lhs, const DenseTensorView::CellsRef & rhsCells, Function &&func)
 {
     for (DenseTensorCellsIterator lhsItr = lhs.cellsIterator(); lhsItr.valid(); lhsItr.next()) {
         combiner.updateLeftAndCommon(lhsItr.address());
-        if (combiner.updateCommonRight()) {
-            combiner.for_each([&func, &builder, &lhsItr](const DenseTensorCellsIterator::Address & combined, double rhsCell) {
+        if (combiner.updateCommon()) {
+            combiner.for_each_right(rhsCells, [&func, &builder, &lhsItr](size_t combined, double rhsCell) {
                 builder.insertCell(combined, func(lhsItr.cell(), rhsCell));
             });
         }
@@ -33,17 +33,18 @@ apply(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
 template <typename Function>
 std::unique_ptr<Tensor>
 apply_no_rightonly_dimensions(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
-                              const DenseTensorView &lhs, Function &&func)  __attribute__((noinline));
+                              const DenseTensorView &lhs, const DenseTensorView::CellsRef & rhsCells,
+                              Function &&func)  __attribute__((noinline));
 
 template <typename Function>
 std::unique_ptr<Tensor>
 apply_no_rightonly_dimensions(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
-                              const DenseTensorView &lhs, Function &&func)
+                              const DenseTensorView &lhs, const DenseTensorView::CellsRef & rhsCells, Function &&func)
 {
     for (DenseTensorCellsIterator lhsItr = lhs.cellsIterator(); lhsItr.valid(); lhsItr.next()) {
         combiner.updateLeftAndCommon(lhsItr.address());
-        if (combiner.updateCommonRight()) {
-            builder.insertCell(combiner.address(), func(lhsItr.cell(), combiner.rightCell()));
+        if (combiner.updateCommon()) {
+            builder.insertCell(combiner.address(), func(lhsItr.cell(), rhsCells[combiner.rightCellIndex()]));
         }
     }
     return builder.build();
@@ -53,12 +54,13 @@ template <typename Function>
 std::unique_ptr<Tensor>
 apply(const DenseTensorView &lhs, const DenseTensorView &rhs, Function &&func)
 {
-    DenseTensorAddressCombiner combiner(lhs.fast_type(), rhs.fast_type(), rhs.cellsRef());
-    DirectDenseTensorBuilder builder(DenseTensorAddressCombiner::combineDimensions(lhs.fast_type(), rhs.fast_type()));
+    eval::ValueType resultType = DenseTensorAddressCombiner::combineDimensions(lhs.fast_type(), rhs.fast_type());
+    DenseTensorAddressCombiner combiner(resultType, lhs.fast_type(), rhs.fast_type());
+    DirectDenseTensorBuilder builder(resultType);
     if (combiner.hasAnyRightOnlyDimensions()) {
-        return apply(combiner, builder, lhs, std::move(func));
+        return apply(combiner, builder, lhs, rhs.cellsRef(), std::move(func));
     } else {
-        return apply_no_rightonly_dimensions(combiner, builder, lhs, std::move(func));
+        return apply_no_rightonly_dimensions(combiner, builder, lhs, rhs.cellsRef(), std::move(func));
     }
 }