A better way to get profiler state from JIT

python_bindings/halide/src/halide_/PyEnums.cpp

@@ -18,6 +18,16 @@ void define_enums(py::module &m) {
         .value("Input", Internal::ArgInfoDirection::Input)
         .value("Output", Internal::ArgInfoDirection::Output);
 
+    py::enum_<halide_profiler_func_kind>(m, "ProfilerFuncKind")
+        .value("Func", halide_profiler_func_kind_func)
+        .value("Overhead", halide_profiler_func_kind_overhead)
+        .value("ThreadIdle", halide_profiler_func_kind_thread_idle)
+        .value("Malloc", halide_profiler_func_kind_malloc)
+        .value("Free", halide_profiler_func_kind_free)
+        .value("CopyToHost", halide_profiler_func_kind_copy_to_host)
+        .value("CopyToDevice", halide_profiler_func_kind_copy_to_device)
+        .value("Allocation", halide_profiler_func_kind_allocation);
+
     py::enum_<DeviceAPI>(m, "DeviceAPI")
         .value("None", DeviceAPI::None)
         .value("Host", DeviceAPI::Host)

python_bindings/halide/src/halide_/PyPipeline.cpp

@@ -20,6 +20,65 @@ py::object realization_to_object(const Realization &r) {
     return to_python_tuple(r);
 }
 
+// Python-owned snapshots of the profiler's stats, so that they stay valid
+// after the ProfilerScope exits and the profiler resets.
+struct ProfilerFuncStats {
+    std::string name;
+    halide_profiler_func_stats stats;
+};
+
+struct ProfilerPipelineStats {
+    std::string name;
+    halide_profiler_pipeline_stats stats;
+    std::vector<ProfilerFuncStats> funcs;
+};
+
+ProfilerFuncStats snapshot(const halide_profiler_func_stats &f) {
+    return {f.name, f};
+}
+
+std::optional<ProfilerPipelineStats> snapshot(const halide_profiler_pipeline_stats *p) {
+    if (!p) {
+        return std::nullopt;
+    }
+    ProfilerPipelineStats result{p->name, *p, {}};
+    for (int i = 0; i < p->num_funcs; i++) {
+        result.funcs.push_back(snapshot(p->funcs[i]));
+    }
+    return result;
+}
+
+std::optional<ProfilerFuncStats> snapshot(const halide_profiler_func_stats *f) {
+    if (!f) {
+        return std::nullopt;
+    }
+    return snapshot(*f);
+}
+
+// Owns a ProfilerScope that a with-statement can end early, via exit(),
+// rather than waiting on garbage collection.
+struct PyProfilerScope {
+    std::unique_ptr<ProfilerScope> scope;
+
+    explicit PyProfilerScope(Pipeline p)
+        : scope(std::make_unique<ProfilerScope>(std::move(p))) {
+    }
+    explicit PyProfilerScope(Func &f)
+        : scope(std::make_unique<ProfilerScope>(f)) {
+    }
+
+    const ProfilerScope &get() const {
+        if (!scope) {
+            throw std::runtime_error("This ProfilerScope has already exited");
+        }
+        return *scope;
+    }
+
+    void exit() {
+        scope.reset();
+    }
+};
+
 }  // namespace
 
 void define_pipeline(py::module &m) {
@@ -305,6 +364,83 @@ void define_pipeline(py::module &m) {
             return create_callable_from_generator(target, name, generator_params);
         },
         py::arg("target"), py::arg("name"), py::arg("generator_params") = std::map<std::string, std::string>{});
