add DenseDimensionCombiner class

3 files changed, 246 insertions, 0 deletions

diff --git a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
index c2638466de6..3fcf51da839 100644
--- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
+++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
@@ -1,6 +1,7 @@
 # Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 vespa_add_library(eval_tensor_dense OBJECT
     SOURCES
+    dense_dimension_combiner.cpp
     dense_add_dimension_optimizer.cpp
     dense_dot_product_function.cpp
     dense_fast_rename_optimizer.cpp
diff --git a/eval/src/vespa/eval/tensor/dense/dense_dimension_combiner.cpp b/eval/src/vespa/eval/tensor/dense/dense_dimension_combiner.cpp
new file mode 100644
index 00000000000..1d73eea6ddc
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_dimension_combiner.cpp
@@ -0,0 +1,110 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "dense_dimension_combiner.h"
+#include <cassert>
+
+namespace vespalib::tensor {
+
+DenseDimensionCombiner::DenseDimensionCombiner(const eval::ValueType &lhs,
+                                               const eval::ValueType &rhs)
+  : _leftDims(), _rightDims(), _commonDims(),
+    _leftIndex(0), _rightIndex(0), _outputIndex(0),
+    _leftOnlySize(1u), _rightOnlySize(1u), _outputSize(1u)
+{
+    eval::ValueType outputType = eval::ValueType::join(lhs, rhs);
+
+    assert(lhs.is_dense());
+    assert(rhs.is_dense());
+    assert(outputType.is_dense());
+
+    const auto &lDims = lhs.dimensions();
+    const auto &rDims = rhs.dimensions();
+    const auto &oDims = outputType.dimensions();
+
+    size_t i = lDims.size();
+    size_t j = rDims.size();
+    size_t k = oDims.size();
+
+    uint32_t lMul = 1;
+    uint32_t rMul = 1;
+    uint32_t oMul = 1;
+
+    while (k-- > 0) {
+        if ((i > 0) && (lDims[i-1].name == oDims[k].name)) {
+            --i;
+            // left dim match
+            if ((j > 0) && (rDims[j-1].name == oDims[k].name)) {
+                // both dim match
+                --j;
+                CommonDim cd;
+                cd.idx = 0;
+                cd.leftMultiplier = lMul;
+                cd.rightMultiplier = rMul;
+                cd.outputMultiplier = oMul;
+                assert(lDims[i].size == oDims[k].size);
+                assert(rDims[j].size == oDims[k].size);
+                cd.size = oDims[k].size;
+                lMul *= cd.size;
+                rMul *= cd.size;
+                oMul *= cd.size;
+                _leftOnlySize *= cd.size;
+                _rightOnlySize *= cd.size;
+                _outputSize *= cd.size;
+                _commonDims.push_back(cd);
+            } else {
+                LeftDim ld;
+                ld.idx = 0;
+                ld.leftMultiplier = lMul;
+                ld.outputMultiplier = oMul;
+                assert(lDims[i].size == oDims[k].size);
+                ld.size = oDims[k].size;
+                lMul *= ld.size;
+                oMul *= ld.size;
+                _leftOnlySize *= ld.size;
+                _outputSize *= ld.size;
+                _leftDims.push_back(ld);
+            }
+        } else {
+            // right dim match
+            assert(j > 0);
+            assert(rDims[j-1].name == oDims[k].name);
+            --j;
+            RightDim rd;
+            rd.idx = 0;
+            rd.rightMultiplier = rMul;
+            rd.outputMultiplier = oMul;
+            assert(rDims[j].size == oDims[k].size);
+            rd.size = oDims[k].size;
+            rMul *= rd.size;
+            oMul *= rd.size;
+            _rightOnlySize *= rd.size;
+            _outputSize *= rd.size;
+            _rightDims.push_back(rd);
+        }
+    }
+}
+
+void
+DenseDimensionCombiner::dump() const
+{
+    fprintf(stderr, "DenseDimensionCombiner: %u * %u -> %u\n", _leftOnlySize, _rightOnlySize, _outputSize);
+    for (const LeftDim& ld : _leftDims) {
+        fprintf(stderr, "ld curidx=%u (of %u) leftmul=%u outmul=%u\n",
+                ld.idx, ld.size, ld.leftMultiplier, ld.outputMultiplier);
+    }
+    for (const RightDim& rd : _rightDims) {
+        fprintf(stderr, "rd curidx=%u (of %u) rightmul=%u outmul=%u\n",
+                rd.idx, rd.size, rd.rightMultiplier, rd.outputMultiplier);
+    }
+    for (const CommonDim& cd : _commonDims) {
+        fprintf(stderr, "cd curidx=%u (of %u) leftmul=%u rightmul=%u outmul=%u\n",
+                cd.idx, cd.size, cd.leftMultiplier, cd.rightMultiplier, cd.outputMultiplier);
+    }
+    fprintf(stderr, "Left Index: %u\n", _leftIndex);
+    fprintf(stderr, "Right Index: %u\n", _rightIndex);
+    fprintf(stderr, "Output Index: %u\n", _outputIndex);
+}
+
+
+} // namespace
+
diff --git a/eval/src/vespa/eval/tensor/dense/dense_dimension_combiner.h b/eval/src/vespa/eval/tensor/dense/dense_dimension_combiner.h
new file mode 100644
index 00000000000..b0539b830db
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_dimension_combiner.h
@@ -0,0 +1,135 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/eval/tensor/tensor.h>
+#include <vespa/eval/tensor/types.h>
+#include <vespa/eval/eval/value_type.h>
+
+namespace vespalib::tensor {
+
+class DenseDimensionCombiner {
+
+    struct LeftDim {
+        uint32_t idx;
+        uint32_t size;
+        uint32_t leftMultiplier;
+        uint32_t outputMultiplier;
+    };
+    struct RightDim {
+        uint32_t idx;
+        uint32_t size;
+        uint32_t rightMultiplier;
+        uint32_t outputMultiplier;
+    };
+    struct CommonDim {
+        uint32_t idx;
+        uint32_t size;
+        uint32_t leftMultiplier;
+        uint32_t rightMultiplier;
+        uint32_t outputMultiplier;
+    };
+
+    std::vector<LeftDim> _leftDims;
+    std::vector<RightDim> _rightDims;
+    std::vector<CommonDim> _commonDims;
+
+    uint32_t _leftIndex;
+    uint32_t _rightIndex;
+    uint32_t _outputIndex;
+
+    uint32_t _leftOnlySize;
+    uint32_t _rightOnlySize;
+    uint32_t _outputSize;
+
+public:
+    size_t leftIdx() const { return _leftIndex; }
+    size_t rightIdx() const { return _rightIndex; }
+    size_t outputIdx() const { return _outputIndex; }
+
+    bool leftInRange() const { return _leftIndex < _leftOnlySize; }
+    bool rightInRange() const { return _rightIndex < _rightOnlySize; }
+    bool commonInRange() const { return _outputIndex < _outputSize; }
+
+    void leftReset() {
+        for (LeftDim& ld : _leftDims) {
+            _leftIndex -= ld.idx * ld.leftMultiplier;
+            _outputIndex -= ld.idx * ld.outputMultiplier;
+            ld.idx = 0;
+        }
+        if (_leftIndex >= _leftOnlySize) _leftIndex -= _leftOnlySize;
+    }
+
+    void stepLeft() {
+        size_t lim = _leftDims.size();
+        for (size_t i = 0; i < lim; ++i) {
+            LeftDim& ld = _leftDims[i];
+            ld.idx++;
+            _leftIndex += ld.leftMultiplier;
+            _outputIndex += ld.outputMultiplier;
+            if (ld.idx < ld.size) return;
+            _leftIndex -= ld.idx * ld.leftMultiplier;
+            _outputIndex -= ld.idx * ld.outputMultiplier;
+            ld.idx = 0;
+        }
+        _leftIndex += _leftOnlySize;
+    }
+
+
+    void rightReset() {
+        for (RightDim& rd : _rightDims) {
+            _rightIndex -= rd.idx * rd.rightMultiplier;
+            _outputIndex -= rd.idx * rd.outputMultiplier;
+            rd.idx = 0;
+        }
+        if (_rightIndex >= _rightOnlySize) _rightIndex -= _rightOnlySize;
+    }
+
+    void stepRight() {
+        size_t lim = _rightDims.size();
+        for (size_t i = 0; i < lim; ++i) {
+            RightDim& rd = _rightDims[i];
+            rd.idx++;
+            _rightIndex += rd.rightMultiplier;
+            _outputIndex += rd.outputMultiplier;
+            if (rd.idx < rd.size) return;
+            _rightIndex -= rd.idx * rd.rightMultiplier;
+            _outputIndex -= rd.idx * rd.outputMultiplier;
+            rd.idx = 0;
+        }
+        _rightIndex += _rightOnlySize;
+    }
+
+    void commonReset() {
+        for (CommonDim& cd : _commonDims) {
+            _leftIndex -= cd.idx * cd.leftMultiplier;
+            _rightIndex -= cd.idx * cd.rightMultiplier;
+            _outputIndex -= cd.idx * cd.outputMultiplier;
+            cd.idx = 0;
+        }
+        if (_outputIndex >= _outputSize) _outputIndex -= _outputSize;
+    }
+
+    void stepCommon() {
+        size_t lim = _commonDims.size();
+        for (size_t i = 0; i < lim; ++i) {
+            CommonDim &cd = _commonDims[i];
+            cd.idx++;
+            _leftIndex += cd.leftMultiplier;
+            _rightIndex += cd.rightMultiplier;
+            _outputIndex += cd.outputMultiplier;
+            if (cd.idx < cd.size) return;
+            _leftIndex -= cd.idx * cd.leftMultiplier;
+            _rightIndex -= cd.idx * cd.rightMultiplier;
+            _outputIndex -= cd.idx * cd.outputMultiplier;
+            cd.idx = 0;
+        }
+        _outputIndex += _outputSize;
+    }
+
+    DenseDimensionCombiner(const eval::ValueType &lhs, const eval::ValueType &rhs);
+
+    void dump() const;
+};
+
+}