Merge pull request #18814 from vespa-engine/havardpe/collect-compile-meta-data

collect meta-data from compile_tensor_function

5 files changed, 88 insertions, 7 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 ba73a578f6f..bcb2e29472c 100644
--- a/eval/src/tests/eval/interpreted_function/interpreted_function_test.cpp
+++ b/eval/src/tests/eval/interpreted_function/interpreted_function_test.cpp
@@ -3,8 +3,10 @@
 #include <vespa/eval/eval/fast_value.h>
 #include <vespa/eval/eval/function.h>
 #include <vespa/eval/eval/tensor_spec.h>
+#include <vespa/eval/eval/tensor_function.h>
 #include <vespa/eval/eval/operation.h>
 #include <vespa/eval/eval/interpreted_function.h>
+#include <vespa/eval/eval/compile_tensor_function.h>
 #include <vespa/eval/eval/test/eval_spec.h>
 #include <vespa/eval/eval/basic_nodes.h>
 #include <vespa/eval/eval/simple_value.h>
@@ -168,4 +170,23 @@ TEST("require that functions with non-compilable lambdas cannot be interpreted")
 
 //-----------------------------------------------------------------------------
 
+TEST("require that compilation meta-data can be collected") {
+    Stash stash;
+    const auto &x2 = tensor_function::inject(ValueType::from_spec("tensor(x[2])"), 0, stash);
+    const auto &x3 = tensor_function::inject(ValueType::from_spec("tensor(x[3])"), 1, stash);
+    const auto &concat_x5 = tensor_function::concat(x3, x2, "x", stash);
+    const auto &x5 = tensor_function::inject(ValueType::from_spec("tensor(x[5])"), 2, stash);
+    const auto &mapped_x5 = tensor_function::map(x5, operation::Relu::f, stash);
+    const auto &flag = tensor_function::inject(ValueType::from_spec("double"), 0, stash);
+    const auto &root = tensor_function::if_node(flag, concat_x5, mapped_x5, stash);
+    CTFMetaData meta;
+    InterpretedFunction ifun(FastValueBuilderFactory::get(), root, meta);
+    fprintf(stderr, "compilation meta-data:\n");
+    for (const auto &step: meta.steps) {
+        fprintf(stderr, "  %s -> %s\n", step.class_name.c_str(), step.symbol_name.c_str());        
+    }
+}
+
+//-----------------------------------------------------------------------------
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/eval/src/vespa/eval/eval/compile_tensor_function.cpp b/eval/src/vespa/eval/eval/compile_tensor_function.cpp
index 4f45aa731b5..aea2ad350b9 100644
--- a/eval/src/vespa/eval/eval/compile_tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/compile_tensor_function.cpp
@@ -3,10 +3,13 @@
 #include "compile_tensor_function.h"
 #include "tensor_function.h"
 
