From 70788a0a520d9826a195449ca1023ef02670dd27 Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Sun, 17 May 2026 08:41:50 -0500 Subject: [PATCH 1/6] Refactor Transform to avoid memory allocation Use a re-entrancy safe thread local storage for intermediate results. --- src/correction.cc | 42 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/src/correction.cc b/src/correction.cc index 2a6b126..5e853e5 100644 --- a/src/correction.cc +++ b/src/correction.cc @@ -146,6 +146,44 @@ namespace { const std::vector& values; }; + // Per-thread scratch storage for Transform::evaluate. + // Depth indexing keeps nested Transform evaluations re-entrant-safe. + class TransformScratch { + public: + explicit TransformScratch(const std::vector& values): + slot_(acquire_slot()) + { + slot_ = values; + } + + ~TransformScratch() { + release_slot(); + } + + std::vector& values() { + return slot_; + } + + private: + static std::vector& acquire_slot() { + if (depth_ == slots_.size()) { + slots_.emplace_back(); + } + return slots_[depth_++]; + } + + static void release_slot() { + depth_--; + } + + std::vector& slot_; + static thread_local std::vector> slots_; + static thread_local std::size_t depth_; + }; + + thread_local std::vector> TransformScratch::slots_; + thread_local std::size_t TransformScratch::depth_ = 0; + std::size_t find_bin_idx(Variable::Type value_variant, const std::variant<_UniformBins, _NonUniformBins> &bins_, const _FlowBehavior &flow, @@ -392,7 +430,9 @@ Transform::Transform(const JSONObject& json, const Correction& context) { } double Transform::evaluate(const std::vector& values) const { - std::vector new_values(values); + TransformScratch scratch(values); + auto& new_values = scratch.values(); + double vnew = std::visit(node_evaluate{values}, *rule_); auto& v = new_values[variableIdx_]; if ( std::holds_alternative(v) ) { From 6a6cb134486f0d36a3db6fbc0def08ea5d9fdb85 Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Sun, 17 May 2026 09:15:48 -0500 Subject: [PATCH 2/6] Add dedicated nested transform regression test --- tests/test_transform_nested.py | 70 ++++++++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) create mode 100644 tests/test_transform_nested.py diff --git a/tests/test_transform_nested.py b/tests/test_transform_nested.py new file mode 100644 index 0000000..a581584 --- /dev/null +++ b/tests/test_transform_nested.py @@ -0,0 +1,70 @@ +import correctionlib._core as core +from correctionlib import schemav2 as schema + + +def wrap(*corrs): + cset = schema.CorrectionSet( + schema_version=schema.VERSION, + corrections=list(corrs), + ) + return core.CorrectionSet.from_string(cset.model_dump_json()) + + +def test_transform_nested_rule_and_content(): + # Build a nested transform tree with transforms in both the rule and + # content paths to exercise recursive evaluation and scratch-buffer reuse. + cset = wrap( + schema.Correction( + name="nested_transform", + version=2, + inputs=[ + schema.Variable(name="x", type="real"), + schema.Variable(name="y", type="real"), + ], + output=schema.Variable(name="out", type="real"), + data=schema.Transform( + nodetype="transform", + input="x", + rule=schema.Transform( + nodetype="transform", + input="y", + rule=1.0, + content=schema.Formula( + nodetype="formula", + expression="x + y", + parser="TFormula", + variables=["x", "y"], + ), + ), + content=schema.Transform( + nodetype="transform", + input="y", + rule=schema.Transform( + nodetype="transform", + input="x", + rule=2.0, + content=schema.Formula( + nodetype="formula", + expression="x + y", + parser="TFormula", + variables=["x", "y"], + ), + ), + content=schema.Formula( + nodetype="formula", + expression="x + y", + parser="TFormula", + variables=["x", "y"], + ), + ), + ), + ) + ) + + corr = cset["nested_transform"] + # x=3, y=4 + # outer rule: transform y->1 then x+y => 3+1 = 4, so x becomes 4 + # outer content: transform y with rule: + # inner rule transform x->2 then x+y => 2+4 = 6, so y becomes 6 + # final content: x+y = 4+6 = 10 + assert corr.evaluate(3.0, 4.0) == 10.0 From 815b421d96df5a7bfffa8d862a844f0c7fd8cd12 Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Sun, 17 May 2026 09:41:21 -0500 Subject: [PATCH 3/6] Use deque for slot storage Otherwise, when the vector grows beyond capacity, all existing references become invalid --- src/correction.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/correction.cc b/src/correction.cc index 5e853e5..ca29401 100644 --- a/src/correction.cc +++ b/src/correction.cc @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include // std::abort @@ -177,11 +178,11 @@ namespace { } std::vector& slot_; - static thread_local std::vector> slots_; + static thread_local std::deque> slots_; static thread_local std::size_t depth_; }; - thread_local std::vector> TransformScratch::slots_; + thread_local std::deque> TransformScratch::slots_; thread_local std::size_t TransformScratch::depth_ = 0; std::size_t find_bin_idx(Variable::Type value_variant, From da964843923c78df1058d59ea3d1b35ac4ee965c Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Mon, 15 Jun 2026 14:24:28 -0500 Subject: [PATCH 4/6] Use thread-local storage for CompoundCorrection as well --- src/correction.cc | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/correction.cc b/src/correction.cc index ca29401..9043d17 100644 --- a/src/correction.cc +++ b/src/correction.cc @@ -806,6 +806,11 @@ size_t CompoundCorrection::input_index(const std::string_view name) const { } double CompoundCorrection::evaluate(const std::vector& values) const { + // Per-thread scratch storage. This call site is not re-entrant so we + // can use a simpler implementation than for TransformScratch + static thread_local std::vector ivalues; + static thread_local std::vector cvalues; + if ( values.size() != inputs_.size() ) { throw std::invalid_argument("Incorrect number of inputs (got " + std::to_string(values.size()) + ", expected " + std::to_string(inputs_.size()) + ")"); @@ -813,9 +818,9 @@ double CompoundCorrection::evaluate(const std::vector& values) c for (size_t i=0; i < inputs_.size(); ++i) { inputs_[i].validate(values[i]); } - std::vector ivalues(values); - std::vector cvalues; + ivalues = values; cvalues.reserve(values.size()); + double out = 0.; double sf = 0.; bool start{true}; From c7ec0ed5f99814befb5894d2bfe70bd106cbd2ce Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Mon, 15 Jun 2026 16:28:31 -0500 Subject: [PATCH 5/6] Maybe fix bins --- src/correction.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/correction.cc b/src/correction.cc index 9043d17..2a4a950 100644 --- a/src/correction.cc +++ b/src/correction.cc @@ -223,7 +223,7 @@ namespace { // otherwise we have non-uniform binning using namespace std::string_literals; - const auto bins = std::get<_NonUniformBins>(bins_); + const auto& bins = std::get<_NonUniformBins>(bins_); if ( flow == _FlowBehavior::wrap ) { double low = bins[0]; double high = bins[bins.size() - 1]; From e892fbef0e18dbed4a3573d36a57ad8a8d93f5dc Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Mon, 15 Jun 2026 16:38:56 -0500 Subject: [PATCH 6/6] Squash one more allocation --- src/correction.cc | 34 ++++++++++++++++------------------ 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/src/correction.cc b/src/correction.cc index 2a4a950..24b7374 100644 --- a/src/correction.cc +++ b/src/correction.cc @@ -682,27 +682,25 @@ Category::Category(const JSONObject& json, const Correction& context) double Category::evaluate(const std::vector& values) const { const Content* child = nullptr; if ( auto pval = std::get_if(&values[variableIdx_]) ) { - try { - child = &std::get(map_).at(*pval); - } catch (std::out_of_range& ex) { - if ( default_ ) { - child = default_.get(); - } - else { - throw std::out_of_range("Index not available in Category for input argument " + std::to_string(variableIdx_) + " val: " + *pval); - } + const auto& m = std::get(map_); + auto it = m.find(*pval); + if ( it != m.end() ) { + child = &it->second; + } else if ( default_ ) { + child = default_.get(); + } else { + throw std::out_of_range("Index not available in Category for input argument " + std::to_string(variableIdx_) + " val: " + *pval); } } else if ( auto pval = std::get_if(&values[variableIdx_]) ) { - try { - child = &std::get(map_).at(*pval); - } catch (std::out_of_range& ex) { - if ( default_ ) { - child = default_.get(); - } - else { - throw std::out_of_range("Index not available in Category for input argument " + std::to_string(variableIdx_) + " val: " + std::to_string(*pval)); - } + const auto& m = std::get(map_); + auto it = m.find(*pval); + if ( it != m.end() ) { + child = &it->second; + } else if ( default_ ) { + child = default_.get(); + } else { + throw std::out_of_range("Index not available in Category for input argument " + std::to_string(variableIdx_) + " val: " + std::to_string(*pval)); } } else { throw std::runtime_error("Invalid variable type");