+
+    auto func_stats_class = py::class_<ProfilerFuncStats>(m, "ProfilerFuncStats")
+                                .def_readonly("name", &ProfilerFuncStats::name)
+                                .def("__repr__", [](const ProfilerFuncStats &s) -> std::string {
+                                    return "<halide.ProfilerFuncStats " + s.name + ">";
+                                });
+#define HALIDE_PROFILER_FUNC_FIELD(field) \
+    func_stats_class.def_property_readonly(#field, [](const ProfilerFuncStats &s) { return s.stats.field; })
+    HALIDE_PROFILER_FUNC_FIELD(parent);
+    HALIDE_PROFILER_FUNC_FIELD(canonical_id);
+    HALIDE_PROFILER_FUNC_FIELD(kind);
+    HALIDE_PROFILER_FUNC_FIELD(buffer_func_id);
+    HALIDE_PROFILER_FUNC_FIELD(counters_approximated);
+    HALIDE_PROFILER_FUNC_FIELD(time);
+    HALIDE_PROFILER_FUNC_FIELD(memory_current);
+    HALIDE_PROFILER_FUNC_FIELD(memory_peak);
+    HALIDE_PROFILER_FUNC_FIELD(stack_peak);
+    HALIDE_PROFILER_FUNC_FIELD(memory_total);
+    HALIDE_PROFILER_FUNC_FIELD(active_threads_numerator);
+    HALIDE_PROFILER_FUNC_FIELD(active_threads_denominator);
+    HALIDE_PROFILER_FUNC_FIELD(num_allocs);
+    HALIDE_PROFILER_FUNC_FIELD(parallel_loops);
+    HALIDE_PROFILER_FUNC_FIELD(parallel_tasks);
+    HALIDE_PROFILER_FUNC_FIELD(points_required_at_root);
+    HALIDE_PROFILER_FUNC_FIELD(points_computed);
+    HALIDE_PROFILER_FUNC_FIELD(scalar_loads);
+    HALIDE_PROFILER_FUNC_FIELD(vector_loads);
+    HALIDE_PROFILER_FUNC_FIELD(gathers);
+    HALIDE_PROFILER_FUNC_FIELD(bytes_loaded);
+    HALIDE_PROFILER_FUNC_FIELD(scalar_stores);
+    HALIDE_PROFILER_FUNC_FIELD(vector_stores);
+    HALIDE_PROFILER_FUNC_FIELD(scatters);
+    HALIDE_PROFILER_FUNC_FIELD(bytes_stored);
+    HALIDE_PROFILER_FUNC_FIELD(realizations);
+    HALIDE_PROFILER_FUNC_FIELD(productions);
+    HALIDE_PROFILER_FUNC_FIELD(points_required_at_realization);
+    HALIDE_PROFILER_FUNC_FIELD(points_required_at_production);
+    HALIDE_PROFILER_FUNC_FIELD(points_required_inwards);
+    HALIDE_PROFILER_FUNC_FIELD(productions_if_inwards);
+#undef HALIDE_PROFILER_FUNC_FIELD
+
+    auto pipeline_stats_class = py::class_<ProfilerPipelineStats>(m, "ProfilerPipelineStats")
+                                    .def_readonly("name", &ProfilerPipelineStats::name)
+                                    .def_readonly("funcs", &ProfilerPipelineStats::funcs)
+                                    .def("__repr__", [](const ProfilerPipelineStats &s) -> std::string {
+                                        return "<halide.ProfilerPipelineStats " + s.name + ">";
+                                    });
+#define HALIDE_PROFILER_PIPELINE_FIELD(field) \
+    pipeline_stats_class.def_property_readonly(#field, [](const ProfilerPipelineStats &s) { return s.stats.field; })
+    HALIDE_PROFILER_PIPELINE_FIELD(time);
+    HALIDE_PROFILER_PIPELINE_FIELD(memory_current);
+    HALIDE_PROFILER_PIPELINE_FIELD(memory_peak);
+    HALIDE_PROFILER_PIPELINE_FIELD(memory_total);
+    HALIDE_PROFILER_PIPELINE_FIELD(active_threads_numerator);
+    HALIDE_PROFILER_PIPELINE_FIELD(active_threads_denominator);
+    HALIDE_PROFILER_PIPELINE_FIELD(native_vector_bytes);
+    HALIDE_PROFILER_PIPELINE_FIELD(runs);
+    HALIDE_PROFILER_PIPELINE_FIELD(billed_runs);
+    HALIDE_PROFILER_PIPELINE_FIELD(samples);
+    HALIDE_PROFILER_PIPELINE_FIELD(num_allocs);
+#undef HALIDE_PROFILER_PIPELINE_FIELD
+
+    py::class_<PyProfilerScope>(m, "ProfilerScope")
+        .def(py::init<Pipeline>(), py::arg("pipeline"))
+        .def(py::init<Func &>(), py::arg("func"))
+        .def("__enter__", [](PyProfilerScope &s) -> PyProfilerScope & { return s; })
+        .def("__exit__", [](PyProfilerScope &s, const py::object &exc_type, const py::object &exc_value, const py::object &exc_traceback) -> bool {
+            s.exit();
+            return false;
+        })
+        .def("exit", &PyProfilerScope::exit)
+        .def("pipeline_stats", [](const PyProfilerScope &s) {
+            return snapshot(s.get().pipeline_stats());
+        })
+        .def("func_stats", [](const PyProfilerScope &s, const Func &f) { return snapshot(s.get().func_stats(f)); }, py::arg("func"))
+        .def("func_stats", [](const PyProfilerScope &s, const std::string &name) { return snapshot(s.get().func_stats(name)); }, py::arg("name"))
+        .def("__repr__", [](const PyProfilerScope &s) -> std::string { return "<halide.ProfilerScope>"; });
 }
 
 }  // namespace PythonBindings