+#include <vespa/vespalib/util/classname.h>
+
 namespace vespalib::eval {
 
 namespace {
 
+using vespalib::getClassName;
 using State = InterpretedFunction::State;
 using Instruction = InterpretedFunction::Instruction;
 
@@ -36,7 +39,16 @@ struct ProgramCompiler {
     Stash &stash;
     std::vector<Frame> stack;
     std::vector<Instruction> prog;
-    ProgramCompiler(const ValueBuilderFactory &factory_in, Stash &stash_in) : factory(factory_in), stash(stash_in), stack(), prog() {}
+    CTFMetaData *meta;
+    ProgramCompiler(const ValueBuilderFactory &factory_in, Stash &stash_in, CTFMetaData *meta_in)
+      : factory(factory_in), stash(stash_in), stack(), prog(), meta(meta_in) {}
+    ~ProgramCompiler();
+
+    void maybe_add_meta(const TensorFunction &node, const Instruction &instr) {
+        if (meta != nullptr) {
+            meta->steps.emplace_back(getClassName(node), instr.resolve_symbol());
+        }
+    }
 
     void append(const std::vector<Instruction> &other_prog) {
         prog.insert(prog.end(), other_prog.begin(), other_prog.end());
@@ -44,9 +56,11 @@ struct ProgramCompiler {
 
     void open(const TensorFunction &node) {
         if (auto if_node = as<tensor_function::If>(node)) {
-            append(compile_tensor_function(factory, if_node->cond(), stash));
-            auto true_prog = compile_tensor_function(factory, if_node->true_child(), stash);
-            auto false_prog = compile_tensor_function(factory, if_node->false_child(), stash);
+            append(compile_tensor_function(factory, if_node->cond(), stash, meta));
+            maybe_add_meta(node, Instruction(op_skip_if_false));
+            auto true_prog = compile_tensor_function(factory, if_node->true_child(), stash, meta);
+            maybe_add_meta(node, Instruction(op_skip));
+            auto false_prog = compile_tensor_function(factory, if_node->false_child(), stash, meta);
             true_prog.emplace_back(op_skip, false_prog.size());
             prog.emplace_back(op_skip_if_false, true_prog.size());
             append(true_prog);
@@ -58,6 +72,7 @@ struct ProgramCompiler {
 
     void close(const TensorFunction &node) {
         prog.push_back(node.compile_self(factory, stash));
+        maybe_add_meta(node, prog.back());
     }
 
     std::vector<Instruction> compile(const TensorFunction &function) {
@@ -73,11 +88,14 @@ struct ProgramCompiler {
         return std::move(prog);
     }
 };
+ProgramCompiler::~ProgramCompiler() = default;
 
 } // namespace vespalib::eval::<unnamed>
 
-std::vector<Instruction> compile_tensor_function(const ValueBuilderFactory &factory, const TensorFunction &function, Stash &stash) {
-    ProgramCompiler compiler(factory, stash);
+CTFMetaData::~CTFMetaData() = default;
+
+std::vector<Instruction> compile_tensor_function(const ValueBuilderFactory &factory, const TensorFunction &function, Stash &stash, CTFMetaData *meta) {
+    ProgramCompiler compiler(factory, stash, meta);
     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 b1526f21c5e..cae30805b08 100644
--- a/eval/src/vespa/eval/eval/compile_tensor_function.h
+++ b/eval/src/vespa/eval/eval/compile_tensor_function.h
@@ -12,6 +12,27 @@ namespace vespalib::eval {
 struct ValueBuilderFactory;
 struct TensorFunction;
 
-std::vector<InterpretedFunction::Instruction> compile_tensor_function(const ValueBuilderFactory &factory, const TensorFunction &function, Stash &stash);
+/**
+ * Meta-data related to the compilation of a tensor function that may
+ * be optionally collected. Each tensor function tree node will be
+ * represented by a single 'Step' containing the class name of the
+ * corresponding tree node and the symbol name of the low-level
+ * function it compiles to. Steps are ordered according to the
+ * instructions of the final program. Note that each 'If' node will
+ * produce 2 steps; one for the conditional jump after the 'if'
+ * condition has been calculated and one for the unconditional jump
+ * after the 'true' branch.
+ **/
+struct CTFMetaData {
+    struct Step {
+        vespalib::string class_name;
+        vespalib::string symbol_name;
+    };
+    std::vector<Step> steps;
+    ~CTFMetaData();
+};
+
+std::vector<InterpretedFunction::Instruction> compile_tensor_function(const ValueBuilderFactory &factory, const TensorFunction &function, Stash &stash,
+                                                                      CTFMetaData *meta = nullptr);
 
 } // namespace vespalib::eval
diff --git a/eval/src/vespa/eval/eval/interpreted_function.cpp b/eval/src/vespa/eval/eval/interpreted_function.cpp
index 10801f276d7..a5368566aa4 100644
--- a/eval/src/vespa/eval/eval/interpreted_function.cpp
+++ b/eval/src/vespa/eval/eval/interpreted_function.cpp
@@ -9,6 +9,7 @@
 #include "compile_tensor_function.h"
 #include <vespa/vespalib/util/classname.h>
 #include <vespa/eval/eval/llvm/compile_cache.h>
+#include <vespa/eval/eval/llvm/addr_to_symbol.h>
 #include <vespa/vespalib/util/benchmark_timer.h>
 #include <set>
 
@@ -60,6 +61,15 @@ InterpretedFunction::Context::Context(const InterpretedFunction &ifun)
 {
 }
 
+vespalib::string
+InterpretedFunction::Instruction::resolve_symbol() const
+{
+    if (function == nullptr) {
+        return "<inject_param>";
+    }
+    return addr_to_symbol((const void *)function);
+}
+
 InterpretedFunction::Instruction
 InterpretedFunction::Instruction::nop()
 {
@@ -74,6 +84,14 @@ InterpretedFunction::InterpretedFunction(const ValueBuilderFactory &factory, con
     _program = compile_tensor_function(factory, function, _stash);
 }
 
+InterpretedFunction::InterpretedFunction(const ValueBuilderFactory &factory, const TensorFunction &function, CTFMetaData &meta)
+    : _program(),
+      _stash(),
+      _factory(factory)
+{
+    _program = compile_tensor_function(factory, function, _stash, &meta);
+}
+
 InterpretedFunction::InterpretedFunction(const ValueBuilderFactory &factory, const nodes::Node &root, const NodeTypes &types)
     : _program(),
       _stash(),
diff --git a/eval/src/vespa/eval/eval/interpreted_function.h b/eval/src/vespa/eval/eval/interpreted_function.h
index 829ae6e3b12..be7c37c3ab4 100644
--- a/eval/src/vespa/eval/eval/interpreted_function.h
+++ b/eval/src/vespa/eval/eval/interpreted_function.h
@@ -13,6 +13,7 @@ namespace vespalib::eval {
 namespace nodes { struct Node; }
 struct TensorFunction;
 class TensorSpec;
+class CTFMetaData;
 
 /**
  * A Function that has been prepared for execution. This will
@@ -72,6 +73,7 @@ public:
             : function(function_in), param(0) {}
         Instruction(op_function function_in, uint64_t param_in) noexcept
             : function(function_in), param(param_in) {}
+        vespalib::string resolve_symbol() const;
         void perform(State &state) const {
             if (function == nullptr) {
                 state.stack.push_back(state.params->resolve(param, state.stash));
@@ -94,6 +96,7 @@ public:
     typedef std::unique_ptr<InterpretedFunction> UP;
     // for testing; use with care; the tensor function must be kept alive
     InterpretedFunction(const ValueBuilderFactory &factory, const TensorFunction &function);
+    InterpretedFunction(const ValueBuilderFactory &factory, const TensorFunction &function, CTFMetaData &meta);
     InterpretedFunction(const ValueBuilderFactory &factory, const nodes::Node &root, const NodeTypes &types);
     InterpretedFunction(const ValueBuilderFactory &factory, const Function &function, const NodeTypes &types)
         : InterpretedFunction(factory, function.root(), types) {}