Put the address back in the combiner.

3 files changed, 24 insertions, 25 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 d6d0a1abf4a..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
@@ -10,7 +10,8 @@ DenseTensorAddressCombiner::~DenseTensorAddressCombiner() = default;
 
 DenseTensorAddressCombiner::DenseTensorAddressCombiner(const eval::ValueType &combined, const eval::ValueType &lhs,
                                                        const eval::ValueType &rhs)
-    : _combinedAddress(combined),
+    : _rightAddress(rhs),
+      _combinedAddress(combined),
       _left(),
       _commonRight(),
       _right()
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 af7c545b3fd..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
@@ -61,6 +61,7 @@ private:
     using CellsRef = vespalib::ConstArrayRef<double>;
     using size_type = eval::ValueType::Dimension::size_type;
 
+    AddressContext         _rightAddress;
     AddressContext         _combinedAddress;
 
     Mapping                _left;
@@ -71,42 +72,42 @@ public:
     DenseTensorAddressCombiner(const eval::ValueType &combined, const eval::ValueType &lhs, const eval::ValueType &rhs);
     ~DenseTensorAddressCombiner();
     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._address; }
-    const Mapping & commonRight() const { return _commonRight; }
+    size_t rightCellIndex() const { return _rightAddress.index(); }
 
     template <typename Func>
-    void for_each(const AddressContext & rightAddress, const CellsRef & rhsCells, 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 rightCellIdx = rightAddress.index();
+        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._address[cdim];
             if (curDimension == lastDimension) {
-                for (cindex = 0; cindex < rightAddress.dimSize(rdim); cindex++) {
+                for (cindex = 0; cindex < _rightAddress.dimSize(rdim); cindex++) {
                     func(combinedCellIdx, rhsCells[rightCellIdx]);
-                    rightCellIdx += rightAddress._accumulatedSize[rdim];
+                    rightCellIdx += _rightAddress._accumulatedSize[rdim];
                     combinedCellIdx += _combinedAddress._accumulatedSize[cdim];
                 }
                 cindex = 0;
-                rightCellIdx -= rightAddress.wholeDimStep(rdim);
+                rightCellIdx -= _rightAddress.wholeDimStep(rdim);
                 combinedCellIdx -= _combinedAddress.wholeDimStep(cdim);
                 curDimension--;
             } else {
-                if (cindex < rightAddress.dimSize(rdim)) {
+                if (cindex < _rightAddress.dimSize(rdim)) {
                     cindex++;
-                    rightCellIdx += rightAddress._accumulatedSize[rdim];
+                    rightCellIdx += _rightAddress._accumulatedSize[rdim];
                     combinedCellIdx += _combinedAddress._accumulatedSize[cdim];
                     curDimension++;
                 } else {
-                    rightCellIdx -= rightAddress.wholeDimStep(rdim);
+                    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 f61dbd05def..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,18 +11,17 @@ namespace vespalib::tensor::dense {
 template <typename Function>
 std::unique_ptr<Tensor>
 apply(DenseTensorAddressCombiner & combiner, DirectDenseTensorBuilder & builder,
-      const DenseTensorView &lhs, AddressContext & rhsAddr, const DenseTensorView::CellsRef & rhsCells,
-      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, AddressContext & rhsAddr, const DenseTensorView::CellsRef & rhsCells, 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 (rhsAddr.updateCommon(combiner.address(), combiner.commonRight())) {
-            combiner.for_each(rhsAddr, rhsCells, [&func, &builder, &lhsItr](size_t 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));
             });
         }
@@ -34,19 +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, AddressContext & rhsAddr,
-                              const DenseTensorView::CellsRef & rhsCells, 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, AddressContext & rhsAddr,
-                              const DenseTensorView::CellsRef & rhsCells, 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 (rhsAddr.updateCommon(combiner.address(), combiner.commonRight())) {
-            builder.insertCell(combiner.address(), func(lhsItr.cell(), rhsCells[rhsAddr.index()]));
+        if (combiner.updateCommon()) {
+            builder.insertCell(combiner.address(), func(lhsItr.cell(), rhsCells[combiner.rightCellIndex()]));
         }
     }
     return builder.build();
@@ -59,11 +57,10 @@ apply(const DenseTensorView &lhs, const DenseTensorView &rhs, Function &&func)
     eval::ValueType resultType = DenseTensorAddressCombiner::combineDimensions(lhs.fast_type(), rhs.fast_type());
     DenseTensorAddressCombiner combiner(resultType, lhs.fast_type(), rhs.fast_type());
     DirectDenseTensorBuilder builder(resultType);
-    AddressContext rhsAddress(rhs.fast_type());
     if (combiner.hasAnyRightOnlyDimensions()) {
-        return apply(combiner, builder, lhs, rhsAddress, rhs.cellsRef(), std::move(func));
+        return apply(combiner, builder, lhs, rhs.cellsRef(), std::move(func));
     } else {
-        return apply_no_rightonly_dimensions(combiner, builder, lhs, rhsAddress, rhs.cellsRef(), std::move(func));
+        return apply_no_rightonly_dimensions(combiner, builder, lhs, rhs.cellsRef(), std::move(func));
     }
 }