perf(tensilelite): recompute the CPU reference only when its inputs were rewritten

projects/hipblaslt/tensilelite/client/include/DataInitialization.hpp

@@ -38,6 +38,7 @@
 #include <mxDataGen.hpp>
 
 #include <cstddef>
+#include <stdexcept>
 #include <map>
 #include <random>
 
@@ -193,6 +194,33 @@ namespace TensileLite
                 return needsSolutionDependentMXPreswizzle(problem, solution);
             }
 
+            // Version of the host-side inputs (A, B, MX scales, ...) that the CPU reference
+            // reads. It is incremented at every point where those buffers are actually rewritten,
+            // so it is the single source of truth for "have the reference inputs changed?".
+            // Consumers (ReferenceValidator) remember the epoch they computed with and only
+            // recompute when it moved; they must not re-derive the answer from their own
+            // conditions (arch, solution format, ...), which is how the two checks used to drift.
+            using InputEpoch = size_t;
+
+            // Problem-level epoch, for use right after prepareCPUInputs().
+            InputEpoch inputEpoch() const
+            {
+                return m_inputEpoch;
+            }
+
+            // Per-solution epoch. Only valid once preSolution() has run for `solution`, because
+            // that is where the inputs are refreshed. Throws otherwise so a listener-ordering
+            // mistake is a loud error instead of a stale reference.
+            InputEpoch inputEpoch(ContractionSolution const* solution) const
+            {
+                if(solution == nullptr || solution != m_currentSolution)
+                    throw std::logic_error(
+                        "DataInitialization::inputEpoch(solution) called before "
+                        "DataInitialization::preSolution() ran for that solution; "
+                        "DataInitialization must be registered before its consumers.");
+                return m_inputEpoch;
+            }
+
             /**
              * Returns a ContractionInputs object with pointers to CPU memory,
              * suitable for using to calculate reference results.
@@ -967,16 +995,25 @@ namespace TensileLite
                 if(needsMXReinit)
                 {
                     ScopedDevice device(m_deviceId);
+                    // The steps below run for every solution that passes the gate above. Report
+                    // them separately depending on whether initializeMXData really regenerated
+                    // the MX data this time (it skips that when its cache is valid).
+                    InputEpoch const epochBefore = m_inputEpoch;
                     {
                         ScopedTimer timer("pre_solution_mx_init_data");
                         initializeMXData(*m_currentGemmProblem);
                     }
+                    bool const regenerated = m_inputEpoch != epochBefore;
                     {
-                        ScopedTimer timer("pre_solution_mx_copy_valid_to_gpu");
+                        ScopedTimer timer(regenerated
+                                              ? "pre_solution_mx_copy_valid_to_gpu_after_regen"
+                                              : "pre_solution_mx_copy_valid_to_gpu_no_regen");
                         copyValidToGPUBuffer(*m_currentGemmProblem);
                     }
                     {
-                        ScopedTimer timer("pre_solution_mx_copy_inputs_d2d");
+                        ScopedTimer timer(regenerated
+                                              ? "pre_solution_mx_copy_inputs_d2d_after_regen"
+                                              : "pre_solution_mx_copy_inputs_d2d_no_regen");
                         copyInputs(m_gpuPtrs,
                                    m_gpuBatchPtrs,
                                    m_maxElements,
@@ -985,7 +1022,8 @@ namespace TensileLite
                                    hipMemcpyDeviceToDevice);
                     }
                     // Sync CPU current buffers so the reference matches GPU data.
-                    ScopedTimer cpuSyncTimer("pre_solution_mx_cpu_sync");
+                    ScopedTimer cpuSyncTimer(regenerated ? "pre_solution_mx_cpu_sync_after_regen"
+                                                         : "pre_solution_mx_cpu_sync_no_regen");
                     for(int ti : {ContractionProblemGemm::TENSOR::A,
                                   ContractionProblemGemm::TENSOR::B,
                                   ContractionProblemGemm::TENSOR::MXSA,
@@ -1265,6 +1303,7 @@ namespace TensileLite
             // Scale descriptor most recently swizzled into gpuInput.valid, per
             // tensor index. The MX equivalent of g_swizzleCache.
             std::map<size_t, TensorDescriptor> m_mxSwizzledDescriptor;
+            InputEpoch                         m_inputEpoch = 0;
         };
 
         template <>

projects/hipblaslt/tensilelite/client/include/ReferenceValidator.hpp

@@ -141,6 +141,17 @@ namespace TensileLite
             std::vector<std::shared_ptr<ProblemInputs>> m_referenceInputs;
             std::vector<std::vector<uint8_t>>           m_referenceRecvs;
 
+            // Per device: DataInitialization::inputEpoch() the reference result was last computed with.
+            // The reference is recomputed only when the epoch has moved.
+            std::vector<DataInitialization::InputEpoch> m_referenceInputEpoch;
+
+            // Debug/CI self-check: when the epoch says the inputs are unchanged, recompute the
+            // reference anyway and require a bitwise-identical result. Catches a missed epoch
+            // bump (stale reference) instead of letting it pass silently.
+            bool m_verifyReferenceReuse = false;
+
+            void verifyReferenceReuse(ContractionProblemGemm const& problem, size_t device);
+
             size_t                   m_cpuResultBufferSize = 0;
             std::shared_ptr<uint8_t> m_cpuResultBuffer;
 

projects/hipblaslt/tensilelite/client/include/TimingInstrumentation.hpp

@@ -23,9 +23,43 @@ namespace TensileLite
         inline double g_calibratedPerCallOverheadMs = 0;
 
         // Set while the per-solution CPU reference recompute runs, so the SolveCPU timers
-        // ("solve_cpu_fast"/"solve_cpu_slow") are reported under *_per_solution names and
-        // do not get attributed to the per-problem cpu_reference_gemm timer.
-        inline thread_local bool g_inPerSolutionReference = false;
+        // ("solve_cpu_fast"/"solve_cpu_slow") are reported under per-solution names and
+        // do not get attributed to the per-problem cpu_reference_gemm timer. The kind says
+        // whether the MX inputs actually changed since the reference was last computed, which
+        // separates necessary recomputes from redundant ones.
+        enum class PerSolutionRefKind
+        {
+            None,
+            InputsChanged,
+            InputsUnchanged,
+        };
+        inline thread_local PerSolutionRefKind g_perSolutionRefKind = PerSolutionRefKind::None;
+
+        inline const char* solveCpuFastCategory()
+        {
+            switch(g_perSolutionRefKind)
+            {
+            case PerSolutionRefKind::InputsChanged:
+                return "solve_cpu_fast_per_solution_inputs_changed";
+            case PerSolutionRefKind::InputsUnchanged:
+                return "solve_cpu_fast_per_solution_inputs_unchanged";
+            default:
+                return "solve_cpu_fast";
+            }
+        }
+
+        inline const char* solveCpuSlowCategory()
+        {
+            switch(g_perSolutionRefKind)
+            {
+            case PerSolutionRefKind::InputsChanged:
+                return "solve_cpu_slow_per_solution_inputs_changed";
+            case PerSolutionRefKind::InputsUnchanged:
+                return "solve_cpu_slow_per_solution_inputs_unchanged";
+            default:
+                return "solve_cpu_slow";
+            }
+        }
 
         // Set by main() around the per-run "gpu_input_reset" call of prepareGPUInputs() so the
         // sub-timers inside it are reported as gpu_input_reset_* instead of gpu_input_prep_*.

projects/hipblaslt/tensilelite/client/main.cpp

@@ -268,6 +268,7 @@ namespace TensileLite
                 ("pristine-on-gpu",          po::value<bool>()->default_value(true), "Keep a pristine copy of inputs on GPU for performance")
                 ("c-equal-d",                po::value<bool>()->default_value(false), "C equals D")
                 ("print-valids",             po::value<bool>()->default_value(false), "Print values that pass validation")
+                ("verify-reference-reuse",   po::value<bool>()->default_value(false), "Recompute the CPU reference even when its inputs are unchanged and fail if the result differs (also: TENSILELITE_VERIFY_REFERENCE_REUSE=1)")
                 ("print-max",                po::value<int>()->default_value(-1), "Max number of values to print")
                 ("num-elements-to-validate", po::value<int>()->default_value(0), "Number of elements to validate")
                 ("check-synchronizer",       po::value<bool>()->default_value(true),

projects/hipblaslt/tensilelite/client/src/DataInitialization.cpp

@@ -1772,6 +1772,8 @@ namespace TensileLite
             if(m_mxDataGenerator == 1)
                 throw std::runtime_error("mx-data-generator=1 is not supported for grouped GEMM");
 
+            // May rewrite the host inputs the reference reads; over-reporting is safe.
+            ++m_inputEpoch;
             for(size_t i = 0; i < m_vdata.size(); i++)
             {
                 if(m_problemDependentData)
@@ -1849,6 +1851,8 @@ namespace TensileLite
 
         void DataInitialization::initializeCPUInputs(ContractionProblemGemm const& problem)
         {
+            // May rewrite the host inputs the reference reads; over-reporting is safe.
+            ++m_inputEpoch;
             // Always drive mxDataGenerator for MX sides so data and scales stay
             // coordinated. initializeMXData also produces the arch-appropriate
             // swizzled scale into gpuInput.valid (when m_mxScaleLayout selects
@@ -2099,6 +2103,7 @@ namespace TensileLite
                                   m_vdata[i].init,
                                   p.second.cpuInput.valid.get(),
                                   tensors[i]);
+                        ++m_inputEpoch; // host input rewritten
                     }
                 }
             };
@@ -2173,6 +2178,20 @@ namespace TensileLite
                           "mx-data-generator=0 supports only Float8/BFloat8/Float4 MX data, got "
                           + TensileLite::ToString(dataDesc.dataType()));
 
+                  // Per-solution calls report whether this side reused its cached inputs
+                  // (cache_hit) or really regenerated data + scales (regenerate). The
+                  // problem-level calls (no solution yet) are covered by cpu_data_init.
+                  struct SideOutcome
+                  {
+                      const char*                          name = nullptr;
+                      TimingClock::time_point              start;
+                      bool                                 enabled;
+                      ~SideOutcome()
+                      {
+                          if(name && enabled)
+                              recordTimingSince(name, start);
+                      }
+                  } sideOutcome{nullptr, TimingClock::now(), m_currentSolution != nullptr};
                   auto         rows       = dataDesc.sizes()[0];
                   auto         cols       = dataDesc.sizes()[1];
                   auto         stride     = dataDesc.strides()[1];
@@ -2199,6 +2218,7 @@ namespace TensileLite
                   // preswizzle choice also depends on the current solution's MI shape.
                   if(m_mxScaleLayout == MXScaleLayout::GFX1250 && descriptorsUnchanged)
                   {
+                      sideOutcome.name = "pre_solution_mx_side_cache_hit";
                       return;
                   }
 
@@ -2223,9 +2243,14 @@ namespace TensileLite
                   {
                       if(hasMatchingSwizzledScale)
                           *preswizzledFlag = true;
+                      sideOutcome.name = "pre_solution_mx_side_cache_hit";
                       return;
                   }
 
+                  // Past the cache checks: A/B data and scales are really regenerated below.
+                  sideOutcome.name = "pre_solution_mx_side_regenerate";
+                  ++m_inputEpoch;
+
                   if(pristineData.initDescriptor.size() == 1)
                       pristineData.initDescriptor[0] = dataDesc;
                   if(pristineScale.initDescriptor.size() == 1)

projects/hipblaslt/tensilelite/client/src/Reference.cpp

@@ -3318,8 +3318,7 @@ namespace TensileLite
 
             if(tryFastPath && isDenseEnoughForFastPath && isFastPathEligible(problem))
             {
-                ScopedTimer timer(g_inPerSolutionReference ? "solve_cpu_fast_per_solution"
-                                                           : "solve_cpu_fast");
+                ScopedTimer timer(solveCpuFastCategory());
                 bool isDouble = (problem.d().dataType() == rocisa::DataType::Double);
                 if(isDouble)
                     solveCPUFast<double>(problem, inputs);
@@ -3331,8 +3330,7 @@ namespace TensileLite
             }
 
             {
-                ScopedTimer timer(g_inPerSolutionReference ? "solve_cpu_slow_per_solution"
-                                                           : "solve_cpu_slow");
+                ScopedTimer timer(solveCpuSlowCategory());
                 auto contractionInputsTypeId = getInputContractionInputsTypeId(problem);
                 SolveCPUTemplates(contractionInputsTypeId, problem, inputs, elementsToValidate);
             }

projects/hipblaslt/tensilelite/client/src/ReferenceValidator.cpp

@@ -39,17 +39,41 @@
 #include <cstring>
 #include <iostream>
 #include <optional>
+#include <cstdlib>
 #include <sstream>
+#include <stdexcept>
+#include <vector>
 
 namespace TensileLite
 {
     namespace Client
     {
+        namespace
+        {
+            // Routes the SolveCPU timers to the per-solution categories for its lifetime.
+            struct PerSolutionRefScope
+            {
+                explicit PerSolutionRefScope(PerSolutionRefKind kind)
+                {
+                    g_perSolutionRefKind = kind;
+                }
+                ~PerSolutionRefScope()
+                {
+                    g_perSolutionRefKind = PerSolutionRefKind::None;
+                }
+            };
+        }
+
         ReferenceValidator::ReferenceValidator(po::variables_map const& args)
         {
             m_elementsToValidate = args["num-elements-to-validate"].as<int>();
             m_printValids        = args["print-valids"].as<bool>();
             m_printMax           = args["print-max"].as<int>();
+            m_verifyReferenceReuse = args.count("verify-reference-reuse")
+                                     && args["verify-reference-reuse"].as<bool>();
+            // Also settable from the environment so the pytest-driven client runs can enable it.
+            if(char const* env = std::getenv("TENSILELITE_VERIFY_REFERENCE_REUSE"))
+                m_verifyReferenceReuse = m_verifyReferenceReuse || std::string(env) == "1";
 
             m_printTensorA             = args["print-tensor-a"].as<bool>();
             m_printTensorB             = args["print-tensor-b"].as<bool>();
@@ -138,6 +162,11 @@ namespace TensileLite
                 if(auto const* gemm = dynamic_cast<ContractionProblemGemm const*>(problem);
                    gemm != nullptr && gemm->fusedGemmA2A())
                     referenceRecvs(*gemm);
+
+                // The references now match the inputs as of these epochs.
+                m_referenceInputEpoch.clear();
+                for(auto const& ctx : *m_devices)
+                    m_referenceInputEpoch.push_back(ctx->dataInit->inputEpoch());
             }
         }
 
@@ -155,29 +184,58 @@ namespace TensileLite
 
             if(auto* gemm = dynamic_cast<ContractionProblemGemm*>(m_problem))
             {
-                // Match DataInitialization MX gate.
-                if(!isMXProblem(*gemm))
-                    return;
-                // Only recompute when DataInitialization actually refreshes MX
-                // inputs for this solution (solution-dependent HostPreSwizzle).
-                // Otherwise the preProblem reference is unchanged, so reuse it
-                // instead of paying a full dense reference GEMM per solution.
+                // DataInitialization::preSolution() has already refreshed the inputs for this
+                // solution (it is registered before this listener; inputEpoch() throws if not).
+                // The reference only needs recomputing if that actually rewrote them. The epoch
+                // is the single source of truth for that; do not re-derive it from arch/format.
                 for(size_t d = 0; d < m_devices->size(); d++)
                 {
-                    if(!(*m_devices)[d]->dataInit->referenceNeedsPerSolutionRecompute(*gemm,
-                                                                                      solution))
-                        continue;
-                    ScopedTimer timer("cpu_reference_gemm_per_solution");
-                    struct PerSolutionScope
+                    auto const epoch = (*m_devices)[d]->dataInit->inputEpoch(solution);
+                    if(epoch == m_referenceInputEpoch[d])
                     {
-                        PerSolutionScope() { g_inPerSolutionReference = true; }
-                        ~PerSolutionScope() { g_inPerSolutionReference = false; }
-                    } perSolutionScope;
+                        if(m_verifyReferenceReuse)
+                            verifyReferenceReuse(*gemm, d);
+                        else
+                        {
+                            // Inputs unchanged: the existing reference is still correct
+                            // (count-only marker).
+                            ScopedTimer timer("pre_solution_ref_skipped");
+                        }
+                        continue;
+                    }
+                    m_referenceInputEpoch[d] = epoch;
+
+                    ScopedTimer timer("cpu_reference_gemm_per_solution_inputs_changed");
+                    PerSolutionRefScope perSolutionScope(PerSolutionRefKind::InputsChanged);
                     SolveCPU(m_problem, m_referenceInputs[d].get(), m_elementsToValidate);
                 }
             }
         }
 
+        void ReferenceValidator::verifyReferenceReuse(ContractionProblemGemm const& problem,
+                                                      size_t                        device)
+        {
+            auto* ref = dynamic_cast<ContractionInputs*>(m_referenceInputs[device].get());
+            if(ref == nullptr || ref->d == nullptr)
+                return;
+
+            size_t const         bytes = problem.d().totalAllocatedBytes();
+            std::vector<uint8_t> reused(bytes);
+            std::memcpy(reused.data(), ref->d, bytes);
+
+            {
+                ScopedTimer timer("cpu_reference_gemm_per_solution_inputs_unchanged");
+                PerSolutionRefScope perSolutionScope(PerSolutionRefKind::InputsUnchanged);
+                SolveCPU(m_problem, m_referenceInputs[device].get(), m_elementsToValidate);
+            }
+
+            if(std::memcmp(reused.data(), ref->d, bytes) != 0)
+                throw std::runtime_error(
+                    "Stale CPU reference: DataInitialization reported unchanged inputs "
+                    "(input epoch not advanced) but recomputing the reference gave a "
+                    "different result. A writer of the host inputs is missing an epoch bump.");
+        }
+
         bool ReferenceValidator::needMoreRunsInSolution() const
         {
             if(m_enabled && !m_validatedSolution)