make tensor engine available when compiling tensor functions

24 files changed, 67 insertions, 55 deletions

diff --git a/eval/src/tests/tensor/dense_replace_type_function/dense_replace_type_function_test.cpp b/eval/src/tests/tensor/dense_replace_type_function/dense_replace_type_function_test.cpp
index 0533b1c92b7..732fc9c3e69 100644
--- a/eval/src/tests/tensor/dense_replace_type_function/dense_replace_type_function_test.cpp
+++ b/eval/src/tests/tensor/dense_replace_type_function/dense_replace_type_function_test.cpp
@@ -23,7 +23,7 @@ struct ChildMock : Leaf {
     bool is_mutable;
     ChildMock(const ValueType &type) : Leaf(type), is_mutable(true) {}
     bool result_is_mutable() const override { return is_mutable; }
-    InterpretedFunction::Instruction compile_self(Stash &) const override { abort(); }
+    InterpretedFunction::Instruction compile_self(const TensorEngine &, Stash &) const override { abort(); }
 };
 
 struct Fixture {
@@ -43,7 +43,7 @@ struct Fixture {
     {
         my_fun.push_children(children);
         state.stack.push_back(*my_value);
-        my_fun.compile_self(state.stash).perform(state);
+        my_fun.compile_self(engine, state.stash).perform(state);
         ASSERT_EQUAL(children.size(), 1u);
         ASSERT_EQUAL(state.stack.size(), 1u);
         ASSERT_TRUE(!new_type.is_error());
diff --git a/eval/src/vespa/eval/eval/compile_tensor_function.cpp b/eval/src/vespa/eval/eval/compile_tensor_function.cpp
index ac36720895f..18ef59507d8 100644
--- a/eval/src/vespa/eval/eval/compile_tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/compile_tensor_function.cpp
@@ -32,10 +32,11 @@ struct Frame {
 };
 
 struct ProgramCompiler {
+    const TensorEngine &engine;
     Stash &stash;
     std::vector<Frame> stack;
     std::vector<Instruction> prog;
-    ProgramCompiler(Stash &stash_in) : stash(stash_in), stack(), prog() {}
+    ProgramCompiler(const TensorEngine &engine_in, Stash &stash_in) : engine(engine_in), stash(stash_in), stack(), prog() {}
 
     void append(const std::vector<Instruction> &other_prog) {
         prog.insert(prog.end(), other_prog.begin(), other_prog.end());
@@ -43,9 +44,9 @@ struct ProgramCompiler {
 
     void open(const TensorFunction &node) {
         if (auto if_node = as<tensor_function::If>(node)) {
-            append(compile_tensor_function(if_node->cond(), stash));
-            auto true_prog = compile_tensor_function(if_node->true_child(), stash);
-            auto false_prog = compile_tensor_function(if_node->false_child(), stash);
+            append(compile_tensor_function(engine, if_node->cond(), stash));
+            auto true_prog = compile_tensor_function(engine, if_node->true_child(), stash);
+            auto false_prog = compile_tensor_function(engine, if_node->false_child(), stash);
             true_prog.emplace_back(op_skip, false_prog.size());
             prog.emplace_back(op_skip_if_false, true_prog.size());
             append(true_prog);
@@ -56,7 +57,7 @@ struct ProgramCompiler {
     }
 
     void close(const TensorFunction &node) {
-        prog.push_back(node.compile_self(stash));
+        prog.push_back(node.compile_self(engine, stash));
     }
 
     std::vector<Instruction> compile(const TensorFunction &function) {
@@ -75,8 +76,8 @@ struct ProgramCompiler {
 
 } // namespace vespalib::eval::<unnamed>
 
-std::vector<Instruction> compile_tensor_function(const TensorFunction &function, Stash &stash) {   
-    ProgramCompiler compiler(stash);
+std::vector<Instruction> compile_tensor_function(const TensorEngine &engine, const TensorFunction &function, Stash &stash) {
+    ProgramCompiler compiler(engine, stash);
     return compiler.compile(function);
 }
 
diff --git a/eval/src/vespa/eval/eval/compile_tensor_function.h b/eval/src/vespa/eval/eval/compile_tensor_function.h
index 013d228c2f9..63f00fde053 100644
--- a/eval/src/vespa/eval/eval/compile_tensor_function.h
+++ b/eval/src/vespa/eval/eval/compile_tensor_function.h
@@ -10,7 +10,8 @@ namespace vespalib { class Stash; }
 namespace vespalib::eval {
 
 struct TensorFunction;
+struct TensorEngine;
 
-std::vector<InterpretedFunction::Instruction> compile_tensor_function(const TensorFunction &function, Stash &stash);
+std::vector<InterpretedFunction::Instruction> compile_tensor_function(const TensorEngine &engine, const TensorFunction &function, Stash &stash);
 
 } // namespace vespalib::eval
diff --git a/eval/src/vespa/eval/eval/interpreted_function.cpp b/eval/src/vespa/eval/eval/interpreted_function.cpp
index f69893ad014..121db4ffb6e 100644
--- a/eval/src/vespa/eval/eval/interpreted_function.cpp
+++ b/eval/src/vespa/eval/eval/interpreted_function.cpp
@@ -63,7 +63,7 @@ InterpretedFunction::InterpretedFunction(const TensorEngine &engine, const Tenso
       _stash(),
       _tensor_engine(engine)
 {
-    _program = compile_tensor_function(function, _stash);
+    _program = compile_tensor_function(engine, function, _stash);
 }
 
 InterpretedFunction::InterpretedFunction(const TensorEngine &engine, const nodes::Node &root, const NodeTypes &types)
@@ -73,7 +73,7 @@ InterpretedFunction::InterpretedFunction(const TensorEngine &engine, const nodes
 {
     const TensorFunction &plain_fun = make_tensor_function(engine, root, types, _stash);
     const TensorFunction &optimized = engine.optimize(plain_fun, _stash);
-    _program = compile_tensor_function(optimized, _stash);
+    _program = compile_tensor_function(engine, optimized, _stash);
 }
 
 InterpretedFunction::~InterpretedFunction() = default;
diff --git a/eval/src/vespa/eval/eval/tensor_function.cpp b/eval/src/vespa/eval/eval/tensor_function.cpp
index 889738a201d..40c4c614f20 100644
--- a/eval/src/vespa/eval/eval/tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/tensor_function.cpp
@@ -243,7 +243,7 @@ Op2::visit_children(vespalib::ObjectVisitor &visitor) const
 //-----------------------------------------------------------------------------
 
 Instruction
-ConstValue::compile_self(Stash &) const
+ConstValue::compile_self(const TensorEngine &, Stash &) const
 {
     return Instruction(op_load_const, wrap_param<Value>(_value));
 }
@@ -262,7 +262,7 @@ ConstValue::visit_self(vespalib::ObjectVisitor &visitor) const
 //-----------------------------------------------------------------------------
 
 Instruction
-Inject::compile_self(Stash &) const
+Inject::compile_self(const TensorEngine &, Stash &) const
 {
     return Instruction::fetch_param(_param_idx);
 }
@@ -277,7 +277,7 @@ Inject::visit_self(vespalib::ObjectVisitor &visitor) const
 //-----------------------------------------------------------------------------
 
 Instruction
-Reduce::compile_self(Stash &stash) const
+Reduce::compile_self(const TensorEngine &, Stash &stash) const
 {
     ReduceParams &params = stash.create<ReduceParams>(_aggr, _dimensions);
     return Instruction(op_tensor_reduce, wrap_param<ReduceParams>(params));
@@ -294,7 +294,7 @@ Reduce::visit_self(vespalib::ObjectVisitor &visitor) const
 //-----------------------------------------------------------------------------
 
 Instruction
-Map::compile_self(Stash &) const
+Map::compile_self(const TensorEngine &, Stash &) const
 {
     if (result_type().is_double()) {
         return Instruction(op_double_map, to_param(_function));
@@ -312,7 +312,7 @@ Map::visit_self(vespalib::ObjectVisitor &visitor) const
 //-----------------------------------------------------------------------------
 
 Instruction
-Join::compile_self(Stash &) const
+Join::compile_self(const TensorEngine &, Stash &) const
 {
     if (result_type().is_double()) {
         if (_function == operation::Mul::f) {
@@ -336,7 +336,7 @@ Join::visit_self(vespalib::ObjectVisitor &visitor) const
 //-----------------------------------------------------------------------------
 
 Instruction
-Merge::compile_self(Stash &) const
+Merge::compile_self(const TensorEngine &, Stash &) const
 {
     return Instruction(op_tensor_merge, to_param(_function));
 }
@@ -351,7 +351,7 @@ Merge::visit_self(vespalib::ObjectVisitor &visitor) const
 //-----------------------------------------------------------------------------
 
 Instruction
-Concat::compile_self(Stash &) const
+Concat::compile_self(const TensorEngine &, Stash &) const
 {
     return Instruction(op_tensor_concat, wrap_param<vespalib::string>(_dimension));
 }
@@ -374,7 +374,7 @@ Create::push_children(std::vector<Child::CREF> &children) const
 }
 
 Instruction
-Create::compile_self(Stash &) const
+Create::compile_self(const TensorEngine &, Stash &) const
 {
     return Instruction(op_tensor_create, wrap_param<Create>(*this));
 }
@@ -436,7 +436,7 @@ Lambda::create_spec_impl(const ValueType &type, const LazyParams &params, const
 }
 
 InterpretedFunction::Instruction
-Lambda::compile_self(Stash &) const
+Lambda::compile_self(const TensorEngine &, Stash &) const
 {
     return Instruction(op_tensor_lambda, wrap_param<Lambda>(*this));
 }
@@ -471,7 +471,7 @@ Peek::push_children(std::vector<Child::CREF> &children) const
 }
 
 Instruction
-Peek::compile_self(Stash &) const
+Peek::compile_self(const TensorEngine &, Stash &) const
 {
     return Instruction(op_tensor_peek, wrap_param<Peek>(*this));
 }
@@ -500,7 +500,7 @@ Peek::visit_children(vespalib::ObjectVisitor &visitor) const
 //-----------------------------------------------------------------------------
 
 Instruction
-Rename::compile_self(Stash &stash) const
+Rename::compile_self(const TensorEngine &, Stash &stash) const
 {
     RenameParams &params = stash.create<RenameParams>(_from, _to);
     return Instruction(op_tensor_rename, wrap_param<RenameParams>(params));
@@ -524,7 +524,7 @@ If::push_children(std::vector<Child::CREF> &children) const
 }
 
 Instruction
-If::compile_self(Stash &) const
+If::compile_self(const TensorEngine &, Stash &) const
 {
     // 'if' is handled directly by compile_tensor_function to enable
     // lazy-evaluation of true/false sub-expressions.
diff --git a/eval/src/vespa/eval/eval/tensor_function.h b/eval/src/vespa/eval/eval/tensor_function.h
index c4e7384abcd..3aedf56affc 100644
--- a/eval/src/vespa/eval/eval/tensor_function.h
+++ b/eval/src/vespa/eval/eval/tensor_function.h
@@ -101,9 +101,10 @@ struct TensorFunction
      * the value stack during execution.
      *
      * @return instruction representing the operation of this node
+     * @param engine the tensor engine used for evaluation
      * @param stash heterogeneous object store
      **/
-    virtual InterpretedFunction::Instruction compile_self(Stash &stash) const = 0;
+    virtual InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const = 0;
 
     // for debug dumping
     vespalib::string as_string() const;
@@ -185,7 +186,7 @@ public:
     ConstValue(const Value &value_in) : Leaf(value_in.type()), _value(value_in) {}
     const Value &value() const { return _value; }
     bool result_is_mutable() const override { return false; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const final override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
 };
 
@@ -201,7 +202,7 @@ public:
         : Leaf(result_type_in), _param_idx(param_idx_in) {}
     size_t param_idx() const { return _param_idx; }
     bool result_is_mutable() const override { return false; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const final override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
 };
 
@@ -222,7 +223,7 @@ public:
     Aggr aggr() const { return _aggr; }
     const std::vector<vespalib::string> &dimensions() const { return _dimensions; }
     bool result_is_mutable() const override { return true; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const final override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
 };
 
@@ -240,7 +241,7 @@ public:
         : Op1(result_type_in, child_in), _function(function_in) {}
     map_fun_t function() const { return _function; }
     bool result_is_mutable() const override { return true; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
 };
 
@@ -259,7 +260,7 @@ public:
         : Op2(result_type_in, lhs_in, rhs_in), _function(function_in) {}
     join_fun_t function() const { return _function; }
     bool result_is_mutable() const override { return true; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
 };
 
@@ -278,7 +279,7 @@ public:
         : Op2(result_type_in, lhs_in, rhs_in), _function(function_in) {}
     join_fun_t function() const { return _function; }
     bool result_is_mutable() const override { return true; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
 };
 
@@ -297,7 +298,7 @@ public:
         : Op2(result_type_in, lhs_in, rhs_in), _dimension(dimension_in) {}
     const vespalib::string &dimension() const { return _dimension; }
     bool result_is_mutable() const override { return true; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const final override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
 };
 
@@ -318,7 +319,7 @@ public:
     }
     const std::map<TensorSpec::Address, Child> &spec() const { return _spec; }
     bool result_is_mutable() const override { return true; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const final override;
     void push_children(std::vector<Child::CREF> &children) const final override;
     void visit_children(vespalib::ObjectVisitor &visitor) const final override;
 };
@@ -337,7 +338,7 @@ public:
     static TensorSpec create_spec_impl(const ValueType &type, const LazyParams &params, const std::vector<size_t> &bind, const InterpretedFunction &fun);
     TensorSpec create_spec(const LazyParams &params) const { return create_spec_impl(result_type(), params, _bindings, _lambda); }
     bool result_is_mutable() const override { return true; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const final override;
     void push_children(std::vector<Child::CREF> &children) const final override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
 };
@@ -372,7 +373,7 @@ public:
     const std::map<vespalib::string, MyLabel> &spec() const { return _spec; }
     const ValueType &param_type() const { return _param.get().result_type(); }
     bool result_is_mutable() const override { return true; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const final override;
     void push_children(std::vector<Child::CREF> &children) const final override;
     void visit_children(vespalib::ObjectVisitor &visitor) const final override;
 };
@@ -394,7 +395,7 @@ public:
     const std::vector<vespalib::string> &from() const { return _from; }
     const std::vector<vespalib::string> &to() const { return _to; }
     bool result_is_mutable() const override { return true; }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const final override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
 };
 
@@ -420,7 +421,7 @@ public:
         return (true_child().result_is_mutable() &&
                 false_child().result_is_mutable());
     }
-    InterpretedFunction::Instruction compile_self(Stash &stash) const final override;
+    InterpretedFunction::Instruction compile_self(const TensorEngine &engine, Stash &stash) const final override;
     void visit_children(vespalib::ObjectVisitor &visitor) const final override;
 };
 
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 2fe89861e9f..c9ff57e4a65 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
@@ -11,6 +11,7 @@ namespace vespalib::tensor {
 
 using eval::ValueType;
 using eval::TensorFunction;
+using eval::TensorEngine;
 using eval::as;
 using eval::Aggr;
 using namespace eval::tensor_function;
@@ -71,7 +72,7 @@ DenseDotProductFunction::DenseDotProductFunction(const eval::TensorFunction &lhs
 }
 
 eval::InterpretedFunction::Instruction
-DenseDotProductFunction::compile_self(Stash &) const
+DenseDotProductFunction::compile_self(const TensorEngine &, Stash &) const
 {
     auto op = my_select(lhs().result_type().cell_type(), rhs().result_type().cell_type());
     return eval::InterpretedFunction::Instruction(op);
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 1d8f749689b..1ee6baff2a5 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
@@ -16,7 +16,7 @@ private:
 public:
     DenseDotProductFunction(const eval::TensorFunction &lhs_in,
                             const eval::TensorFunction &rhs_in);
-    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
     bool result_is_mutable() const override { return true; }
     static bool compatible_types(const ValueType &res, const ValueType &lhs, const ValueType &rhs);
     static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
diff --git a/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.cpp
index 990215cbbe8..2107c7661f2 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.cpp
@@ -10,6 +10,7 @@ namespace vespalib::tensor {
 using eval::Value;
 using eval::ValueType;
 using eval::TensorFunction;
+using eval::TensorEngine;
 using eval::as;
 using namespace eval::tensor_function;
 
@@ -74,7 +75,7 @@ DenseInplaceJoinFunction::~DenseInplaceJoinFunction()
 }
 
 eval::InterpretedFunction::Instruction
-DenseInplaceJoinFunction::compile_self(Stash &) const
+DenseInplaceJoinFunction::compile_self(const TensorEngine &, Stash &) const
 {
     auto op = my_select(lhs().result_type().cell_type(),
                         rhs().result_type().cell_type(), _write_left);
diff --git a/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.h b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.h
index 83acee6d31f..acd1a2d716b 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_join_function.h
@@ -25,7 +25,7 @@ public:
     ~DenseInplaceJoinFunction();
     bool write_left() const { return _write_left; }
     bool result_is_mutable() const override { return true; }
-    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
     static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
 };
diff --git a/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.cpp
index e086eab0b13..62434073f8e 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.cpp
@@ -9,6 +9,7 @@ namespace vespalib::tensor {
 using eval::Value;
 using eval::ValueType;
 using eval::TensorFunction;
+using eval::TensorEngine;
 using eval::as;
 using namespace eval::tensor_function;
 
@@ -42,7 +43,7 @@ DenseInplaceMapFunction::~DenseInplaceMapFunction()
 }
 
 eval::InterpretedFunction::Instruction
-DenseInplaceMapFunction::compile_self(Stash &) const
+DenseInplaceMapFunction::compile_self(const TensorEngine &, Stash &) const
 {
     auto op = select_1<MyInplaceMapOp>(result_type().cell_type());
     return eval::InterpretedFunction::Instruction(op, (uint64_t)function());
diff --git a/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.h b/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.h
index acc4504176a..52122f4d8dc 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_inplace_map_function.h
@@ -18,7 +18,7 @@ public:
                             map_fun_t function_in);
     ~DenseInplaceMapFunction();
     bool result_is_mutable() const override { return true; }
-    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
     static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
 };
 
diff --git a/eval/src/vespa/eval/tensor/dense/dense_matmul_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_matmul_function.cpp
index 1feecc31e51..695e0fddd08 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_matmul_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_matmul_function.cpp
@@ -13,6 +13,7 @@ namespace vespalib::tensor {
 
 using eval::ValueType;
 using eval::TensorFunction;
+using eval::TensorEngine;
 using eval::as;
 using eval::Aggr;
 using namespace eval::tensor_function;
@@ -194,7 +195,7 @@ DenseMatMulFunction::DenseMatMulFunction(const eval::ValueType &result_type,
 DenseMatMulFunction::~DenseMatMulFunction() = default;
 
 eval::InterpretedFunction::Instruction
-DenseMatMulFunction::compile_self(Stash &stash) const
+DenseMatMulFunction::compile_self(const TensorEngine &, Stash &stash) const
 {
     Self &self = stash.create<Self>(result_type(), _lhs_size, _common_size, _rhs_size);
     auto op = my_select(lhs().result_type().cell_type(), rhs().result_type().cell_type(),
diff --git a/eval/src/vespa/eval/tensor/dense/dense_matmul_function.h b/eval/src/vespa/eval/tensor/dense/dense_matmul_function.h
index f0b6d8b6c19..88d7b9f37e0 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_matmul_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_matmul_function.h
@@ -50,7 +50,7 @@ public:
     bool lhs_common_inner() const { return _lhs_common_inner; }
     bool rhs_common_inner() const { return _rhs_common_inner; }
 
-    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
     static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
 };
diff --git a/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.cpp
index b81b0f2c876..e2dc2241555 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.cpp
@@ -9,6 +9,7 @@ namespace vespalib::tensor {
 using eval::Value;
 using eval::ValueType;
 using eval::TensorFunction;
+using eval::TensorEngine;
 using eval::as;
 using namespace eval::tensor_function;
 
@@ -38,7 +39,7 @@ DenseReplaceTypeFunction::~DenseReplaceTypeFunction()
 }
 
 eval::InterpretedFunction::Instruction
-DenseReplaceTypeFunction::compile_self(Stash &) const
+DenseReplaceTypeFunction::compile_self(const TensorEngine &, Stash &) const
 {
     return eval::InterpretedFunction::Instruction(my_replace_type_op, (uint64_t)&(result_type()));
 }
diff --git a/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.h b/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.h
index 4ad1f4c1cee..22c3886022d 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.h
@@ -16,7 +16,7 @@ public:
     DenseReplaceTypeFunction(const eval::ValueType &result_type,
                              const eval::TensorFunction &child);
     ~DenseReplaceTypeFunction();
-    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
     bool result_is_mutable() const override { return child().result_is_mutable(); }
     static const DenseReplaceTypeFunction &create_compact(const eval::ValueType &result_type,
             const eval::TensorFunction &child,
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.cpp
index b072f334f13..3533ab20175 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.cpp
@@ -12,6 +12,7 @@ using eval::DoubleValue;
 using eval::ValueType;
 using eval::TensorSpec;
 using eval::TensorFunction;
+using eval::TensorEngine;
 using Child = eval::TensorFunction::Child;
 using eval::as;
 using namespace eval::tensor_function;
@@ -68,7 +69,7 @@ DenseTensorCreateFunction::push_children(std::vector<Child::CREF> &target) const
 }
 
 eval::InterpretedFunction::Instruction
-DenseTensorCreateFunction::compile_self(Stash &) const
+DenseTensorCreateFunction::compile_self(const TensorEngine &, Stash &) const
 {
     static_assert(sizeof(uint64_t) == sizeof(&_self));
     auto op = select_1<MyTensorCreateOp>(result_type().cell_type());
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.h b/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.h
index 6d262f48aa6..d471658fba0 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_create_function.h
@@ -25,7 +25,7 @@ public:
     ~DenseTensorCreateFunction();
     const eval::ValueType &result_type() const override { return _self.result_type; }
     void push_children(std::vector<Child::CREF> &children) const override;
-    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
     bool result_is_mutable() const override { return true; }
     static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
 };
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.cpp
index 07a2f08c423..6dc081cba1c 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.cpp
@@ -12,6 +12,7 @@ using eval::DoubleValue;
 using eval::ValueType;
 using eval::TensorSpec;
 using eval::TensorFunction;
+using eval::TensorEngine;
 using Child = eval::TensorFunction::Child;
 using eval::as;
 using namespace eval::tensor_function;
@@ -67,7 +68,7 @@ DenseTensorPeekFunction::push_children(std::vector<Child::CREF> &target) const
 }
 
 eval::InterpretedFunction::Instruction
-DenseTensorPeekFunction::compile_self(Stash &) const
+DenseTensorPeekFunction::compile_self(const TensorEngine &, Stash &) const
 {
     static_assert(sizeof(uint64_t) == sizeof(&_spec));
     auto op = select_1<MyTensorPeekOp>(_children[0].get().result_type().cell_type());
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.h b/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.h
index 150bab5211f..8ed672f95b0 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_peek_function.h
@@ -26,7 +26,7 @@ public:
     ~DenseTensorPeekFunction();
     const eval::ValueType &result_type() const override { return eval::DoubleValue::double_type(); }
     void push_children(std::vector<Child::CREF> &children) const override;
-    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
     bool result_is_mutable() const override { return true; }
     static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
 };
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 fb126b03d19..a0d63a1ce1e 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
@@ -13,6 +13,7 @@ namespace vespalib::tensor {
 
 using eval::ValueType;
 using eval::TensorFunction;
+using eval::TensorEngine;
 using eval::as;
 using eval::Aggr;
 using namespace eval::tensor_function;
@@ -156,7 +157,7 @@ DenseXWProductFunction::DenseXWProductFunction(const eval::ValueType &result_typ
 }
 
 eval::InterpretedFunction::Instruction
-DenseXWProductFunction::compile_self(Stash &stash) const
+DenseXWProductFunction::compile_self(const TensorEngine &, Stash &stash) const
 {
     Self &self = stash.create<Self>(result_type(), _vector_size, _result_size);
     auto op = my_select(lhs().result_type().cell_type(),
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 d7c39fa45a2..9f05222fff6 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
@@ -44,7 +44,7 @@ public:
     size_t result_size() const { return _result_size; }
     bool common_inner() const { return _common_inner; }
 
-    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
     void visit_self(vespalib::ObjectVisitor &visitor) const override;
     static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
 };
diff --git a/eval/src/vespa/eval/tensor/dense/vector_from_doubles_function.cpp b/eval/src/vespa/eval/tensor/dense/vector_from_doubles_function.cpp
index 4ef678a1de9..7a4b5917f00 100644
--- a/eval/src/vespa/eval/tensor/dense/vector_from_doubles_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/vector_from_doubles_function.cpp
@@ -9,6 +9,7 @@ namespace vespalib::tensor {
 using eval::Value;
 using eval::ValueType;
 using eval::TensorFunction;
+using eval::TensorEngine;
 using Child = eval::TensorFunction::Child;
 using eval::as;
 using namespace eval::tensor_function;
@@ -90,7 +91,7 @@ VectorFromDoublesFunction::push_children(std::vector<Child::CREF> &target) const
 }
 
 eval::InterpretedFunction::Instruction
-VectorFromDoublesFunction::compile_self(Stash &) const
+VectorFromDoublesFunction::compile_self(const TensorEngine &, Stash &) const
 {
     return eval::InterpretedFunction::Instruction(my_vector_from_doubles_op, (uint64_t)&_self);
 }
diff --git a/eval/src/vespa/eval/tensor/dense/vector_from_doubles_function.h b/eval/src/vespa/eval/tensor/dense/vector_from_doubles_function.h
index 378c9026f84..28346c4cb3b 100644
--- a/eval/src/vespa/eval/tensor/dense/vector_from_doubles_function.h
+++ b/eval/src/vespa/eval/tensor/dense/vector_from_doubles_function.h
@@ -30,7 +30,7 @@ public:
         return _self.resultType.dimensions()[0].name;
     }
     size_t size() const { return _self.resultSize; }
-    eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    eval::InterpretedFunction::Instruction compile_self(const eval::TensorEngine &engine, Stash &stash) const override;
     bool result_is_mutable() const override { return true; }
     static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
 };