feat: add host range profiling

CMakeLists.txt

@@ -46,6 +46,8 @@ option(GENERATE_OPERATOR_CALL_INSTANTIATIONS
     "Generate explicit operator call instantiations" ON)
 option(GENERATE_PYTHON_BINDINGS "Generate Python bindings" OFF)
 option(INFINI_OPS_BUILD_DOCS "Build InfiniOps documentation" OFF)
+option(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING
+    "Enable host-side range profiling" OFF)
 
 set(_DEFAULT_HYGON_DTK_ROOT "/opt/dtk")
 

scripts/generate_public_operator_header.py

@@ -133,6 +133,7 @@ def _hide_operator_call_definition(text: str) -> str:
 
 def generate_public_operator_header(source: Path, output: Path) -> None:
     public_text = _hide_operator_call_definition(source.read_text())
+    public_text = public_text.replace('#include "host_range_profiler.h"\n', "")
     output.parent.mkdir(parents=True, exist_ok=True)
     output.write_text(public_text)
 

scripts/generate_wrappers.py

@@ -677,6 +677,8 @@ def _generate_call(op_name, call, method=True):
 
             return (
                 f'  m.def("{op_name}", []({params}) {{\n'
+                f"    [[maybe_unused]] HostRangeScope host_range_binding_body{{\n"
+                f"        HostRangeLayer::kBindingBody}};\n"
                 f"    Handle handle;\n"
                 f"    if (stream) {{\n"
                 f"      handle.set_stream(reinterpret_cast<void*>(stream));\n"
@@ -746,6 +748,7 @@ def _overload_order_key(node):
 #include "config.h"
 #include "generated/bindings/generated_dispatch.h"
 #include "handle.h"
+#include "host_range_profiler.h"
 #include "pybind11_utils.h"
 
 namespace py = pybind11;
@@ -1078,9 +1081,11 @@ def _append_optional_params(prefix, params):
 
     declarations.append(f"void ClearCacheFor{symbol_name}();")
     definitions.append(
-        f"""void ClearCacheFor{symbol_name}() {{
+        f"""#if !defined(INFINI_OPS_USE_OPERATOR_CALL_INSTANTIATIONS)
+void ClearCacheFor{symbol_name}() {{
   Operator<{op_type}>::clear_cache();
-}}"""
+}}
+#endif"""
     )
 
     emitted_make_params = set()
@@ -1129,6 +1134,8 @@ def _append_optional_params(prefix, params):
         )
         definitions.append(
             f"""void Call{symbol_name}(const Handle& handle, {_append_optional_params("const Config& config", params)}) {{
+  [[maybe_unused]] HostRangeScope host_range_dispatch_call{{
+      HostRangeLayer::kDispatchCall}};
   return Operator<{op_type}>::Call({_append_optional_args("handle, config", args)});
 }}"""
         )
@@ -1175,6 +1182,7 @@ def _generate_generated_dispatch_source(impl_paths, definitions):
     impl_includes = "\n".join(f'#include "{impl_path}"' for impl_path in impl_paths)
 
     return f"""#include "generated_dispatch.h"
+#include "host_range_profiler.h"
 
 // clang-format off
 {impl_includes}
@@ -1188,6 +1196,27 @@ def _generate_generated_dispatch_source(impl_paths, definitions):
 """
 
 
+def _generate_cache_clear_dispatch_source(op_names):
+    base_includes = "\n".join(f'#include "base/{name}.h"' for name in op_names)
+    definitions = "\n\n".join(
+        f"""void ClearCacheFor{_op_symbol_name(name)}() {{
+  Operator<{_op_cpp_type(name)}>::clear_cache();
+}}"""
+        for name in op_names
+    )
+
+    return f"""#include "operator.h"
+
+{base_includes}
+
+namespace infini::ops::generated_dispatch {{
+
+{definitions}
+
+}}  // namespace infini::ops::generated_dispatch
+"""
+
+
 def _strip_top_level_const(type_spelling):
     type_spelling = " ".join(type_spelling.split())
 
@@ -1605,6 +1634,80 @@ def _use_monolithic_bindings():
     return value.upper() in {"1", "ON", "TRUE"}
 
 
