optimized dot product now has children

also remove instruction trampoline

6 files changed, 35 insertions, 26 deletions

diff --git a/eval/src/tests/eval/interpreted_function/interpreted_function_test.cpp b/eval/src/tests/eval/interpreted_function/interpreted_function_test.cpp
index f0306e99a91..e0de483fde2 100644
--- a/eval/src/tests/eval/interpreted_function/interpreted_function_test.cpp
+++ b/eval/src/tests/eval/interpreted_function/interpreted_function_test.cpp
@@ -193,7 +193,7 @@ struct InnerProduct {
           types(function, {ValueType::from_spec(a.type()), ValueType::from_spec(b.type())}),
           interpreted(engine, function, types) {}
     void verify_optimized() const {
-        EXPECT_EQUAL(1u, interpreted.program_size());
+        EXPECT_LESS(interpreted.program_size(), 4u);
         InterpretedFunction::Context ctx(interpreted);
         Value::UP va = engine.from_spec(a);
         Value::UP vb = engine.from_spec(b);
diff --git a/eval/src/tests/tensor/dense_dot_product_function/dense_dot_product_function_test.cpp b/eval/src/tests/tensor/dense_dot_product_function/dense_dot_product_function_test.cpp
index 8b4b1497243..80e994d8c0b 100644
--- a/eval/src/tests/tensor/dense_dot_product_function/dense_dot_product_function_test.cpp
+++ b/eval/src/tests/tensor/dense_dot_product_function/dense_dot_product_function_test.cpp
@@ -73,8 +73,10 @@ public:
 
 struct Fixture
 {
-    DenseDotProductFunction function;
     FunctionInput input;
+    tensor_function::Inject a;
+    tensor_function::Inject b;
+    DenseDotProductFunction function;
     Fixture(size_t lhsNumCells, size_t rhsNumCells);
     ~Fixture();
     double eval() const {
@@ -90,8 +92,10 @@ struct Fixture
 };
 
 Fixture::Fixture(size_t lhsNumCells, size_t rhsNumCells)
-    : function(0, 1),
-      input(lhsNumCells, rhsNumCells)
+    : input(lhsNumCells, rhsNumCells),
+      a(input.param(0).type(), 0),
+      b(input.param(1).type(), 1),
+      function(a, b)
 { }
 
 Fixture::~Fixture() { }
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 57d03c09686..1c8527a3eb4 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
@@ -26,16 +26,22 @@ optimizeDotProduct(const vespalib::string &lhsType,
     return DenseTensorFunctionOptimizer::optimize(reduceNode, stash);
 }
 
+void assertParam(const TensorFunction &node, size_t expect_idx) {
+    auto inject = as<Inject>(node);
+    ASSERT_TRUE(inject);
+    EXPECT_EQUAL(inject->param_idx(), expect_idx);
+}
+
 void
 assertOptimizedDotProduct(const vespalib::string &lhsType,
-                         const vespalib::string &rhsType)
+                          const vespalib::string &rhsType)
 {
     Stash stash;
     const TensorFunction &func = optimizeDotProduct(lhsType, rhsType, stash);
     const DenseDotProductFunction *dotProduct = as<DenseDotProductFunction>(func);
     ASSERT_TRUE(dotProduct);
-    EXPECT_EQUAL(1u, dotProduct->lhsTensorId());
-    EXPECT_EQUAL(3u, dotProduct->rhsTensorId());
+    TEST_DO(assertParam(dotProduct->lhs(), 1));
+    TEST_DO(assertParam(dotProduct->rhs(), 3));
 }
 
 void
diff --git a/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.cpp
index 0f395bd353b..48bc6e3b110 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.cpp
@@ -10,9 +10,9 @@ namespace vespalib::tensor {
 
 using CellsRef = DenseTensorView::CellsRef;
 
-DenseDotProductFunction::DenseDotProductFunction(size_t lhsTensorId_, size_t rhsTensorId_)
-    : _lhsTensorId(lhsTensorId_),
-      _rhsTensorId(rhsTensorId_),
+DenseDotProductFunction::DenseDotProductFunction(const eval::TensorFunction &lhs_in,
+                                                 const eval::TensorFunction &rhs_in)
+    : eval::tensor_function::Op2(eval::ValueType::double_type(), lhs_in, rhs_in),
       _hwAccelerator(hwaccelrated::IAccelrated::getAccelrator())
 {
 }
@@ -26,9 +26,13 @@ getCellsRef(const eval::Value &value)
     return denseTensor.cellsRef();
 }
 
-void op_call_leaf_eval(eval::InterpretedFunction::State &state, uint64_t param) {
-    DenseDotProductFunction *self = (DenseDotProductFunction *)(param);
-    state.stack.push_back(self->eval(state.engine, *state.params, state.stash));
+void my_op(eval::InterpretedFunction::State &state, uint64_t param) {
+    auto *hw_accelerator = (hwaccelrated::IAccelrated *)(param);
+    DenseTensorView::CellsRef lhsCells = getCellsRef(state.peek(1));
+    DenseTensorView::CellsRef rhsCells = getCellsRef(state.peek(0));
+    size_t numCells = std::min(lhsCells.size(), rhsCells.size());
+    double result = hw_accelerator->dotProduct(lhsCells.cbegin(), rhsCells.cbegin(), numCells);
+    state.pop_pop_push(state.stash.create<eval::DoubleValue>(result));
 }
 
 }
@@ -36,14 +40,14 @@ void op_call_leaf_eval(eval::InterpretedFunction::State &state, uint64_t param)
 eval::InterpretedFunction::Instruction
 DenseDotProductFunction::compile_self(Stash &) const
 {
-    return eval::InterpretedFunction::Instruction(op_call_leaf_eval, (uint64_t)(this));
+    return eval::InterpretedFunction::Instruction(my_op, (uint64_t)(_hwAccelerator.get()));
 }
 
 const eval::Value &
-DenseDotProductFunction::eval(const eval::TensorEngine &, const eval::LazyParams &params, Stash &stash) const
+DenseDotProductFunction::eval(const eval::TensorEngine &engine, const eval::LazyParams &params, Stash &stash) const
 {
-    DenseTensorView::CellsRef lhsCells = getCellsRef(params.resolve(_lhsTensorId, stash));
-    DenseTensorView::CellsRef rhsCells = getCellsRef(params.resolve(_rhsTensorId, stash));
+    DenseTensorView::CellsRef lhsCells = getCellsRef(lhs().eval(engine, params, stash));
+    DenseTensorView::CellsRef rhsCells = getCellsRef(rhs().eval(engine, params, stash));
     size_t numCells = std::min(lhsCells.size(), rhsCells.size());
     double result = _hwAccelerator->dotProduct(lhsCells.cbegin(), rhsCells.cbegin(), numCells);
     return stash.create<eval::DoubleValue>(result);
diff --git a/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.h b/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.h
index d313602bd53..b496986453f 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.h
@@ -10,19 +10,14 @@ namespace vespalib::tensor {
 /**
  * Tensor function for a dot product between two 1-dimensional dense tensors.
  */
-class DenseDotProductFunction : public eval::TensorFunction
+class DenseDotProductFunction : public eval::tensor_function::Op2
 {
 private:
-    size_t _lhsTensorId;
-    size_t _rhsTensorId;
     hwaccelrated::IAccelrated::UP _hwAccelerator;
 
 public:
-    DenseDotProductFunction(size_t lhsTensorId_, size_t rhsTensorId_);
-    size_t lhsTensorId() const { return _lhsTensorId; }
-    size_t rhsTensorId() const { return _rhsTensorId; }
-    const eval::ValueType &result_type() const override { return eval::DoubleValue::double_type(); }
-    void push_children(std::vector<Child::CREF> &) const override {}
+    DenseDotProductFunction(const eval::TensorFunction &lhs_in,
+                            const eval::TensorFunction &rhs_in);
     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;
 };
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 7e167a92819..17ae7763dd4 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
@@ -66,7 +66,7 @@ struct InnerProductFunctionOptimizer
                 const Inject *rhs = as<Inject>(join->rhs());
                 if (lhs && rhs) {
                     if (isDenseDotProduct(result_type, lhs->result_type(), rhs->result_type())) {
-                        return stash.create<DenseDotProductFunction>(lhs->param_idx(), rhs->param_idx());
+                        return stash.create<DenseDotProductFunction>(*lhs, *rhs);
                     }
                     if (isDenseXWProduct(result_type, lhs->result_type(), rhs->result_type())) {
                         return createDenseXWProduct(result_type, *lhs, *rhs, stash);