python_bindings/halide/test/correctness/CMakeLists.txt

@@ -24,6 +24,7 @@ set(tests
     memoize.py
     multi_method_module_test.py
     multipass_constraints.py
+    profiler_scope.py
     pystub.py
     rdom.py
     realize_warnings.py

python_bindings/halide/test/correctness/profiler_scope.py

@@ -0,0 +1,62 @@
+import halide as hl
+
+
+def test_profiler_scope():
+    target = hl.get_jit_target_from_environment()
+    if target.arch == hl.TargetArch.WebAssembly:
+        print("[SKIP] Profiler state is not accessible under WebAssembly.")
+        return
+    target = target.with_feature(hl.TargetFeature.Profile)
+
+    x, y = hl.Var("x"), hl.Var("y")
+    g = hl.Func("g_profiled")
+    f = hl.Func("f_profiled")
+    g[x, y] = x + y
+    f[x, y] = g[x, y] * 2
+    g.compute_root()
+
+    size = 256
+    with hl.ProfilerScope(f) as scope:
+        assert scope.pipeline_stats() is None
+
+        for _ in range(3):
+            f.realize([size, size], target)
+
+        p = scope.pipeline_stats()
+        assert p is not None
+        assert p.runs == 3
+        assert p.name == f.name()
+        assert any(fs.name == g.name() for fs in p.funcs)
+
+        gs = scope.func_stats(g)
+        assert gs is not None
+        assert gs.kind == hl.ProfilerFuncKind.Func
+        assert gs.num_allocs == 3
+        assert gs.memory_peak == size * size * 4
+        assert gs.memory_total == 3 * size * size * 4
+        assert scope.func_stats(g.name()).num_allocs == gs.num_allocs
+        assert scope.func_stats("nonexistent") is None
+
+    # Snapshots outlive the scope.
+    assert gs.num_allocs == 3
+
+    try:
+        scope.pipeline_stats()
+        raise AssertionError("Expected an error after the scope exited")
+    except RuntimeError:
+        pass
+
+    # Explicitly constructing a Pipeline works too, and the scope's exit
+    # reset the stats.
+    pipe = hl.Pipeline(f)
+    with hl.ProfilerScope(pipe) as scope:
+        pipe.realize([size, size], target)
+        assert scope.pipeline_stats().runs == 1
+
+
+def main():
+    test_profiler_scope()
+
+
+if __name__ == "__main__":
+    main()

src/Func.h

@@ -794,6 +794,8 @@ class Func {
      * creating it (and freezing the Func) if necessary. */
     Pipeline pipeline();
 
+    friend class ProfilerScope;
+
     // Helper function for recursive reordering support
     Func &reorder_storage(const std::vector<Var> &dims, size_t start);
 

src/Pipeline.cpp

@@ -151,6 +151,10 @@ struct PipelineContents {
 
     bool trace_pipeline = false;
 
+    /** The number of live ProfilerScopes for this pipeline. While
+     * nonzero, realize leaves the profiler's statistics in place. */
+    int profiler_scopes = 0;
+
     /** Optional prefixes used to rename halide_-prefixed runtime symbols.
      * Empty unless set via Pipeline::apply_runtime_prefixes(). */
     RuntimePrefixParams runtime_prefixes_params;
@@ -837,7 +841,9 @@ Realization Pipeline::realize(JITUserContext *context,
     }
 
     // If we're profiling, report runtimes and reset profiler stats.
-    contents->jit_cache.finish_profiling(context);
+    if (contents->profiler_scopes == 0) {
+        contents->jit_cache.finish_profiling(context);
+    }
     jit_context.finalize(exit_status);
 
     // Crop back to the requested size if necessary
@@ -900,6 +906,88 @@ void Pipeline::trace_pipeline() {
     contents->trace_pipeline = true;
 }
 
+ProfilerScope::ProfilerScope(Pipeline p)
+    : pipeline(std::move(p)) {
+    user_assert(pipeline.defined()) << "Pipeline is undefined\n";
+    pipeline.contents->profiler_scopes++;
+}
+
+ProfilerScope::ProfilerScope(Func &f)
+    : ProfilerScope(f.pipeline()) {
+}
+
+ProfilerScope::~ProfilerScope() {
+    if (--pipeline.contents->profiler_scopes > 0) {
+        return;
+    }
+    // Report and reset as a realize outside of any scope would have.
+    JITUserContext context{};
+    JITFuncCallContext jit_context(&context, pipeline.jit_handlers());
+    pipeline.contents->jit_cache.finish_profiling(&context);
+    jit_context.finalize(0);
+}
+
+const halide_profiler_pipeline_stats *ProfilerScope::pipeline_stats() const {
+    const JITCache &cache = pipeline.contents->jit_cache;
+    if (!cache.jit_target.has_feature(Target::Profile) &&
+        !cache.jit_target.has_feature(Target::ProfileByTimer)) {
+        return nullptr;
+    }
+    // The profiler lives in the shared JIT runtime, which the wasm
+    // module does not link against, so the symbols may not exist.
+    using GetStateFn = halide_profiler_state *(*)();
+    using LockFn = void (*)(halide_profiler_state *);
+    auto find = [&](const char *symbol) {
+        return cache.jit_module.find_symbol_by_name(symbol).address;
+    };
+    auto get_state = (GetStateFn)find("halide_profiler_get_state");
+    auto lock = (LockFn)find("halide_profiler_lock");
+    auto unlock = (LockFn)find("halide_profiler_unlock");
+    if (!get_state || !lock || !unlock) {
+        return nullptr;
+    }
+
+    // halide_profiler_get_pipeline_state compares names by pointer, so
+    // walk the list comparing by string instead. Recompiling the
+    // pipeline produces a new entry with the same name; the newest is
+    // at the head of the list.
+    const std::string name = pipeline.generate_function_name();
+    halide_profiler_state *state = get_state();
+    const halide_profiler_pipeline_stats *result = nullptr;
+    lock(state);
+    for (const halide_profiler_pipeline_stats *p = state->pipelines; p;
+         p = (const halide_profiler_pipeline_stats *)p->next) {
+        if (name == p->name) {
+            result = p;
+            break;
+        }
+    }
+    unlock(state);
+    return result;
+}
+
+const halide_profiler_func_stats *ProfilerScope::func_stats(const std::string &name) const {
+    const halide_profiler_pipeline_stats *p = pipeline_stats();
+    if (!p) {
+        return nullptr;
+    }
+    for (int i = 0; i < p->num_funcs; i++) {
+        const halide_profiler_func_stats &f = p->funcs[i];
+        if (f.kind == halide_profiler_func_kind_func &&
+            f.canonical_id == i &&
+            name == f.name) {
+            return &f;
+        }
+    }
+    return nullptr;
+}
+
+const halide_profiler_func_stats *ProfilerScope::func_stats(const Func &f) const {
+    // The profiler reports a Func under its display name if it has one.
+    const std::string &display_name = f.function().profiler_display_name();
+    return func_stats(display_name.empty() ? f.name() : display_name);
+}
+
 // Make a vector of void *'s to pass to the jit call using the
 // currently bound value for all of the params and image
 // params.
@@ -1114,7 +1202,9 @@ void Pipeline::realize(JITUserContext *context,
     debug(2) << "Back from jitted function. Exit status was " << exit_status << "\n";
 
     // If we're profiling, report runtimes and reset profiler stats.
-    contents->jit_cache.finish_profiling(context);
+    if (contents->profiler_scopes == 0) {
+        contents->jit_cache.finish_profiling(context);
+    }
 
     jit_call_context.finalize(exit_status);
 }