Go back to using a single combiner.

3 files changed, 58 insertions, 98 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 1de7253e186..5ad4181bbf5 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
@@ -6,11 +6,15 @@
 
 namespace vespalib::tensor {
 
-DenseTensorAddressCombiner::~DenseTensorAddressCombiner() { }
+DenseTensorAddressCombiner::~DenseTensorAddressCombiner() = default;
 
-DenseTensorAddressCombiner::DenseTensorAddressCombiner(const eval::ValueType &lhs,
-                                                       const eval::ValueType &rhs)
-    : _combinedAddress(),
+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()),
       _left(),
       _commonRight(),
       _right()
@@ -37,11 +41,15 @@ DenseTensorAddressCombiner::DenseTensorAddressCombiner(const eval::ValueType &lh
         ++rhsItr;
     }
     _combinedAddress.resize(numDimensions);
+    size_t multiplier = 1;
+    for (int32_t i(_rightAddress.size() - 1); i >= 0; i--) {
+        _accumulatedSize[i] = multiplier;
+        multiplier *= _rightType.dimensions()[i].size;
+    }
 }
 
 eval::ValueType
-DenseTensorAddressCombiner::combineDimensions(const eval::ValueType &lhs,
-                                              const eval::ValueType &rhs)
+DenseTensorAddressCombiner::combineDimensions(const eval::ValueType &lhs, const eval::ValueType &rhs)
 {
     // NOTE: both lhs and rhs are sorted according to dimension names.
     std::vector<eval::ValueType::Dimension> result;
@@ -50,8 +58,7 @@ DenseTensorAddressCombiner::combineDimensions(const eval::ValueType &lhs,
     while (lhsItr != lhs.dimensions().end() &&
            rhsItr != rhs.dimensions().end()) {
         if (lhsItr->name == rhsItr->name) {
-            result.emplace_back(lhsItr->name,
-                                std::min(lhsItr->size, rhsItr->size));
+            result.emplace_back(lhsItr->name, std::min(lhsItr->size, rhsItr->size));
             ++lhsItr;
             ++rhsItr;
         } else if (lhsItr->name < rhsItr->name) {
@@ -71,24 +78,4 @@ DenseTensorAddressCombiner::combineDimensions(const eval::ValueType &lhs,
             eval::ValueType::tensor_type(std::move(result)));
 }
 
-
-CommonDenseTensorCellsIterator::CommonDenseTensorCellsIterator(const Mapping & common,
-                                                               const Mapping & right,
-                                                               const eval::ValueType &type_in,
-                                                               CellsRef cells)
-    : _type(type_in),
-      _cells(cells),
-      _address(type_in.dimensions().size(), 0),
-      _common(common),
-      _right(right),
-      _accumulatedSize(_address.size())
-{
-    size_t multiplier = 1;
-    for (int32_t i(_address.size() - 1); i >= 0; i--) {
-        _accumulatedSize[i] = multiplier;
-        multiplier *= type_in.dimensions()[i].size;
-    }
-}
-CommonDenseTensorCellsIterator::~CommonDenseTensorCellsIterator() = default;
-
 }
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 2650d1df47d..fdd8da467cc 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
@@ -22,8 +22,14 @@ public:
 
 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;
     Mapping                _left;
     Mapping                _commonRight;
     Mapping                _right;
@@ -32,97 +38,67 @@ private:
             _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);
+    DenseTensorAddressCombiner(const eval::ValueType &lhs, const eval::ValueType &rhs, CellsRef rhsCells);
     ~DenseTensorAddressCombiner();
     void updateLeftAndCommon(const Address & addr) { update(addr, _left); }
-    const Mapping & getCommonRight() const { return _commonRight; }
-    const Mapping & getRight() const { return _right; }
 
     bool hasAnyRightOnlyDimensions() const { return ! _right.empty(); }
 
     const Address &address() const { return _combinedAddress; }
-    Address &address() { return _combinedAddress; }
-
-    static eval::ValueType combineDimensions(const eval::ValueType &lhs, const eval::ValueType &rhs);
-};
-
 
