From 6e05d4232a61f55587def5fd929c2f218ec7fb85 Mon Sep 17 00:00:00 2001 From: Jacques Liao Date: Fri, 5 Jun 2026 14:57:15 -0700 Subject: [PATCH] n/a PiperOrigin-RevId: 927494385 --- BUILD | 76 ++- hwy/cache_control.h | 38 +- hwy/contrib/dot/one_to_many-inl.h | 266 +++++++++ hwy/contrib/dot/one_to_many_kernels-inl.h | 139 +++++ hwy/contrib/dot/one_to_many_test.cc | 510 ++++++++++++++++++ hwy/contrib/pipeline/README.md | 22 + hwy/contrib/pipeline/prefetch_pipeline.h | 316 +++++++++++ hwy/contrib/pipeline/prefetch_pipeline_2d.h | 205 +++++++ .../pipeline/prefetch_pipeline_2d_test.cc | 188 +++++++ .../pipeline/prefetch_pipeline_test.cc | 336 ++++++++++++ .../pipeline/prefetch_pipeline_types.h | 322 +++++++++++ .../pipeline/prefetch_pipeline_types_test.cc | 130 +++++ hwy/contrib/pipeline/prefetch_tuner.h | 199 +++++++ .../pipeline/prefetch_tuner_registry.h | 187 +++++++ .../pipeline/prefetch_tuner_registry_test.cc | 249 +++++++++ hwy/contrib/pipeline/prefetch_tuner_test.cc | 82 +++ hwy_tests.bzl | 32 ++ 17 files changed, 3274 insertions(+), 23 deletions(-) create mode 100644 hwy/contrib/dot/one_to_many-inl.h create mode 100644 hwy/contrib/dot/one_to_many_kernels-inl.h create mode 100644 hwy/contrib/dot/one_to_many_test.cc create mode 100644 hwy/contrib/pipeline/README.md create mode 100644 hwy/contrib/pipeline/prefetch_pipeline.h create mode 100644 hwy/contrib/pipeline/prefetch_pipeline_2d.h create mode 100644 hwy/contrib/pipeline/prefetch_pipeline_2d_test.cc create mode 100644 hwy/contrib/pipeline/prefetch_pipeline_test.cc create mode 100644 hwy/contrib/pipeline/prefetch_pipeline_types.h create mode 100644 hwy/contrib/pipeline/prefetch_pipeline_types_test.cc create mode 100644 hwy/contrib/pipeline/prefetch_tuner.h create mode 100644 hwy/contrib/pipeline/prefetch_tuner_registry.h create mode 100644 hwy/contrib/pipeline/prefetch_tuner_registry_test.cc create mode 100644 hwy/contrib/pipeline/prefetch_tuner_test.cc diff --git a/BUILD b/BUILD index 6ebb89f484..d18c114ba4 100644 --- a/BUILD +++ b/BUILD @@ -4,6 +4,7 @@ load("@rules_cc//cc:cc_test.bzl", "cc_test") load("@bazel_skylib//lib:selects.bzl", "selects") load("//:hwy_tests.bzl", "HWY_TESTS") load("@rules_license//rules:license.bzl", "license") +load("//tools/build_defs/testing:bzl_library.bzl", "bzl_library") package( default_applicable_licenses = [":license"], @@ -118,6 +119,27 @@ COPTS = select({ ], }) +HWY_TEST_COPTS = select({ + ":compiler_msvc": [], + "//conditions:default": [ + # gTest triggers this warning (which is enabled by the + # extra-semi in COPTS), so we need to disable it here, + # but it's still enabled for :hwy. + "-Wno-c++98-compat-extra-semi", + ], +}) + +# Common to all tests. +HWY_TEST_DEPS = [ + ":hwy_test_util", + ":hwy", + ":nanobenchmark", + ":timer", +] + select({ + ":compiler_msvc": [], + "//conditions:default": ["@com_google_googletest//:gtest_main"], +}) + DEFINES = select({ ":compiler_msvc": ["HWY_SHARED_DEFINE"], ":compiler_clangcl": ["HWY_SHARED_DEFINE"], @@ -344,9 +366,16 @@ cc_library( copts = COPTS, textual_hdrs = [ "hwy/contrib/dot/dot-inl.h", + # copybara:strip_begin(internal) + "hwy/contrib/dot/one_to_many-inl.h", + "hwy/contrib/dot/one_to_many_kernels-inl.h", + # copybara:strip_end ], deps = [ ":hwy", + # copybara:strip_begin(internal) + ":prefetch_pipeline", + # copybara:strip_end ], ) @@ -607,6 +636,25 @@ cc_library( ], ) +# copybara:strip_begin(internal) +cc_library( + name = "prefetch_pipeline", + hdrs = [ + "hwy/contrib/pipeline/prefetch_pipeline.h", + "hwy/contrib/pipeline/prefetch_pipeline_2d.h", + "hwy/contrib/pipeline/prefetch_pipeline_types.h", + "hwy/contrib/pipeline/prefetch_tuner.h", + "hwy/contrib/pipeline/prefetch_tuner_registry.h", + ], + compatible_with = [], + copts = COPTS, + deps = [ + ":hwy", + ":timer", + ], +) +# copybara:strip_end + cc_test( name = "list_targets", size = "small", @@ -628,27 +676,6 @@ cc_test( ], ) -HWY_TEST_COPTS = select({ - ":compiler_msvc": [], - "//conditions:default": [ - # gTest triggers this warning (which is enabled by the - # extra-semi in COPTS), so we need to disable it here, - # but it's still enabled for :hwy. - "-Wno-c++98-compat-extra-semi", - ], -}) - -# Common to all tests. -HWY_TEST_DEPS = [ - ":hwy_test_util", - ":hwy", - ":nanobenchmark", - ":timer", -] + select({ - ":compiler_msvc": [], - "//conditions:default": ["@com_google_googletest//:gtest_main"], -}) - [ [ cc_test( @@ -714,3 +741,10 @@ test_suite( name = "hwy_ops_tests", tags = ["hwy_ops_test"], ) + +bzl_library( + name = "hwy_tests_bzl", + srcs = ["hwy_tests.bzl"], + parse_tests = False, + visibility = ["//visibility:private"], +) diff --git a/hwy/cache_control.h b/hwy/cache_control.h index 90743cd3f2..10cabaf46e 100644 --- a/hwy/cache_control.h +++ b/hwy/cache_control.h @@ -92,8 +92,10 @@ HWY_INLINE HWY_ATTR_CACHE void FlushStream() { #endif } -// Optionally begins loading the cache line containing "p" to reduce latency of -// subsequent actual loads. +// Optionally begins loading the cache line containing "p" into all cache +// levels, including L1, to reduce latency of subsequent actual loads. This +// corresponds to the T0 temporal locality hint on x86, which is ideal when data +// is about to be directly consumed. template HWY_INLINE HWY_ATTR_CACHE void Prefetch(const T* p) { (void)p; @@ -109,6 +111,38 @@ HWY_INLINE HWY_ATTR_CACHE void Prefetch(const T* p) { #endif // HWY_DISABLE_CACHE_CONTROL } +// Begins loading the cache line containing "p" into the L1 cache only, passing +// a Non-Temporal Access (NTA) hint. This minimizes pollution of outer memory +// caches (L2/L3) and is ideal for data accessed exactly once. +template +HWY_INLINE HWY_ATTR_CACHE void PrefetchForImmediateuse(const T* p) { + (void)p; +#ifndef HWY_DISABLE_CACHE_CONTROL +#if HWY_ARCH_X86 && !(HWY_COMPILER_CLANGCL && !defined(__MMX__)) + _mm_prefetch(reinterpret_cast(p), _MM_HINT_NTA); +#elif HWY_COMPILER_GCC || HWY_COMPILER_CLANGCL // includes clang + // Hint=0 specifically sets Non-Temporal local locality + __builtin_prefetch(p, /*write=*/0, /*hint=*/0); +#endif +#endif // HWY_DISABLE_CACHE_CONTROL +} + +// Attempts to stage the cache line containing "p" into the L3/L2 outer caches +// without aggressively staging it immediately into the L1. This restricts L1 +// and LFB thrashing on architectures like Intel when hiding massive DRAM delay. +template +HWY_INLINE HWY_ATTR_CACHE void PrefetchForFutureUse(const T* p) { + (void)p; +#ifndef HWY_DISABLE_CACHE_CONTROL +#if HWY_ARCH_X86 && !(HWY_COMPILER_CLANGCL && !defined(__MMX__)) + _mm_prefetch(reinterpret_cast(p), _MM_HINT_T2); +#elif HWY_COMPILER_GCC || HWY_COMPILER_CLANGCL // includes clang + // Hint=1 requests Moderate degrees of temporal locality (L2/L3 bounds) + __builtin_prefetch(p, /*write=*/0, /*hint=*/1); +#endif +#endif // HWY_DISABLE_CACHE_CONTROL +} + // Invalidates and flushes the cache line containing "p", if possible. HWY_INLINE HWY_ATTR_CACHE void FlushCacheline(const void* p) { #if HWY_ARCH_X86 && !defined(HWY_DISABLE_CACHE_CONTROL) diff --git a/hwy/contrib/dot/one_to_many-inl.h b/hwy/contrib/dot/one_to_many-inl.h new file mode 100644 index 0000000000..49b0a5d6a9 --- /dev/null +++ b/hwy/contrib/dot/one_to_many-inl.h @@ -0,0 +1,266 @@ +// Generic One-to-Many 2D-Tiled SIMD Scoring Framework. +// +// This header provides a generic Highway-based 2D block-tiling architecture +// for One-To-Many (OTM) dot product / distance operations. It abstracts the +// pipeline mechanics so that any downstream application can leverage +// high-throughput vector search strategies. +// +// NOTE ON HEADER ARCHITECTURE: While the outer pipeline driver is structurally +// generic, we anticipate that the underlying compute kernels (ScorerKernel) +// will predominantly invoke per-target SIMD vector instructions. By packaging +// this framework as an `-inl.h` header included after `foreach_target.h`, we +// ensure flawless per-target SIMD kernel inlining and Clang AST target +// adaptation for maximum performance and architectural simplicity. +// +// 1. Separation of Concerns: +// - MemoryLayout: Defines how memory is accessed, prefetching rules, and +// final score calibration (e.g., metric adjustments). +// - ScorerKernel: The raw compute abstraction (e.g. BF16/F32 FMA loops). +// - Pipeline: The nested-loop driver that manages registers and cache. + +#include "hwy/contrib/pipeline/prefetch_pipeline_2d.h" +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/contrib/pipeline/prefetch_tuner_registry.h" +#if defined(HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_INL_H_) == \ + defined(HWY_TARGET_TOGGLE) +#ifdef HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_INL_H_ +#undef HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_INL_H_ +#else +#define HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_INL_H_ +#endif + +#include +#include // IWYU pragma: keep +#include +#include + +#include "hwy/base.h" +#include "hwy/highway.h" + +HWY_BEFORE_NAMESPACE(); +namespace hwy { +namespace HWY_NAMESPACE { +namespace dot { + +namespace low_level { + +// ============================================================================= +// Generic Concept Definitions +// ============================================================================= + +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L + +// Concept: Accessor +// --------------------------------------------------------------------------- +// Represents the mapping interface between the flat iteration index (0 to N) +// and the underlying database indices. It also serves as the final sink where +// MemoryLayout writes calibrated scores. +template +concept OTM_Accessor = requires(const A c_accessor, size_t i) { + { c_accessor.size() } -> std::convertible_to; + { c_accessor.Index(i) } -> std::convertible_to; +}; + +// Concept: MemoryLayout +// --------------------------------------------------------------------------- +// Defines the contract for fetching dimensions, strides, and dynamically +// prefetching inner loops, independent of precision type (Float, BF16, Int8, or +// other quantization types). +// +// - Layout is the concrete type satisfying this contract (e.g. an +// implementation that fetches dimensions and strides for a specific +// quantization type and layout). +// - Accessor: The type of the accessor, which must satisfy OTM_Accessor. +// - AccumT: The type of the accumulator. +// - B_dp: The number of datapoints in a block. +template +concept OTM_MemoryLayout = + OTM_Accessor && + requires(const Layout layout, size_t dp_idx, size_t dp_end, + size_t dim_start, size_t dim_end, const AccumT* raw_accumulators, + Accessor& accessor, + hwy::pipeline::CachelineBundle& collector) { + // Yields the discrete memory cachelines to prefetch for a chunk + { layout.PopulateCachelines(dp_idx, dim_start, dim_end, collector) }; + { layout.CalibrateScores(raw_accumulators, dp_idx, dp_end, accessor) }; + // GetBasePtr MUST return a valid readable memory pointer even if dp_idx + // is out-of-bounds (e.g., mapping invalid indices to index 0). This + // guarantees that the ScorerKernel never segfaults and avoids branch + // divergence during unrolled SIMD loads. + { layout.GetBasePtr(dp_idx) } -> std::convertible_to; + }; + +// Concept: ScorerKernel +// --------------------------------------------------------------------------- +// The innermost compute block dealing exclusively with raw vectorized loops. +// +// - Kernel is the concrete type satisfying this contract. +template +concept OTM_ScorerKernel = + requires(Kernel& kernel, size_t dp_idx, const Layout& layout, + size_t dim_offset, size_t inner_end, size_t dim_end) { + typename Kernel::AccumT; + + // Pipeline Configuration Traits + // ----------------------------------------------------------------------- + // - kBlockDimensions: Query dimension slicing scale suitable for FMA + // loops. + { Kernel::kBlockDimensions } -> std::convertible_to; + + // Hook designed for executing pre-block operations (like broadcasting + // sub-centroid codebooks into SIMD registers prior to inner loop scans). + { kernel.PrepareDimensionBlock(dim_offset, dim_end) }; + // Core compute evaluating datapoint vector comparisons, accumulating + // into the requested accum reference. + { + kernel.ScoreBlock(dp_idx, layout, dim_offset, inner_end, dim_end, + std::declval()) + }; + }; + +#endif // __cpp_concepts + +// =========================================================================== +// WARNING: Low-Level Execution API +// =========================================================================== +// Do NOT call this struct's methods directly in production code! Bypassing the +// tuning framework prevents fleet-wide continuous profiling and telemetry +// collection. Prefer `hwy::OneToMany2DTiledPipeline::Run` below. +// +// A generic pipeline for executing 2D block-tiled one-to-many distance scoring. +// This splits scoring out into localized blocks that prevent cache evictions +// and FPU instruction bottlenecks. +// +// This generic pipeline retrieves hardware cache-bound constants directly from +// the execution kernel and policy descriptors. +// +// Required Traits: +// ScorerKernelType::kBlockDimensions: The number of dimensions scored per +// inner phase. This pins active query segments inside SIMD vector +// registers / L1 cache. Often 256 or 512, representing a pipelined +// saturation threshold. +struct OneToMany2DTiledPipeline { + private: + template + struct Callbacks { + const MemoryLayout& layout; + Accessor& accessor; + ScorerKernelType& kernel; + size_t simd_end; + size_t outer_block; + std::vector& raw_accumulators; + + void OnOuterBlockStart(size_t /*outer_idx*/, size_t /*outer_end*/) { + std::fill(raw_accumulators.begin(), raw_accumulators.end(), AccumT{0}); + } + + void PrepareInnerBlock(size_t /*outer_idx*/, size_t /*outer_end*/, + size_t inner_idx, size_t inner_end) { + kernel.PrepareDimensionBlock(inner_idx, inner_end); + } + + template + void PopulateCachelines(size_t outer_i, size_t inner_idx, size_t inner_end, + hwy::pipeline::CachelineBundle& collector) { + layout.PopulateCachelines(accessor.Index(outer_i), inner_idx, inner_end, + collector); + } + + void ComputeTask(size_t outer_i, size_t inner_idx, size_t inner_end) { + const size_t mapped_idx = accessor.Index(outer_i); + kernel.ScoreBlock(mapped_idx, layout, inner_idx, + std::min(inner_end, simd_end), inner_end, + raw_accumulators[outer_i % outer_block]); + } + + void OnOuterBlockFinish(size_t outer_idx, size_t outer_end) { + layout.CalibrateScores(raw_accumulators.data(), outer_idx, outer_end, + accessor); + } + }; + + public: + template + struct OTM_Policy : hwy::pipeline::Default2DTiledPrefetchPolicy { + static constexpr size_t kMaxCachelinesPerIter = + MemoryLayout::kMaxCachelinesPerIter; + }; + + // Executes the 2D-tiled scoring pipeline with explicitly provided arguments. + // Primarily useful for isolated micro-benchmarks or rigid environments that + // must bypass dynamic tuning. + template +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L + requires OTM_MemoryLayout && + OTM_ScorerKernel +#endif + static HWY_INLINE HWY_ATTR void Run( + const MemoryLayout& layout, Accessor& accessor, size_t total_dims, + size_t simd_end, ScorerKernelType& kernel, + hwy::pipeline::Tiling2DArgs tiling, + const hwy::pipeline::PrefetchArgs& prefetch_args) { + using AccumT = typename ScorerKernelType::AccumT; + + const size_t outer_block = tiling.outer_block; + std::vector raw_accumulators(outer_block); + + Callbacks cb{ + layout, accessor, kernel, simd_end, outer_block, raw_accumulators}; + + hwy::pipeline::low_level::PrefetchPipeline2DTiledLoop< + OTM_Policy>(accessor.size(), total_dims, cb, tiling, + prefetch_args); + } +}; + +} // namespace low_level + +// =========================================================================== +// Public Context-Aware Wrapper +// =========================================================================== +// OneToMany2DTiledPipeline +// --------------------------------------------------------------------------- +// A generic pipeline for executing 2D block-tiled one-to-many distance scoring. +// This splits scoring out into localized blocks that prevent cache evictions +// and FPU instruction bottlenecks. +struct OneToMany2DTiledPipeline { + // Executes the 2D-tiled scoring pipeline utilizing Context-Aware Tuning. + // This automatically resolves tuning arguments from the global registry, + // preserving caller-specified tiling geometry while injecting auto-tuned + // lookahead velocity. + template +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L + requires hwy::pipeline::low_level::IsPrefetchPolicy< + low_level::OneToMany2DTiledPipeline::OTM_Policy> +#endif + static HWY_INLINE HWY_ATTR void Run( + const MemoryLayout& layout, Accessor& accessor, size_t total_dims, + size_t simd_end, ScorerKernelType& kernel, + hwy::pipeline::Tiling2DArgs tiling = hwy::pipeline::Tiling2DArgs(), +#if defined(__clang__) || defined(__GNUC__) + const char* file_loc = __builtin_FILE(), int line_loc = __builtin_LINE() +#else + const char* file_loc = nullptr, int line_loc = 0 +#endif + ) { + auto loop_runner = [&](const hwy::pipeline::PrefetchArgs& args) { + low_level::OneToMany2DTiledPipeline::Run(layout, accessor, total_dims, + simd_end, kernel, tiling, args); + }; + + hwy::pipeline::low_level::DispatchTunedWorkload< + Hint, low_level::OneToMany2DTiledPipeline::OTM_Policy>( + tiling.TotalElements(accessor.size(), total_dims), file_loc, line_loc, + loop_runner); + } +}; + +} // namespace dot +} // namespace HWY_NAMESPACE +} // namespace hwy +HWY_AFTER_NAMESPACE(); +#endif // HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_INL_H_ diff --git a/hwy/contrib/dot/one_to_many_kernels-inl.h b/hwy/contrib/dot/one_to_many_kernels-inl.h new file mode 100644 index 0000000000..67f5fd0bfe --- /dev/null +++ b/hwy/contrib/dot/one_to_many_kernels-inl.h @@ -0,0 +1,139 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#if defined(HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_KERNELS_INL_H_) == \ + defined(HWY_TARGET_TOGGLE) +#ifdef HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_KERNELS_INL_H_ +#undef HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_KERNELS_INL_H_ +#else +#define HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_KERNELS_INL_H_ +#endif + +#include +#include + +#include "hwy/base.h" +#include "hwy/contrib/dot/dot-inl.h" + +HWY_BEFORE_NAMESPACE(); +namespace hwy { +namespace HWY_NAMESPACE { + +// ============================================================================= +// Reusable Kernels which use Highway's overloaded Dot::Compute. +// ============================================================================= + +struct OTM_DefaultScorerKernel { + class FloatToBFloat16 { + public: + using AccumT = float; + + static constexpr size_t kBlockDimensions = 256; + + explicit FloatToBFloat16(const float* query_data, size_t query_dims) + : query_data_(query_data) { + (void)query_data_; // Silence unused private field warnings +#if HWY_NATIVE_DOT_BF16 || HWY_IDE + // If the CPU supports native bfloat16 dot-products (e.g., AVX512_BF16 or + // ARM Neon BFDOT), it is drastically faster to demote the query to + // bfloat16 once and execute pure bf16 * bf16 loops instead of promoting + // every single database element to f32 on the fly! + query_bf16_data_.resize(query_dims); + for (size_t d = 0; d < query_dims; ++d) { + query_bf16_data_[d] = BF16FromF32(query_data[d]); + } +#endif + } + + HWY_INLINE void PrepareDimensionBlock(size_t /*dim_offset*/, + size_t /*dim_end*/) {} + + template + HWY_INLINE void ScoreBlock(size_t dp_idx, const Policy& policy, + size_t dim_offset, size_t /*inner_end*/, + size_t dim_end, float& accum) { + const bfloat16_t* bfloat_ptr = + static_cast(policy.GetBasePtr(dp_idx)) + + dim_offset; + size_t num_elements = dim_end - dim_offset; + +#if HWY_NATIVE_DOT_BF16 || HWY_IDE + const hwy::HWY_NAMESPACE::ScalableTag dbf; + const bfloat16_t* bfloat_query_ptr = query_bf16_data_.data() + dim_offset; + // Dispatches natively to CPU BFloat16 FMA instructions (AVX512_BF16 / + // AMX) + accum += Dot::Compute<0>(dbf, bfloat_query_ptr, bfloat_ptr, num_elements); +#else + const hwy::HWY_NAMESPACE::ScalableTag df; + const float* float_ptr = query_data_ + dim_offset; + // Falls back to highway-abstracted upcast-and-multiply: f32 * bf16 -> f32 + accum += Dot::Compute<0>(df, float_ptr, bfloat_ptr, num_elements); +#endif + } + + private: + HWY_MAYBE_UNUSED const float* query_data_; +#if HWY_NATIVE_DOT_BF16 || HWY_IDE + // std::vector should be included where FloatToBFloat16ScorerKernel is used. + std::vector query_bf16_data_; +#endif + }; + + // SameType. + // --------------------------------------------------------------------------- + // A universal ScorerKernel for computing dot products when the query and + // database elements are natively stored in the same data type. By leveraging + // Highway's overloaded Dot::Compute architecture, this struct handles + // execution and proper accumulator typing automatically (e.g. + // float*float->float, int16*int16->int32). + template + class SameType { + public: + using AccumT = AccumType; // i.e. float for f32/bf16, int32_t for int16 + + static constexpr size_t kBlockDimensions = 256; + + explicit SameType(const T* query_data, size_t query_dims) + : query_data_(query_data) { + (void)query_dims; // Silence unused variable warning uniformly + } + + HWY_INLINE void PrepareDimensionBlock(size_t /*dim_offset*/, + size_t /*dim_end*/) {} + + template + HWY_INLINE void ScoreBlock(size_t dp_idx, const Policy& policy, + size_t dim_offset, size_t /*inner_end*/, + size_t dim_end, AccumT& accum) { + const T* t_dp_ptr = + static_cast(policy.GetBasePtr(dp_idx)) + dim_offset; + const T* t_query_ptr = query_data_ + dim_offset; + size_t num_elements = dim_end - dim_offset; + + const hwy::HWY_NAMESPACE::ScalableTag d; + accum += Dot::Compute<0>(d, t_query_ptr, t_dp_ptr, num_elements); + } + + private: + const T* query_data_; + }; +}; // struct OTM_ReusableKernels + +} // namespace HWY_NAMESPACE +} // namespace hwy +HWY_AFTER_NAMESPACE(); +#endif // HIGHWAY_HWY_CONTRIB_DOT_ONE_TO_MANY_KERNELS_INL_H_ diff --git a/hwy/contrib/dot/one_to_many_test.cc b/hwy/contrib/dot/one_to_many_test.cc new file mode 100644 index 0000000000..059586024d --- /dev/null +++ b/hwy/contrib/dot/one_to_many_test.cc @@ -0,0 +1,510 @@ +#include +#include +#include + +#include "hwy/contrib/pipeline/prefetch_pipeline_2d.h" +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/contrib/pipeline/prefetch_tuner.h" +#include "hwy/contrib/pipeline/prefetch_tuner_registry.h" + +// clang-format off +#undef HWY_TARGET_INCLUDE +#define HWY_TARGET_INCLUDE "hwy/contrib/dot/one_to_many_test.cc" // NOLINT +#include "hwy/foreach_target.h" // IWYU pragma: keep +#include "hwy/highway.h" +#include "hwy/contrib/dot/one_to_many_kernels-inl.h" +#include "hwy/contrib/dot/one_to_many-inl.h" +#include "hwy/tests/test_util-inl.h" +// clang-format on + +HWY_BEFORE_NAMESPACE(); +namespace hwy { +namespace HWY_NAMESPACE { +namespace { + +// ========================================================================= +// Example 1: Float32 Query against BFloat16 Dataset +// ========================================================================= + +// An accessor handles mapping logic from loop iteration `i` to database `idx`, +// and stores the final calibrated scores into output vectors. +class SimpleAccessor { + public: + explicit SimpleAccessor(size_t size) : size_(size), scores_(size, 0.0f) {} + + size_t size() const { return size_; } + size_t Index(size_t i) const { return i; } // Direct 1:1 mapped indices + + void SetScore(size_t i, float score) { scores_[i] = score; } + float GetScore(size_t i) const { return scores_[i]; } + + private: + size_t size_; + std::vector scores_; +}; + +// LayoutPolicy dictating how to iterate over our flat BFloat16 array and +// prefetch. +class BFloat16LayoutPolicy { + public: + HWY_MAYBE_UNUSED static constexpr size_t kMaxCachelinesPerIter = 1; + HWY_MAYBE_UNUSED static constexpr size_t kBlockDatapoints = 128; + HWY_MAYBE_UNUSED static constexpr size_t kPrefetchLookaheadL1 = 4; + HWY_MAYBE_UNUSED static constexpr size_t kPrefetchLookaheadL3 = 32; + + BFloat16LayoutPolicy(const bfloat16_t* dataset_ptr, size_t dimensions) + : dataset_ptr_(dataset_ptr), dimensions_(dimensions) {} + + HWY_INLINE void PopulateCachelines(size_t /*dp_idx*/, size_t /*dim_start*/, + size_t /*dim_end*/, + auto& /*collector*/) const {} + + // Gets the exact starting address in memory for a target datapoint + HWY_INLINE const void* GetBasePtr(size_t dp_idx) const { + return dataset_ptr_ + dp_idx * dimensions_; + } + + // Translates intermediate pipeline blocks into final stored outputs. + // In a real framework, this might apply scaling factors and subtract biases. + HWY_INLINE void CalibrateScores(const float* raw_accumulators, size_t dp_idx, + size_t dp_end, + SimpleAccessor& accessor) const { + for (size_t i = dp_idx; i < dp_end; ++i) { + accessor.SetScore(i, raw_accumulators[i - dp_idx]); + } + } + + private: + const bfloat16_t* dataset_ptr_; + size_t dimensions_; +}; + +template +class GenericLayoutPolicy { + public: + HWY_MAYBE_UNUSED static constexpr size_t kMaxCachelinesPerIter = 1; + HWY_MAYBE_UNUSED static constexpr size_t kBlockDatapoints = 128; + HWY_MAYBE_UNUSED static constexpr size_t kPrefetchLookaheadL1 = 4; + HWY_MAYBE_UNUSED static constexpr size_t kPrefetchLookaheadL3 = 32; + + GenericLayoutPolicy(const T* dataset_ptr, size_t dimensions) + : dataset_ptr_(dataset_ptr), dimensions_(dimensions) {} + + HWY_INLINE void PopulateCachelines(size_t /*dp_idx*/, size_t /*dim_start*/, + size_t /*dim_end*/, + auto& /*collector*/) const {} + + HWY_INLINE const void* GetBasePtr(size_t dp_idx) const { + return dataset_ptr_ + dp_idx * dimensions_; + } + + template + HWY_INLINE void CalibrateScores(const AccumT* raw_accumulators, size_t dp_idx, + size_t dp_end, + SimpleAccessor& accessor) const { + for (size_t i = dp_idx; i < dp_end; ++i) { + accessor.SetScore(i, static_cast(raw_accumulators[i - dp_idx])); + } + } + + private: + const T* dataset_ptr_; + size_t dimensions_; +}; + +// ========================================================================= +// Example 3: LUT16 with PrepareDimensionBlock (Fake Codebook) +// ========================================================================= + +// LayoutPolicy for 4-bit packed dataset stored as uint8_t (2 dimensions per +// byte) +class LUT16LayoutPolicy { + public: + HWY_MAYBE_UNUSED static constexpr size_t kMaxCachelinesPerIter = 1; + // Since LUT16 stores 2 dimensions per byte (0.5 bytes per dimension), + // we can heavily scale the datapoint blocking and lookahead arrays + // while comfortably fitting in L1D cache constraints! + HWY_MAYBE_UNUSED static constexpr size_t kBlockDatapoints = 512; + HWY_MAYBE_UNUSED static constexpr size_t kPrefetchLookaheadL1 = 8; + HWY_MAYBE_UNUSED static constexpr size_t kPrefetchLookaheadL3 = 64; + + LUT16LayoutPolicy(const uint8_t* dataset_ptr, size_t dimensions) + : dataset_ptr_(dataset_ptr), dimensions_(dimensions) {} + + HWY_INLINE void PopulateCachelines(size_t /*dp_idx*/, size_t /*dim_start*/, + size_t /*dim_end*/, + auto& /*collector*/) const {} + + HWY_INLINE const void* GetBasePtr(size_t dp_idx) const { + return dataset_ptr_ + dp_idx * (dimensions_ / 2); + } + + HWY_INLINE void CalibrateScores(const float* raw_accumulators, size_t dp_idx, + size_t dp_end, + SimpleAccessor& accessor) const { + for (size_t i = dp_idx; i < dp_end; ++i) { + accessor.SetScore(i, raw_accumulators[i - dp_idx]); + } + } + + private: + const uint8_t* dataset_ptr_; + size_t dimensions_; +}; + +class FakeLut16ScorerKernel { + public: + using AccumT = float; + HWY_MAYBE_UNUSED static constexpr size_t kBlockDimensions = 512; + + explicit FakeLut16ScorerKernel(const std::vector& fake_codebooks) + : fake_codebooks_(fake_codebooks) {} + + HWY_INLINE void PrepareDimensionBlock(size_t dim_offset, size_t dim_end) { + // Simulates loading a codebook (e.g. centroids) for the current dimensions + // into hot state/registers. + active_codebook_.clear(); + for (size_t d = dim_offset; d < dim_end; ++d) { + // For this fake test, each dimension has 16 float values + for (size_t c = 0; c < 16; ++c) { + active_codebook_.push_back(fake_codebooks_[d * 16 + c]); + } + } + } + + template + HWY_INLINE void ScoreBlock(size_t dp_idx, const Policy& policy, + size_t dim_offset, size_t /*inner_end*/, + size_t dim_end, float& accum) { + const void* dp_ptr = policy.GetBasePtr(dp_idx); + // Read the 4-bit packed data + const uint8_t* u8_ptr = + static_cast(dp_ptr) + (dim_offset / 2); + size_t num_elements = dim_end - dim_offset; + size_t num_bytes = num_elements / 2; + float sum0 = 0.0f; + float sum1 = 0.0f; + float sum2 = 0.0f; + float sum3 = 0.0f; + size_t i = 0; + + // Unroll by 4 bytes (8 dimensions) per iteration to maximize FPU pipeline + // throughput and reduce loop branch conditions drastically. + // + // We also use 4 separate local accumulators so the CPU can calculate the + // addition instructions in parallel without waiting on a single `sum` + // execution dependency chain! + for (; i + 3 < num_bytes; i += 4) { + uint8_t p0 = u8_ptr[i + 0]; + uint8_t p1 = u8_ptr[i + 1]; + uint8_t p2 = u8_ptr[i + 2]; + uint8_t p3 = u8_ptr[i + 3]; + + // Evaluate distances dynamically from the active block cache + sum0 += active_codebook_[((i + 0) * 2) * 16 + (p0 & 0x0F)]; + sum0 += active_codebook_[((i + 0) * 2 + 1) * 16 + (p0 >> 4)]; + + sum1 += active_codebook_[((i + 1) * 2) * 16 + (p1 & 0x0F)]; + sum1 += active_codebook_[((i + 1) * 2 + 1) * 16 + (p1 >> 4)]; + + sum2 += active_codebook_[((i + 2) * 2) * 16 + (p2 & 0x0F)]; + sum2 += active_codebook_[((i + 2) * 2 + 1) * 16 + (p2 >> 4)]; + + sum3 += active_codebook_[((i + 3) * 2) * 16 + (p3 & 0x0F)]; + sum3 += active_codebook_[((i + 3) * 2 + 1) * 16 + (p3 >> 4)]; + } + + // Combine accumulators using an addition tree + sum0 += sum1; + sum2 += sum3; + sum0 += sum2; + + // Scalar cleanup loop + for (; i < num_bytes; ++i) { + uint8_t packed = u8_ptr[i]; + sum0 += active_codebook_[(i * 2) * 16 + (packed & 0x0F)]; + sum0 += active_codebook_[(i * 2 + 1) * 16 + (packed >> 4)]; + } + accum += sum0; + } + + private: + const std::vector& fake_codebooks_; + // Reusable hot state initialized per block by PrepareDimensionBlock + std::vector active_codebook_; +}; + +void TestOneToManyFloatFloat() { + const size_t num_dps = 1000; + const size_t dims = 512; + + std::vector dataset(num_dps * dims); + for (size_t i = 0; i < dataset.size(); ++i) { + dataset[i] = static_cast(i % 5) * 0.1f; + } + + std::vector query(dims); + for (size_t i = 0; i < dims; ++i) { + query[i] = static_cast(i % 3) * 0.2f; + } + + SimpleAccessor accessor(num_dps); + GenericLayoutPolicy policy(dataset.data(), dims); + OTM_DefaultScorerKernel::SameType kernel(query.data(), dims); + dot::low_level::OneToMany2DTiledPipeline::Run( + policy, accessor, dims, dims, kernel, hwy::pipeline::Tiling2DArgs(), + hwy::pipeline::PrefetchArgs::DefaultSequential()); + + for (size_t i = 0; i < num_dps; ++i) { + float expected = 0.0f; + for (size_t d = 0; d < dims; ++d) { + expected += dataset[i * dims + d] * query[d]; + } + + float actual = accessor.GetScore(i); + HWY_ASSERT(std::abs(expected - actual) < 1e-4); + } +} + +void TestOneToManyFloatBfloat16() { + const size_t num_dps = 1000; + const size_t dims = 512; + + std::vector dataset(num_dps * dims); + for (size_t i = 0; i < dataset.size(); ++i) { + dataset[i] = BF16FromF32(static_cast(i % 5) * 0.1f); + } + + std::vector query(dims); + for (size_t i = 0; i < dims; ++i) { + query[i] = static_cast(i % 3) * 0.2f; + } + + SimpleAccessor accessor(num_dps); + BFloat16LayoutPolicy policy(dataset.data(), dims); + OTM_DefaultScorerKernel::FloatToBFloat16 kernel(query.data(), dims); + + dot::low_level::OneToMany2DTiledPipeline::Run( + policy, accessor, dims, dims, kernel, hwy::pipeline::Tiling2DArgs(), + hwy::pipeline::PrefetchArgs::DefaultSequential()); + + // Validate mathematical computations. Note that BFloat16 only has 8 bits of + // precision total (7 bits mantissa). Accumulating 512 elements can drift + // significantly around truncation limits across CPU architectures + // (e.g. native AVX512_BF16 vs emulated floats), so we must scale tolerance + // appropriately against the machine epsilon! + for (size_t i = 0; i < num_dps; ++i) { + double expected = 0.0; + for (size_t d = 0; d < dims; ++d) { + expected += + static_cast(F32FromBF16(dataset[i * dims + d])) * query[d]; + } + + float actual = accessor.GetScore(i); + HWY_ASSERT(std::abs(expected - actual) < 1.0); + } +} + +void TestOneToManyLUT16() { + const size_t num_dps = 1000; + const size_t dims = 512; + + // 4-bit LUT16 packs 2 elements per byte. + std::vector dataset(num_dps * (dims / 2)); + for (size_t i = 0; i < dataset.size(); ++i) { + dataset[i] = static_cast(i % 256); + } + + // Codebooks contain 16 float centroids per dimension slice. + std::vector fake_codebooks(dims * 16); + for (size_t i = 0; i < fake_codebooks.size(); ++i) { + fake_codebooks[i] = static_cast(i % 7) * 0.1f; + } + + SimpleAccessor accessor(num_dps); + LUT16LayoutPolicy policy(dataset.data(), dims); + FakeLut16ScorerKernel kernel(fake_codebooks); + dot::low_level::OneToMany2DTiledPipeline::Run( + policy, accessor, dims, dims, kernel, hwy::pipeline::Tiling2DArgs(), + hwy::pipeline::PrefetchArgs::DefaultSequential()); + + for (size_t i = 0; i < num_dps; ++i) { + float expected = 0.0f; + for (size_t d = 0; d < dims / 2; ++d) { + uint8_t packed = dataset[i * (dims / 2) + d]; + uint8_t c1 = packed & 0x0F; + uint8_t c2 = packed >> 4; + expected += fake_codebooks[(d * 2) * 16 + c1]; + expected += fake_codebooks[(d * 2 + 1) * 16 + c2]; + } + + float actual = accessor.GetScore(i); + HWY_ASSERT(std::abs(expected - actual) < 1e-4); + } +} + +// ========================================================================= +// Example 4: Public Wrapper Test Coverage +// ========================================================================= + +void TestOneToManyPublicWrapperCorrectness() { + const size_t num_dps = 1000; + const size_t dims = 512; + + std::vector dataset(num_dps * dims); + for (size_t i = 0; i < dataset.size(); ++i) { + dataset[i] = static_cast(i % 5) * 0.1f; + } + + std::vector query(dims); + for (size_t i = 0; i < dims; ++i) { + query[i] = static_cast(i % 3) * 0.2f; + } + + SimpleAccessor accessor(num_dps); + GenericLayoutPolicy policy(dataset.data(), dims); + OTM_DefaultScorerKernel::SameType kernel(query.data(), dims); + + // Call the PUBLIC wrapper instead of low_level + dot::OneToMany2DTiledPipeline::Run< + hwy::pipeline::PrefetchTuningHint::kSequential>(policy, accessor, dims, + dims, kernel); + + for (size_t i = 0; i < num_dps; ++i) { + float expected = 0.0f; + for (size_t d = 0; d < dims; ++d) { + expected += dataset[i * dims + d] * query[d]; + } + float actual = accessor.GetScore(i); + HWY_ASSERT(std::abs(expected - actual) < 1e-4); + } +} + +struct TestOtmMetricContext { + bool called = false; + std::string captured_file; + int captured_line = 0; + size_t captured_total = 0; + bool captured_tiny = false; + hwy::pipeline::PrefetchTuningHint captured_hint = + hwy::pipeline::PrefetchTuningHint::kAuto; +}; + +void FakeOtmMetricCb(void* user_data, float) { + static_cast(user_data)->called = true; +} + +class MockOtmTuner : public hwy::pipeline::low_level::PrefetchTuner { + public: + explicit MockOtmTuner(TestOtmMetricContext* ctx) : ctx_(ctx) {} + hwy::pipeline::low_level::PrefetchTuningScope CreateScope( + const hwy::pipeline::low_level::PrefetchTuningContext& context) + const override { + if (context.file_loc != nullptr) { + ctx_->captured_file = context.file_loc; + } + ctx_->captured_line = context.line_loc; + ctx_->captured_total = context.total_elements; + ctx_->captured_tiny = context.is_ultra_tiny; + ctx_->captured_hint = context.hint; + + hwy::pipeline::PrefetchArgs args{.deep_lookahead = 12, + .shallow_lookahead = 2}; + return hwy::pipeline::low_level::PrefetchTuningScope(args, FakeOtmMetricCb, + ctx_, 1.0f); + } + + hwy::pipeline::CallsiteId RegisterContext( + const hwy::pipeline::low_level::PrefetchTuningContext& context) + const override { + return 0; + } + + hwy::pipeline::low_level::PrefetchTuningScope CreateScopeByCallsiteId( + hwy::pipeline::CallsiteId callsite_id, + const hwy::pipeline::low_level::PrefetchTuningContext& context) + const override { + return CreateScope(context); + } + + private: + TestOtmMetricContext* ctx_; +}; + +void TestOneToManyPublicWrapperTunerIntegration() { + TestOtmMetricContext ctx; + MockOtmTuner tuner(&ctx); + hwy::pipeline::low_level::GetGlobalPrefetchTunerRegistry() = &tuner; + + const size_t num_dps = 128000; + const size_t dims = 128; // total_elements = 128,000 (Huge bucket!) + + std::vector dataset(num_dps * dims, 0.1f); + std::vector query(dims, 0.2f); + SimpleAccessor accessor(num_dps); + GenericLayoutPolicy policy(dataset.data(), dims); + OTM_DefaultScorerKernel::SameType kernel(query.data(), dims); + + const int expected_base_line = __LINE__ + 1; + dot::OneToMany2DTiledPipeline::Run< + hwy::pipeline::PrefetchTuningHint::kRandom>(policy, accessor, dims, dims, + kernel); + + HWY_ASSERT(ctx.called); + HWY_ASSERT(ctx.captured_total == 128000); + HWY_ASSERT(!ctx.captured_tiny); + HWY_ASSERT(ctx.captured_hint == hwy::pipeline::PrefetchTuningHint::kRandom); + HWY_ASSERT(ctx.captured_file.find("one_to_many_test.cc") != + std::string::npos); // NOLINT + // Random Huge (128,000 <= 65536 is false -> Huge Random -> +600000) + HWY_ASSERT(ctx.captured_line == expected_base_line + 600000); + + hwy::pipeline::low_level::GetGlobalPrefetchTunerRegistry() = nullptr; +} + +void TestOneToManyPublicWrapperUltraTinyShortCircuit() { + TestOtmMetricContext ctx; + MockOtmTuner tuner(&ctx); + hwy::pipeline::low_level::GetGlobalPrefetchTunerRegistry() = &tuner; + + const size_t num_dps = 2; + const size_t dims = 8; // total_elements = 16 (< 32 UltraTiny!) + + std::vector dataset(num_dps * dims, 0.1f); + std::vector query(dims, 0.2f); + SimpleAccessor accessor(num_dps); + GenericLayoutPolicy policy(dataset.data(), dims); + OTM_DefaultScorerKernel::SameType kernel(query.data(), dims); + + dot::OneToMany2DTiledPipeline::Run< + hwy::pipeline::PrefetchTuningHint::kSequential>(policy, accessor, dims, + dims, kernel); + + // Tuner plugin should be bypassed entirely! + HWY_ASSERT(!ctx.called); + + hwy::pipeline::low_level::GetGlobalPrefetchTunerRegistry() = nullptr; +} + +} // namespace +} // namespace HWY_NAMESPACE +} // namespace hwy +HWY_AFTER_NAMESPACE(); + +#if HWY_ONCE +namespace hwy { +namespace { +HWY_BEFORE_TEST(OneToManyTest); +HWY_EXPORT_AND_TEST_P(OneToManyTest, TestOneToManyLUT16); +HWY_EXPORT_AND_TEST_P(OneToManyTest, TestOneToManyFloatFloat); +HWY_EXPORT_AND_TEST_P(OneToManyTest, TestOneToManyFloatBfloat16); +HWY_EXPORT_AND_TEST_P(OneToManyTest, TestOneToManyPublicWrapperCorrectness); +HWY_EXPORT_AND_TEST_P(OneToManyTest, + TestOneToManyPublicWrapperTunerIntegration); +HWY_EXPORT_AND_TEST_P(OneToManyTest, + TestOneToManyPublicWrapperUltraTinyShortCircuit); +HWY_AFTER_TEST(); +} // namespace +} // namespace hwy +HWY_TEST_MAIN(); +#endif // HWY_ONCE diff --git a/hwy/contrib/pipeline/README.md b/hwy/contrib/pipeline/README.md new file mode 100644 index 0000000000..bb7a328b33 --- /dev/null +++ b/hwy/contrib/pipeline/README.md @@ -0,0 +1,22 @@ +# Highway Prefetch Pipelining (`//hwy/contrib/pipeline`) + +## Overview + +This directory provides abstractions for **Dual-Tier Software Prefetching**, +designed to maximize execution throughput on workloads bound by memory latency +and hardware prefetcher limitations (such as L1 Line Fill Buffer starvation). + +By decoupling memory traversals into two separate hardware horizons (a Deep L3 +fetch and a Shallow L1 fetch), pipelines can hide massive DRAM latencies while +safely avoiding cache churn and buffer exhaustion. + +## Status + +**Experimental / Alpha** + +The dual tier prefetch loop abstractions are actively in development. While the +core algorithms deliver significant performance enhancements on memory-bound +workloads, the **API is subject to iteration and change**. + +Users should expect potential breaking changes in parameter structures, callback +signatures in future releases. diff --git a/hwy/contrib/pipeline/prefetch_pipeline.h b/hwy/contrib/pipeline/prefetch_pipeline.h new file mode 100644 index 0000000000..3c8d5b01fc --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_pipeline.h @@ -0,0 +1,316 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_PIPELINE_H_ +#define HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_PIPELINE_H_ + +#include + +#include "hwy/base.h" +#include "hwy/cache_control.h" +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/contrib/pipeline/prefetch_tuner_registry.h" + +namespace hwy { +namespace pipeline { +namespace low_level { + +// =========================================================================== +// WARNING: Low-Level Execution API +// =========================================================================== +// Do NOT call this function directly in production code! Bypassing the tuning +// framework prevents fleet-wide continuous profiling and telemetry collection. +// Prefer `hwy::PrefetchPipelineLoop` in `prefetch_pipeline.h`. +// +// Design Philosophy: +// While this helper aims to deeply accelerate predictable memory-bound loops, +// its core tenet is "do no harm in the worst case". By carefully staging +// memory into L3 before migrating it to L1, and rigorously capping the active +// footprint against hardware Line Fill Buffer (LFB) limits, it ensures that +// memory bandwidth is maximized without accidentally thrashing the cache or +// stalling the processor pipelines. +// +// Evaluates a pipelined loop over the range [start, end) with two stages of +// rolling prefetch lookahead: a deep prefetch (e.g. L3) and a shallow prefetch +// (e.g. L1). This correctly stages data transitions from main memory -> L3 -> +// L1 to maximize memory bandwidth and keep CPU Line Fill Buffers from stalling. +// +// Policy: A struct dictating cache limits and loop constants. Defaults to +// `DefaultPrefetchPolicy`. +// CachelinesProvider: A callable type matching the signature documented +// in `PrefetchPipelineCachelineProvider` concept. It +// resolves cacheline pointers for a given index `i`. +// TaskFn: A callable type matching the signature documented in +// `PrefetchPipelineTask` concept. It represents the core evaluation +// logic for index `i`. +// args: A `PrefetchArgs` configuration that controls the deep (L3) and +// shallow (L1) lookahead pipeline distances. Passing a reasonable +template +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L + requires low_level::IsPrefetchPolicy && + PrefetchPipelineCachelineProvider && + PrefetchPipelineTask +#endif +inline void PrefetchPipelineLoop(size_t start, size_t end, + const CachelinesProvider& populate_cachelines, + const TaskFn& task, const PrefetchArgs& args) { + // Gracefully degrade inverted configurations if AutoTune generates them. + size_t actual_shallow = args.shallow_lookahead; + size_t actual_deep = args.deep_lookahead; + HWY_DASSERT(actual_deep == 0 || actual_deep > actual_shallow); + if (actual_shallow >= actual_deep) { + actual_deep = 0; + } + + const size_t initial_shallow_prefetch_end = + HWY_MIN(start + actual_shallow, end); + const size_t initial_deep_prefetch_end = HWY_MIN(start + actual_deep, end); + + // Reusable cache injection loops: + auto execute_deep_prefetch = [&](const CachelineBundle& bundle) + HWY_ATTR_CACHE { + if (actual_deep == 0) return; + for (size_t r = 0; r < bundle.count; ++r) { + DeepPrefetchFn(bundle.ptrs[r]); + } + }; + auto execute_shallow_prefetch = + [&](const CachelineBundle& bundle) HWY_ATTR_CACHE { +#if HWY_IS_DEBUG_BUILD + // Safeguard: The active L1 hardware queue footprint should not + // exceed the physical Miss Status Holding Registers (MSHRs). + // (e.g. 10-12 Line Fill Buffers on legacy Intel architectures). + // Exceeding this generates catastrophic silent memory stalls. + HWY_DASSERT(bundle.count * actual_shallow <= Policy::kNumMSHRs); +#endif + for (size_t r = 0; r < bundle.count; ++r) { + ShallowPrefetchFn(bundle.ptrs[r]); + } + }; + + // Hoisted state variables to avoid tight-loop reallocation thrashing: + CachelineBundle bundle; + + // ------------------------------------------------------------------------ + // Branchless loop execution via compile-time strategy + // ------------------------------------------------------------------------ + + // NOTE: Use a strict `if constexpr ... else if constexpr ... else` chain to + // discard not matched branches at compile time. + if constexpr (Policy::kStrategy == PrefetchStrategy::kNoPrefetch) { + for (size_t i = start; i < end; ++i) { + task(i); + } + } else if constexpr (Policy::kStrategy == + PrefetchStrategy::kShallowLookaheadOnly) { + const size_t limit_shallow = + (end > start + actual_shallow) ? end - actual_shallow : start; + // Startup prefetching + for (size_t i = start; i < initial_shallow_prefetch_end; ++i) { + bundle.count = 0; + populate_cachelines(i, bundle); + execute_shallow_prefetch(bundle); + } + // Main sliding loop + for (size_t i = start; i < limit_shallow; ++i) { + bundle.count = 0; + populate_cachelines(i + actual_shallow, bundle); + execute_shallow_prefetch(bundle); + task(i); + } + // Task drain + for (size_t i = limit_shallow; i < end; ++i) { + task(i); + } + } else if constexpr (Policy::kStrategy == + PrefetchStrategy::kDeepLookaheadOnly) { + const size_t limit_deep = + (end > start + actual_deep) ? end - actual_deep : start; + // Startup prefetching + for (size_t i = start; i < initial_deep_prefetch_end; ++i) { + bundle.count = 0; + populate_cachelines(i, bundle); + execute_deep_prefetch(bundle); + } + // Main sliding loop + for (size_t i = start; i < limit_deep; ++i) { + bundle.count = 0; + populate_cachelines(i + actual_deep, bundle); + execute_deep_prefetch(bundle); + task(i); + } + // Task drain + for (size_t i = limit_deep; i < end; ++i) { + task(i); + } + } else if constexpr (Policy::kStrategy == PrefetchStrategy::kMiniBatchDeep || + Policy::kStrategy == + PrefetchStrategy::kMiniBatchShallow) { + const size_t batch_size = + (Policy::kStrategy == PrefetchStrategy::kMiniBatchDeep) + ? actual_deep + : actual_shallow; + + for (size_t b = start; b < end; b += batch_size) { + const size_t b_end = HWY_MIN(b + batch_size, end); + for (size_t p = b; p < b_end; ++p) { + bundle.count = 0; + populate_cachelines(p, bundle); + if constexpr (Policy::kStrategy == PrefetchStrategy::kMiniBatchDeep) { + execute_deep_prefetch(bundle); + } else { + execute_shallow_prefetch(bundle); + } + } + for (size_t p = b; p < b_end; ++p) { + task(p); + } + } + } else { + // Fallback: Default Staggered Pipeline (PrefetchStrategy::kDualTier) + // A meticulously unrolled Dual-Tier pipeline is functionally necessary here + // to stage data into L3 before pulling it into L1, preventing LFB stalls + // on Intel, while providing optimal micro-pipelining on Zen/Maple. + + // Phase 1: Overlapping L1 (shallow) and L3 (deep) startup pipeline horizons + for (size_t i = start; i < initial_shallow_prefetch_end; ++i) { + bundle.count = 0; + populate_cachelines(i, bundle); + execute_deep_prefetch(bundle); + execute_shallow_prefetch(bundle); + } + + // Phase 2: Outstanding L3 (deep) horizons which haven't entered L1 window + for (size_t i = initial_shallow_prefetch_end; i < initial_deep_prefetch_end; + ++i) { + bundle.count = 0; + populate_cachelines(i, bundle); + execute_deep_prefetch(bundle); + } + + // Phase 3: Main execution. + // Instead of a single loop with bounds-checking `if` statements, we split + // it into three branchless phases. This avoids branch mispredictions in the + // hot loop and prevents the user provided `populate_cachelines` from + // ever receiving an out-of-bounds index (requiring them to defensively + // handle it). + + const size_t limit_deep = (actual_deep == 0 || start + actual_deep >= end) + ? start + : end - actual_deep; + const size_t limit_shallow = + (actual_shallow == 0 || start + actual_shallow >= end) + ? start + : end - actual_shallow; + + // 3a: Both deep and shallow prefetches are within bounds. + for (size_t i = start; i < limit_deep; ++i) { + bundle.count = 0; + populate_cachelines(i + actual_deep, bundle); + execute_deep_prefetch(bundle); + + if (actual_shallow > 0) { + bundle.count = 0; + populate_cachelines(i + actual_shallow, bundle); + execute_shallow_prefetch(bundle); + } + + task(i); + } + + // 3b: Only shallow prefetch is within bounds. + for (size_t i = limit_deep; i < limit_shallow; ++i) { + bundle.count = 0; + populate_cachelines(i + actual_shallow, bundle); + execute_shallow_prefetch(bundle); + + task(i); + } + + // 3c: No prefetches are within bounds, just finish the tasks. + const size_t task_drain_start = HWY_MAX(limit_deep, limit_shallow); + for (size_t i = task_drain_start; i < end; ++i) { + task(i); + } + } +} + +} // namespace low_level + +// --------------------------------------------------------------------------- +// Context-Aware Loop Wrapper (Production Entry Point) +// --------------------------------------------------------------------------- +// Evaluates a pipelined loop over [start, end) using dynamic lookahead tuning. +// +// Parameters: +// start, end: The iteration range of the workload. +// populate_cachelines: Callable providing cacheline addresses for index `i`. +// task: Callable executing the compute logic for index `i`. +// file_loc, line_loc: Lexical call-site identifiers used as the primary key +// by the autotuning framework (AutoTune / AutoFDO) to +// isolate and maintain historical tuning profiles. +// +// Call-Site Stability & Autotuning Best Practices: +// By default, `file_loc` and `line_loc` automatically capture the physical +// source file and line number via `__builtin_FILE()` and `__builtin_LINE()`. +// Under normal development, this provides zero-boilerplate profile isolation. +// +// However, the stability requirements for `file_loc` and `line_loc` largely +// depend on the underlying prefetch tuner mechanism in use. +// +// If your callsite shifts frequently and relies on compile-time AutoFDO or +// requires strict zero-warmup latency stability, it is highly recommended to +// pass a fixed `file_loc` and `line_loc` (e.g., a stable virtual file string +// like "my_project_stable_loop" and a fixed dummy line number like 1000). +// +// WARNING: When providing fixed identifiers, ensure they are globally unique +// across your project. Reusing the same fixed `(file_loc, line_loc)` for +// distinct loops will cause profile aliasing, where the tuner attempts to fit +// conflicting access patterns into a single shared profile bucket. +template +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L + requires low_level::IsPrefetchPolicy +#endif +inline void PrefetchPipelineLoop(size_t start, size_t end, + const CachelinesProvider& populate_cachelines, + const TaskFn& task, +#if defined(__clang__) || defined(__GNUC__) + const char* file_loc = __builtin_FILE(), + int line_loc = __builtin_LINE()) { +#else + const char* file_loc = nullptr, + int line_loc = 0) { +#endif + const size_t total_elements = end > start ? end - start : 0; + + auto loop_runner = [&](const PrefetchArgs& args) { + hwy::pipeline::low_level::PrefetchPipelineLoop( + start, end, populate_cachelines, task, args); + }; + + low_level::DispatchTunedWorkload(total_elements, file_loc, + line_loc, loop_runner); +} + +} // namespace pipeline +} // namespace hwy + +#endif // HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_PIPELINE_H_ diff --git a/hwy/contrib/pipeline/prefetch_pipeline_2d.h b/hwy/contrib/pipeline/prefetch_pipeline_2d.h new file mode 100644 index 0000000000..4884708844 --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_pipeline_2d.h @@ -0,0 +1,205 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_PIPELINE_2D_H_ +#define HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_PIPELINE_2D_H_ + +#include + +#include "hwy/base.h" +#include "hwy/contrib/pipeline/prefetch_pipeline.h" +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/contrib/pipeline/prefetch_tuner_registry.h" + +namespace hwy { +namespace pipeline { + +// --------------------------------------------------------------------------- +// 2D Cache Tiling Geometry +// --------------------------------------------------------------------------- +// A dedicated POD struct specifying outer and inner loop blocking dimensions. +// Dictates cache locality boundaries independently of prefetch lookahead +// velocity. +struct Tiling2DArgs { + // In contrast to prefetch lookaheads (which must be aggressively tuned to + // negotiate volatile memory latency boundaries), the 2D block shapes rarely + // require dynamic adjustment. + // + // `outer_block` dictates how many items in the outer loop are amortized + // together to ensure their intermediate working state remains resident inside + // the L1 Data Cache. For most algorithms, an outer block around ~128 items + // consumes minimal memory (e.g., 512 bytes for 32-bit accumulators) while + // fully amortizing outer-loop overhead. + // + // `inner_block` dictates how many iterations of the innermost phase are + // processed continuously to saturate instruction-level parallelism. This is + // usually overridden at compile-time or runtime by the specific compute + // kernel, which intrinsically knows its own ideal unrolling bounds (e.g., to + // perfectly fill pipeline stages with Fused Multiply-Add instructions). + size_t outer_block = 128; + size_t inner_block = 256; + + // Calculates the total number of discrete 1D tile-iterations to be processed + // across the 2D geometry. + // + // NOTE: This is NOT the total number of 2D tiles (which would be + // `num_outer_tiles * num_inner_tiles`). Instead, it represents + // `outer_size * num_inner_tiles`. Because the underlying 1D prefetch pipeline + // unrolls across `outer_size` rows for each inner tile, this value accurately + // reflects the total number of discrete inner loop invocations, ensuring + // correct AutoFDO fission bucket assignment and cost normalization. + HWY_INLINE size_t TotalElements(size_t outer_size, size_t inner_size) const { + const size_t num_tile_per_row = + inner_block > 0 ? (inner_size + inner_block - 1) / inner_block : 1; + return outer_size * num_tile_per_row; + } +}; + +struct Default2DTiledPrefetchPolicy : public DefaultPrefetchPolicy {}; + +namespace low_level { + +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L +// --------------------------------------------------------------------------- +// 2D-Tiled Pipeline Callbacks Concept +// --------------------------------------------------------------------------- +// To use PrefetchPipeline2DTiledLoop, the user must provide a callbacks object +// that adheres to the following signatures: +// +// // Called before evaluating the inner loop phases for an outer block. +// // Useful for allocating or resetting local accumulators on the stack. +// void OnOuterBlockStart(size_t outer_idx, size_t outer_end); +// +// // Called before processing a new inner segment across all items in the +// // current outer block. Useful for broadcasting data into SIMD registers. +// void PrepareInnerBlock(size_t outer_idx, size_t outer_end, +// size_t inner_idx, size_t inner_end); +// +// // Evaluates the sequence index `outer_i` and adds memory pointers to the +// // `collector` collection to be prefetched for the given inner block. +// template +// void PopulateCachelines( +// size_t outer_i, size_t inner_idx, size_t inner_end, +// CachelineBundle& collector); +// +// // The core loop body to execute for a given sequence index `outer_i` and +// // inner segment. +// void ComputeTask(size_t outer_i, size_t inner_idx, size_t inner_end); +// +// // Called after all inner segments have been processed for the outer block. +// // Useful for committing accumulated data. +// void OnOuterBlockFinish(size_t outer_idx, size_t outer_end); +template +concept PrefetchPipeline2DTiledCallbacks = + requires(T& cb, size_t outer_idx, size_t outer_end, size_t inner_idx, + size_t inner_end, CachelineBundle& collector) { + { cb.OnOuterBlockStart(outer_idx, outer_end) }; + { cb.PrepareInnerBlock(outer_idx, outer_end, inner_idx, inner_end) }; + { cb.PopulateCachelines(outer_idx, inner_idx, inner_end, collector) }; + { cb.ComputeTask(outer_idx, inner_idx, inner_end) }; + { cb.OnOuterBlockFinish(outer_idx, outer_end) }; + }; +#endif + +// =========================================================================== +// WARNING: Low-Level Execution API +// =========================================================================== +// Do NOT call this function directly in production code! Bypassing the tuning +// framework prevents fleet-wide continuous profiling and telemetry collection. +// Prefer `hwy::PrefetchPipeline2DTiledLoop` in `prefetch_pipeline_2d.h`. +// +// A generic pipeline for executing 2D block-tiled computations with +// prefetching. This splits processing into localized blocks that prevent cache +// evictions and FPU instruction bottlenecks. +// +// Template Parameters: +// Policy: A struct extending `Default2DTiledPrefetchPolicy` dictating cache +// limits, tiling dimensions, and loop constants. +// Callbacks: A type matching the signature documented in +// `PrefetchPipeline2DTiledCallbacks` above. Passed by reference to +// allow state mutation. +// tiling: A `Tiling2DArgs` object dictating the outer/inner block dimensions. +// prefetch_args: A `PrefetchArgs` object dictating prefetch lookahead depths. +template +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L + requires low_level::IsPrefetchPolicy && + PrefetchPipeline2DTiledCallbacks +#endif +inline void PrefetchPipeline2DTiledLoop(size_t outer_size, size_t inner_size, + Callbacks& cb, + const Tiling2DArgs& tiling, + const PrefetchArgs& prefetch_args) { + for (size_t outer_idx = 0; outer_idx < outer_size; + outer_idx += tiling.outer_block) { + const size_t outer_end = (outer_size < outer_idx + tiling.outer_block) + ? outer_size + : outer_idx + tiling.outer_block; + cb.OnOuterBlockStart(outer_idx, outer_end); + + for (size_t inner_idx = 0; inner_idx < inner_size; + inner_idx += tiling.inner_block) { + const size_t inner_end = (inner_size < inner_idx + tiling.inner_block) + ? inner_size + : inner_idx + tiling.inner_block; + + cb.PrepareInnerBlock(outer_idx, outer_end, inner_idx, inner_end); + + hwy::pipeline::low_level::PrefetchPipelineLoop( + outer_idx, outer_end, + [&](size_t i, auto& bundle) { + cb.PopulateCachelines(i, inner_idx, inner_end, bundle); + }, + [&](size_t i) { cb.ComputeTask(i, inner_idx, inner_end); }, + prefetch_args); + } + + cb.OnOuterBlockFinish(outer_idx, outer_end); + } +} + +} // namespace low_level + +// --------------------------------------------------------------------------- +// PrefetchPipeline2DTiledLoop (Production Entry Point) +// --------------------------------------------------------------------------- +template +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L + requires low_level::IsPrefetchPolicy +#endif +inline void PrefetchPipeline2DTiledLoop(size_t outer_size, size_t inner_size, + Callbacks& cb, + Tiling2DArgs tiling = Tiling2DArgs(), +#if defined(__clang__) || defined(__GNUC__) + const char* file_loc = __builtin_FILE(), + int line_loc = __builtin_LINE()) { +#else + const char* file_loc = nullptr, + int line_loc = 0) { +#endif + auto loop_runner = [&](const PrefetchArgs& args) { + hwy::pipeline::low_level::PrefetchPipeline2DTiledLoop( + outer_size, inner_size, cb, tiling, args); + }; + + low_level::DispatchTunedWorkload( + tiling.TotalElements(outer_size, inner_size), file_loc, line_loc, + loop_runner); +} + +} // namespace pipeline +} // namespace hwy + +#endif // HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_PIPELINE_2D_H_ diff --git a/hwy/contrib/pipeline/prefetch_pipeline_2d_test.cc b/hwy/contrib/pipeline/prefetch_pipeline_2d_test.cc new file mode 100644 index 0000000000..55ef243b24 --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_pipeline_2d_test.cc @@ -0,0 +1,188 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 + +#include "hwy/contrib/pipeline/prefetch_pipeline_2d.h" + +#include +#include +#include + +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/contrib/pipeline/prefetch_tuner.h" +#include "hwy/contrib/pipeline/prefetch_tuner_registry.h" +#include "hwy/tests/hwy_gtest.h" + +namespace hwy { +namespace pipeline { +namespace { + +template +struct TestPolicy : Default2DTiledPrefetchPolicy { + static constexpr PrefetchStrategy kStrategy = TargetStrategy; + static constexpr size_t kMaxCachelinesPerIter = 1; +}; + +std::vector g_trace; + +struct Fake2DCallbacks { + void OnOuterBlockStart(size_t outer_idx, size_t outer_end) { + g_trace.push_back("OuterStart(" + std::to_string(outer_idx) + "," + + std::to_string(outer_end) + ")"); + } + + void PrepareInnerBlock(size_t outer_idx, size_t outer_end, size_t inner_idx, + size_t inner_end) { + g_trace.push_back("PrepareInner(" + std::to_string(inner_idx) + "," + + std::to_string(inner_end) + ")"); + } + + template + void PopulateCachelines(size_t outer_i, size_t inner_idx, size_t inner_end, + CachelineBundle& collector) { + g_trace.push_back("Populate(" + std::to_string(outer_i) + "," + + std::to_string(inner_idx) + ")"); + collector.Add(reinterpret_cast(outer_i * 1000 + inner_idx)); + } + + void ComputeTask(size_t outer_i, size_t inner_idx, size_t inner_end) { + g_trace.push_back("Compute(" + std::to_string(outer_i) + "," + + std::to_string(inner_idx) + ")"); + } + + void OnOuterBlockFinish(size_t outer_idx, size_t outer_end) { + g_trace.push_back("OuterFinish(" + std::to_string(outer_idx) + "," + + std::to_string(outer_end) + ")"); + } +}; + +class PrefetchPipeline2DTest : public ::testing::Test { + protected: + void SetUp() override { g_trace.clear(); } +}; + +template +void Call2DPipeline(size_t outer_size, size_t inner_size, Tiling2DArgs tiling) { + PrefetchArgs args; + args.deep_lookahead = Deep; + args.shallow_lookahead = Shallow; + Fake2DCallbacks cb; + low_level::PrefetchPipeline2DTiledLoop, + Fake2DCallbacks>( + outer_size, inner_size, cb, tiling, args); +} + +TEST_F(PrefetchPipeline2DTest, NoPrefetchStrategy) { + // Outer size = 4, Inner size = 4. Tiling: outer=2, inner=2. + Call2DPipeline( + 4, 4, Tiling2DArgs{.outer_block = 2, .inner_block = 2}); + + std::vector expected = { + // Outer Block 1 (0 to 2) + "OuterStart(0,2)", "PrepareInner(0,2)", "Compute(0,0)", "Compute(1,0)", + "PrepareInner(2,4)", "Compute(0,2)", "Compute(1,2)", "OuterFinish(0,2)", + // Outer Block 2 (2 to 4) + "OuterStart(2,4)", "PrepareInner(0,2)", "Compute(2,0)", "Compute(3,0)", + "PrepareInner(2,4)", "Compute(2,2)", "Compute(3,2)", "OuterFinish(2,4)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipeline2DTest, EndToEndPublicWrapperTuning) { + struct TestMetricContext { + bool called = false; + std::string captured_file; + int captured_line = 0; + size_t captured_total = 0; + bool captured_tiny = false; + PrefetchTuningHint captured_hint = PrefetchTuningHint::kAuto; + } ctx; + + class Mock2DTuner : public low_level::PrefetchTuner { + public: + explicit Mock2DTuner(TestMetricContext* ctx) : ctx_(ctx) {} + low_level::PrefetchTuningScope CreateScope( + const low_level::PrefetchTuningContext& context) const override { + if (context.file_loc != nullptr) { + ctx_->captured_file = context.file_loc; + } + ctx_->captured_line = context.line_loc; + ctx_->captured_total = context.total_elements; + ctx_->captured_tiny = context.is_ultra_tiny; + ctx_->captured_hint = context.hint; + + PrefetchArgs args{.deep_lookahead = 0, .shallow_lookahead = 0}; + return low_level::PrefetchTuningScope( + args, + [](void* user_data, float) { + static_cast(user_data)->called = true; + }, + ctx_, 1.0f); + } + + CallsiteId RegisterContext( + const low_level::PrefetchTuningContext& context) const override { + return 0; + } + + low_level::PrefetchTuningScope CreateScopeByCallsiteId( + CallsiteId callsite_id, + const low_level::PrefetchTuningContext& context) const override { + return CreateScope(context); + } + + private: + TestMetricContext* ctx_; + }; + + Mock2DTuner tuner(&ctx); + low_level::GetGlobalPrefetchTunerRegistry() = &tuner; + + g_trace.clear(); + Fake2DCallbacks cb; + + // Execute on 8x8 matrix. total_elements = 32 (>= 32). kAuto hint. kNoPrefetch + // strategy. Capture the exact line number where the wrapper is invoked! + const int expected_base_line = __LINE__ + 1; + PrefetchPipeline2DTiledLoop, + Fake2DCallbacks>( + 8, 8, cb, Tiling2DArgs{.outer_block = 2, .inner_block = 2}); + + // 1. Verify RAII Telemetry Firing + EXPECT_TRUE(ctx.called); + + // 2. Verify Context Propagation + EXPECT_EQ(ctx.captured_total, 32); + EXPECT_FALSE(ctx.captured_tiny); + EXPECT_EQ(ctx.captured_hint, PrefetchTuningHint::kAuto); + EXPECT_NE(ctx.captured_file.find("prefetch_pipeline_2d_test.cc"), + std::string::npos); + + // 3. Verify Fission Line Enrichment (64 items <= 256 -> Tiny kAuto -> + // +700000) + EXPECT_EQ(ctx.captured_line, expected_base_line + 700000); + + // 4. Verify Low-Level Pipeline Execution Sequence + EXPECT_FALSE(g_trace.empty()); + EXPECT_EQ(g_trace[0], "OuterStart(0,2)"); + + low_level::GetGlobalPrefetchTunerRegistry() = nullptr; +} + +TEST_F(PrefetchPipeline2DTest, Tiling2DArgsTotalElements) { + Tiling2DArgs tiling{.outer_block = 128, .inner_block = 256}; + + // Case 1: inner_size <= inner_block (1 tile per row) + EXPECT_EQ(tiling.TotalElements(10, 100), 10); + + // Case 2: inner_size > inner_block (multiple tiles per row) + // 300 / 256 -> 2 tiles per row. 10 outer * 2 = 20. + EXPECT_EQ(tiling.TotalElements(10, 300), 20); + + // Case 3: inner_block == 0 fallback safety (1 tile per row) + Tiling2DArgs zero_tiling{.outer_block = 128, .inner_block = 0}; + EXPECT_EQ(zero_tiling.TotalElements(5, 500), 5); +} + +} // namespace +} // namespace pipeline +} // namespace hwy diff --git a/hwy/contrib/pipeline/prefetch_pipeline_test.cc b/hwy/contrib/pipeline/prefetch_pipeline_test.cc new file mode 100644 index 0000000000..1eede5b2cc --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_pipeline_test.cc @@ -0,0 +1,336 @@ +#include "hwy/contrib/pipeline/prefetch_pipeline.h" + +#include +#include +#include +#include +#include + +#include "hwy/base.h" +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/contrib/pipeline/prefetch_tuner.h" +#include "hwy/contrib/pipeline/prefetch_tuner_registry.h" +#include "hwy/tests/hwy_gtest.h" + +namespace hwy { +namespace pipeline { +namespace { + +// A generic policy builder to parameterize the loop strategy for testing. +template +struct TestPolicy : DefaultPrefetchPolicy { + static constexpr PrefetchStrategy kStrategy = TargetStrategy; + static constexpr size_t kMaxCachelinesPerIter = 1; +}; + +// Global trace vector to track the exact chronological sequence of operations. +std::vector g_trace; + +void FakeDeepPrefetch(const void* ptr) { + // "PD" --> Deep Prefetch. + g_trace.push_back("PD(" + std::to_string(reinterpret_cast(ptr)) + + ")"); +} + +void FakeShallowPrefetch(const void* ptr) { + // "PS" --> Shallow Prefetch. + g_trace.push_back("PS(" + std::to_string(reinterpret_cast(ptr)) + + ")"); +} + +// A fake provider that pushes string events into the global trace. +struct FakeCachelinesProvider { + template + void operator()(size_t i, CachelineBundle& collector) const { + // "S" --> Supply cachelines. + g_trace.push_back("S(" + std::to_string(i) + ")"); + // Add a dummy pointer so the `Prefetch` Assembly compiles cleanly natively. + collector.Add(reinterpret_cast(i)); + } +}; + +// A fake task to assert chronological execution order. +struct FakeTask { + void operator()(size_t i) const { + // "T" --> Task. + g_trace.push_back("T(" + std::to_string(i) + ")"); + } +}; + +class PrefetchPipelineTest : public ::testing::Test { + protected: + void SetUp() override { g_trace.clear(); } +}; + +template +void CallPipeline(size_t start, size_t end) { + PrefetchArgs args; + args.deep_lookahead = Deep; + args.shallow_lookahead = Shallow; + low_level::PrefetchPipelineLoop, + FakeCachelinesProvider, FakeTask, + FakeDeepPrefetch, FakeShallowPrefetch>( + start, end, FakeCachelinesProvider(), FakeTask(), args); +} + +TEST_F(PrefetchPipelineTest, NoPrefetchStrategy) { + CallPipeline(0, 5); + + std::vector expected = {"T(0)", "T(1)", "T(2)", "T(3)", "T(4)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipelineTest, DualTierStrategy) { + // Dual-Tier relies on a deeply staggered Phase 1, Phase 2, Phase 3 pipeline. + // We test on an array of length 6. Lookaheads: kShallow = 2, kDeep = 4. + CallPipeline(0, 6); + + std::vector expected = { + // Phase 1: Overlapping limits. L1/L3 horizons are primed (i = 0 to 1). + // Note: get_cachelines is called once per `i` and then distributed to + // deep/shallow. + "S(0)", "PD(0)", "PS(0)", "S(1)", "PD(1)", "PS(1)", + // Phase 2: Outstanding L3. L1 window is exhausted (i = 2 to 3). + "S(2)", "PD(2)", "S(3)", "PD(3)", + // Phase 3a: Main Sequence (i = 0 to 1). + "S(4)", "PD(4)", "S(2)", "PS(2)", "T(0)", "S(5)", "PD(5)", "S(3)", + "PS(3)", "T(1)", + // Phase 3b: Limit Deep (i = 2 to 3). Deep lookahead has reached array + // bounds. + "S(4)", "PS(4)", "T(2)", "S(5)", "PS(5)", "T(3)", + // Phase 3c: Drain (i = 4 to 5). No fetches remaining. + "T(4)", "T(5)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipelineTest, DualTierStrategy_DifferentLookaheads) { + // Test with kShallow = 1, kDeep = 3 on an array of length 5. + CallPipeline(0, 5); + + std::vector expected = { + // Phase 1: Overlapping limits. L1/L3 horizons are primed (i = 0 to 0). + "S(0)", "PD(0)", "PS(0)", + // Phase 2: Outstanding L3. L1 window is exhausted (i = 1 to 2). + "S(1)", "PD(1)", "S(2)", "PD(2)", + // Phase 3a: Main Sequence (i = 0 to 1). + // i=0 triggers deep+3 (3) and shallow+1 (1) + "S(3)", "PD(3)", "S(1)", "PS(1)", "T(0)", + // i=1 triggers deep+3 (4) and shallow+1 (2) + "S(4)", "PD(4)", "S(2)", "PS(2)", "T(1)", + // Phase 3b: Limit Deep (i = 2 to 3). Deep lookahead ends. + // i=2 shallow+1 (3) + "S(3)", "PS(3)", "T(2)", + // i=3 shallow+1 (4) + "S(4)", "PS(4)", "T(3)", + // Phase 3c: Drain (i = 4 to 4). No fetches remaining. + "T(4)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipelineTest, ShallowRollingLookaheadStrategy) { + // Tests the 1D rolling array. Only shallow lookahead is active (kShallow=2). + CallPipeline(0, 6); + + std::vector expected = {// Startup Phase (i = 0 to 1) for L1 + "S(0)", "PS(0)", "S(1)", "PS(1)", + // Main Sliding Loop (i = 0 to 3) + "S(2)", "PS(2)", "T(0)", "S(3)", "PS(3)", + "T(1)", "S(4)", "PS(4)", "T(2)", "S(5)", + "PS(5)", "T(3)", + // Drain (i = 4 to 5) + "T(4)", "T(5)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipelineTest, DeepRollingLookaheadStrategy) { + // Tests the 1D rolling array using the Deep boundary (kDeep=4). + CallPipeline(0, 6); + + std::vector expected = { + // Startup Phase (i = 0 to 3) for L3 (kDeep = 4) + "S(0)", "PD(0)", "S(1)", "PD(1)", "S(2)", "PD(2)", "S(3)", "PD(3)", + // Main Sliding Loop (i = 0 to 1) + "S(4)", "PD(4)", "T(0)", "S(5)", "PD(5)", "T(1)", + // Drain (i = 2 to 5) + "T(2)", "T(3)", "T(4)", "T(5)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipelineTest, MiniBatchShallowStrategy) { + // Tests the blocked/batch prefetch array. Shallow lookahead = 2 = Batch size. + CallPipeline(0, 5); + + std::vector expected = { + // Block 1 (i = 0 to 1) + "S(0)", "PS(0)", "S(1)", "PS(1)", "T(0)", "T(1)", + // Block 2 (i = 2 to 3) + "S(2)", "PS(2)", "S(3)", "PS(3)", "T(2)", "T(3)", + // Block 3 (remainder, i = 4 to 4) + "S(4)", "PS(4)", "T(4)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipelineTest, MiniBatchDeepStrategy) { + // Tests the blocked/batch prefetch array. Deep lookahead = 4 = Batch size. + CallPipeline(0, 5); + + std::vector expected = {// Block 1 (i = 0 to 3) + "S(0)", "PD(0)", "S(1)", "PD(1)", "S(2)", + "PD(2)", "S(3)", "PD(3)", "T(0)", "T(1)", + "T(2)", "T(3)", + // Block 2 (remainder, i = 4 to 4) + "S(4)", "PD(4)", "T(4)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipelineTest, ZeroShallowZeroDeepFallback) { + // Disabling both tiers should degrade down to NoPrefetch behavior entirely. + CallPipeline(0, 5); + + std::vector expected = {"T(0)", "T(1)", "T(2)", "T(3)", "T(4)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipelineTest, ZeroShallowFallback) { + // If shallow is 0, we degrade to pure DeepOnly logic (e.g. Rolling L3 Only) + CallPipeline(0, 4); + + std::vector expected = { + // Startup Deep + "S(0)", "PD(0)", "S(1)", "PD(1)", "S(2)", "PD(2)", + // Sliding Loop + "S(3)", "PD(3)", "T(0)", + // Drain + "T(1)", "T(2)", "T(3)"}; + EXPECT_EQ(g_trace, expected); +} + +TEST_F(PrefetchPipelineTest, ShallowGreaterOrEqualDeepFallback) { + // If shallow >= deep, the L3 tier is bypassed entirely protecting LFBs. + // Tests DualTier logic natively degrading to Rolling L1-only behavior. + +#if !HWY_IS_DEBUG_BUILD + // Note: This is only tested in release builds because the check `actual_deep + // > actual_shallow` is an `assert`. + CallPipeline(0, 5); +#else + CallPipeline(0, 5); +#endif + + std::vector expected = { + // Startup Shallow (for 4 steps) + "S(0)", "PS(0)", "S(1)", "PS(1)", "S(2)", "PS(2)", "S(3)", "PS(3)", + // Sliding Loop + "S(4)", "PS(4)", "T(0)", + // Drain + "T(1)", "T(2)", "T(3)", "T(4)"}; + EXPECT_EQ(g_trace, expected); +} + +struct TestMetricContext { + bool called = false; + float ticks = 0; + std::string captured_file; + int captured_line = 0; + size_t captured_total = 0; + bool captured_tiny = false; + PrefetchTuningHint captured_hint = PrefetchTuningHint::kAuto; + uint32_t captured_dense_id = 9999; +}; + +void FakeMetricCollectorCb(void* user_data, float elapsed_ticks) { + auto* ctx = static_cast(user_data); + ctx->called = true; + ctx->ticks = elapsed_ticks; +} + +class MockMetricTuner : public low_level::PrefetchTuner { + public: + explicit MockMetricTuner(TestMetricContext* ctx) : ctx_(ctx) {} + low_level::PrefetchTuningScope CreateScope( + const low_level::PrefetchTuningContext& context) const override { + if (context.file_loc != nullptr) { + ctx_->captured_file = context.file_loc; + } + ctx_->captured_line = context.line_loc; + ctx_->captured_total = context.total_elements; + ctx_->captured_tiny = context.is_ultra_tiny; + ctx_->captured_hint = context.hint; + + // Return custom lookahead depths: deep=3, shallow=1 + PrefetchArgs args{.deep_lookahead = 3, .shallow_lookahead = 1}; + return low_level::PrefetchTuningScope(args, FakeMetricCollectorCb, ctx_, + 1.0f); + } + + CallsiteId RegisterContext( + const low_level::PrefetchTuningContext& context) const override { + CallsiteId id = next_id_++; + registered_line_locs_[id] = context.line_loc; + return id; + } + + low_level::PrefetchTuningScope CreateScopeByCallsiteId( + CallsiteId callsite_id, + const low_level::PrefetchTuningContext& context) const override { + ctx_->captured_dense_id = callsite_id; + EXPECT_TRUE(registered_line_locs_.find(callsite_id) != + registered_line_locs_.end()); + EXPECT_EQ(registered_line_locs_.at(callsite_id), context.line_loc); + return CreateScope(context); + } + + mutable CallsiteId next_id_ = 0; + mutable std::map registered_line_locs_; + + private: + TestMetricContext* ctx_; +}; + +TEST_F(PrefetchPipelineTest, EndToEndPublicWrapperTuning) { + TestMetricContext ctx; + MockMetricTuner tuner(&ctx); + low_level::GetGlobalPrefetchTunerRegistry() = &tuner; + + g_trace.clear(); + + // Helper lambda encapsulates a single physical call site in the codebase. + auto run_callsite_a = [&]() { + PrefetchPipelineLoop, + FakeCachelinesProvider, FakeTask>( + 0, 35, FakeCachelinesProvider(), FakeTask()); + }; + + // 1. First call site execution (Dense ID 0 assigned and cached) + run_callsite_a(); + + EXPECT_TRUE(ctx.called); + EXPECT_EQ(ctx.captured_total, 35); + EXPECT_EQ(ctx.captured_dense_id, 0); + EXPECT_EQ(tuner.next_id_, 1); + + // 2. Execute the FIRST call site again. Reuses cached Dense ID 0! + ctx.called = false; + run_callsite_a(); + + EXPECT_TRUE(ctx.called); + EXPECT_EQ(ctx.captured_dense_id, 0); + EXPECT_EQ(tuner.next_id_, 1); // No new registration! + + // 3. Execute a SECOND distinct physical call site. Gets Dense ID 1! + ctx.called = false; + PrefetchPipelineLoop, + FakeCachelinesProvider, FakeTask>( + 0, 35, FakeCachelinesProvider(), FakeTask()); + + EXPECT_TRUE(ctx.called); + EXPECT_EQ(ctx.captured_dense_id, 1); + EXPECT_EQ(tuner.next_id_, 2); + + low_level::GetGlobalPrefetchTunerRegistry() = nullptr; +} +} // namespace +} // namespace pipeline +} // namespace hwy diff --git a/hwy/contrib/pipeline/prefetch_pipeline_types.h b/hwy/contrib/pipeline/prefetch_pipeline_types.h new file mode 100644 index 0000000000..e57697f85d --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_pipeline_types.h @@ -0,0 +1,322 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_PIPELINE_TYPES_H_ +#define HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_PIPELINE_TYPES_H_ + +#include + +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L +#include // IWYU pragma: keep +#include +#endif + +#include + +#include "hwy/base.h" +namespace hwy { +namespace pipeline { + +// --------------------------------------------------------------------------- +// PrefetchStrategy +// --------------------------------------------------------------------------- +// Enumerates the structural looping algorithm that the pipeline will compile +// down into. Used to decouple hardware limits from loop execution mechanics. +enum class PrefetchStrategy { + // Bypasses all explicit software prefetch instructions entirely. Used as a + // control baseline during A/B testing or when hardware stream trackers are + // already fully saturating memory bandwidth. + kNoPrefetch, + + // Issues prefetch instructions exclusively for the deep lookahead distance + // (targeting L3 cache or DRAM). Ideal for pointer-chasing or irregular + // scatter-gather workloads where hiding massive DRAM latency is the primary + // bottleneck. + kDeepLookaheadOnly, + + // Issues prefetch instructions exclusively for the shallow lookahead distance + // (targeting L1/L2 cache). Ideal for linear memory scans where L3 is warmed + // up by hardware prefetchers, but explicit hints are needed to bridge the + // final L3->L1 latency gap without thrashing Line Fill Buffers (LFBs). + kShallowLookaheadOnly, + + // Aggregates deep prefetch instructions into discrete mini-batches executed + // before compute phases. Decouples prefetch dispatch from inner loop compute + // unrolling, preventing instruction-cache bloat and front-end bottlenecks. + kMiniBatchDeep, + + // Aggregates shallow prefetch instructions into discrete mini-batches. + // Ensures L1 cachelines are fetched in tight bursts right before SIMD vector + // execution, maximizing instruction-level parallelism (ILP). + kMiniBatchShallow, + + // The flagship, production-grade execution strategy. Simultaneously + // orchestrates both deep (L3/DRAM) and shallow (L1/L2) prefetching in a + // coordinated, two-tier pipeline. Keeps the memory hierarchy perfectly + // saturated across all latency boundaries. + kDualTier +}; + +// Identifier for a prefetch call site. It is recommended to be unique for +// different call sites, but it is functionally safe even if collisions occur. +using CallsiteId = uint64_t; + +// --------------------------------------------------------------------------- +// PrefetchArgs +// --------------------------------------------------------------------------- +// Bit-field struct encapsulating the two prefetch lookahead parameters for a +// single prefetch call site. Total size is exactly 2 bytes (16 bits). +struct PrefetchArgs { + // The iteration distance (in loop iterations) to look ahead for the deep L3 + // prefetch (max 511). If 0, no deep prefetch will be issued. + uint16_t deep_lookahead : 9; + + // The iteration distance (in loop iterations) to look ahead for the shallow + // L1 prefetch (max 127). If 0, no shallow prefetch will be issued. + uint16_t shallow_lookahead : 7; + + // ------------------------------------------------------------------------- + // Safe Default Factories + // ------------------------------------------------------------------------- + // Tuning memory prefetching is notoriously difficult because lookahead bounds + // change dramatically depending on the spatial distribution of the workload. + + // Random Access / Scatter-Gather (e.g. Hash Table Probing, Graph Walks) + // + // Random array accesses constantly suffer TLB (Translation Lookaside Buffer) + // misses, resulting in massive Page Walk delays. To absorb these colossal + // ~300-cycle stalls natively inside the L3 queue, the deep lookahead must + // aggressively stretch out by large margins (e.g. 32-48 iterations). + static constexpr PrefetchArgs DefaultRandom() { +#if HWY_ARCH_ARM_A64 + return PrefetchArgs{.deep_lookahead = 64, .shallow_lookahead = 8}; +#else + return PrefetchArgs{.deep_lookahead = 32, .shallow_lookahead = 4}; +#endif + } + + // Sequential Scans / Linear Memory (e.g. Matrix Vector, Filter Scans) + // + // Linear accesses benefit intimately from native CPU stream-trackers (which + // already mask bulk DRAM latency). Here, a heavy L3 lookahead is + // counter-productive; it merely crowds the queue. Instead, we tighten the + // lookaheads down to safely bridge the narrower L3 -> L1 latency gap + // (~40 cycles) without overflowing LFBs during heavy SIMD evaluation. + static constexpr PrefetchArgs DefaultSequential() { +#if HWY_ARCH_ARM_A64 + return PrefetchArgs{.deep_lookahead = 32, .shallow_lookahead = 4}; +#else + return PrefetchArgs{.deep_lookahead = 8, .shallow_lookahead = 2}; +#endif + } +}; + +// --------------------------------------------------------------------------- +// PrefetchTuningHint +// --------------------------------------------------------------------------- +// Hint for the expected memory access pattern. +// +// These hints are primarily used to establish safe baseline configuration +// limits when the tuning framework executes without a registered telemetry +// plugin, or to artificially restrict the scope of candidate grids when +// auto-tuning is enabled. +enum class PrefetchTuningHint { + // The workload access pattern is unknown or highly variable. + // - Without Tuner: Safely defaults to conservative, tight lookahead bounds. + // - With Tuner: Evaluates the maximum exhaustive search grid for optimal + // configuration. + kAuto, + + // The workload scatters reads across wide memory distributions (e.g., hash + // tables, graph node aggregations) resulting in TLB/Page Walk delays. + // - Without Tuner: Defaults to aggressive, deep lookahead bounds. + // - With Tuner: Artificially restricts candidate grids to deep lookahead + // bounds only to optimize profiling. + kRandom, + + // The workload linearly scans contiguous memory blocks (e.g., Matrix-Vector + // multiplication). + // - Without Tuner: Safely defaults to conservative, tight lookahead bounds. + // - With Tuner: Artificially restricts candidate grids to shallow lookahead + // bounds only to optimize profiling. + kSequential +}; + +// --------------------------------------------------------------------------- +// PrefetchPolicy +// --------------------------------------------------------------------------- +struct DefaultPrefetchPolicy { + // The specific execution strategy the pipeline will adopt. + // + // Note this is mainly used for benchmarking and testing purposes. + // For the production code, just use `kDualTier` and let the tuning framework + // optimize the lookahead values, e.g., when deep lookahead is zero, it will + // automatically fall back to shallow-only. + static constexpr PrefetchStrategy kStrategy = PrefetchStrategy::kDualTier; + + // Threshold below which the workload short-circuits dynamic tuning queries + // overhead, and run the loop with the default prefetch strategy and lookahead + // values per the available hints. + // + // Note that the optimal value may vary, depending on the specific hardware, + // workload, and the tuner capabilities. + static constexpr size_t kUltraTinyThreshold = 32; + + // The maximum number of explicit cachelines that should be prefetched per + // iteration (i.e., cachelines containing the data to be used by the upcoming + // compute task). + static constexpr size_t kMaxCachelinesPerIter = 4; + + // ------------------------------------------------------------------------- + // Hardware Architecture Matrix + // ------------------------------------------------------------------------- + // Resolved automatically within the struct using HWY_ARCH_* + +#if HWY_ARCH_ARM || HWY_ARCH_ARM_A64 + // ARM platforms vary by their natures, we use 24 here to be conservative. + static constexpr size_t kNumMSHRs = 24; + +#elif HWY_ARCH_X86 + // On x86, we cannot differentiate Intel vs AMD reliably at compile time. + // Because hitting the edge of the Intel LFB pool (10-12) causes a hard CPU + // stall, we MUST gracefully pick the most restrictive denominator (12) to + // ensure safety. AMD Zen architectures uniquely "absorb" the excess prefetch + // instructions into their 124 MABs, suffering zero penalty from this wrapper. + static constexpr size_t kNumMSHRs = 12; + +#else + // Safe, generic bounds. + static constexpr size_t kNumMSHRs = 12; +#endif +}; + +// --------------------------------------------------------------------------- +// CachelineBundle +// --------------------------------------------------------------------------- +// A lightweight, stack-allocated, fixed-capacity container for collecting +// memory addresses to be prefetched. Provides bulletproof bounds checking in +// production while maintaining strict debug capacity assertions. +template +struct CachelineBundle { + // Array of explicit memory addresses to prefetch. + const void* ptrs[Policy::kMaxCachelinesPerIter]; + + // The number of valid pointers currently registered in the array. + size_t count = 0; + + // Registers a discrete, individual memory address to be prefetched. + // USE CASE: Non-contiguous, scatter-gather accesses (e.g., hash tables, graph + // walks). + // Use `AddContiguousRange` for linear memory buffers. + HWY_INLINE void Add(const void* ptr) { + HWY_DASSERT(count < Policy::kMaxCachelinesPerIter); + if (HWY_LIKELY(count < Policy::kMaxCachelinesPerIter)) { + ptrs[count] = ptr; + ++count; + } + } + + // Registers a dense, contiguous memory range to be prefetched. + // USE CASE: Linear memory buffers (e.g., embedding vectors, matrix tiles, + // image spans). Automatically calculates 64-byte cacheline strides and + // short-circuits if full. + HWY_INLINE void AddContiguousRange(const void* base_ptr, size_t num_bytes) { + const char* base = static_cast(base_ptr); + for (size_t offset = 0; offset < num_bytes; offset += 64) { + if (HWY_LIKELY(count < Policy::kMaxCachelinesPerIter)) { + Add(base + offset); + } else { + break; // Container is full; short-circuit remaining iterations + } + } + } +}; + +namespace low_level { + +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L +// --------------------------------------------------------------------------- +// PrefetchPolicy Concept +// --------------------------------------------------------------------------- +// Mandates that any user-provided Policy struct defines the required hardware +// and execution policy constants. +template +concept IsPrefetchPolicy = requires { + requires ::std::is_integral_v; + requires ::std::is_integral_v; + requires ::std::is_same_v<::std::remove_const_t, + PrefetchStrategy>; + requires ::std::is_integral_v; +}; + +// --------------------------------------------------------------------------- +// Cachelines Provider Concept +// --------------------------------------------------------------------------- +// To use PrefetchPipelineLoop, the user must provide a callable that adheres +// to the following signature: +// +// template +// void operator()(size_t i, CachelineBundle& +// bundle) const; +// +// Parameters: +// - i: The current sequence index evaluated by the loop. +// This is guaranteed to be in the range `[start, end)`. +// - bundle: Output collection. Call `bundle.Add(ptr)` to +// register cachelines. For invalid or conditional +// indices, simply do not add anything. +// +// Execution constraints: +// - Depending on the architecture, prefetching optimizations, or runtime +// auto-tuning limits (see Policy configurations below), there is no +// hardware or programmatic guarantee how many times this function will be +// called for a given `i` (including zero times if prefetches are stripped +// via `constexpr`). Therefore, this callable MUST be completely pure and +// strictly side-effect free! +// - Missing or zero-length entries are natively skipped by the pipeline +// without incurring branching overhead. +template +concept PrefetchPipelineCachelineProvider = + requires(const T& provider, size_t i, CachelineBundle& bundle) { + { provider(i, bundle) }; + }; + +// --------------------------------------------------------------------------- +// Pipeline Task Concept +// --------------------------------------------------------------------------- +// To use PrefetchPipelineLoop, the user must provide a callable that adheres +// to the following signature: +// +// void operator()(size_t i) const; +// +// Parameters: +// - i: The current sequence index being evaluated by the pipeline. +// +// Execution constraints: +// - Guaranteed to be called exactly once for each index `i` in the range +// `[start, end)`. Furthermore, it is guaranteed to be invoked purely +// sequentially (i.e. `i`, `i+1`, `i+2`), preserving any cross-iteration +// dependencies or internal accumulator state. +template +concept PrefetchPipelineTask = requires(const T& task, size_t i) { + { task(i) }; +}; +#endif + +} // namespace low_level +} // namespace pipeline +} // namespace hwy + +#endif // HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_PIPELINE_TYPES_H_ diff --git a/hwy/contrib/pipeline/prefetch_pipeline_types_test.cc b/hwy/contrib/pipeline/prefetch_pipeline_types_test.cc new file mode 100644 index 0000000000..f152f8af1c --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_pipeline_types_test.cc @@ -0,0 +1,130 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 + +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" + +#include + +#include "hwy/base.h" +#include "hwy/tests/hwy_gtest.h" + +namespace hwy { +namespace pipeline { +namespace { + +TEST(PrefetchPipelineTypesTest, PrefetchArgsFactories) { + PrefetchArgs r = PrefetchArgs::DefaultRandom(); + PrefetchArgs s = PrefetchArgs::DefaultSequential(); + +#if HWY_ARCH_ARM_A64 + EXPECT_EQ(r.deep_lookahead, 64); + EXPECT_EQ(r.shallow_lookahead, 8); + EXPECT_EQ(s.deep_lookahead, 32); + EXPECT_EQ(s.shallow_lookahead, 4); +#else + EXPECT_EQ(r.deep_lookahead, 32); + EXPECT_EQ(r.shallow_lookahead, 4); + EXPECT_EQ(s.deep_lookahead, 8); + EXPECT_EQ(s.shallow_lookahead, 2); +#endif +} + +TEST(PrefetchPipelineTypesTest, PrefetchArgsBitFieldLayout) { + // 1. Ensure sizeof is exactly 2 bytes (16 bits), this is crucial for + // minimizing the tuner overhead. + EXPECT_EQ(sizeof(PrefetchArgs), 2); + + // 2. Ensure proper value range (deep: 9 bits -> max 511, shallow: 7 bits -> + // max 127) + PrefetchArgs max_args{.deep_lookahead = 511, .shallow_lookahead = 127}; + EXPECT_EQ(max_args.deep_lookahead, 511); + EXPECT_EQ(max_args.shallow_lookahead, 127); + + // 3. Verify unsigned wrap-around behavior at exact bit boundaries + PrefetchArgs wrap_args; + uint16_t deep_val = 512; + uint16_t shallow_val = 128; + wrap_args.deep_lookahead = deep_val; + wrap_args.shallow_lookahead = shallow_val; + EXPECT_EQ(wrap_args.deep_lookahead, 0); + EXPECT_EQ(wrap_args.shallow_lookahead, 0); +} + +struct TinyPolicy : DefaultPrefetchPolicy { + static constexpr size_t kMaxCachelinesPerIter = 2; +}; + +TEST(PrefetchPipelineTypesTest, CachelineBundleAdd) { + CachelineBundle bundle; + EXPECT_EQ(bundle.count, 0); + + bundle.Add(reinterpret_cast(0x1000)); + EXPECT_EQ(bundle.count, 1); + EXPECT_EQ(bundle.ptrs[0], reinterpret_cast(0x1000)); + + bundle.Add(reinterpret_cast(0x2000)); + EXPECT_EQ(bundle.count, 2); + EXPECT_EQ(bundle.ptrs[1], reinterpret_cast(0x2000)); + + // Adding a 3rd pointer should be safely ignored or asserted in debug builds. + // In production builds, it short-circuits cleanly without buffer overflow. +#if !HWY_IS_DEBUG_BUILD + bundle.Add(reinterpret_cast(0x3000)); + EXPECT_EQ(bundle.count, 2); +#endif +} + +TEST(PrefetchPipelineTypesTest, CachelineBundleContiguousRange) { + CachelineBundle bundle; + char buffer[256]; + + // Adding 128 bytes (2 cachelines at 64B each) perfectly fills TinyPolicy + // (capacity 2). + bundle.AddContiguousRange(buffer, 128); + EXPECT_EQ(bundle.count, 2); + EXPECT_EQ(bundle.ptrs[0], buffer); + EXPECT_EQ(bundle.ptrs[1], buffer + 64); + + // Adding more should short-circuit safely. +#if !HWY_IS_DEBUG_BUILD + bundle.AddContiguousRange(buffer + 128, 64); + EXPECT_EQ(bundle.count, 2); +#endif +} + +#if defined(__cpp_concepts) && __cpp_concepts >= 201907L + +struct ValidCustomPolicy { + static constexpr size_t kMaxCachelinesPerIter = 4; + static constexpr size_t kNumMSHRs = 12; + static constexpr PrefetchStrategy kStrategy = PrefetchStrategy::kDualTier; + static constexpr size_t kUltraTinyThreshold = 32; +}; +static_assert(low_level::IsPrefetchPolicy); +static_assert(low_level::IsPrefetchPolicy); + +struct InvalidMissingMSHRs { + [[maybe_unused]] static constexpr size_t kMaxCachelinesPerIter = 4; + [[maybe_unused]] static constexpr PrefetchStrategy kStrategy = + PrefetchStrategy::kDualTier; + [[maybe_unused]] static constexpr size_t kUltraTinyThreshold = 32; +}; +static_assert(!low_level::IsPrefetchPolicy); + +struct MockCachelineProvider { + template + void operator()(size_t, CachelineBundle&) const {} +}; +static_assert(low_level::PrefetchPipelineCachelineProvider< + MockCachelineProvider, DefaultPrefetchPolicy>); + +struct MockTask { + void operator()(size_t) const {} +}; +static_assert(low_level::PrefetchPipelineTask); + +#endif + +} // namespace +} // namespace pipeline +} // namespace hwy diff --git a/hwy/contrib/pipeline/prefetch_tuner.h b/hwy/contrib/pipeline/prefetch_tuner.h new file mode 100644 index 0000000000..917a60c126 --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_tuner.h @@ -0,0 +1,199 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_TUNER_H_ +#define HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_TUNER_H_ + +#include +#include + +#include "hwy/base.h" +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/timer.h" + +namespace hwy { +namespace pipeline { +namespace low_level { + +struct DefaultPrefetchTimer { + static HWY_INLINE uint64_t Start() { return hwy::timer::Start(); } + static HWY_INLINE uint64_t Stop() { return hwy::timer::Stop(); } +}; + +// --------------------------------------------------------------------------- +// Telemetry & Tuning Scopes +// --------------------------------------------------------------------------- +// An RAII scope manager returned by PrefetchTuner plugins. +// Automatically orchestrates timer allocation on construction and metric +// reporting on destruction, decoupling telemetry from the core POD structs. +template +class PrefetchTuningScopeT { + public: + // Steady-state constructor (Zero overhead, no timer started) + PrefetchTuningScopeT() + : args_{}, + cb_(nullptr), + user_data_(nullptr), + t0_(0), + cost_normalization_factor_(1.0f) {} + + // Active profiling constructor (Starts timer automatically) + PrefetchTuningScopeT(PrefetchArgs args, void (*cb)(void*, float), + void* user_data, float cost_normalization_factor) + : args_(args), + cb_(cb), + user_data_(user_data), + t0_(cb != nullptr ? Timer::Start() : 0), + cost_normalization_factor_(cost_normalization_factor > 0.0f + ? cost_normalization_factor + : 1.0f) {} + + ~PrefetchTuningScopeT() { + if (HWY_UNLIKELY(cb_ != nullptr)) { + float elapsed = static_cast(Timer::Stop() - t0_); + cb_(user_data_, elapsed / cost_normalization_factor_); + } + } + + // Move semantics guarantee exactly-once firing + PrefetchTuningScopeT(PrefetchTuningScopeT&& other) noexcept + : args_(other.args_), + cb_(other.cb_), + user_data_(other.user_data_), + t0_(other.t0_), + cost_normalization_factor_(other.cost_normalization_factor_) { + other.cb_ = nullptr; + } + + PrefetchTuningScopeT& operator=(PrefetchTuningScopeT&& other) noexcept { + if (this != &other) { + if (cb_ != nullptr) { + float elapsed = static_cast(Timer::Stop() - t0_); + cb_(user_data_, elapsed / cost_normalization_factor_); + } + args_ = other.args_; + cb_ = other.cb_; + user_data_ = other.user_data_; + t0_ = other.t0_; + cost_normalization_factor_ = other.cost_normalization_factor_; + other.cb_ = nullptr; + } + return *this; + } + + // Prevent copying + PrefetchTuningScopeT(const PrefetchTuningScopeT&) = delete; + PrefetchTuningScopeT& operator=(const PrefetchTuningScopeT&) = delete; + + // Returns the current prefetch arguments for the scope. + const PrefetchArgs& GetArgs() const { return args_; } + + private: + // The prefetch arguments to use for this scope. If the scope is actively + // profiling, this is the candidate being evaluated. Otherwise, it is the best + // converged configuration (or default if no tuning is performed). + PrefetchArgs args_; + + // Telemetry ingestion callback function pointer invoked upon scope + // destruction. Converts stateless lambdas directly into machine code stubs to + // guarantee zero heap allocations and zero virtual table dispatch overhead. + // + // Parameters: + // - udata: Opaque 64-bit payload (e.g., packed dense ID and + // candidate args) passed through registers. + // - normalized_cost: Raw hardware timer ticks elapsed during active + // profiling normalized by cost_normalization_factor. + void (*cb_)(void* udata, float normalized_cost) = nullptr; + + // Opaque 64-bit payload passed through registers to the callback. + // This could be a pointer to any allocated memory, or even an arbitrary + // packed 64-bit value. + // + // May be nullptr if no callback is installed. + void* user_data_ = nullptr; + + // The unnormalized cost (e.g., elapsed timer ticks) offset at the start of + // profiling. + uint64_t t0_ = 0; + + // The normalization factor to divide the raw cost by before reporting. + float cost_normalization_factor_ = 1.0f; +}; + +// Default RAII scope manager using the default timer. +using PrefetchTuningScope = PrefetchTuningScopeT<>; + +// --------------------------------------------------------------------------- +// Tuning Context & Plugin Interface +// --------------------------------------------------------------------------- + +// Establishes the context around the data geometry and access patterns before +// executing the pipelined loops. +struct PrefetchTuningContext { + // Hint for the underlying memory baseline in absence of a plugin. + PrefetchTuningHint hint = PrefetchTuningHint::kAuto; + + // Flag indicating whether the workload is an ultra-tiny micro-batch. + bool is_ultra_tiny = false; + + // Total number of discrete workload elements to be processed. + // E.g., in 1D workloads, the number of outer loop iterations. In 2D + // workloads, explicitly populated by callers as `num_rows * num_tile_per_row` + // (total tile-iterations) to ensure accurate AutoFDO fission bucket + // assignment. + size_t total_elements = 0; + + // The estimated number of discrete cachelines accessed per workload element. + // E.g., in 1D workloads, cachelines accessed per element. In 2D workloads, + // represents cachelines accessed per tile-iteration (typically + // `Policy::kMaxCachelinesPerIter`). + size_t active_cachelines_per_element = 0; + + // Captures the precise physical source file and line number of the loop + // invocation. + const char* file_loc = nullptr; + int line_loc = 0; + + // Pointer to an array of static string tags representing the active ambient + // context. + const char* const* scope_tags = nullptr; + size_t num_scope_tags = 0; +}; + +// A tuning plugin interface used to instantiate dynamically-optimized pipeline +// arguments based on the runtime context prior to loop execution. +class PrefetchTuner { + public: + virtual ~PrefetchTuner() = default; + + // Derives dynamically-tuned lookahead depths based on the runtime context. + virtual PrefetchTuningScope CreateScope( + const PrefetchTuningContext& context) const = 0; + + // Registers a static call site context out-of-band and returns a callsite ID. + // This is called exactly once at static initialization time per call site. + virtual CallsiteId RegisterContext( + const PrefetchTuningContext& context) const = 0; + + // Fast-path lock-free lookup of winning prefetch arguments by callsite ID. + virtual PrefetchTuningScope CreateScopeByCallsiteId( + CallsiteId callsite_id, const PrefetchTuningContext& context) const = 0; +}; + +} // namespace low_level +} // namespace pipeline +} // namespace hwy + +#endif // HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_TUNER_H_ diff --git a/hwy/contrib/pipeline/prefetch_tuner_registry.h b/hwy/contrib/pipeline/prefetch_tuner_registry.h new file mode 100644 index 0000000000..edf7dee5c8 --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_tuner_registry.h @@ -0,0 +1,187 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_TUNER_REGISTRY_H_ +#define HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_TUNER_REGISTRY_H_ + +#include +#include + +#include + +#include "hwy/base.h" +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/contrib/pipeline/prefetch_tuner.h" + +namespace hwy { +namespace pipeline { +namespace low_level { + +// Link-time hook to retrieve a context-aware prefetch tuner if one is linked. +// Returns nullptr if no plugin is available (falling back to generic defaults). +inline const PrefetchTuner*& GetGlobalPrefetchTunerRegistry() { + static const PrefetchTuner* tuner = nullptr; + return tuner; +} + +// Sets the global prefetch tuner registry. +// This should only be called once at initialization time. +inline void SetGlobalPrefetchTunerRegistry(const PrefetchTuner* tuner) { + HWY_ASSERT(GetGlobalPrefetchTunerRegistry() == nullptr); + GetGlobalPrefetchTunerRegistry() = tuner; +} + +// Registers a prefetch tuning context with the global tuner registry. Returns a +// callsite ID for the context. +inline CallsiteId RegisterPrefetchContext( + const PrefetchTuningContext& context) { + const PrefetchTuner* tuner = GetGlobalPrefetchTunerRegistry(); + if (tuner != nullptr) { + return tuner->RegisterContext(context); + } + return 0; +} + +// Creates a prefetch tuning scope for the given callsite ID and context. +// If no tuner is available or if the workload is ultra-tiny, returns a scope +// with default args and cost normalization factor. +inline PrefetchTuningScope CreatePrefetchTuningScopeByCallsiteId( + CallsiteId callsite_id, const PrefetchTuningContext& context) { + if (HWY_LIKELY(!context.is_ultra_tiny)) { + const PrefetchTuner* tuner = GetGlobalPrefetchTunerRegistry(); + if (HWY_LIKELY(tuner != nullptr)) { + return tuner->CreateScopeByCallsiteId(callsite_id, context); + } + } + + // Use default args and cost normalization factor if no tuner is available or + // if the workload is ultra-tiny (to bypass profiling overhead). + constexpr float kDefaultCostNormalizationFactor = 1.0f; + return PrefetchTuningScope(context.hint == PrefetchTuningHint::kRandom + ? PrefetchArgs::DefaultRandom() + : PrefetchArgs::DefaultSequential(), + nullptr, nullptr, kDefaultCostNormalizationFactor); +} + +// --------------------------------------------------------------------------- +// Dynamic Dispatch Fission Helper +// --------------------------------------------------------------------------- +#if defined(__clang__) || defined(__GNUC__) +#define HWY_NO_CODE_MERGE __attribute__((noinline)) +#else +#define HWY_NO_CODE_MERGE +#endif + +// Dispatches workload execution across distinct physical call sites based on +// the compile-time memory hint and runtime workload size. +// +// Under the hood, this mechanism pairs with forced compiler fission +// (HWY_NO_CODE_MERGE on the caller's lambda) to guarantee that LLVM AutoFDO +// records separate profile buckets and prefetch tuners allocate independent +// lookahead tables for disparate workload categories (Tiny, Medium, Huge) +// originating from the exact same user call site. +// +// Microarchitectural Footprint: Generates exactly one physical machine code +// instance of the unrolled loop body in the .text section (zero code bloat), +// incurring only ~2-3 cycles of fast integer comparison jumping. +template +HWY_INLINE void DispatchWorkloadFission(size_t total_elements, + Executor&& exec) { + if (HWY_UNLIKELY(total_elements < Policy::kUltraTinyThreshold)) { + exec.template operator()<0>(true); + return; + } + + constexpr size_t kTinyThreshold = 256; + constexpr size_t kMediumThreshold = 65536; + + if constexpr (Hint == PrefetchTuningHint::kSequential) { + if (total_elements <= kTinyThreshold) { + exec.template operator()<100000>(false); + } else if (total_elements <= kMediumThreshold) { + exec.template operator()<200000>(false); + } else { + exec.template operator()<300000>(false); + } + } else if constexpr (Hint == PrefetchTuningHint::kRandom) { + if (total_elements <= kTinyThreshold) { + exec.template operator()<400000>(false); + } else if (total_elements <= kMediumThreshold) { + exec.template operator()<500000>(false); + } else { + exec.template operator()<600000>(false); + } + } else { + // kAuto fallback + if (total_elements <= kTinyThreshold) { + exec.template operator()<700000>(false); + } else if (total_elements <= kMediumThreshold) { + exec.template operator()<800000>(false); + } else { + exec.template operator()<900000>(false); + } + } +} + +template +struct UniversalFissionExecutor { + size_t total_elements; + const char* file_loc; + int line_loc; + LoopRunner runner; + + template + HWY_NO_CODE_MERGE void operator()(bool is_ultra_tiny) const { + const int derived_line_loc = line_loc + static_cast(FissionOffset); + static const CallsiteId kCallsiteId = [this, derived_line_loc]() { + PrefetchTuningContext init_ctx; + init_ctx.file_loc = file_loc; + init_ctx.line_loc = derived_line_loc; + init_ctx.hint = Hint; + init_ctx.active_cachelines_per_element = Policy::kMaxCachelinesPerIter; + return RegisterPrefetchContext(init_ctx); + }(); + + PrefetchTuningContext run_ctx; + run_ctx.hint = Hint; + run_ctx.total_elements = total_elements; + run_ctx.active_cachelines_per_element = Policy::kMaxCachelinesPerIter; + run_ctx.file_loc = file_loc; + run_ctx.line_loc = derived_line_loc; + run_ctx.is_ultra_tiny = is_ultra_tiny; + + PrefetchTuningScope tuning_scope = + CreatePrefetchTuningScopeByCallsiteId(kCallsiteId, run_ctx); + + runner(tuning_scope.GetArgs()); + } +}; + +template +HWY_INLINE void DispatchTunedWorkload(size_t total_elements, + const char* file_loc, int line_loc, + LoopRunner&& runner) { + UniversalFissionExecutor::type> + exec{total_elements, file_loc, line_loc, + std::forward(runner)}; + DispatchWorkloadFission(total_elements, exec); +} + +} // namespace low_level +} // namespace pipeline +} // namespace hwy + +#endif // HIGHWAY_HWY_CONTRIB_PIPELINE_PREFETCH_TUNER_REGISTRY_H_ diff --git a/hwy/contrib/pipeline/prefetch_tuner_registry_test.cc b/hwy/contrib/pipeline/prefetch_tuner_registry_test.cc new file mode 100644 index 0000000000..647b3600a5 --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_tuner_registry_test.cc @@ -0,0 +1,249 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 + +#include "hwy/contrib/pipeline/prefetch_tuner_registry.h" + +#include +#include +#include + +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/contrib/pipeline/prefetch_tuner.h" +#include "hwy/detect_compiler_arch.h" +#include "hwy/tests/hwy_gtest.h" + +namespace hwy { +namespace pipeline { +namespace low_level { +namespace { + +struct TestMetricContext { + bool called = false; + float ticks = 0; + uint32_t captured_dense_id = 9999; + std::string captured_file; + int captured_line = 0; + size_t registered_count = 0; +}; + +void FakeMetricCb(void* user_data, float elapsed_ticks) { + auto* ctx = static_cast(user_data); + ctx->called = true; + ctx->ticks = elapsed_ticks; +} + +class MockTunerPlugin : public PrefetchTuner { + public: + explicit MockTunerPlugin(TestMetricContext* ctx) : ctx_(ctx) {} + PrefetchTuningScope CreateScope( + const PrefetchTuningContext& context) const override { + if (context.file_loc != nullptr) ctx_->captured_file = context.file_loc; + ctx_->captured_line = context.line_loc; + PrefetchArgs args{.deep_lookahead = 99, .shallow_lookahead = 9}; + return PrefetchTuningScope(args, FakeMetricCb, ctx_, 1.0f); + } + + CallsiteId RegisterContext( + const PrefetchTuningContext& context) const override { + ctx_->registered_count++; + return 77; + } + + PrefetchTuningScope CreateScopeByCallsiteId( + CallsiteId callsite_id, + const PrefetchTuningContext& context) const override { + ctx_->captured_dense_id = callsite_id; + EXPECT_EQ(callsite_id, 77); + return CreateScope(context); + } + + private: + TestMetricContext* ctx_; +}; + +TEST(PrefetchTunerRegistryTest, SetRegistryAndStandaloneWrappers) { + TestMetricContext ctx; + MockTunerPlugin plugin(&ctx); + + // Verify initial state + EXPECT_EQ(GetGlobalPrefetchTunerRegistry(), nullptr); + + SetGlobalPrefetchTunerRegistry(&plugin); + EXPECT_EQ(GetGlobalPrefetchTunerRegistry(), &plugin); + + PrefetchTuningContext context; + context.hint = PrefetchTuningHint::kRandom; + context.total_elements = 1000; + context.file_loc = "standalone.cc"; + context.line_loc = 100; + + // 1. Test RegisterPrefetchContext + CallsiteId dense_id = RegisterPrefetchContext(context); + EXPECT_EQ(dense_id, 77); + EXPECT_EQ(ctx.registered_count, 1); + + // 2. Test CreatePrefetchTuningScopeByCallsiteId + PrefetchTuningScope scope = + CreatePrefetchTuningScopeByCallsiteId(77, context); + EXPECT_EQ(scope.GetArgs().deep_lookahead, 99); + EXPECT_EQ(ctx.captured_dense_id, 77); + + // 3. Test Fallback when registry is nullptr + GetGlobalPrefetchTunerRegistry() = nullptr; + EXPECT_EQ(RegisterPrefetchContext(context), 0); + + PrefetchTuningScope fallback_scope = + CreatePrefetchTuningScopeByCallsiteId(77, context); +#if HWY_ARCH_ARM_A64 + EXPECT_EQ(fallback_scope.GetArgs().deep_lookahead, 64); +#else + EXPECT_EQ(fallback_scope.GetArgs().deep_lookahead, 32); +#endif +} + +TEST(PrefetchTunerRegistryTest, GlobalRegistryRoutingAndShortCircuit) { + TestMetricContext ctx; + MockTunerPlugin plugin(&ctx); + GetGlobalPrefetchTunerRegistry() = &plugin; + + PrefetchTuningContext context; + context.hint = PrefetchTuningHint::kRandom; + context.total_elements = 1000; + context.is_ultra_tiny = false; + + { + PrefetchTuningScope scope = + CreatePrefetchTuningScopeByCallsiteId(77, context); + EXPECT_EQ(scope.GetArgs().deep_lookahead, 99); + } + EXPECT_TRUE(ctx.called); + + // Test short-circuiting on ultra-tiny workloads + ctx.called = false; + context.is_ultra_tiny = true; + { + PrefetchTuningScope scope = + CreatePrefetchTuningScopeByCallsiteId(77, context); + // Bypasses plugin, returns DefaultRandom() +#if HWY_ARCH_ARM_A64 + EXPECT_EQ(scope.GetArgs().deep_lookahead, 64); +#else + EXPECT_EQ(scope.GetArgs().deep_lookahead, 32); +#endif + } + EXPECT_FALSE(ctx.called); + + GetGlobalPrefetchTunerRegistry() = nullptr; + + // Test fallback when tuner plugin is not available + context.is_ultra_tiny = false; + { + PrefetchTuningScope scope = + CreatePrefetchTuningScopeByCallsiteId(77, context); + // Fallback returns DefaultRandom() +#if HWY_ARCH_ARM_A64 + EXPECT_EQ(scope.GetArgs().deep_lookahead, 64); +#else + EXPECT_EQ(scope.GetArgs().deep_lookahead, 32); +#endif + } +} + +struct TestFissionPolicy : DefaultPrefetchPolicy { + [[maybe_unused]] static constexpr size_t kUltraTinyThreshold = 32; +}; + +struct TestFissionExec { + int* captured_line; + bool* captured_tiny; + int base_line; + + template + void operator()(bool tiny) const { + *captured_line = base_line + static_cast(FissionOffset); + *captured_tiny = tiny; + } +}; + +TEST(PrefetchTunerRegistryTest, WorkloadFissionLineEnrichment) { + int captured_line = 0; + bool captured_tiny = false; + TestFissionExec exec{&captured_line, &captured_tiny, 42}; + + // 1. Ultra-Tiny (< 32) + DispatchWorkloadFission( + 10, exec); + EXPECT_EQ(captured_line, 42); + EXPECT_TRUE(captured_tiny); + + // 2. Sequential Tiny (<= 256) + DispatchWorkloadFission( + 128, exec); + EXPECT_EQ(captured_line, 100042); + EXPECT_FALSE(captured_tiny); + + // 3. Sequential Medium (<= 65536) + DispatchWorkloadFission( + 1000, exec); + EXPECT_EQ(captured_line, 200042); + EXPECT_FALSE(captured_tiny); + + // 4. Sequential Huge (> 65536) + DispatchWorkloadFission( + 100000, exec); + EXPECT_EQ(captured_line, 300042); + EXPECT_FALSE(captured_tiny); + + // 5. Random Tiny (<= 256) + DispatchWorkloadFission(128, + exec); + EXPECT_EQ(captured_line, 400042); + EXPECT_FALSE(captured_tiny); + + // 6. Auto Medium (<= 65536) + DispatchWorkloadFission(1000, + exec); + EXPECT_EQ(captured_line, 800042); + EXPECT_FALSE(captured_tiny); +} + +TEST(PrefetchTunerRegistryTest, UniversalFissionDispatchAndStaticCaching) { + TestMetricContext ctx; + MockTunerPlugin plugin(&ctx); + GetGlobalPrefetchTunerRegistry() = &plugin; + + bool runner_called = false; + PrefetchArgs runner_args; + auto loop_runner = [&](const PrefetchArgs& args) { + runner_called = true; + runner_args = args; + }; + + // 1. First execution (triggers static init RegisterPrefetchContext) + DispatchTunedWorkload( + 1000, "fission_test.cc", 50, loop_runner); + + EXPECT_TRUE(runner_called); + EXPECT_EQ(runner_args.deep_lookahead, 99); + EXPECT_EQ(ctx.registered_count, 1); + EXPECT_EQ(ctx.captured_dense_id, 77); + EXPECT_EQ(ctx.captured_line, 200050); // 50 + 200000 (Medium Sequential) + + // 2. Re-execute the exact same lambda. Confirms RegisterPrefetchContext is + // bypassed! + runner_called = false; + ctx.captured_dense_id = 9999; + DispatchTunedWorkload( + 1000, "fission_test.cc", 50, loop_runner); + + EXPECT_TRUE(runner_called); + EXPECT_EQ(ctx.registered_count, 1); // Still 1! Zero re-registration! + EXPECT_EQ(ctx.captured_dense_id, 77); // Uses statically cached 77! + + GetGlobalPrefetchTunerRegistry() = nullptr; +} + +} // namespace +} // namespace low_level +} // namespace pipeline +} // namespace hwy diff --git a/hwy/contrib/pipeline/prefetch_tuner_test.cc b/hwy/contrib/pipeline/prefetch_tuner_test.cc new file mode 100644 index 0000000000..569e5e4c07 --- /dev/null +++ b/hwy/contrib/pipeline/prefetch_tuner_test.cc @@ -0,0 +1,82 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 + +#include "hwy/contrib/pipeline/prefetch_tuner.h" + +#include +#include +#include + +#include "hwy/contrib/pipeline/prefetch_pipeline_types.h" +#include "hwy/tests/hwy_gtest.h" + +namespace hwy { +namespace pipeline { +namespace low_level { +namespace { + +struct TestMetricContext { + bool called = false; + float ticks = 0; +}; + +void FakeMetricCb(void* user_data, float elapsed_ticks) { + auto* ctx = static_cast(user_data); + ctx->called = true; + ctx->ticks = elapsed_ticks; +} + +struct FakeTimer { + static uint64_t fake_ticks; + static uint64_t Start() { return fake_ticks; } + static uint64_t Stop() { return fake_ticks; } +}; +uint64_t FakeTimer::fake_ticks = 0; + +TEST(PrefetchTunerTest, ScopeLifecycleAndMoveSemantics) { + TestMetricContext ctx; + PrefetchArgs args{.deep_lookahead = 16, .shallow_lookahead = 4}; + + { + PrefetchTuningScope scope1(args, FakeMetricCb, &ctx, 1.0f); + EXPECT_EQ(scope1.GetArgs().deep_lookahead, 16); + EXPECT_FALSE(ctx.called); + + // Move construct into scope2 + PrefetchTuningScope scope2(std::move(scope1)); + EXPECT_EQ(scope2.GetArgs().deep_lookahead, 16); + EXPECT_FALSE(ctx.called); + + // Move assign into scope3 + PrefetchTuningScope scope3; + scope3 = std::move(scope2); + EXPECT_EQ(scope3.GetArgs().deep_lookahead, 16); + EXPECT_FALSE(ctx.called); + } + // Destructor of scope3 fires exactly once + EXPECT_TRUE(ctx.called); +} + +TEST(PrefetchTunerTest, CostNormalizationWithFakeTimer) { + TestMetricContext ctx; + PrefetchArgs args{.deep_lookahead = 32, .shallow_lookahead = 8}; + + FakeTimer::fake_ticks = 1000; + { + // Create scope with cost normalization factor = 10.0f + PrefetchTuningScopeT scope(args, FakeMetricCb, &ctx, 10.0f); + EXPECT_EQ(scope.GetArgs().deep_lookahead, 32); + EXPECT_FALSE(ctx.called); + + FakeTimer::fake_ticks = 2000; // elapsed = 1000 ticks + } + + EXPECT_TRUE(ctx.called); + // normalized cost = 1000 ticks / 10.0f = 100.0f + EXPECT_FLOAT_EQ(ctx.ticks, 100.0f); +} + +} // namespace +} // namespace low_level +} // namespace pipeline +} // namespace hwy diff --git a/hwy_tests.bzl b/hwy_tests.bzl index 660e59a2dc..0608c33f65 100644 --- a/hwy_tests.bzl +++ b/hwy_tests.bzl @@ -67,6 +67,38 @@ HWY_CONTRIB_TESTS = ( "math_hyper_test", [":math"], ), + # copybara:strip_begin(internal) + ( + "hwy/contrib/pipeline/", + "prefetch_pipeline_test", + [":prefetch_pipeline"], + ), + ( + "hwy/contrib/pipeline/", + "prefetch_pipeline_2d_test", + [":prefetch_pipeline"], + ), + ( + "hwy/contrib/pipeline/", + "prefetch_pipeline_types_test", + [":prefetch_pipeline"], + ), + ( + "hwy/contrib/pipeline/", + "prefetch_tuner_test", + [":prefetch_pipeline"], + ), + ( + "hwy/contrib/pipeline/", + "prefetch_tuner_registry_test", + [":prefetch_pipeline"], + ), + ( + "hwy/contrib/dot/", + "one_to_many_test", + [":dot", ":prefetch_pipeline"], + ), + # copybara:strip_end ( "hwy/contrib/math/", "math_tan_test",