diff --git a/CMakeLists.txt b/CMakeLists.txt index ec92727cf..293b15dce 100644 --- a/CMakeLists.txt +++ b/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") diff --git a/scripts/generate_public_operator_header.py b/scripts/generate_public_operator_header.py index 19577a3b2..eb6a3add4 100644 --- a/scripts/generate_public_operator_header.py +++ b/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) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 89e31016f..44044461e 100644 --- a/scripts/generate_wrappers.py +++ b/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(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& 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 + +#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 - -// 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 - -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) diff --git a/scripts/run_host_overhead_control.py b/scripts/run_host_overhead_control.py new file mode 100644 index 000000000..61c8c8e2f --- /dev/null +++ b/scripts/run_host_overhead_control.py @@ -0,0 +1,193 @@ +"""Measure InfiniOps host submission with optional host-range collection.""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import statistics +import time + + +def measure_host_submission( + callback, + synchronize, + *, + warmup_iterations, + iterations, + rounds, + before_round=None, + after_round=None, + timer_ns=time.perf_counter_ns, +): + """Return per-call host timing with synchronization outside each interval.""" + if warmup_iterations < 0: + raise ValueError("warmup_iterations must be non-negative") + + if iterations <= 0: + raise ValueError("iterations must be positive") + + if rounds <= 0: + raise ValueError("rounds must be positive") + + for _ in range(warmup_iterations): + callback() + + synchronize() + + samples = [] + + for _ in range(rounds): + synchronize() + if before_round is not None: + before_round() + + try: + start = timer_ns() + + for _ in range(iterations): + callback() + + elapsed = timer_ns() - start + finally: + try: + if after_round is not None: + after_round() + finally: + synchronize() + + samples.append(elapsed / iterations) + + return { + "warmup_iterations": warmup_iterations, + "iterations_per_round": iterations, + "rounds": rounds, + "unit": "ns", + "mean": statistics.fmean(samples), + "median": statistics.median(samples), + "samples": samples, + } + + +def _nvidia_cases(device_name): + import infini.ops as ops + import torch + + device = torch.device(device_name) + if device.type != "cuda": + raise ValueError("the host-overhead control currently supports CUDA only") + + stream = torch.cuda.current_stream(device).cuda_stream + + add_input = torch.randn((13, 4), dtype=torch.float32, device=device) + add_other = torch.randn_like(add_input) + add_output = torch.empty_like(add_input) + + def add(): + ops.add( + add_input, + add_other, + add_output, + stream=stream, + implementation_index=0, + ) + + gemm_a = torch.randn((4, 48, 64), dtype=torch.float32, device=device) + gemm_b = torch.randn((4, 64, 6), dtype=torch.float32, device=device) + gemm_c = torch.empty((4, 48, 6), dtype=torch.float32, device=device) + + def gemm(): + ops.gemm( + gemm_a, + gemm_b, + 1.0, + 0.0, + False, + False, + gemm_c, + stream=stream, + implementation_index=1, + ) + + compiled = getattr(ops, "_host_range_profile_compiled", lambda: False)() + cases = ( + ( + "add", + { + "shape": [13, 4], + "dtype": "float32", + "implementation_index": 0, + }, + add, + ), + ( + "gemm", + { + "a_shape": [4, 48, 64], + "b_shape": [4, 64, 6], + "c_shape": [4, 48, 6], + "dtype": "float32", + "implementation_index": 1, + }, + gemm, + ), + ) + + return ops, torch, bool(compiled), cases + + +def _parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=pathlib.Path, required=True) + parser.add_argument("--label", required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--activate-profile", action="store_true") + parser.add_argument("--warmup-iterations", type=int, default=2_000) + parser.add_argument("--iterations", type=int, default=5_000) + parser.add_argument("--rounds", type=int, default=7) + + return parser.parse_args() + + +def main(): + args = _parse_args() + ops, torch, profiling_compiled, cases = _nvidia_cases(args.device) + if args.activate_profile and not profiling_compiled: + raise RuntimeError("cannot activate profiling in a profiling-disabled build") + + before_round = ops._host_range_profile_start if args.activate_profile else None + after_round = ops._host_range_profile_stop if args.activate_profile else None + + results = [] + for benchmark, params, callback in cases: + result = measure_host_submission( + callback, + lambda: torch.cuda.synchronize(args.device), + warmup_iterations=args.warmup_iterations, + iterations=args.iterations, + rounds=args.rounds, + before_round=before_round, + after_round=after_round, + ) + results.append({"benchmark": benchmark, "params": params, **result}) + + report = { + "label": args.label, + "profiling_compiled": profiling_compiled, + "device": args.device, + "measurement": ( + "host submission loop timed with time.perf_counter_ns; " + "device synchronization occurs only outside each timed interval" + ), + "profiling_active": args.activate_profile, + "results": results, + } + encoded = json.dumps(report, indent=2, sort_keys=True) + "\n" + + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(encoded, encoding="utf-8") + print(encoded, end="") + + +if __name__ == "__main__": + main() diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 138312707..32b408083 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -6,6 +6,16 @@ target_compile_options(infiniops PRIVATE "$<$:-Wno-deprecated-declarations>" ) +if(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) + target_compile_definitions(infiniops PUBLIC + INFINI_OPS_ENABLE_HOST_RANGE_PROFILING=1) +endif() + +if(GENERATE_OPERATOR_CALL_INSTANTIATIONS) + target_compile_definitions(infiniops PRIVATE + INFINI_OPS_USE_OPERATOR_CALL_INSTANTIATIONS=1) +endif() + function(_infini_ops_write_if_different path content) set(_should_write TRUE) if(EXISTS "${path}") @@ -657,6 +667,10 @@ if(WITH_TORCH) if(WITH_CPU) list(APPEND _torch_extra_flags "-DWITH_CPU=1") endif() + if(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) + list(APPEND _torch_extra_flags + "-DINFINI_OPS_ENABLE_HOST_RANGE_PROFILING=1") + endif() if(DEFINED TORCH_CXX11_ABI) list(APPEND _torch_extra_flags "-D_GLIBCXX_USE_CXX11_ABI=${TORCH_CXX11_ABI}") endif() @@ -731,6 +745,12 @@ target_include_directories(infiniops $ ) +set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/operator.h" + "${PROJECT_SOURCE_DIR}/scripts/generate_public_operator_header.py" + "${PROJECT_SOURCE_DIR}/scripts/generate_wrappers.py" +) + if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS) find_package(Python COMPONENTS Interpreter REQUIRED) # Always regenerate wrappers so emitted call instantiations and bindings @@ -785,6 +805,9 @@ configure_file( ) if(GENERATE_OPERATOR_CALL_INSTANTIATIONS) + target_sources(infiniops PRIVATE + "${PROJECT_SOURCE_DIR}/generated/src/cache_clear_dispatch.cc") + file(GLOB_RECURSE OPERATOR_CALL_INSTANTIATION_SOURCES CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/generated/src/operator_call_instantiations_*.cc") @@ -803,7 +826,13 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS) list(APPEND _iluvatar_call_instantiation_include_flags "-I${_dir}") endforeach() - set(_iluvatar_call_instantiation_defs -DWITH_ILUVATAR=1) + set(_iluvatar_call_instantiation_defs + -DWITH_ILUVATAR=1 + -DINFINI_OPS_USE_OPERATOR_CALL_INSTANTIATIONS=1) + if(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) + list(APPEND _iluvatar_call_instantiation_defs + -DINFINI_OPS_ENABLE_HOST_RANGE_PROFILING=1) + endif() if(WITH_CPU) list(APPEND _iluvatar_call_instantiation_defs -DWITH_CPU=1) endif() @@ -979,6 +1008,14 @@ if(GENERATE_PYTHON_BINDINGS) endforeach() set(_iluvatar_dispatch_defs -DWITH_ILUVATAR=1) + if(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) + list(APPEND _iluvatar_dispatch_defs + -DINFINI_OPS_ENABLE_HOST_RANGE_PROFILING=1) + endif() + if(GENERATE_OPERATOR_CALL_INSTANTIATIONS) + list(APPEND _iluvatar_dispatch_defs + -DINFINI_OPS_USE_OPERATOR_CALL_INSTANTIATIONS=1) + endif() if(WITH_TORCH) list(APPEND _iluvatar_dispatch_defs -DWITH_TORCH=1) endif() @@ -1046,6 +1083,11 @@ if(GENERATE_PYTHON_BINDINGS) "$<$:-Wno-deprecated-declarations>" ) + if(GENERATE_OPERATOR_CALL_INSTANTIATIONS) + target_compile_definitions(ops PRIVATE + INFINI_OPS_USE_OPERATOR_CALL_INSTANTIATIONS=1) + endif() + if(WITH_TORCH AND CMAKE_GENERATOR MATCHES "Ninja") set(INFINI_OPS_BINDING_COMPILE_JOBS "2" CACHE STRING "Maximum concurrent generated pybind11 binding compilations") @@ -1153,7 +1195,8 @@ endif() file(GLOB INFINI_OPS_PUBLIC_CORE_HEADERS CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/*.h") list(REMOVE_ITEM INFINI_OPS_PUBLIC_CORE_HEADERS - "${CMAKE_CURRENT_SOURCE_DIR}/operator.h") + "${CMAKE_CURRENT_SOURCE_DIR}/operator.h" + "${CMAKE_CURRENT_SOURCE_DIR}/host_range_profiler.h") install(FILES ${PROJECT_SOURCE_DIR}/generated/include/operator.h diff --git a/src/host_range_profiler.cc b/src/host_range_profiler.cc new file mode 100644 index 000000000..c18d5f59d --- /dev/null +++ b/src/host_range_profiler.cc @@ -0,0 +1,281 @@ +#include "host_range_profiler.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace infini::ops { +namespace { + +constexpr std::array(HostRangeLayer::kCount)> + kLayerNames{ + "binding.body", + "binding.tensor_conversion", + "binding.device_conversion", + "dispatch.call", + "operator.call", + "cache.key", + "cache.lookup", + "cache.construct", + "operator.invoke", + "backend.submit", + "calibration.depth1", + "calibration.depth2", + "calibration.depth3", + }; + +std::size_t LayerIndex(HostRangeLayer layer) { + const auto index = static_cast(layer); + if (index >= kLayerNames.size()) { + throw std::runtime_error("invalid host range layer"); + } + return index; +} + +} // namespace + +const char* HostRangeLayerName(HostRangeLayer layer) { + return kLayerNames[LayerIndex(layer)]; +} + +#if defined(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) + +namespace { + +using Clock = std::chrono::steady_clock; +using DurationNs = std::uint64_t; + +constexpr std::size_t kInitialStackCapacity = 16; +constexpr std::size_t kInitialSampleCapacity = 256; + +struct Frame { + HostRangeLayer layer; + Clock::time_point start; + DurationNs child_duration{0}; +}; + +struct LayerSamples { + std::vector inclusive; + std::vector self; +}; + +struct CollectorState { + bool active{false}; + bool recording_failed{false}; + std::uint64_t session_id{0}; + std::vector stack; + std::array samples; +}; + +thread_local CollectorState collector_state; +thread_local std::uint64_t next_session_id{0}; + +double Mean(const std::vector& samples) { + long double total = 0; + for (const auto sample : samples) { + total += sample; + } + return static_cast(total / samples.size()); +} + +double Median(std::vector samples) { + std::sort(samples.begin(), samples.end()); + const auto middle = samples.size() / 2; + if (samples.size() % 2 != 0) { + return static_cast(samples[middle]); + } + return static_cast( + (static_cast(samples[middle - 1]) + samples[middle]) / 2.0L); +} + +void ReserveCollectorStorage(CollectorState& state) { + state.stack.reserve(kInitialStackCapacity); + for (auto& samples : state.samples) { + samples.inclusive.reserve(kInitialSampleCapacity); + samples.self.reserve(kInitialSampleCapacity); + } +} + +} // namespace + +bool HostRangeProfiler::IsCompiled() { return true; } + +void HostRangeProfiler::Start() { + if (collector_state.active) { + throw std::runtime_error("host range profiler is already active"); + } + + CollectorState state; + ReserveCollectorStorage(state); + if (next_session_id == std::numeric_limits::max()) { + throw std::runtime_error("host range profiler session IDs are exhausted"); + } + state.session_id = ++next_session_id; + state.active = true; + collector_state = std::move(state); +} + +std::vector HostRangeProfiler::Stop() { + if (!collector_state.active) { + throw std::runtime_error("host range profiler is not active"); + } + + CollectorState completed_state = std::move(collector_state); + collector_state = CollectorState{}; + + if (!completed_state.stack.empty()) { + throw std::runtime_error("host range profiler has unclosed ranges"); + } + if (completed_state.recording_failed) { + throw std::runtime_error("host range profiler failed to record a sample"); + } + + std::vector summaries; + summaries.reserve(kLayerNames.size()); + for (std::size_t index = 0; index < completed_state.samples.size(); ++index) { + const auto& samples = completed_state.samples[index]; + if (samples.inclusive.empty()) { + continue; + } + if (samples.inclusive.size() != samples.self.size()) { + throw std::runtime_error("host range profiler sample counts differ"); + } + + summaries.push_back(HostRangeSummary{ + static_cast(index), + samples.inclusive.size(), + Mean(samples.inclusive), + Median(samples.inclusive), + Mean(samples.self), + Median(samples.self), + }); + } + + return summaries; +} + +std::vector HostRangeProfiler::Calibrate( + std::size_t iterations) { + if (collector_state.active) { + throw std::runtime_error( + "cannot calibrate while the host range profiler is active"); + } + if (iterations == 0) { + throw std::runtime_error("host range calibration requires iterations"); + } + + Start(); + try { + for (std::size_t iteration = 0; iteration < iterations; ++iteration) { + HostRangeScope depth1{HostRangeLayer::kCalibrationDepth1}; + HostRangeScope depth2{HostRangeLayer::kCalibrationDepth2}; + HostRangeScope depth3{HostRangeLayer::kCalibrationDepth3}; + } + return Stop(); + } catch (...) { + collector_state = CollectorState{}; + throw; + } +} + +HostRangeScope::HostRangeScope(HostRangeLayer layer) { + if (!collector_state.active) { + return; + } + + LayerIndex(layer); + collector_state.stack.push_back(Frame{layer, Clock::now(), 0}); + owner_ = &collector_state; + session_id_ = collector_state.session_id; + layer_ = layer; + depth_ = collector_state.stack.size(); + active_ = true; +} + +HostRangeScope::~HostRangeScope() noexcept { + if (!active_) { + return; + } + if (owner_ != &collector_state) { + return; + } + if (!collector_state.active || collector_state.session_id != session_id_) { + return; + } + if (collector_state.stack.empty()) { + collector_state.recording_failed = true; + return; + } + if (collector_state.stack.size() != depth_ || + collector_state.stack.back().layer != layer_) { + collector_state.recording_failed = true; + return; + } + + const auto end = Clock::now(); + const auto frame = collector_state.stack.back(); + collector_state.stack.pop_back(); + const auto elapsed = + std::chrono::duration_cast(end - frame.start) + .count(); + const auto inclusive = + elapsed > 0 ? static_cast(elapsed) : DurationNs{0}; + const auto self = inclusive >= frame.child_duration + ? inclusive - frame.child_duration + : DurationNs{0}; + + if (!collector_state.stack.empty()) { + auto& parent_child_duration = collector_state.stack.back().child_duration; + if (inclusive > + std::numeric_limits::max() - parent_child_duration) { + parent_child_duration = std::numeric_limits::max(); + } else { + parent_child_duration += inclusive; + } + } + + auto& samples = + collector_state.samples[static_cast(frame.layer)]; + try { + samples.inclusive.push_back(inclusive); + try { + samples.self.push_back(self); + } catch (...) { + samples.inclusive.pop_back(); + throw; + } + } catch (...) { + collector_state.recording_failed = true; + } +} + +#else + +namespace { + +[[noreturn]] void ThrowNotCompiled() { + throw std::runtime_error("host range profiling is not compiled"); +} + +} // namespace + +bool HostRangeProfiler::IsCompiled() { return false; } + +void HostRangeProfiler::Start() { ThrowNotCompiled(); } + +std::vector HostRangeProfiler::Stop() { ThrowNotCompiled(); } + +std::vector HostRangeProfiler::Calibrate( + std::size_t iterations) { + (void)iterations; + ThrowNotCompiled(); +} + +#endif + +} // namespace infini::ops diff --git a/src/host_range_profiler.h b/src/host_range_profiler.h new file mode 100644 index 000000000..caf71af4f --- /dev/null +++ b/src/host_range_profiler.h @@ -0,0 +1,79 @@ +#pragma once + +#include +#include +#include + +namespace infini::ops { + +enum class HostRangeLayer { + kBindingBody, + kTensorConversion, + kDeviceConversion, + kDispatchCall, + kOperatorCall, + kCacheKey, + kCacheLookup, + kCacheConstruct, + kOperatorInvoke, + kBackendSubmit, + kCalibrationDepth1, + kCalibrationDepth2, + kCalibrationDepth3, + kCount, +}; + +const char* HostRangeLayerName(HostRangeLayer layer); + +struct HostRangeSummary { + HostRangeLayer layer; + std::size_t count; + double inclusive_mean; + double inclusive_median; + double self_mean; + double self_median; +}; + +class HostRangeProfiler { + public: + static bool IsCompiled(); + static void Start(); + static std::vector Stop(); + static std::vector Calibrate(std::size_t iterations); +}; + +#if defined(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) + +class HostRangeScope { + public: + explicit HostRangeScope(HostRangeLayer layer); + ~HostRangeScope() noexcept; + + HostRangeScope(const HostRangeScope&) = delete; + HostRangeScope& operator=(const HostRangeScope&) = delete; + HostRangeScope(HostRangeScope&&) = delete; + HostRangeScope& operator=(HostRangeScope&&) = delete; + + private: + bool active_{false}; + const void* owner_{nullptr}; + std::uint64_t session_id_{0}; + HostRangeLayer layer_{HostRangeLayer::kCount}; + std::size_t depth_{0}; +}; + +#else + +class HostRangeScope { + public: + explicit constexpr HostRangeScope(HostRangeLayer) noexcept {} + + HostRangeScope(const HostRangeScope&) = delete; + HostRangeScope& operator=(const HostRangeScope&) = delete; + HostRangeScope(HostRangeScope&&) = delete; + HostRangeScope& operator=(HostRangeScope&&) = delete; +}; + +#endif + +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/gemm/cublaslt.h b/src/native/cuda/nvidia/ops/gemm/cublaslt.h index fa83b827c..4728ca04c 100644 --- a/src/native/cuda/nvidia/ops/gemm/cublaslt.h +++ b/src/native/cuda/nvidia/ops/gemm/cublaslt.h @@ -9,6 +9,7 @@ // clang-format on #include "base/gemm.h" +#include "host_range_profiler.h" #include "native/cuda/nvidia/blas_utils.h" #include "native/cuda/nvidia/runtime_.h" @@ -41,6 +42,9 @@ class Operator : public Gemm { void operator()(const Tensor a, const Tensor b, std::optional alpha, std::optional beta, std::optional trans_a, std::optional trans_b, Tensor c) const override { + [[maybe_unused]] HostRangeScope host_range_backend_submit{ + HostRangeLayer::kBackendSubmit}; + const auto alpha_value{alpha.value_or(alpha_)}; const auto beta_value{beta.value_or(beta_)}; const auto trans_a_value{trans_a.value_or(trans_a_)}; diff --git a/src/native/cuda/ops/add/kernel.h b/src/native/cuda/ops/add/kernel.h index 65d0ea7fb..0125cd3dd 100644 --- a/src/native/cuda/ops/add/kernel.h +++ b/src/native/cuda/ops/add/kernel.h @@ -8,6 +8,7 @@ #include "base/add.h" #include "common/generic_utils.h" +#include "host_range_profiler.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/add/kernel.cuh" #include "native/cuda/runtime_utils.h" @@ -62,6 +63,9 @@ class CudaAdd : public Add { void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const override { + [[maybe_unused]] HostRangeScope host_range_backend_submit{ + HostRangeLayer::kBackendSubmit}; + int block_size = RuntimeUtils::GetOptimalBlockSize(); DispatchFunc( {static_cast(out_type_), block_size}, diff --git a/src/operator.h b/src/operator.h index dc34d25bc..9d868bb95 100644 --- a/src/operator.h +++ b/src/operator.h @@ -1,6 +1,7 @@ #ifndef INFINI_OPS_OPERATOR_H_ #define INFINI_OPS_OPERATOR_H_ +#include #include #include #include @@ -13,6 +14,7 @@ #include "config.h" #include "dispatcher.h" #include "handle.h" +#include "host_range_profiler.h" #include "tensor.h" namespace infini::ops::detail { @@ -188,10 +190,12 @@ template class Operator : public OperatorBase { public: - // Invalidate the operator cache. Cached operators are destroyed on the - // next `call()` invocation. Intended for test isolation — production - // code should never call this. - static void clear_cache() { ++cache_generation_; } + // Invalidate the operator cache. Cached operators are destroyed on the next + // `Call()` invocation. Intended for test isolation; production code should + // never call this. + static void clear_cache() { + cache_generation_.fetch_add(1, std::memory_order_relaxed); + } template static std::unique_ptr Make(const Config& config, @@ -224,16 +228,39 @@ class Operator : public OperatorBase { template static void Call(const Handle& handle, const Config& config, const Args&... args) { + [[maybe_unused]] HostRangeScope host_range_operator_call{ + HostRangeLayer::kOperatorCall}; + static thread_local std::unordered_map> cache; static thread_local std::size_t generation{0}; - if (generation != cache_generation_) { + const auto cache_generation = + cache_generation_.load(std::memory_order_relaxed); + if (generation != cache_generation) { cache.clear(); - generation = cache_generation_; + generation = cache_generation; } +#if defined(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) + auto key = [&]() { + HostRangeScope host_range_cache_key{HostRangeLayer::kCacheKey}; + return CacheKeyBuilder{}(config, args...); + }(); + + auto it = [&]() { + HostRangeScope host_range_cache_lookup{HostRangeLayer::kCacheLookup}; + return cache.find(key); + }(); + + if (it == cache.end()) { + HostRangeScope host_range_cache_construct{ + HostRangeLayer::kCacheConstruct}; + auto new_op = Make(config, args...); + it = cache.emplace(std::move(key), std::move(new_op)).first; + } +#else auto key = CacheKeyBuilder{}(config, args...); auto it{cache.find(key)}; @@ -241,9 +268,12 @@ class Operator : public OperatorBase { if (it == cache.end()) { it = cache.emplace(std::move(key), Make(config, args...)).first; } +#endif auto& op{it->second}; + [[maybe_unused]] HostRangeScope host_range_operator_invoke{ + HostRangeLayer::kOperatorInvoke}; return (*op)(handle, args...); } @@ -348,10 +378,7 @@ class Operator : public OperatorBase { return op_ptr; } - // Generation counter for lazy cache invalidation. Bumped by - // `clear_cache()`; the next `call()` detects the mismatch and - // destroys all cached operator instances. - static inline std::size_t cache_generation_{0}; + static inline std::atomic cache_generation_{0}; }; // Maximum number of implementation slots per (operator, device) pair. diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index c60748228..292a2c00b 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -6,6 +6,7 @@ #include +#include "host_range_profiler.h" #include "tensor.h" #include "torch/device_.h" @@ -33,7 +34,9 @@ inline DataType DataTypeFromString(const std::string& name) { return kStringToDataType.at(name); } -inline DataType DataTypeFromPybind11Handle(py::handle obj) { +namespace detail { + +inline DataType DataTypeFromPybind11HandleImpl(py::handle obj) { auto dtype_str{py::str(obj).cast()}; const auto pos{dtype_str.find_last_of('.')}; @@ -41,6 +44,14 @@ inline DataType DataTypeFromPybind11Handle(py::handle obj) { pos == std::string::npos ? dtype_str : dtype_str.substr(pos + 1)); } +} // namespace detail + +inline DataType DataTypeFromPybind11Handle(py::handle obj) { + [[maybe_unused]] HostRangeScope host_range_tensor_conversion{ + HostRangeLayer::kTensorConversion}; + return detail::DataTypeFromPybind11HandleImpl(obj); +} + template inline Device::Type DeviceTypeFromString(const std::string& name) { static const auto kTorchNameToTypes{ @@ -120,7 +131,9 @@ inline std::optional TryDeviceTypeFromString( return std::nullopt; } -inline Device DeviceFromPybind11Handle(py::handle obj) { +namespace detail { + +inline Device DeviceFromPybind11HandleImpl(py::handle obj) { auto device_obj{obj.attr("device")}; auto device_type_str{device_obj.attr("type").cast()}; auto device_index_obj{device_obj.attr("index")}; @@ -130,33 +143,51 @@ inline Device DeviceFromPybind11Handle(py::handle obj) { return Device{DeviceTypeFromString(device_type_str), device_index}; } -inline Tensor TensorFromPybind11Handle(py::handle obj) { +inline Tensor TensorFromPybind11HandleImpl(py::handle obj) { auto data{ reinterpret_cast(obj.attr("data_ptr")().cast())}; auto shape{obj.attr("shape").cast()}; - auto dtype{DataTypeFromPybind11Handle(obj.attr("dtype"))}; + auto dtype{DataTypeFromPybind11HandleImpl(obj.attr("dtype"))}; - auto device{DeviceFromPybind11Handle(obj)}; + auto device{DeviceFromPybind11HandleImpl(obj)}; auto strides{obj.attr("stride")().cast()}; return Tensor{data, std::move(shape), dtype, device, std::move(strides)}; } +} // namespace detail + +inline Device DeviceFromPybind11Handle(py::handle obj) { + [[maybe_unused]] HostRangeScope host_range_device_conversion{ + HostRangeLayer::kDeviceConversion}; + return detail::DeviceFromPybind11HandleImpl(obj); +} + +inline Tensor TensorFromPybind11Handle(py::handle obj) { + [[maybe_unused]] HostRangeScope host_range_tensor_conversion{ + HostRangeLayer::kTensorConversion}; + return detail::TensorFromPybind11HandleImpl(obj); +} + inline std::optional OptionalTensorFromPybind11Handle( const std::optional& obj) { + [[maybe_unused]] HostRangeScope host_range_tensor_conversion{ + HostRangeLayer::kTensorConversion}; if (!obj.has_value() || obj->is_none()) return std::nullopt; - return TensorFromPybind11Handle(*obj); + return detail::TensorFromPybind11HandleImpl(*obj); } inline std::vector VectorTensorFromPybind11Handle( const std::vector& objs) { + [[maybe_unused]] HostRangeScope host_range_tensor_conversion{ + HostRangeLayer::kTensorConversion}; std::vector result; result.reserve(objs.size()); for (const auto& obj : objs) { - result.push_back(TensorFromPybind11Handle(obj)); + result.push_back(detail::TensorFromPybind11HandleImpl(obj)); } return result; } diff --git a/tests/conftest.py b/tests/conftest.py index 1361c1752..d9016815d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,10 +1,12 @@ import hashlib import random +import time import pytest import torch import torch.utils.benchmark as benchmark +from tests import host_range_profile from tests.report import register_reporter from tests.utils import clone_strided, get_available_devices @@ -13,6 +15,12 @@ def pytest_addoption(parser): parser.addoption( "--benchmark", action="store_true", help="Run performance benchmarks." ) + parser.addoption( + "--host-range-profile", + action="store", + default=None, + help="Write host-range benchmark timings as JSON lines.", + ) parser.addoption( "--devices", nargs="+", @@ -30,9 +38,39 @@ def pytest_addoption(parser): ) +def _host_range_profile_compiled(): + import infini.ops as ops + + compiled_fn = getattr(ops, "_host_range_profile_compiled", None) + + return callable(compiled_fn) and compiled_fn() is True + + def pytest_configure(config): torch.backends.fp32_precision = "tf32" + profile_output = config.getoption("--host-range-profile") + config._infini_host_range_profile_output = None + + if profile_output: + try: + host_range_profile.validate_request( + profile_output, + benchmark_enabled=config.getoption("--benchmark"), + compiled=_host_range_profile_compiled, + xdist_workers=getattr(config.option, "numprocesses", None), + ) + except ValueError as error: + raise pytest.UsageError(str(error)) from error + except Exception as error: + raise pytest.UsageError( + "host-range profiling is not compiled into infini.ops" + ) from error + + config._infini_host_range_profile_output = host_range_profile.truncate_output( + profile_output + ) + config.addinivalue_line( "markers", "auto_act_and_assert: automatically perform Act and Assert phases using the return values", @@ -627,22 +665,85 @@ def pytest_pyfunc_call(pyfuncitem): if pyfuncitem.config.getoption("--benchmark"): stmt = "func(*args, **kwargs)" + profile_output = pyfuncitem.config._infini_host_range_profile_output + timer_kwargs = ( + {"timer": time.perf_counter} if profile_output is not None else {} + ) func_timer = benchmark.Timer( stmt=stmt, globals={"func": func, "args": args, "kwargs": kwargs}, label=func.__name__, description="InfiniOps", + **timer_kwargs, ) + if profile_output is None: + func_measurement = func_timer.blocked_autorange() + else: + import infini.ops as ops + + op_cls = _op_class_from_module(pyfuncitem.module) + clear_cache = getattr(op_cls, "clear_cache", None) + + if callable(clear_cache): + clear_cache() + + params = _callspec_params(pyfuncitem) + context = { + "nodeid": pyfuncitem.nodeid, + "operator": host_range_profile.operator_from_module( + pyfuncitem.module + ), + "backend": host_range_profile.backend_from_devices( + pyfuncitem.config.getoption("--devices"), + params.get("device"), + ), + } + + def synchronize(): + host_range_profile.synchronize_torch_device( + torch, params.get("device") + ) + + _, cold_summaries = host_range_profile.collect_ranges( + ops._host_range_profile_start, + ops._host_range_profile_stop, + lambda: func(*args, **kwargs), + synchronize=synchronize, + ) + host_range_profile.append_rows( + profile_output, + host_range_profile.expand_summary_rows( + cold_summaries, phase="cold", **context + ), + ) + + func_measurement = func_timer.blocked_autorange() + warm_summaries = host_range_profile.collect_replayed_ranges( + ops._host_range_profile_start, + ops._host_range_profile_stop, + lambda: func(*args, **kwargs), + func_measurement, + synchronize=synchronize, + ) + warm_rows = host_range_profile.expand_summary_rows( + warm_summaries, phase="warm", **context + ) + warm_rows.append( + host_range_profile.measurement_row( + func_measurement, phase="warm", **context + ) + ) + host_range_profile.append_rows(profile_output, warm_rows) + ref_timer = benchmark.Timer( stmt=stmt, globals={"func": ref, "args": ref_args, "kwargs": ref_kwargs}, label=func.__name__, description="Reference", + **timer_kwargs, ) - - func_measurement = func_timer.blocked_autorange() ref_measurement = ref_timer.blocked_autorange() benchmark.Compare((func_measurement, ref_measurement)).print() diff --git a/tests/host_range_profile.py b/tests/host_range_profile.py new file mode 100644 index 000000000..bb93940e9 --- /dev/null +++ b/tests/host_range_profile.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +import json +import pathlib + + +ROW_FIELDS = ( + "nodeid", + "operator", + "backend", + "phase", + "range", + "metric", + "count", + "unit", + "mean", + "median", +) + +_PLATFORM_TO_TORCH_DEVICE = { + "nvidia": "cuda", + "metax": "cuda", + "iluvatar": "cuda", + "hygon": "cuda", + "moore": "musa", + "cambricon": "mlu", + "ascend": "npu", +} + + +def validate_request(output_path, *, benchmark_enabled, compiled, xdist_workers=None): + if not output_path: + return + + if not benchmark_enabled: + raise ValueError("--host-range-profile requires --benchmark") + + if xdist_workers: + raise ValueError("host range profiling does not support pytest-xdist") + + if callable(compiled): + compiled = compiled() + + if compiled is not True: + raise ValueError("host-range profiling is not compiled into infini.ops") + + +def truncate_output(output_path): + path = pathlib.Path(output_path).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("", encoding="utf-8") + + return path + + +def append_rows(output_path, rows): + path = pathlib.Path(output_path).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + + with path.open("a", encoding="utf-8") as output: + for row in rows: + if set(row) != set(ROW_FIELDS): + raise ValueError( + "host-range profile rows must contain exactly: " + + ", ".join(ROW_FIELDS) + ) + + ordered = {field: row[field] for field in ROW_FIELDS} + output.write(json.dumps(ordered, separators=(",", ":")) + "\n") + + +def expand_summary_rows(summaries, *, nodeid, operator, backend, phase): + rows = [] + + for summary in summaries: + for metric in ("inclusive", "self"): + rows.append( + { + "nodeid": nodeid, + "operator": operator, + "backend": backend, + "phase": phase, + "range": summary["range"], + "metric": metric, + "count": summary["count"], + "unit": summary["unit"], + "mean": summary[f"{metric}_mean"], + "median": summary[f"{metric}_median"], + } + ) + + return rows + + +def measurement_row(measurement, *, nodeid, operator, backend, phase): + return { + "nodeid": nodeid, + "operator": operator, + "backend": backend, + "phase": phase, + "range": "end_to_end", + "metric": "inclusive", + "count": len(measurement.raw_times) * measurement.number_per_run, + "unit": "ns", + "mean": measurement.mean * 1e9, + "median": measurement.median * 1e9, + } + + +def collect_ranges(start, stop, callback, synchronize=None): + if synchronize is not None: + synchronize() + + try: + start() + + try: + result = callback() + except BaseException: + try: + stop() + except BaseException: + pass + + raise + + try: + summaries = stop() + except BaseException: + try: + stop() + except BaseException: + pass + + raise + + return result, summaries + finally: + if synchronize is not None: + synchronize() + + +def collect_replayed_ranges(start, stop, callback, measurement, synchronize=None): + call_count = len(measurement.raw_times) * measurement.number_per_run + + def replay_calls(): + for _ in range(call_count): + callback() + + _, summaries = collect_ranges(start, stop, replay_calls, synchronize=synchronize) + + return summaries + + +def synchronize_torch_device(torch_module, torch_device): + if torch_device is None: + return + + object_device_type = getattr(torch_device, "type", None) + device_type = object_device_type + if device_type is None: + device_type = str(torch_device).split(":", maxsplit=1)[0] + + device_type = _PLATFORM_TO_TORCH_DEVICE.get(device_type, device_type) + if device_type == "cpu": + return + + runtime = getattr(torch_module, device_type, None) + synchronize = getattr(runtime, "synchronize", None) + if not callable(synchronize): + raise RuntimeError(f"cannot synchronize PyTorch device type {device_type!r}") + + device_name = str(torch_device) + has_explicit_index = ":" in device_name + if object_device_type is not None: + synchronize(torch_device) + elif has_explicit_index: + _, device_index = device_name.split(":", maxsplit=1) + synchronize(f"{device_type}:{device_index}") + else: + synchronize() + + +def operator_from_module(module): + module_name = module.__name__.rsplit(".", 1)[-1] + + if module_name.startswith("test_"): + return module_name[len("test_") :] + + return module_name + + +def backend_from_devices(devices, torch_device=None): + devices = tuple(devices or ()) + + if len(devices) == 1: + return devices[0] + + if torch_device in devices: + return torch_device + + matching_platforms = tuple( + device + for device in devices + if _PLATFORM_TO_TORCH_DEVICE.get(device) == torch_device + ) + + if len(matching_platforms) == 1: + return matching_platforms[0] + + return torch_device or "unknown" diff --git a/tests/test_generate_public_operator_header.py b/tests/test_generate_public_operator_header.py new file mode 100644 index 000000000..e7e2656bb --- /dev/null +++ b/tests/test_generate_public_operator_header.py @@ -0,0 +1,31 @@ +import importlib.util +import pathlib + + +def _load_generator_module(): + path = ( + pathlib.Path(__file__).resolve().parents[1] + / "scripts" + / "generate_public_operator_header.py" + ) + spec = importlib.util.spec_from_file_location( + "generate_public_operator_header_under_test", path + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_public_operator_header_omits_host_range_profiler(tmp_path): + root = pathlib.Path(__file__).resolve().parents[1] + output = tmp_path / "operator.h" + module = _load_generator_module() + + module.generate_public_operator_header(root / "src" / "operator.h", output) + + public_header = output.read_text() + assert '#include "host_range_profiler.h"' not in public_header + assert ( + "static void Call(const Handle& handle, const Config& config," in public_header + ) diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 7c033f8cc..443ad6943 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -1,5 +1,6 @@ import importlib.util import pathlib +import re import sys @@ -263,6 +264,163 @@ class DtypeOp : public Operator { """ +def _make_profile_operator(module, tmp_path, monkeypatch): + base_header = tmp_path / "profile_op.h" + base_header.write_text( + """ +class ProfileOp { + public: + ProfileOp(const Tensor input, Tensor out); + virtual void operator()(const Tensor input, Tensor out) const = 0; + virtual void operator()(const Tensor input, const double alpha, + Tensor out) const = 0; +}; +""" + ) + monkeypatch.setattr(module, "_find_base_header", lambda op_name: base_header) + + arguments = [ + module._ParsedArgument("const Tensor", "input"), + module._ParsedArgument("Tensor", "out"), + ] + scaled_arguments = [ + module._ParsedArgument("const Tensor", "input"), + module._ParsedArgument("const double", "alpha"), + module._ParsedArgument("Tensor", "out"), + ] + + return module._Operator( + "profile_op", + constructors=[module._ParsedFunction(arguments)], + calls=[ + module._ParsedFunction(arguments), + module._ParsedFunction(scaled_arguments), + ], + ) + + +def test_generated_free_calls_start_with_binding_body_profile_scope( + tmp_path, monkeypatch +): + module = _load_generator_module() + operator = _make_profile_operator(module, tmp_path, monkeypatch) + + text = module._generate_pybind11(operator) + + scope = ( + "[[maybe_unused]] HostRangeScope host_range_binding_body{\n" + " HostRangeLayer::kBindingBody};" + ) + assert '#include "host_range_profiler.h"' in text + free_call_bodies = re.findall( + r'm\.def\("profile_op", \[\]\([^)]*\) \{\n(.*?)\n \},', + text, + flags=re.DOTALL, + ) + assert len(free_call_bodies) == 2 + + for body in free_call_bodies: + assert body.lstrip().startswith(scope) + assert body.count(scope) == 1 + assert body.index(scope) < body.index("Handle handle;") + assert body.index(scope) < body.index("Config config;") + + +def test_generated_ops_module_exposes_host_range_profile_controls_once(): + module = _load_generator_module() + + text = module._generate_ops_module_source(["BindProfileOp"]) + + assert '#include "host_range_profiler.h"' in text + expected_targets = { + "_host_range_profile_compiled": "HostRangeProfiler::IsCompiled", + "_host_range_profile_start": "HostRangeProfiler::Start", + "_host_range_profile_stop": "HostRangeProfiler::Stop", + "_host_range_profile_calibrate": "HostRangeProfiler::Calibrate", + } + + for name, target in expected_targets.items(): + marker = f'm.def("{name}"' + assert text.count(marker) == 1 + binding = text.split(marker, maxsplit=1)[1].split("\n m.def(", maxsplit=1)[0] + assert target in binding + + +def test_iluvatar_custom_compilers_receive_host_range_profile_definition(): + cmake = (pathlib.Path(__file__).parents[1] / "src" / "CMakeLists.txt").read_text( + encoding="utf-8" + ) + + for definitions in ( + "_iluvatar_call_instantiation_defs", + "_iluvatar_dispatch_defs", + ): + assert ( + f"list(APPEND {definitions}\n" + " -DINFINI_OPS_ENABLE_HOST_RANGE_PROFILING=1)" + ) in cmake + + +def test_torch_system_compiler_receives_host_range_profile_definition(): + cmake = (pathlib.Path(__file__).parents[1] / "src" / "CMakeLists.txt").read_text( + encoding="utf-8" + ) + + assert ( + "list(APPEND _torch_extra_flags\n" + ' "-DINFINI_OPS_ENABLE_HOST_RANGE_PROFILING=1")' + ) in cmake + + +def test_generated_dispatch_calls_start_with_dispatch_profile_scope( + tmp_path, monkeypatch +): + module = _load_generator_module() + operator = _make_profile_operator(module, tmp_path, monkeypatch) + _, definitions = module._generate_generated_dispatch_entries(operator) + + text = module._generate_generated_dispatch_source([], definitions) + + scope = ( + "[[maybe_unused]] HostRangeScope host_range_dispatch_call{\n" + " HostRangeLayer::kDispatchCall};" + ) + assert '#include "host_range_profiler.h"' in text + dispatch_call_bodies = re.findall( + r"void CallProfileOp\([^)]*\) \{\n(.*?)\n\}", + text, + flags=re.DOTALL, + ) + assert len(dispatch_call_bodies) == 2 + + for body in dispatch_call_bodies: + assert body.lstrip().startswith(scope) + assert body.count(scope) == 1 + assert body.index(scope) < body.index( + "Operator<::infini::ops::ProfileOp>::Call" + ) + + +def test_generated_cache_clear_has_a_backend_independent_core_source( + tmp_path, monkeypatch +): + module = _load_generator_module() + operator = _make_profile_operator(module, tmp_path, monkeypatch) + + _, definitions = module._generate_generated_dispatch_entries(operator) + dispatch_text = "\n".join(definitions) + core_text = module._generate_cache_clear_dispatch_source(["profile_op"]) + + assert ( + "#if !defined(INFINI_OPS_USE_OPERATOR_CALL_INSTANTIATIONS)\n" + "void ClearCacheForProfileOp()" + ) in dispatch_text + assert "INFINI_OPS_BUILD_CORE_DISPATCH" not in dispatch_text + assert core_text.count("void ClearCacheForProfileOp()") == 1 + assert '#include "base/profile_op.h"' in core_text + assert "Operator<::infini::ops::ProfileOp>::clear_cache();" in core_text + + def test_pybind_converts_data_type_arguments_from_torch_dtype(tmp_path, monkeypatch): text = _generate_binding("dtype_op", tmp_path, monkeypatch, _DTYPE_OP_SOURCE) diff --git a/tests/test_host_overhead_control.py b/tests/test_host_overhead_control.py new file mode 100644 index 000000000..5f09d5ff9 --- /dev/null +++ b/tests/test_host_overhead_control.py @@ -0,0 +1,107 @@ +import importlib.util +import pathlib + + +def _load_control_module(): + path = ( + pathlib.Path(__file__).resolve().parents[1] + / "scripts" + / "run_host_overhead_control.py" + ) + spec = importlib.util.spec_from_file_location( + "run_host_overhead_control_under_test", path + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + return module + + +def test_measure_host_submission_keeps_synchronization_outside_timed_interval(): + control = _load_control_module() + events = [] + clock = iter((100, 160, 200, 300)) + + def synchronize(): + events.append("synchronize") + + def callback(): + events.append("call") + + def timer_ns(): + events.append("timer") + + return next(clock) + + def before_round(): + events.append("profile_start") + + def after_round(): + events.append("profile_stop") + + result = control.measure_host_submission( + callback, + synchronize, + warmup_iterations=2, + iterations=2, + rounds=2, + before_round=before_round, + after_round=after_round, + timer_ns=timer_ns, + ) + + assert events == [ + "call", + "call", + "synchronize", + "synchronize", + "profile_start", + "timer", + "call", + "call", + "timer", + "profile_stop", + "synchronize", + "synchronize", + "profile_start", + "timer", + "call", + "call", + "timer", + "profile_stop", + "synchronize", + ] + assert result == { + "warmup_iterations": 2, + "iterations_per_round": 2, + "rounds": 2, + "unit": "ns", + "mean": 40.0, + "median": 40.0, + "samples": [30.0, 50.0], + } + + +def test_measure_host_submission_rejects_empty_measurements(): + control = _load_control_module() + + for name, arguments in ( + ( + "warmup_iterations", + {"warmup_iterations": -1, "iterations": 1, "rounds": 1}, + ), + ( + "iterations", + {"warmup_iterations": 1, "iterations": 0, "rounds": 1}, + ), + ( + "rounds", + {"warmup_iterations": 1, "iterations": 1, "rounds": 0}, + ), + ): + try: + control.measure_host_submission(lambda: None, lambda: None, **arguments) + except ValueError as error: + assert name in str(error) + else: + raise AssertionError(f"expected invalid {name} to fail") diff --git a/tests/test_host_range_profile.py b/tests/test_host_range_profile.py new file mode 100644 index 000000000..b959d8bb3 --- /dev/null +++ b/tests/test_host_range_profile.py @@ -0,0 +1,209 @@ +import math +import os + +import pytest + +import infini.ops as ops + + +_PROFILE_API = ( + "_host_range_profile_compiled", + "_host_range_profile_start", + "_host_range_profile_stop", + "_host_range_profile_calibrate", +) +_SUMMARY_KEYS = { + "range", + "count", + "unit", + "inclusive_mean", + "inclusive_median", + "self_mean", + "self_median", +} +_CALIBRATION_RANGES = ( + "calibration.depth1", + "calibration.depth2", + "calibration.depth3", +) +_TIMING_KEYS = ( + "inclusive_mean", + "inclusive_median", + "self_mean", + "self_median", +) +_EXPECT_PROFILE_ENV = "INFINI_OPS_EXPECT_HOST_RANGE_PROFILING" +_COLLECTOR_DRAIN_LIMIT = 8 + + +def _restore_inactive_profile_collector(): + for _ in range(_COLLECTOR_DRAIN_LIMIT): + try: + ops._host_range_profile_stop() + except RuntimeError: + try: + ops._host_range_profile_start() + except RuntimeError: + continue + + try: + ops._host_range_profile_stop() + except RuntimeError as error: + pytest.fail(f"host-range profile collector cleanup failed: {error}") + + try: + ops._host_range_profile_stop() + except RuntimeError: + return + + pytest.fail( + "host-range profile collector did not become inactive after " + f"{_COLLECTOR_DRAIN_LIMIT} cleanup attempts" + ) + + +@pytest.fixture(autouse=True) +def _isolate_profile_collector(): + if any(not callable(getattr(ops, name, None)) for name in _PROFILE_API): + yield + + return + + compiled = getattr(ops, "_host_range_profile_compiled", None) + if compiled() is not True: + yield + + return + + _restore_inactive_profile_collector() + yield + _restore_inactive_profile_collector() + + +def _require_profile_build(): + missing = [name for name in _PROFILE_API if not callable(getattr(ops, name, None))] + assert not missing, "infini.ops is missing host-range profiling API: " + ", ".join( + missing + ) + + compiled = ops._host_range_profile_compiled() + assert isinstance(compiled, bool) + if not compiled: + if os.environ.get(_EXPECT_PROFILE_ENV) == "1": + pytest.fail( + f"{_EXPECT_PROFILE_ENV}=1 but InfiniOps reports host-range " + "profiling is not compiled" + ) + + pytest.skip("InfiniOps was built without host-range profiling") + + +def test_host_range_profile_start_rejects_an_active_collector(): + _require_profile_build() + + ops._host_range_profile_start() + try: + with pytest.raises(RuntimeError): + ops._host_range_profile_start() + finally: + ops._host_range_profile_stop() + + +def test_host_range_profile_stop_rejects_an_inactive_collector(): + _require_profile_build() + + with pytest.raises(RuntimeError): + ops._host_range_profile_stop() + + +def test_host_range_profile_calibration_reports_nested_summary_schema(): + _require_profile_build() + iterations = 7 + + rows = ops._host_range_profile_calibrate(iterations) + + assert [row["range"] for row in rows] == list(_CALIBRATION_RANGES) + assert all(set(row) == _SUMMARY_KEYS for row in rows) + assert all(row["count"] == iterations for row in rows) + assert all(row["unit"] == "ns" for row in rows) + + timing_values = [] + for row in rows: + for key in _TIMING_KEYS: + value = row[key] + assert isinstance(value, (int, float)) and not isinstance(value, bool) + assert math.isfinite(value) + assert value >= 0 + timing_values.append(value) + + assert row["inclusive_mean"] >= row["self_mean"] >= 0 + assert row["inclusive_median"] >= row["self_median"] >= 0 + + assert any(value > 0 for value in timing_values) + + depth1, depth2, depth3 = rows + assert depth1["inclusive_mean"] == pytest.approx( + depth1["self_mean"] + depth2["inclusive_mean"] + ) + assert depth2["inclusive_mean"] == pytest.approx( + depth2["self_mean"] + depth3["inclusive_mean"] + ) + assert depth3["inclusive_mean"] == depth3["self_mean"] + assert depth3["inclusive_median"] == depth3["self_median"] + + inclusive_medians = [row["inclusive_median"] for row in rows] + assert inclusive_medians == sorted(inclusive_medians, reverse=True) + + +def test_host_range_profile_clear_cache_forces_next_call_to_construct(): + _require_profile_build() + + import torch + + from tests.utils import get_stream + + add_cls = getattr(ops, "Add", None) + if add_cls is None: + pytest.skip("Add is not included in this profiling build") + + if 0 in add_cls.active_implementation_indices("cpu"): + device = "cpu" + elif torch.cuda.is_available() and 0 in add_cls.active_implementation_indices( + "nvidia" + ): + device = "cuda" + else: + pytest.skip("Add implementation 0 is not active on CPU or NVIDIA") + + input = torch.randn((13, 4), device=device) + other = torch.randn_like(input) + out = torch.empty_like(input) + kwargs = { + "stream": get_stream(input.device), + "implementation_index": 0, + } + + ops.add(input, other, out, **kwargs) + add_cls.clear_cache() + + ops._host_range_profile_start() + try: + ops.add(input, other, out, **kwargs) + finally: + rows = ops._host_range_profile_stop() + + construct = [row for row in rows if row["range"] == "cache.construct"] + assert len(construct) == 1 + assert construct[0]["count"] == 1 + + by_range = {row["range"]: row for row in rows} + assert by_range["binding.tensor_conversion"]["count"] == 3 + assert by_range["binding.device_conversion"]["count"] == 1 + + ops._host_range_profile_start() + try: + ops.add(input, other, out, **kwargs) + finally: + rows = ops._host_range_profile_stop() + + assert all(row["range"] != "cache.construct" for row in rows) diff --git a/tests/test_host_range_profile_writer.py b/tests/test_host_range_profile_writer.py new file mode 100644 index 000000000..46c362d20 --- /dev/null +++ b/tests/test_host_range_profile_writer.py @@ -0,0 +1,251 @@ +import json +from types import SimpleNamespace + +import pytest + +from tests import host_range_profile + + +_CONTEXT = { + "nodeid": "tests/test_add.py::test_case[nvidia]", + "operator": "add", + "backend": "nvidia", + "phase": "warm", +} + + +def test_expand_summaries_writes_inclusive_and_self_rows(): + summaries = [ + { + "range": "operator.call", + "count": 3, + "unit": "ns", + "inclusive_mean": 17.5, + "inclusive_median": 16.0, + "self_mean": 4.5, + "self_median": 4.0, + } + ] + + rows = host_range_profile.expand_summary_rows(summaries, **_CONTEXT) + + assert rows == [ + { + **_CONTEXT, + "range": "operator.call", + "metric": "inclusive", + "count": 3, + "unit": "ns", + "mean": 17.5, + "median": 16.0, + }, + { + **_CONTEXT, + "range": "operator.call", + "metric": "self", + "count": 3, + "unit": "ns", + "mean": 4.5, + "median": 4.0, + }, + ] + assert all(tuple(row) == host_range_profile.ROW_FIELDS for row in rows) + + +def test_truncate_then_append_json_lines(tmp_path): + output = tmp_path / "profile.jsonl" + output.write_text("stale\n", encoding="utf-8") + row = { + **_CONTEXT, + "range": "operator.call", + "metric": "inclusive", + "count": 3, + "unit": "ns", + "mean": 17.5, + "median": 16.0, + } + + host_range_profile.truncate_output(output) + host_range_profile.append_rows(output, [row]) + host_range_profile.append_rows(output, [{**row, "phase": "cold"}]) + + lines = output.read_text(encoding="utf-8").splitlines() + assert [json.loads(line) for line in lines] == [row, {**row, "phase": "cold"}] + + +@pytest.mark.parametrize( + ( + "output_path", + "benchmark_enabled", + "compiled", + "xdist_workers", + "message", + ), + [ + ("profile.jsonl", False, True, None, "requires --benchmark"), + ("profile.jsonl", True, False, None, "not compiled"), + ("profile.jsonl", True, True, "auto", "does not support pytest-xdist"), + ], +) +def test_validate_request_rejects_invalid_profile_configuration( + output_path, benchmark_enabled, compiled, xdist_workers, message +): + with pytest.raises(ValueError, match=message): + host_range_profile.validate_request( + output_path, + benchmark_enabled=benchmark_enabled, + compiled=compiled, + xdist_workers=xdist_workers, + ) + + +def test_validate_request_does_not_restrict_xdist_without_profile_output(): + host_range_profile.validate_request( + None, + benchmark_enabled=False, + compiled=False, + xdist_workers="logical", + ) + + +def test_measurement_row_converts_seconds_to_nanoseconds_and_counts_calls(): + measurement = SimpleNamespace( + raw_times=[0.000_001, 0.000_003], + number_per_run=4, + mean=0.000_002, + median=0.000_001_5, + ) + + row = host_range_profile.measurement_row(measurement, **_CONTEXT) + + assert row == { + **_CONTEXT, + "range": "end_to_end", + "metric": "inclusive", + "count": 8, + "unit": "ns", + "mean": 2000.0, + "median": 1500.0, + } + + +def test_backend_and_operator_names_preserve_nvidia_and_snake_case(): + module = SimpleNamespace(__name__="tests.test_causal_softmax") + + assert host_range_profile.backend_from_devices(["nvidia"], "cuda") == "nvidia" + assert host_range_profile.operator_from_module(module) == "causal_softmax" + + +def test_collect_ranges_stops_after_profiled_call_raises(): + state = SimpleNamespace(active=False) + + def start(): + state.active = True + + def stop(): + state.active = False + + return [] + + def fail(): + raise RuntimeError("profiled call failed") + + with pytest.raises(RuntimeError, match="profiled call failed"): + host_range_profile.collect_ranges(start, stop, fail) + + assert state.active is False + + +def test_collect_ranges_synchronizes_only_outside_profile_window(): + state = SimpleNamespace(active=False) + events = [] + + def synchronize(): + assert state.active is False + events.append("synchronize") + + def start(): + state.active = True + events.append("start") + + def callback(): + assert state.active is True + events.append("callback") + + return "result" + + def stop(): + assert state.active is True + state.active = False + events.append("stop") + + return ["summary"] + + result, summaries = host_range_profile.collect_ranges( + start, stop, callback, synchronize=synchronize + ) + + assert result == "result" + assert summaries == ["summary"] + assert events == ["synchronize", "start", "callback", "stop", "synchronize"] + + +def test_collect_replayed_ranges_profiles_exact_measurement_call_count(): + state = SimpleNamespace(active=False, calls=0) + measurement = SimpleNamespace(raw_times=[0.1, 0.2, 0.3], number_per_run=4) + + def start(): + state.active = True + + def stop(): + state.active = False + + return [{"range": "binding.body", "count": state.calls}] + + def callback(): + assert state.active is True + state.calls += 1 + + summaries = host_range_profile.collect_replayed_ranges( + start, stop, callback, measurement + ) + + assert state.active is False + assert state.calls == 12 + assert summaries == [{"range": "binding.body", "count": 12}] + + +@pytest.mark.parametrize( + ("torch_device", "expected_args"), + [ + ("cuda", ()), + ("cuda:1", ("cuda:1",)), + ("nvidia", ()), + ("nvidia:1", ("cuda:1",)), + ], +) +def test_synchronize_torch_device_uses_platform_runtime(torch_device, expected_args): + calls = [] + torch_module = SimpleNamespace( + cuda=SimpleNamespace(synchronize=lambda *args: calls.append(args)) + ) + + host_range_profile.synchronize_torch_device(torch_module, torch_device) + + assert calls == [expected_args] + + +def test_synchronize_torch_device_preserves_device_object(): + calls = [] + torch_device = SimpleNamespace(type="cuda", index=1) + torch_module = SimpleNamespace( + cuda=SimpleNamespace(synchronize=lambda *args: calls.append(args)) + ) + + host_range_profile.synchronize_torch_device(torch_module, torch_device) + + assert calls == [(torch_device,)] + + +def test_synchronize_torch_device_skips_cpu(): + host_range_profile.synchronize_torch_device(SimpleNamespace(), "cpu")