+def _generate_ops_module_source(bind_func_names, op_includes=(), monolithic=False):
+    bind_func_names = list(bind_func_names)
+    bind_func_calls = "\n".join(
+        f"{bind_func_name}(m);" for bind_func_name in bind_func_names
+    )
+    profile_bindings = """namespace {
+
+pybind11::list HostRangeSummariesToPybind11(
+    const std::vector<HostRangeSummary>& summaries) {
+  pybind11::list result;
+  for (const auto& summary : summaries) {
+    pybind11::dict row;
+    row["range"] = HostRangeLayerName(summary.layer);
+    row["count"] = summary.count;
+    row["unit"] = "ns";
+    row["inclusive_mean"] = summary.inclusive_mean;
+    row["inclusive_median"] = summary.inclusive_median;
+    row["self_mean"] = summary.self_mean;
+    row["self_median"] = summary.self_median;
+    result.append(row);
+  }
+  return result;
+}
+
+void BindHostRangeProfileControls(pybind11::module& m) {
+  m.def("_host_range_profile_compiled", &HostRangeProfiler::IsCompiled);
+  m.def("_host_range_profile_start", &HostRangeProfiler::Start);
+  m.def("_host_range_profile_stop", []() {
+    return HostRangeSummariesToPybind11(HostRangeProfiler::Stop());
+  });
+  m.def("_host_range_profile_calibrate", [](std::size_t iterations) {
+    return HostRangeSummariesToPybind11(
+        HostRangeProfiler::Calibrate(iterations));
+  }, pybind11::arg("iterations"));
+}
+
+}  // namespace
+"""
+
+    if monolithic:
+        op_includes = "\n".join(op_includes)
+        pre_namespace = f"""// Generated with `INFINI_OPS_MONOLITHIC_BINDINGS=1`.
+{op_includes}"""
+        bind_func_declarations = ""
+    else:
+        pre_namespace = ""
+        bind_func_declarations = "\n".join(
+            f"void {bind_func_name}(pybind11::module& m);"
+            for bind_func_name in bind_func_names
+        )
+
+    module_calls = "BindHostRangeProfileControls(m);"
+    if bind_func_calls:
+        module_calls = f"{module_calls}\n{bind_func_calls}"
+
+    return f"""#include <pybind11/pybind11.h>
+
+#include "host_range_profiler.h"
+
+{pre_namespace}
+
+namespace infini::ops {{
+
+{bind_func_declarations}
+
+{profile_bindings}
+PYBIND11_MODULE(ops, m) {{
+{textwrap.indent(module_calls, _INDENTATION)}
+}}
+
+}}  // namespace infini::ops
+"""
+
+
 def _dispatch_gen_batch_size():
     raw = os.environ.get("INFINI_OPS_DISPATCH_BATCH_SIZE")
 
@@ -1745,6 +1848,13 @@ def _dispatch_gen_batch_size():
     _write_text_if_changed(call_instantiation_header_path, call_instantiation_header)
     expected_include_files.add(call_instantiation_header_path)
 
+    cache_clear_dispatch_source = _generate_cache_clear_dispatch_source(op_names)
+    cache_clear_dispatch_source_path = _GENERATED_SRC_DIR / "cache_clear_dispatch.cc"
+    _write_text_if_changed(
+        cache_clear_dispatch_source_path, cache_clear_dispatch_source
+    )
+    expected_generated_src_files.add(cache_clear_dispatch_source_path)
+
     dispatch_batch_size = _dispatch_gen_batch_size()
 
     for dispatch_batch_index, start in enumerate(
@@ -1787,43 +1897,12 @@ def _dispatch_gen_batch_size():
         )
         expected_generated_src_files.add(call_instantiation_source_path)
 
-    bind_func_calls = "\n".join(
-        f"{bind_func_name}(m);" for bind_func_name in bind_func_names
+    ops_source = _generate_ops_module_source(
+        bind_func_names,
+        op_includes=op_includes,
+        monolithic=use_monolithic_bindings,
     )
 
-    if use_monolithic_bindings:
-        op_includes = "\n".join(op_includes)
-        ops_source = f"""#include <pybind11/pybind11.h>
-
-// Generated with `INFINI_OPS_MONOLITHIC_BINDINGS=1`.
-{op_includes}
-
-namespace infini::ops {{
-
-PYBIND11_MODULE(ops, m) {{
-{textwrap.indent(bind_func_calls, _INDENTATION)}
-}}
-
-}}  // namespace infini::ops
-"""
-    else:
-        bind_func_declarations = "\n".join(
-            f"void {bind_func_name}(pybind11::module& m);"
-            for bind_func_name in bind_func_names
-        )
-        ops_source = f"""#include <pybind11/pybind11.h>
-
-namespace infini::ops {{
-
-{bind_func_declarations}
-
-PYBIND11_MODULE(ops, m) {{
-{textwrap.indent(bind_func_calls, _INDENTATION)}
-}}
-
-}}  // namespace infini::ops
-"""
-
     ops_source_path = _BINDINGS_DIR / "ops.cc"
     _write_text_if_changed(ops_source_path, ops_source)
     expected_binding_files.add(ops_source_path)