optimized XW product now has children

also remove instruction trampoline

5 files changed, 99 insertions, 75 deletions

diff --git a/eval/src/tests/tensor/dense_tensor_function_optimizer/dense_tensor_function_optimizer_test.cpp b/eval/src/tests/tensor/dense_tensor_function_optimizer/dense_tensor_function_optimizer_test.cpp
index 1c8527a3eb4..c09b218508e 100644
--- a/eval/src/tests/tensor/dense_tensor_function_optimizer/dense_tensor_function_optimizer_test.cpp
+++ b/eval/src/tests/tensor/dense_tensor_function_optimizer/dense_tensor_function_optimizer_test.cpp
@@ -85,10 +85,10 @@ assertOptimizedXWProduct(const vespalib::string &vecTypeStr,
     ASSERT_TRUE(xwProduct);
     ASSERT_TRUE(inv_xwProduct);
     ASSERT_TRUE(common_idx != ValueType::Dimension::npos);
-    EXPECT_EQUAL(xwProduct->vectorId(), 1u);
-    EXPECT_EQUAL(inv_xwProduct->vectorId(), 3u);
-    EXPECT_EQUAL(xwProduct->matrixId(), 3u);
-    EXPECT_EQUAL(inv_xwProduct->matrixId(), 1u);
+    TEST_DO(assertParam(xwProduct->lhs(), 1));
+    TEST_DO(assertParam(inv_xwProduct->lhs(), 3));
+    TEST_DO(assertParam(xwProduct->rhs(), 3));
+    TEST_DO(assertParam(inv_xwProduct->rhs(), 1));
     EXPECT_EQUAL(xwProduct->vectorSize(), vecType.dimensions()[0].size);
     EXPECT_EQUAL(inv_xwProduct->vectorSize(), vecType.dimensions()[0].size);
     EXPECT_EQUAL(xwProduct->resultSize(), matType.dimensions()[1 - common_idx].size);
diff --git a/eval/src/tests/tensor/dense_xw_product_function/dense_xw_product_function_test.cpp b/eval/src/tests/tensor/dense_xw_product_function/dense_xw_product_function_test.cpp
index 5e18df10921..b45ccd5ceba 100644
--- a/eval/src/tests/tensor/dense_xw_product_function/dense_xw_product_function_test.cpp
+++ b/eval/src/tests/tensor/dense_xw_product_function/dense_xw_product_function_test.cpp
@@ -20,6 +20,7 @@ LOG_SETUP("dense_dot_product_function_test");
 using namespace vespalib;
 using namespace vespalib::eval;
 using namespace vespalib::tensor;
+using namespace vespalib::eval::tensor_function;
 
 const TensorEngine &ref_engine = SimpleTensorEngine::ref();
 const TensorEngine &prod_engine = DefaultTensorEngine::ref();
@@ -48,14 +49,20 @@ void verify_result(const TensorSpec &v, const TensorSpec &m, bool happy) {
     Value::UP prod_vec = prod_engine.from_spec(v);
     Value::UP prod_mat = prod_engine.from_spec(m);
 
-    DenseXWProductFunction fun1(expect.type(), 0, 1,
+    Inject vec_first(prod_vec->type(), 0);
+    Inject mat_last(prod_mat->type(), 1);
+
+    DenseXWProductFunction fun1(expect.type(), vec_first, mat_last,
                                 prod_vec->type().dimensions()[0].size,
                                 expect.type().dimensions()[0].size,
                                 happy);
     const Value &actual1 = fun1.eval(prod_engine, wrap({*prod_vec, *prod_mat}), stash);
     TEST_DO(verify_equal(expect, actual1));
 
-    DenseXWProductFunction fun2(expect.type(), 1, 0,
+    Inject vec_last(prod_vec->type(), 1);
+    Inject mat_first(prod_mat->type(), 0);
+
+    DenseXWProductFunction fun2(expect.type(), vec_last, mat_first,
                                 prod_vec->type().dimensions()[0].size,
                                 expect.type().dimensions()[0].size,
                                 happy);
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.cpp b/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.cpp
index 17ae7763dd4..c5ebe837151 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.cpp
@@ -48,7 +48,7 @@ bool isDenseXWProduct(const ValueType &res, const ValueType &vec, const ValueTyp
 
 const TensorFunction &createDenseXWProduct(const ValueType &res, const Inject &vec, const Inject &mat, Stash &stash) {
     bool common_is_inner = (mat.result_type().dimension_index(vec.result_type().dimensions()[0].name) == 1);
-    return stash.create<DenseXWProductFunction>(res, vec.param_idx(), mat.param_idx(),
+    return stash.create<DenseXWProductFunction>(res, vec, mat,
                                                 vec.result_type().dimensions()[0].size,
                                                 res.dimensions()[0].size,
                                                 common_is_inner);
diff --git a/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.cpp
index a62dafb6831..4684116e413 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.cpp
@@ -10,91 +10,107 @@
 
 namespace vespalib::tensor {
 
-DenseXWProductFunction::DenseXWProductFunction(const eval::ValueType &resultType,
-                                               size_t vectorId,
-                                               size_t matrixId,
-                                               size_t vectorSize,
-                                               size_t resultSize,
-                                               bool matrixHasCommonDimensionInnermost)
-    : _resultType(resultType),
-      _vectorId(vectorId),
-      _matrixId(matrixId),
-      _vectorSize(vectorSize),
-      _resultSize(resultSize),
-      _commonDimensionInnermost(matrixHasCommonDimensionInnermost),
-      _hwAccelerator(hwaccelrated::IAccelrated::getAccelrator())
-{}
+namespace {
 
-void
-DenseXWProductFunction::multiDotProduct(const XWInput &vectorCells,
-                                        const XWInput &matrixCells,
-                                        XWOutput &result) const
+DenseTensorView::CellsRef
+getCellsRef(const eval::Value &value)
+{
+    const DenseTensorView &denseTensor = static_cast<const DenseTensorView &>(value);
+    return denseTensor.cellsRef();
+}
+
+void multiDotProduct(const DenseXWProductFunction::Self &self,
+                     const XWInput &vectorCells, const XWInput &matrixCells, XWOutput &result)
 {
     double *out = result.begin();
     const double *matrixP = matrixCells.cbegin();
     const double * const vectorP = vectorCells.cbegin();
-    for (size_t row = 0; row < _resultSize; ++row) {
-        double cell = _hwAccelerator->dotProduct(vectorP, matrixP, _vectorSize);
+    for (size_t row = 0; row < self._resultSize; ++row) {
+        double cell = self._hwAccelerator->dotProduct(vectorP, matrixP, self._vectorSize);
         *out++ = cell;
-        matrixP += _vectorSize;
+        matrixP += self._vectorSize;
     }
     assert(out == result.end());
     assert(matrixP == matrixCells.cend());
 }
 
-void
-DenseXWProductFunction::transposedProduct(const XWInput &vectorCells,
-                                          const XWInput &matrixCells,
-                                          XWOutput &result) const
+void transposedProduct(const DenseXWProductFunction::Self &self,
+                       const XWInput &vectorCells, const XWInput &matrixCells, XWOutput &result)
 {
     double *out = result.begin();
     const double * const matrixP = matrixCells.cbegin();
     const double * const vectorP = vectorCells.cbegin();
-    for (size_t row = 0; row < _resultSize; ++row) {
+    for (size_t row = 0; row < self._resultSize; ++row) {
         double cell = 0;
-        for (size_t col = 0; col < _vectorSize; ++col) {
-            cell += matrixP[col*_resultSize + row] * vectorP[col];
+        for (size_t col = 0; col < self._vectorSize; ++col) {
+            cell += matrixP[col*self._resultSize + row] * vectorP[col];
         }
         *out++ = cell;
     }
     assert(out == result.end());
 }
 
-namespace {
+void my_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    DenseXWProductFunction::Self *self = (DenseXWProductFunction::Self *)(param);
 
-DenseTensorView::CellsRef
-getCellsRef(const eval::Value &value)
-{
-    const DenseTensorView &denseTensor = static_cast<const DenseTensorView &>(value);
-    return denseTensor.cellsRef();
-}
+    DenseTensorView::CellsRef vectorCells = getCellsRef(state.peek(1));
+    DenseTensorView::CellsRef matrixCells = getCellsRef(state.peek(0));
 
-void op_call_leaf_eval(eval::InterpretedFunction::State &state, uint64_t param) {
-    DenseXWProductFunction *self = (DenseXWProductFunction *)(param);
-    state.stack.push_back(self->eval(state.engine, *state.params, state.stash));
+    ArrayRef<double> outputCells = state.stash.create_array<double>(self->_resultSize);
+    if (self->_commonDimensionInnermost) {
+        multiDotProduct(*self, vectorCells, matrixCells, outputCells);
+    } else {
+        transposedProduct(*self, vectorCells, matrixCells, outputCells);
+    }
+    state.pop_pop_push(state.stash.create<DenseTensorView>(self->_resultType, outputCells));
 }
 
+} // namespace vespalib::tensor::<unnamed>
+
+DenseXWProductFunction::Self::Self(const eval::ValueType &resultType,
+                                   size_t vectorSize,
+                                   size_t resultSize,
+                                   bool matrixHasCommonDimensionInnermost)
+    : _resultType(resultType),
+      _vectorSize(vectorSize),
+      _resultSize(resultSize),
+      _commonDimensionInnermost(matrixHasCommonDimensionInnermost),
+      _hwAccelerator(hwaccelrated::IAccelrated::getAccelrator())
+{}
+
+DenseXWProductFunction::DenseXWProductFunction(const eval::ValueType &resultType,
+                                               const eval::TensorFunction &vector_in,
+                                               const eval::TensorFunction &matrix_in,
+                                               size_t vectorSize,
+                                               size_t resultSize,
+                                               bool matrixHasCommonDimensionInnermost)
+    : eval::tensor_function::Op2(resultType, vector_in, matrix_in),
+      _self(resultType, vectorSize, resultSize, matrixHasCommonDimensionInnermost)
+{}
+
+namespace {
+
 } // namespace <unnamed>
 
 eval::InterpretedFunction::Instruction
 DenseXWProductFunction::compile_self(Stash &) const
 {
-    return eval::InterpretedFunction::Instruction(op_call_leaf_eval, (uint64_t)(this));
+    return eval::InterpretedFunction::Instruction(my_op, (uint64_t)(&_self));
 }
 
 const eval::Value &
-DenseXWProductFunction::eval(const eval::TensorEngine &, const eval::LazyParams &params, Stash &stash) const
+DenseXWProductFunction::eval(const eval::TensorEngine &engine, const eval::LazyParams &params, Stash &stash) const
 {
-    DenseTensorView::CellsRef vectorCells = getCellsRef(params.resolve(_vectorId, stash));
-    DenseTensorView::CellsRef matrixCells = getCellsRef(params.resolve(_matrixId, stash));
+    DenseTensorView::CellsRef vectorCells = getCellsRef(lhs().eval(engine, params, stash));
+    DenseTensorView::CellsRef matrixCells = getCellsRef(rhs().eval(engine, params, stash));
 
-    ArrayRef<double> outputCells = stash.create_array<double>(_resultSize);
-    if (_commonDimensionInnermost) {
-        multiDotProduct(vectorCells, matrixCells, outputCells);
+    ArrayRef<double> outputCells = stash.create_array<double>(_self._resultSize);
+    if (_self._commonDimensionInnermost) {
+        multiDotProduct(_self, vectorCells, matrixCells, outputCells);
     } else {
-        transposedProduct(vectorCells, matrixCells, outputCells);
+        transposedProduct(_self, vectorCells, matrixCells, outputCells);
     }
-    return stash.create<DenseTensorView>(_resultType, outputCells);
+    return stash.create<DenseTensorView>(result_type(), outputCells);
 }
 
 }
diff --git a/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.h b/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.h
index 4d2a85d96f7..9655bc1c386 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.h
@@ -14,39 +14,40 @@ using XWOutput = ArrayRef<double>;
 /**
  * Tensor function for product of one 1-dimensional and one 2-dimensional dense tensor.
  */
-class DenseXWProductFunction : public eval::TensorFunction
+class DenseXWProductFunction : public eval::tensor_function::Op2
 {
+public:
+    struct Self {
+        const eval::ValueType _resultType;
+        const size_t _vectorSize;
+        const size_t _resultSize;
+        bool _commonDimensionInnermost;
+        hwaccelrated::IAccelrated::UP _hwAccelerator;
+        Self(const eval::ValueType &resultType,
+             size_t vectorSize,
+             size_t resultSize,
+             bool matrixHasCommonDimensionInnermost);
+        ~Self() {}
+    };
+
 private:
-    const eval::ValueType _resultType;
-    const size_t _vectorId;
-    const size_t _matrixId;
-    const size_t _vectorSize;
-    const size_t _resultSize;
-    bool _commonDimensionInnermost;
-    hwaccelrated::IAccelrated::UP _hwAccelerator;
-
-    void multiDotProduct(const XWInput &v, const XWInput &m, XWOutput &r) const;
-    void transposedProduct(const XWInput &v, const XWInput &m, XWOutput &r) const;
+    Self _self;
+
 public:
     DenseXWProductFunction(const eval::ValueType &resultType,
-                           size_t vectorId,
-                           size_t matrixId,
+                           const eval::TensorFunction &vector_in,
+                           const eval::TensorFunction &matrix_in,
                            size_t vectorSize,
                            size_t resultSize,
                            bool matrixHasCommonDimensionInnermost);
 
     ~DenseXWProductFunction() {}
 
-    size_t vectorId() const { return _vectorId; }
-    size_t matrixId() const { return _matrixId; }
-
-    size_t vectorSize() const { return _vectorSize; }
-    size_t resultSize() const { return _resultSize; }
+    size_t vectorSize() const { return _self._vectorSize; }
+    size_t resultSize() const { return _self._resultSize; }
 
-    bool matrixHasCommonDimensionInnermost() const { return _commonDimensionInnermost; }
+    bool matrixHasCommonDimensionInnermost() const { return _self._commonDimensionInnermost; }
 
-    const eval::ValueType &result_type() const override { return _resultType; }
-    void push_children(std::vector<Child::CREF> &) const override {}
     eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
     const eval::Value &eval(const eval::TensorEngine &engine, const eval::LazyParams &params, Stash &stash) const override;
 };