-/**
- * Utility class to iterate over common cells in a dense tensor.
- */
-class CommonDenseTensorCellsIterator
-{
-public:
-    using size_type = eval::ValueType::Dimension::size_type;
-    using Address = std::vector<size_type>;
-    using Mapping = DenseTensorAddressCombiner::Mapping;
-private:
-    using Dims = std::vector<uint32_t>;
-    using CellsRef = vespalib::ConstArrayRef<double>;
-    const eval::ValueType &_type;
-    CellsRef               _cells;
-    Address                _address;
-    const Mapping         &_common;
-    const Mapping         &_right;
-    std::vector<size_t>    _accumulatedSize;
-
-    double cell(size_t cellIdx) const { return _cells[cellIdx]; }
-    size_t index(const Address &address) const {
-        size_t cellIdx(0);
-        for (uint32_t i(0); i < address.size(); i++) {
-            cellIdx += address[i]*_accumulatedSize[i];
+    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 cellIdx;
+        return true;
     }
-public:
-    CommonDenseTensorCellsIterator(const Mapping & common, const Mapping & right,
-                                   const eval::ValueType &type_in, CellsRef cells);
-    ~CommonDenseTensorCellsIterator();
+    double rightCell() { return rightCell(rightIndex(_rightAddress)); }
+
     template <typename Func>
-    void for_each(Address & combined, Func && func) const {
+    void for_each(Func && func) {
         const int32_t lastDimension = _right.size() - 1;
         int32_t curDimension = lastDimension;
-        size_t cellIdx = index(_address);
+        size_t cellIdx = rightIndex(_rightAddress);
         while (curDimension >= 0) {
             const uint32_t rdim = _right[curDimension].second;
             const uint32_t cdim = _right[curDimension].first;
-            size_type & cindex = combined[cdim];
+            size_type & cindex = _combinedAddress[cdim];
             if (curDimension == lastDimension) {
-                for (cindex = 0; cindex < _type.dimensions()[rdim].size; cindex++) {
-                    func(combined, cell(cellIdx));
+                for (cindex = 0; cindex < _rightType.dimensions()[rdim].size; cindex++) {
+                    func(_combinedAddress, rightCell(cellIdx));
                     cellIdx += _accumulatedSize[rdim];
                 }
                 cindex = 0;
-                cellIdx -= _accumulatedSize[rdim] * _type.dimensions()[rdim].size;
+                cellIdx -= _accumulatedSize[rdim] * _rightType.dimensions()[rdim].size;
                 curDimension--;
             } else {
-                if (cindex < _type.dimensions()[rdim].size) {
+                if (cindex < _rightType.dimensions()[rdim].size) {
                     cindex++;
                     cellIdx += _accumulatedSize[rdim];
                     curDimension++;
                 } else {
-                    cellIdx -= _accumulatedSize[rdim] * _type.dimensions()[rdim].size;
+                    cellIdx -= _accumulatedSize[rdim] * _rightType.dimensions()[rdim].size;
                     cindex = 0;
                     curDimension--;
                 }
             }
         }
     }
-    bool updateCommon(const Address & combined) {
-        for (const auto & m : _common) {
-            if (combined[m.first] >= _type.dimensions()[m.second].size) {
-                return false;
-            }
-            _address[m.second] = combined[m.first];
-        }
-        return true;
-    }
-    double cell() const {
-        return cell(index(_address));
-    }
 
-    const eval::ValueType &fast_type() const { return _type; }
+    static eval::ValueType combineDimensions(const eval::ValueType &lhs, const eval::ValueType &rhs);
 };
 
+
 }
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 0603e1e80ee..80720bd10e1 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,
-      CommonDenseTensorCellsIterator & rhsIter, const DenseTensorView &lhs, Function &&func) __attribute__((noinline));
+      const DenseTensorView &lhs, Function &&func) __attribute__((noinline));
 
 template <typename Function>
 std::unique_ptr<Tensor>
 apply(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
-      CommonDenseTensorCellsIterator & rhsIter, const DenseTensorView &lhs, Function &&func)
+      const DenseTensorView &lhs, Function &&func)
 {
     for (DenseTensorCellsIterator lhsItr = lhs.cellsIterator(); lhsItr.valid(); lhsItr.next()) {
         combiner.updateLeftAndCommon(lhsItr.address());
-        if (rhsIter.updateCommon(combiner.address())) {
-            rhsIter.for_each(combiner.address(), [&func, &builder, &lhsItr](const DenseTensorCellsIterator::Address & combined, double rhsCell) {
+        if (combiner.updateCommonRight()) {
+            combiner.for_each([&func, &builder, &lhsItr](const DenseTensorCellsIterator::Address & combined, double rhsCell) {
                 builder.insertCell(combined, func(lhsItr.cell(), rhsCell));
             });
         }
@@ -33,19 +33,17 @@ apply(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
 template <typename Function>
 std::unique_ptr<Tensor>
 apply_no_rightonly_dimensions(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
-                              CommonDenseTensorCellsIterator & rhsIter,
                               const DenseTensorView &lhs, Function &&func)  __attribute__((noinline));
 
 template <typename Function>
 std::unique_ptr<Tensor>
 apply_no_rightonly_dimensions(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
-                              CommonDenseTensorCellsIterator & rhsIter,
                               const DenseTensorView &lhs, Function &&func)
 {
     for (DenseTensorCellsIterator lhsItr = lhs.cellsIterator(); lhsItr.valid(); lhsItr.next()) {
         combiner.updateLeftAndCommon(lhsItr.address());
-        if (rhsIter.updateCommon(combiner.address())) {
-            builder.insertCell(combiner.address(), func(lhsItr.cell(), rhsIter.cell()));
+        if (combiner.updateCommonRight()) {
+            builder.insertCell(combiner.address(), func(lhsItr.cell(), combiner.rightCell()));
         }
     }
     return builder.build();
@@ -55,13 +53,12 @@ template <typename Function>
 std::unique_ptr<Tensor>
 apply(const DenseTensorView &lhs, const DenseTensorView &rhs, Function &&func)
 {
-    DenseTensorAddressCombiner combiner(lhs.fast_type(), rhs.fast_type());
+    DenseTensorAddressCombiner combiner(lhs.fast_type(), rhs.fast_type(), rhs.cellsRef());
     DirectDenseTensorBuilder builder(DenseTensorAddressCombiner::combineDimensions(lhs.fast_type(), rhs.fast_type()));
-    CommonDenseTensorCellsIterator rhsIter(combiner.getCommonRight(), combiner.getRight(), rhs.fast_type(), rhs.cellsRef());
     if (combiner.hasAnyRightOnlyDimensions()) {
-        return apply(combiner, builder, rhsIter, lhs, std::move(func));
+        return apply(combiner, builder, lhs, std::move(func));
     } else {
-        return apply_no_rightonly_dimensions(combiner, builder, rhsIter, lhs, std::move(func));
+        return apply_no_rightonly_dimensions(combiner, builder, lhs, std::move(func));
     }
 }