Merge pull request #23765 from vespa-engine/havardpe/rank-program-profiling-support

enable profiling of rank programs

3 files changed, 72 insertions, 4 deletions

diff --git a/searchlib/src/tests/fef/rank_program/rank_program_test.cpp b/searchlib/src/tests/fef/rank_program/rank_program_test.cpp
index 40e1cee3629..327f88ae22b 100644
--- a/searchlib/src/tests/fef/rank_program/rank_program_test.cpp
+++ b/searchlib/src/tests/fef/rank_program/rank_program_test.cpp
@@ -13,10 +13,14 @@
 #include <vespa/searchlib/fef/test/plugin/double.h>
 #include <vespa/searchlib/fef/rank_program.h>
 #include <vespa/searchlib/fef/test/test_features.h>
+#include <vespa/vespalib/util/execution_profiler.h>
+#include <vespa/vespalib/data/slime/slime.h>
 
 using namespace search::fef;
 using namespace search::fef::test;
 using namespace search::features;
+using vespalib::ExecutionProfiler;
+using vespalib::Slime;
 
 uint32_t default_docid = 1;
 
@@ -109,12 +113,12 @@ struct Fixture {
         overrides.add(feature, vespalib::make_string("%g", value));
         return *this;
     }
-    Fixture &compile() {
+    Fixture &compile(ExecutionProfiler *profiler = nullptr) {
         ASSERT_TRUE(resolver->compile());
         MatchDataLayout mdl;
         QueryEnvironment queryEnv(&indexEnv);
         match_data = mdl.createMatchData();
-        program.setup(*match_data, queryEnv, overrides);
+        program.setup(*match_data, queryEnv, overrides, profiler);
         return *this;
     }
     vespalib::string final_executor_name() const {
@@ -391,4 +395,31 @@ TEST_F("require that fast-forest gbdt evaluation is pure", Fixture()) {
     EXPECT_EQUAL(f1.final_executor_name(), "search::features::FastForestExecutor");
 }
 
+TEST_F("require that rank program can be profiled", Fixture()) {
+    ExecutionProfiler profiler(64);
+    f1.add("mysum(value(10),ivalue(5))").compile(&profiler);
+    EXPECT_EQUAL(3u, f1.program.num_executors());
+    EXPECT_EQUAL(3u, count_features(f1.program));
+    EXPECT_EQUAL(1u, count_const_features(f1.program));
+    EXPECT_EQUAL(15.0, f1.get(1));
+    EXPECT_EQUAL(15.0, f1.get(2));
+    EXPECT_EQUAL(15.0, f1.get(3));
+    Slime slime;
+    profiler.report(slime.setObject());
+    fprintf(stderr, "%s", slime.toString().c_str());
+    EXPECT_EQUAL(slime["roots"].entries(), 2u);
+    auto *a = &slime["roots"][0];
+    auto *b = &slime["roots"][1];
+    if ((*b)["count"].asLong() > (*a)["count"].asLong()) {
+        std::swap(a, b);
+    }
+    EXPECT_EQUAL((*a)["name"].asString().make_string(), vespalib::string("mysum(value(10),ivalue(5))"));
+    EXPECT_EQUAL((*a)["count"].asLong(), 3);
+    EXPECT_EQUAL((*a)["children"].entries(), 1u);
+    EXPECT_EQUAL((*a)["children"][0]["name"].asString().make_string(), vespalib::string("ivalue(5)"));
+    EXPECT_EQUAL((*a)["children"][0]["count"].asLong(), 3);
+    EXPECT_EQUAL((*b)["name"].asString().make_string(), vespalib::string("value(10)"));
+    EXPECT_EQUAL((*b)["count"].asLong(), 1);
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/searchlib/src/vespa/searchlib/fef/rank_program.cpp b/searchlib/src/vespa/searchlib/fef/rank_program.cpp
index c8c2d59e872..ec4a22ea10f 100644
--- a/searchlib/src/vespa/searchlib/fef/rank_program.cpp
+++ b/searchlib/src/vespa/searchlib/fef/rank_program.cpp
@@ -8,6 +8,7 @@
 #include <vespa/vespalib/stllike/hash_set.hpp>
 #include <vespa/vespalib/util/size_literals.h>
 #include <vespa/vespalib/util/issue.h>
+#include <vespa/vespalib/util/execution_profiler.h>
 #include <algorithm>
 #include <cassert>
 
@@ -16,6 +17,7 @@ LOG_SETUP(".fef.rankprogram");
 
 using vespalib::Stash;
 using vespalib::Issue;
+using vespalib::ExecutionProfiler;
 using vespalib::eval::Value;
 using vespalib::eval::ValueType;
 using vespalib::eval::FastValueBuilderFactory;
@@ -101,6 +103,33 @@ struct UnboxingExecutor : FeatureExecutor {
     }
 };
 
+struct ProfiledExecutor : FeatureExecutor {
+    ExecutionProfiler &profiler;
+    FeatureExecutor &executor;
+    ExecutionProfiler::TaskId self;
+    ProfiledExecutor(ExecutionProfiler &profiler_in,
+                     FeatureExecutor &executor_in,
+                     const vespalib::string &name)
+      : profiler(profiler_in), executor(executor_in), self(profiler.resolve(name)) {}
+    void handle_bind_match_data(const MatchData &md) override {
+        executor.bind_match_data(md);
+    }
+    void handle_bind_inputs(vespalib::ConstArrayRef<LazyValue> inputs) override {
+        executor.bind_inputs(inputs);
+    }
+    void handle_bind_outputs(vespalib::ArrayRef<NumberOrObject> outputs) override {
+        executor.bind_outputs(outputs);
+    }
+    bool isPure() override {
+        return executor.isPure();
+    }
+    void execute(uint32_t docId) override {
+        profiler.start(self);
+        executor.lazy_execute(docId);
+        profiler.complete();
+    }
+};
+
 class StashSelector {
 private:
     Stash &_primary;
@@ -203,7 +232,8 @@ RankProgram::~RankProgram() = default;
 void
 RankProgram::setup(const MatchData &md,
                    const IQueryEnvironment &queryEnv,
-                   const Properties &featureOverrides)
+                   const Properties &featureOverrides,
+                   ExecutionProfiler *profiler)
 {
     const auto &specs = _resolver->getExecutorSpecs();
     assert(_executors.empty());
@@ -239,6 +269,10 @@ RankProgram::setup(const MatchData &md,
             FeatureExecutor *tmp = executor;
             executor = &(stash.get().create<FeatureOverrider>(*tmp, override->ref.output, override->number, std::move(override->object)));
         }
+        if (profiler) {
+            FeatureExecutor *tmp = executor;
+            executor = &(stash.get().create<ProfiledExecutor>(*profiler, *tmp, specs[i].blueprint->getName()));
+        }
         executor->bind_inputs(inputs);
         executor->bind_outputs(outputs);
         executor->bind_match_data(md);
diff --git a/searchlib/src/vespa/searchlib/fef/rank_program.h b/searchlib/src/vespa/searchlib/fef/rank_program.h
index f3c7541adbd..aa6446e44ab 100644
--- a/searchlib/src/vespa/searchlib/fef/rank_program.h
+++ b/searchlib/src/vespa/searchlib/fef/rank_program.h
@@ -12,6 +12,8 @@
 #include <vespa/vespalib/util/stash.h>
 #include <vespa/vespalib/stllike/hash_set.h>
 
+namespace vespalib { class ExecutionProfiler; }
+
 namespace search::fef {
 
 /**
@@ -68,7 +70,8 @@ public:
      **/
     void setup(const MatchData &md,
                const IQueryEnvironment &queryEnv,
-               const Properties &featureOverrides = Properties());
+               const Properties &featureOverrides = Properties(),
+               vespalib::ExecutionProfiler *profiler = nullptr);
 
     /**
      * Obtain the names and storage locations of all seed features for