https://github.com/Xazax-hun updated https://github.com/llvm/llvm-project/pull/209807
From f66ca1bf1ea79d2790530023bde3deda7b44cc5b Mon Sep 17 00:00:00 2001 From: Gabor Horvath <[email protected]> Date: Wed, 15 Jul 2026 15:46:16 +0100 Subject: [PATCH] [ADT] Add single-pass merge for ImmutableSet/ImmutableMap and use it in dataflow joins The immutable-set/map dataflow joins merged two containers by inserting the elements of one into the other one at a time, costing O(|B| * log|A|) and re-copying shared spine nodes on every insert. Add a single-pass, structure-sharing tree merge on ImutAVLFactory (mergeTrees), exposed as ImmutableSet::Factory::unionSets and ImmutableMap::Factory::mergeWith. It recurses over the larger operand and splits the smaller at each key, returning non-overlapping subtrees unchanged, so each spine node is copied at most once: O(|B| * log(|A|/|B| + 1)). Two flags tune it for the different joins: * KeepUnmatched - keep keys unique to one side (set union / a lattice join with an identity) vs. pass them through the combiner (a symmetric join, e.g. liveness Must->Maybe). * SkipShared - return a pointer-identical subtree unchanged in O(1). Valid for an idempotent merge (Combine(a, a) == a), which holds for set union and every lattice join. Since dataflow states are path-copied from one another, a join's operands usually share most subtrees by pointer, so a join that touches only a few keys becomes nearly O(diff). Wire it into the Clang lifetime-safety analysis (set and map joins) and the LiveVariables merge, and add ImmutableMergeBM to track it. Speedups on the lifetime-safety dataflow benchmark (LoanPropagation phase): pointer cycle in a loop ~2.6x (1.59 s -> 0.61 s) CFG merges ~190x (1.83 s -> 9.5 ms) switch fan-out ~1.9x (LiveOrigins; disjoint states, no skip) LiveVariables is ~2-3x faster on liveness-heavy inputs and neutral on typical code (the join is a small fraction of real-code analysis time). ImmutableMergeBM confirms the merge beats the add loop for small and near-identical operands, except for independently-built sets differing by exactly one element; when one operand is derived from the other -- the real dataflow case -- the pointer-skip makes even that a large win. New randomized ADT stress tests check the merge against std::set/std::map oracles for both canonicalizing and non-canonicalizing factories. Assisted by: Claude Opus 4.8 --- .../Analysis/Analyses/LifetimeSafety/Utils.h | 48 +-- clang/lib/Analysis/LiveVariables.cpp | 50 +-- llvm/benchmarks/CMakeLists.txt | 1 + llvm/benchmarks/ImmutableMergeBM.cpp | 379 ++++++++++++++++++ llvm/include/llvm/ADT/ImmutableMap.h | 21 + llvm/include/llvm/ADT/ImmutableSet.h | 167 ++++++++ llvm/unittests/ADT/ImmutableMapTest.cpp | 82 ++++ llvm/unittests/ADT/ImmutableSetTest.cpp | 54 +++ 8 files changed, 735 insertions(+), 67 deletions(-) create mode 100644 llvm/benchmarks/ImmutableMergeBM.cpp diff --git a/clang/include/clang/Analysis/Analyses/LifetimeSafety/Utils.h b/clang/include/clang/Analysis/Analyses/LifetimeSafety/Utils.h index b2b19dd3030bb..62eec670c54a2 100644 --- a/clang/include/clang/Analysis/Analyses/LifetimeSafety/Utils.h +++ b/clang/include/clang/Analysis/Analyses/LifetimeSafety/Utils.h @@ -47,13 +47,7 @@ using MapTy = llvm::ImmutableMap<KeyT, ValT, llvm::ImutKeyValueInfo<KeyT, ValT>, /// Computes the union of two ImmutableSets. template <typename T> SetTy<T> join(SetTy<T> A, SetTy<T> B, typename SetTy<T>::Factory &F) { - if (A.getRootWithoutRetain() == B.getRootWithoutRetain()) - return A; - if (A.getHeight() < B.getHeight()) - std::swap(A, B); - for (const auto &E : B) - A = F.add(A, E); - return A; + return F.unionSets(A, B); } /// Describes the strategy for joining two `ImmutableMap` instances, primarily @@ -72,30 +66,36 @@ enum class JoinKind { Asymmetric, }; -/// Computes the key-wise union of two ImmutableMaps. -// TODO(opt): This key-wise join is a performance bottleneck. A more -// efficient merge could be implemented using a Patricia Trie or HAMT -// instead of the current AVL-tree-based ImmutableMap. +/// Computes the key-wise union of two ImmutableMaps in a single traversal +/// (see ImmutableMap::Factory::mergeWith), sharing subtrees the two maps do +/// not overlap. This assumes -- as the swap below already does -- that +/// JoinValues is commutative with a left identity, which holds for the +/// lifetime lattices. template <typename KeyT, typename ValT, typename Joiner> -MapTy<KeyT, ValT> join(const MapTy<KeyT, ValT> &A, const MapTy<KeyT, ValT> &B, +MapTy<KeyT, ValT> join(MapTy<KeyT, ValT> A, MapTy<KeyT, ValT> B, typename MapTy<KeyT, ValT>::Factory &F, Joiner JoinValues, JoinKind Kind) { if (A.getRootWithoutRetain() == B.getRootWithoutRetain()) return A; + // Drive the merge with the taller map so the shorter one is the one split. if (A.getHeight() < B.getHeight()) - return join(B, A, F, JoinValues, Kind); + std::swap(A, B); - // For each element in B, join it with the corresponding element in A - // (or with an empty value if it doesn't exist in A). - MapTy<KeyT, ValT> Res = A; - for (const auto &[Key, ValB] : B) - Res = F.add(Res, Key, JoinValues(A.lookup(Key), &ValB)); - if (Kind == JoinKind::Symmetric) { - for (const auto &[Key, ValA] : A) - if (!B.contains(Key)) - Res = F.add(Res, Key, JoinValues(&ValA, nullptr)); - } - return Res; + using ValueTy = typename MapTy<KeyT, ValT>::value_type; + auto Combine = [&JoinValues](const ValueTy *AElem, + const ValueTy *BElem) -> std::pair<KeyT, ValT> { + const KeyT &Key = AElem ? AElem->first : BElem->first; + return std::pair<KeyT, ValT>(Key, + JoinValues(AElem ? &AElem->second : nullptr, + BElem ? &BElem->second : nullptr)); + }; + // Asymmetric keeps keys unique to either map as-is (valid because JoinValues + // has a left identity); symmetric passes unmatched keys through JoinValues. + // The lifetime joins are idempotent lattice joins, so pointer-identical + // subtrees (common once one state is derived from the other) can be shared. + return F.mergeWith(A, B, Combine, + /*KeepUnmatched=*/Kind == JoinKind::Asymmetric, + /*SkipShared=*/true); } } // namespace clang::lifetimes::internal::utils diff --git a/clang/lib/Analysis/LiveVariables.cpp b/clang/lib/Analysis/LiveVariables.cpp index 3f5ff9ec33c22..c9e547ac85380 100644 --- a/clang/lib/Analysis/LiveVariables.cpp +++ b/clang/lib/Analysis/LiveVariables.cpp @@ -30,9 +30,6 @@ namespace { class LiveVariablesImpl { public: template <typename T> using SetTy = LiveVariables::SetTy<T>; - template <typename T> - using SetRefTy = llvm::ImmutableSetRef<T, llvm::ImutContainerInfo<T>, - /*Canonicalize=*/false>; AnalysisDeclContext &analysisContext; SetTy<const Expr *>::Factory ESetFact; @@ -89,51 +86,18 @@ bool LiveVariables::LivenessValues::isLive(const VarDecl *D) const { return liveDecls.contains(D); } -namespace { -template <typename SET> SET mergeSets(SET A, SET B) { - if (A.getRootWithoutRetain() == B.getRootWithoutRetain()) - return A; - - if (A.getHeight() < B.getHeight()) - std::swap(A, B); - - for (const auto *Elem : B) { - A = A.add(Elem); - } - return A; -} -} // namespace - void LiveVariables::Observer::anchor() { } LiveVariables::LivenessValues LiveVariablesImpl::merge(LiveVariables::LivenessValues valsA, LiveVariables::LivenessValues valsB) { - - SetRefTy<const Expr *> SSetRefA(valsA.liveExprs.getRootWithoutRetain(), - ESetFact.getTreeFactory()), - SSetRefB(valsB.liveExprs.getRootWithoutRetain(), - ESetFact.getTreeFactory()); - - SetRefTy<const VarDecl *> DSetRefA(valsA.liveDecls.getRootWithoutRetain(), - DSetFact.getTreeFactory()), - DSetRefB(valsB.liveDecls.getRootWithoutRetain(), - DSetFact.getTreeFactory()); - - SetRefTy<const BindingDecl *> BSetRefA( - valsA.liveBindings.getRootWithoutRetain(), BSetFact.getTreeFactory()), - BSetRefB(valsB.liveBindings.getRootWithoutRetain(), - BSetFact.getTreeFactory()); - - SSetRefA = mergeSets(SSetRefA, SSetRefB); - DSetRefA = mergeSets(DSetRefA, DSetRefB); - BSetRefA = mergeSets(BSetRefA, BSetRefB); - - // Finalize the merged builder refs into immutable sets. These are not - // canonicalized; LivenessValues::operator== compares them structurally. - return LiveVariables::LivenessValues(SSetRefA.asImmutableSet(), - DSetRefA.asImmutableSet(), - BSetRefA.asImmutableSet()); + // Liveness at a merge point is the union of the successors' live sets. These + // sets are not canonicalized; LivenessValues::operator== compares them + // structurally. + return LiveVariables::LivenessValues( + ESetFact.unionSets(valsA.liveExprs, valsB.liveExprs), + DSetFact.unionSets(valsA.liveDecls, valsB.liveDecls), + BSetFact.unionSets(valsA.liveBindings, valsB.liveBindings)); } bool LiveVariables::LivenessValues::operator==(const LivenessValues &V) const { diff --git a/llvm/benchmarks/CMakeLists.txt b/llvm/benchmarks/CMakeLists.txt index d8c4503e23b92..51c47e5ada602 100644 --- a/llvm/benchmarks/CMakeLists.txt +++ b/llvm/benchmarks/CMakeLists.txt @@ -17,6 +17,7 @@ add_benchmark(DWARFVerifierBM DWARFVerifierBM.cpp PARTIAL_SOURCES_INTENDED) add_benchmark(PointerUnionBM PointerUnionBM.cpp PARTIAL_SOURCES_INTENDED) add_benchmark(ImmutableSetIteratorBM ImmutableSetIteratorBM.cpp PARTIAL_SOURCES_INTENDED) add_benchmark(ImmutableSetBuildBM ImmutableSetBuildBM.cpp PARTIAL_SOURCES_INTENDED) +add_benchmark(ImmutableMergeBM ImmutableMergeBM.cpp PARTIAL_SOURCES_INTENDED) add_benchmark(RuntimeLibcallsBench RuntimeLibcalls.cpp PARTIAL_SOURCES_INTENDED) add_benchmark(writeToOutputInParallelBench writeToOutputInParallel.cpp PARTIAL_SOURCES_INTENDED) diff --git a/llvm/benchmarks/ImmutableMergeBM.cpp b/llvm/benchmarks/ImmutableMergeBM.cpp new file mode 100644 index 0000000000000..a638f3e1e453a --- /dev/null +++ b/llvm/benchmarks/ImmutableMergeBM.cpp @@ -0,0 +1,379 @@ +//===- ImmutableMergeBM.cpp - Benchmark ImmutableSet/Map merges -----------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Compares the single-pass, structure-sharing merge (ImmutableSet::unionSets / +// ImmutableMap::mergeWith) against the per-element add loop it replaced, on the +// operand shapes that dominate dataflow joins: +// +// * Small - both operands tiny (constant-factor sensitive). +// * NearId - both large but differing by only a few keys (the fixpoint +// case, where most subtrees are shared). +// +// It covers a plain set, a scalar-valued map, and the map-of-sets shape used by +// the Clang lifetime analysis (ImmutableMap<Origin, ImmutableSet<Loan>>, whose +// join key-wise unions the inner loan sets). The point of these benchmarks is +// to confirm the bulk merge is not slower than the add loop for small or nearly +// identical inputs. Run before/after and compare (e.g. llvm/utils/compare.py). +// +//===----------------------------------------------------------------------===// + +#include "benchmark/benchmark.h" +#include "llvm/ADT/ImmutableMap.h" +#include "llvm/ADT/ImmutableSet.h" +#include <algorithm> +#include <memory> +#include <numeric> +#include <vector> + +using namespace llvm; + +namespace { + +using IntSet = + ImmutableSet<int, ImutContainerInfo<int>, /*Canonicalize=*/false>; +using IntMap = + ImmutableMap<int, int, ImutKeyValueInfo<int, int>, /*Canonicalize=*/false>; +// The lifetime analysis' OriginLoanMap shape: a map whose values are sets. +using SetMap = ImmutableMap<int, IntSet, ImutKeyValueInfo<int, IntSet>, + /*Canonicalize=*/false>; + +//===----------------------------------------------------------------------===// +// The per-element merges the bulk path replaced. +//===----------------------------------------------------------------------===// + +static IntSet setUnionAdd(IntSet::Factory &F, IntSet A, IntSet B) { + if (A.getRootWithoutRetain() == B.getRootWithoutRetain() || B.isEmpty()) + return A; + if (A.isEmpty()) + return B; + if (A.getHeight() < B.getHeight()) + std::swap(A, B); + for (int E : B) + A = F.add(A, E); + return A; +} + +static IntMap mapUnionAdd(IntMap::Factory &F, IntMap A, IntMap B) { + if (A.getRootWithoutRetain() == B.getRootWithoutRetain()) + return A; + if (A.getHeight() < B.getHeight()) + std::swap(A, B); + IntMap Res = A; + for (IntMap::iterator I = B.begin(), E = B.end(); I != E; ++I) { + const int *AV = A.lookup(I.getKey()); + Res = F.add(Res, I.getKey(), AV ? std::max(*AV, I.getData()) : I.getData()); + } + return Res; +} + +static SetMap setMapUnionAdd(SetMap::Factory &MF, IntSet::Factory &SF, SetMap A, + SetMap B) { + if (A.getRootWithoutRetain() == B.getRootWithoutRetain()) + return A; + if (A.getHeight() < B.getHeight()) + std::swap(A, B); + SetMap Res = A; + for (SetMap::iterator I = B.begin(), E = B.end(); I != E; ++I) { + const IntSet *AV = A.lookup(I.getKey()); + IntSet Merged = AV ? setUnionAdd(SF, *AV, I.getData()) : I.getData(); + Res = MF.add(Res, I.getKey(), Merged); + } + return Res; +} + +//===----------------------------------------------------------------------===// +// The bulk merges (mirroring the caller-side short-circuit/swap the analyses +// perform around them, so both sides are compared on equal footing). +//===----------------------------------------------------------------------===// + +static IntSet setUnionNew(IntSet::Factory &F, IntSet A, IntSet B) { + return F.unionSets(A, B); +} + +static IntMap mapUnionNew(IntMap::Factory &F, IntMap A, IntMap B) { + if (A.getRootWithoutRetain() == B.getRootWithoutRetain()) + return A; + if (A.getHeight() < B.getHeight()) + std::swap(A, B); + auto Combine = [](const IntMap::value_type *AE, + const IntMap::value_type *BE) -> std::pair<int, int> { + int Key = AE ? AE->first : BE->first; + return {Key, std::max(AE ? AE->second : BE->second, + BE ? BE->second : AE->second)}; + }; + return F.mergeWith(A, B, Combine, /*KeepUnmatched=*/true); +} + +static SetMap setMapUnionNew(SetMap::Factory &MF, IntSet::Factory &SF, SetMap A, + SetMap B) { + if (A.getRootWithoutRetain() == B.getRootWithoutRetain()) + return A; + if (A.getHeight() < B.getHeight()) + std::swap(A, B); + auto Combine = [&SF](const SetMap::value_type *AE, + const SetMap::value_type *BE) -> std::pair<int, IntSet> { + int Key = AE ? AE->first : BE->first; + IntSet AS = AE ? AE->second : SF.getEmptySet(); + IntSet BS = BE ? BE->second : SF.getEmptySet(); + return {Key, SF.unionSets(AS, BS)}; + }; + return MF.mergeWith(A, B, Combine, /*KeepUnmatched=*/true); +} + +//===----------------------------------------------------------------------===// +// Key generators. Small: two tiny half-overlapping sets. NearId: two size-N +// sets sharing all but D keys (D unique to each side). +//===----------------------------------------------------------------------===// + +static void makeKeys(bool NearId, int N, int D, std::vector<int> &KA, + std::vector<int> &KB) { + KA.resize(N); + std::iota(KA.begin(), KA.end(), 0); + KB = KA; + if (NearId) { + // Replace the first D keys of B with fresh keys, so A and B share N-D keys + // and each has D unique ones. + for (int I = 0; I < D && I < N; ++I) + KB[I] = N + I; + } else { + // Half-overlap: shift B up by N/2. + for (int &V : KB) + V += N / 2; + } +} + +// Bound allocator growth: results accumulate in the factory, so periodically +// rebuild it (excluded from timing) rather than let it grow unboundedly. +static constexpr int ResetEvery = 4096; + +static IntSet buildSet(IntSet::Factory &F, const std::vector<int> &Keys) { + IntSet S = F.getEmptySet(); + for (int V : Keys) + S = F.add(S, V); + return S; +} + +//===----------------------------------------------------------------------===// +// Benchmarks. +//===----------------------------------------------------------------------===// + +static void BM_Set(benchmark::State &State, bool UseNew, bool NearId, int N, + int D) { + std::vector<int> KA, KB; + makeKeys(NearId, N, D, KA, KB); + + std::unique_ptr<IntSet::Factory> F; + IntSet A{nullptr}, B{nullptr}; + auto reset = [&] { + A = IntSet(nullptr); + B = IntSet(nullptr); + F = std::make_unique<IntSet::Factory>(); + A = buildSet(*F, KA); + B = buildSet(*F, KB); + }; + reset(); + + int Since = 0; + for (auto _ : State) { + // Rebuild before creating this iteration's result, so no live result + // references the factory being torn down. + if (++Since == ResetEvery) { + State.PauseTiming(); + reset(); + Since = 0; + State.ResumeTiming(); + } + IntSet U = UseNew ? setUnionNew(*F, A, B) : setUnionAdd(*F, A, B); + benchmark::DoNotOptimize(U.getRootWithoutRetain()); + } +} + +// Like BM_Set, but B is derived from A by adding D fresh keys, so B shares all +// of A's untouched subtrees by pointer -- the realistic "state after a small +// edit" shape that the pointer-equality skip in unionSets is meant to exploit. +static void BM_SetShared(benchmark::State &State, bool UseNew, int N, int D) { + std::vector<int> KA(N); + std::iota(KA.begin(), KA.end(), 0); + + std::unique_ptr<IntSet::Factory> F; + IntSet A{nullptr}, B{nullptr}; + auto reset = [&] { + A = IntSet(nullptr); + B = IntSet(nullptr); + F = std::make_unique<IntSet::Factory>(); + A = buildSet(*F, KA); + B = A; + for (int I = 0; I < D; ++I) + B = F->add(B, N + I); + }; + reset(); + + int Since = 0; + for (auto _ : State) { + if (++Since == ResetEvery) { + State.PauseTiming(); + reset(); + Since = 0; + State.ResumeTiming(); + } + IntSet U = UseNew ? setUnionNew(*F, A, B) : setUnionAdd(*F, A, B); + benchmark::DoNotOptimize(U.getRootWithoutRetain()); + } +} + +static void BM_Map(benchmark::State &State, bool UseNew, bool NearId, int N, + int D) { + std::vector<int> KA, KB; + makeKeys(NearId, N, D, KA, KB); + + std::unique_ptr<IntMap::Factory> F; + IntMap A{nullptr}, B{nullptr}; + auto reset = [&] { + A = IntMap(nullptr); + B = IntMap(nullptr); + F = std::make_unique<IntMap::Factory>(); + IntMap MA = F->getEmptyMap(), MB = F->getEmptyMap(); + for (int V : KA) + MA = F->add(MA, V, V); + for (int V : KB) + MB = F->add(MB, V, V); + A = MA; + B = MB; + }; + reset(); + + int Since = 0; + for (auto _ : State) { + if (++Since == ResetEvery) { + State.PauseTiming(); + reset(); + Since = 0; + State.ResumeTiming(); + } + IntMap U = UseNew ? mapUnionNew(*F, A, B) : mapUnionAdd(*F, A, B); + benchmark::DoNotOptimize(U.getRootWithoutRetain()); + } +} + +// Map<int, Set<int>>: each key maps to a small set; the merge key-wise unions +// the value sets (the lifetime OriginLoanMap join). +static void BM_MapOfSets(benchmark::State &State, bool UseNew, bool NearId, + int N, int D, int SetSize) { + std::vector<int> KA, KB; + makeKeys(NearId, N, D, KA, KB); + + std::unique_ptr<SetMap::Factory> MF; + std::unique_ptr<IntSet::Factory> SF; + SetMap A{nullptr}, B{nullptr}; + auto valueSet = [&](int Key) { + // A deterministic small set derived from the key. + IntSet S = SF->getEmptySet(); + for (int I = 0; I < SetSize; ++I) + S = SF->add(S, Key * 31 + I); + return S; + }; + auto reset = [&] { + A = SetMap(nullptr); + B = SetMap(nullptr); + MF = std::make_unique<SetMap::Factory>(); + SF = std::make_unique<IntSet::Factory>(); + SetMap MA = MF->getEmptyMap(), MB = MF->getEmptyMap(); + for (int V : KA) + MA = MF->add(MA, V, valueSet(V)); + for (int V : KB) + MB = MF->add(MB, V, valueSet(V)); + A = MA; + B = MB; + }; + reset(); + + int Since = 0; + for (auto _ : State) { + if (++Since == ResetEvery) { + State.PauseTiming(); + reset(); + Since = 0; + State.ResumeTiming(); + } + SetMap U = UseNew ? setMapUnionNew(*MF, *SF, A, B) + : setMapUnionAdd(*MF, *SF, A, B); + benchmark::DoNotOptimize(U.getRootWithoutRetain()); + } +} + +} // namespace + +// Small: both operands tiny. NearId: size N, differ by D keys each side. +BENCHMARK_CAPTURE(BM_Set, "Set/Small/Add", false, false, 8, 0); +BENCHMARK_CAPTURE(BM_Set, "Set/Small/Union", true, false, 8, 0); +BENCHMARK_CAPTURE(BM_Set, "Set/Small16/Add", false, false, 16, 0); +BENCHMARK_CAPTURE(BM_Set, "Set/Small16/Union", true, false, 16, 0); +// Small and nearly identical (share all but one key each side) -- the common +// dataflow live-set shape, and the case most exposed to split overhead. +BENCHMARK_CAPTURE(BM_Set, "Set/SmallNearId8/Add", false, true, 8, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/SmallNearId8/Union", true, true, 8, 1); +// Very small sets, delta 2 (and a couple of neighbours) to see if the Δ1 +// regression persists at tiny sizes once the diff is at least 2. +BENCHMARK_CAPTURE(BM_Set, "Set/N8D2/Add", false, true, 8, 2); +BENCHMARK_CAPTURE(BM_Set, "Set/N8D2/Union", true, true, 8, 2); +BENCHMARK_CAPTURE(BM_Set, "Set/N16D2/Add", false, true, 16, 2); +BENCHMARK_CAPTURE(BM_Set, "Set/N16D2/Union", true, true, 16, 2); +BENCHMARK_CAPTURE(BM_Set, "Set/N32D2/Add", false, true, 32, 2); +BENCHMARK_CAPTURE(BM_Set, "Set/N32D2/Union", true, true, 32, 2); +BENCHMARK_CAPTURE(BM_Set, "Set/SmallNearId64/Add", false, true, 64, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/SmallNearId64/Union", true, true, 64, 1); +// Delta sweep at N=64 (where the near-identical regression peaked): how does +// the gap behave as the two sets differ by more keys? +BENCHMARK_CAPTURE(BM_Set, "Set/N64D2/Add", false, true, 64, 2); +BENCHMARK_CAPTURE(BM_Set, "Set/N64D2/Union", true, true, 64, 2); +BENCHMARK_CAPTURE(BM_Set, "Set/N64D4/Add", false, true, 64, 4); +BENCHMARK_CAPTURE(BM_Set, "Set/N64D4/Union", true, true, 64, 4); +BENCHMARK_CAPTURE(BM_Set, "Set/N64D8/Add", false, true, 64, 8); +BENCHMARK_CAPTURE(BM_Set, "Set/N64D8/Union", true, true, 64, 8); +BENCHMARK_CAPTURE(BM_Set, "Set/N64D16/Add", false, true, 64, 16); +BENCHMARK_CAPTURE(BM_Set, "Set/N64D16/Union", true, true, 64, 16); +BENCHMARK_CAPTURE(BM_Set, "Set/N64D32/Add", false, true, 64, 32); +BENCHMARK_CAPTURE(BM_Set, "Set/N64D32/Union", true, true, 64, 32); +// Δ1 but B derived from A (shares subtrees by pointer): the case the +// pointer-equality skip targets. Compare against the independent-build +// Set/*NearId* numbers above. +BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N8D1/Add", false, 8, 1); +BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N8D1/Union", true, 8, 1); +BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N64D1/Add", false, 64, 1); +BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N64D1/Union", true, 64, 1); +BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N1024D1/Add", false, 1024, 1); +BENCHMARK_CAPTURE(BM_SetShared, "SetShared/N1024D1/Union", true, 1024, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/NearId128/Add", false, true, 128, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/NearId128/Union", true, true, 128, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/NearId256/Add", false, true, 256, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/NearId256/Union", true, true, 256, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/NearId512/Add", false, true, 512, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/NearId512/Union", true, true, 512, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/NearId/Add", false, true, 1024, 1); +BENCHMARK_CAPTURE(BM_Set, "Set/NearId/Union", true, true, 1024, 1); + +BENCHMARK_CAPTURE(BM_Map, "Map/Small/Add", false, false, 8, 0); +BENCHMARK_CAPTURE(BM_Map, "Map/Small/Union", true, false, 8, 0); +BENCHMARK_CAPTURE(BM_Map, "Map/SmallNearId8/Add", false, true, 8, 1); +BENCHMARK_CAPTURE(BM_Map, "Map/SmallNearId8/Union", true, true, 8, 1); +BENCHMARK_CAPTURE(BM_Map, "Map/SmallNearId64/Add", false, true, 64, 1); +BENCHMARK_CAPTURE(BM_Map, "Map/SmallNearId64/Union", true, true, 64, 1); +BENCHMARK_CAPTURE(BM_Map, "Map/NearId/Add", false, true, 1024, 1); +BENCHMARK_CAPTURE(BM_Map, "Map/NearId/Union", true, true, 1024, 1); + +BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/Small/Add", false, false, 8, 0, 4); +BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/Small/Union", true, false, 8, 0, 4); +BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/SmallNearId8/Add", false, true, 8, 1, + 4); +BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/SmallNearId8/Union", true, true, 8, + 1, 4); +BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/NearId/Add", false, true, 256, 1, 8); +BENCHMARK_CAPTURE(BM_MapOfSets, "MapOfSets/NearId/Union", true, true, 256, 1, + 8); + +BENCHMARK_MAIN(); diff --git a/llvm/include/llvm/ADT/ImmutableMap.h b/llvm/include/llvm/ADT/ImmutableMap.h index 6362d25cc2bd0..f95eda31dde79 100644 --- a/llvm/include/llvm/ADT/ImmutableMap.h +++ b/llvm/include/llvm/ADT/ImmutableMap.h @@ -103,6 +103,27 @@ class ImmutableMap { return ImmutableMap(T); } + /// Merges \p A and \p B in a single traversal (see + /// ImutAVLFactory::mergeTrees), sharing subtrees the two maps do not + /// overlap. \p Combine(AElem, BElem) yields the stored element for a key; + /// \p KeepUnmatched governs keys unique to one side. \p SkipShared returns + /// a pointer-identical subtree unchanged in O(1); only pass true when the + /// merge is idempotent (Combine(a, a) == a, e.g. a lattice join). This is + /// more efficient than repeatedly adding \p B's entries to \p A when \p B + /// is large. Only valid for non-canonicalizing factories (the bulk path + /// does not canonicalize the nodes it creates). Does not short-circuit + /// equal/empty operands or reorder by size; callers wanting those apply + /// them first. + template <typename CombineFn> + [[nodiscard]] ImmutableMap mergeWith(ImmutableMap A, ImmutableMap B, + CombineFn Combine, bool KeepUnmatched, + bool SkipShared = false) { + static_assert(!Canonicalize, + "mergeWith does not canonicalize the nodes it creates"); + return ImmutableMap(F.mergeTrees(A.Root.get(), B.Root.get(), Combine, + KeepUnmatched, SkipShared)); + } + [[nodiscard]] ImmutableMap remove(ImmutableMap Old, key_type_ref K) { TreeTy *T = F.remove(Old.Root.get(), K); if constexpr (Canonicalize) diff --git a/llvm/include/llvm/ADT/ImmutableSet.h b/llvm/include/llvm/ADT/ImmutableSet.h index 7e94b75d17b3c..23d66caafce33 100644 --- a/llvm/include/llvm/ADT/ImmutableSet.h +++ b/llvm/include/llvm/ADT/ImmutableSet.h @@ -404,6 +404,7 @@ class ImutAVLFactory friend class ImutAVLTree<ImutInfo, Canonicalize>; using TreeTy = ImutAVLTree<ImutInfo, Canonicalize>; + using value_type = typename TreeTy::value_type; using value_type_ref = typename TreeTy::value_type_ref; using key_type_ref = typename TreeTy::key_type_ref; @@ -440,6 +441,39 @@ class ImutAVLFactory return T; } + /// Merges \p A and \p B in a single traversal, sharing every subtree that the + /// two operands do not overlap. \p Combine(AElem, BElem) produces the element + /// stored for a key present in both; \p KeepUnmatched governs keys unique to + /// one side (see merge_internal). For merging |B| entries into |A| + /// (|B| <= |A|) this costs O(|B| * log(|A|/|B| + 1)) and copies each spine + /// node at most once, versus O(|B| * log|A|) repeated \ref add descents. + /// \p A and \p B must be immutable. This does not short-circuit equal or + /// empty operands (merge_internal handles them correctly but not specially); + /// callers that want those fast paths, or size-driven operand ordering, + /// should apply them first (see ImmutableSet::Factory::unionSets). + template <typename CombineFn> + TreeTy *mergeTrees(TreeTy *A, TreeTy *B, CombineFn Combine, + bool KeepUnmatched, bool SkipShared = false) { + TreeTy *T = merge_internal(A, B, Combine, KeepUnmatched, SkipShared); + recoverNodes(T); + return T; + } + + /// Returns the set union of \p A and \p B (keeping \p A's element on matching + /// keys). Shorthand for the fully sharing \ref mergeTrees. + TreeTy *unionTrees(TreeTy *A, TreeTy *B) { + // With KeepUnmatched=true, unmatched elements are shared as-is and Combine + // is invoked only for keys present in both, where it keeps A's element. + auto KeepFirst = [](const value_type *L, + const value_type *R) -> const value_type & { + return L ? *L : *R; + }; + // Set union is idempotent, so identical (pointer-equal) subtrees -- common + // once one operand is derived from the other -- can be shared in O(1). + return mergeTrees(A, B, KeepFirst, /*KeepUnmatched=*/true, + /*SkipShared=*/true); + } + TreeTy* remove(TreeTy* T, key_type_ref V) { T = remove_internal(V,T); recoverNodes(T); @@ -561,6 +595,115 @@ class ImutAVLFactory return createNode(L,V,R); } + /// Combines \p L and \p R with the value \p V (every key in \p L less than + /// \p V, every key in \p R greater) into one balanced tree. Unlike + /// balanceTree this tolerates an arbitrary height difference between \p L and + /// \p R: it descends the taller side's spine and rebalances on the way back + /// up, exactly as an insertion would. + TreeTy *joinTrees(TreeTy *L, value_type_ref V, TreeTy *R) { + if (getHeight(L) > getHeight(R) + 2) + return balanceTree(getLeft(L), getValue(L), joinTrees(getRight(L), V, R)); + if (getHeight(R) > getHeight(L) + 2) + return balanceTree(joinTrees(L, V, getLeft(R)), getValue(R), getRight(R)); + return createNode(L, V, R); + } + + /// Splits \p T into \p L (all keys less than \p K) and \p R (all keys greater + /// than \p K). If \p K is present in \p T, \p Match is set to point at its + /// element (which is dropped from \p L and \p R); otherwise \p Match is null. + void splitLookup(TreeTy *T, key_type_ref K, TreeTy *&L, + const value_type *&Match, TreeTy *&R) { + if (isEmpty(T)) { + L = R = getEmptyTree(); + Match = nullptr; + return; + } + key_type_ref KCurrent = ImutInfo::KeyOfValue(getValue(T)); + if (ImutInfo::isEqual(K, KCurrent)) { + L = getLeft(T); + R = getRight(T); + // Use the tree accessor, which returns a reference to the stored element + // (the factory's getValue returns value_type_ref, which is by value for + // pointer-like element types). + Match = &T->getValue(); + } else if (ImutInfo::isLess(K, KCurrent)) { + TreeTy *LR; + splitLookup(getLeft(T), K, L, Match, LR); + R = joinTrees(LR, getValue(T), getRight(T)); + } else { + TreeTy *RL; + splitLookup(getRight(T), K, RL, Match, R); + L = joinTrees(getLeft(T), getValue(T), RL); + } + } + + /// Rebuilds \p T with the same shape but each element replaced by + /// \p Combine applied to it. \p FromB selects which side of \p Combine the + /// element is passed on (it is the sole non-null argument). + template <typename CombineFn> + TreeTy *transformTree(TreeTy *T, CombineFn &Combine, bool FromB) { + if (isEmpty(T)) + return T; + TreeTy *L = transformTree(getLeft(T), Combine, FromB); + TreeTy *R = transformTree(getRight(T), Combine, FromB); + const value_type &E = getValue(T); + return createNode(L, FromB ? Combine(nullptr, &E) : Combine(&E, nullptr), + R); + } + + /// Merges \p A and \p B by recursing over \p A's structure and splitting \p B + /// at each of \p A's keys. For a key in both, the stored element is + /// Combine(AElem, BElem). \p KeepUnmatched controls keys unique to one side: + /// when true, such elements (and whole non-overlapping subtrees) are taken + /// unchanged and shared, and \p Combine is invoked only on keys present in + /// both (valid when \p Combine is an identity for a missing side, e.g. a set + /// union or a lattice join with an identity element); when false every key is + /// passed through \p Combine with the absent side null (needed for a join + /// that transforms unmatched keys, e.g. liveness downgrading Must to Maybe). + template <typename CombineFn> + TreeTy *merge_internal(TreeTy *A, TreeTy *B, CombineFn &Combine, + bool KeepUnmatched, bool SkipShared) { + // When A and B are the same tree (which happens all the time once B is + // derived from A by a small edit, since the untouched side is shared by + // pointer), an idempotent merge returns it unchanged in O(1). Only valid + // when merge(x, x) == x, so the caller opts in via SkipShared. + if (SkipShared && A == B) + return A; + if (isEmpty(A)) + return KeepUnmatched ? B : transformTree(B, Combine, /*FromB=*/true); + if (isEmpty(B)) + return KeepUnmatched ? A : transformTree(A, Combine, /*FromB=*/false); + + const value_type &AElem = getValue(A); + TreeTy *BL, *BR; + const value_type *BMatch; + splitLookup(B, ImutInfo::KeyOfValue(AElem), BL, BMatch, BR); + + TreeTy *NewL = + merge_internal(getLeft(A), BL, Combine, KeepUnmatched, SkipShared); + TreeTy *NewR = + merge_internal(getRight(A), BR, Combine, KeepUnmatched, SkipShared); + + if (!BMatch) { + // Key present only in A. + if (KeepUnmatched) { + if (NewL == getLeft(A) && NewR == getRight(A)) + return A; + return joinTrees(NewL, AElem, NewR); + } + return joinTrees(NewL, Combine(&AElem, nullptr), NewR); + } + // Key present in both: combine the two elements. Preserve sharing when the + // combined value is unchanged and neither subtree moved, so that a join + // that only touches a few keys does not rebuild the whole spine. + auto NewElem = Combine(&AElem, BMatch); + if (NewL == getLeft(A) && NewR == getRight(A) && + ImutInfo::isDataEqual(ImutInfo::DataOfValue(NewElem), + ImutInfo::DataOfValue(AElem))) + return A; + return joinTrees(NewL, NewElem, NewR); + } + /// add_internal - Creates a new tree that includes the specified /// data and the data from the original tree. If the original tree /// already contained the data item, the original tree is returned. @@ -994,6 +1137,30 @@ class ImmutableSet { return ImmutableSet(NewT); } + /// Returns the union of \p A and \p B, computed in a single traversal that + /// shares subtrees of both operands wherever possible (see + /// ImutAVLFactory::unionTrees). This is more efficient than repeatedly + /// adding \p B's elements to \p A when \p B is large. + [[nodiscard]] ImmutableSet unionSets(ImmutableSet A, ImmutableSet B) { + if (A.Root.get() == B.Root.get() || B.isEmpty()) + return A; + if (A.isEmpty()) + return B; + // Drive the recursion with the taller tree so the shorter one is the one + // being split. + if (A.getHeight() < B.getHeight()) + std::swap(A, B); + if constexpr (Canonicalize) { + // The bulk path does not canonicalize the nodes it creates, so fall + // back to per-element insertion for canonicalizing factories. + for (value_type_ref V : B) + A = add(A, V); + return A; + } else { + return ImmutableSet(F.unionTrees(A.Root.get(), B.Root.get())); + } + } + /// Creates a new immutable set that contains all of the values /// of the original set with the exception of the specified value. If /// the original set did not contain the value, the original set is diff --git a/llvm/unittests/ADT/ImmutableMapTest.cpp b/llvm/unittests/ADT/ImmutableMapTest.cpp index 1217718826f77..a2ed6df1e8198 100644 --- a/llvm/unittests/ADT/ImmutableMapTest.cpp +++ b/llvm/unittests/ADT/ImmutableMapTest.cpp @@ -8,9 +8,18 @@ #include "llvm/ADT/ImmutableMap.h" #include "gtest/gtest.h" +#include <map> +#include <random> +#include <utility> using namespace llvm; +// Non-canonicalizing map, matching how the clang lifetime analyses use it and +// the only configuration for which the bulk mergeWith path is valid. +template <typename K, typename V> +using NCMap = + ImmutableMap<K, V, ImutKeyValueInfo<K, V>, /*Canonicalize=*/false>; + namespace { TEST(ImmutableMapTest, EmptyIntMapTest) { @@ -87,4 +96,77 @@ TEST(ImmutableMapTest, MultiElemIntMapRefTest) { EXPECT_TRUE(f.getEmptyMap() == f.getEmptyMap()); } + // Exercises Factory::mergeWith against a std::map oracle across the + // KeepA/KeepB configurations used by the set union (true/true), an asymmetric + // map join (true/false), and a symmetric map join (false/false), with a + // combiner that is order- and side-sensitive so the flag semantics are + // actually tested. + TEST(ImmutableMapTest, MergeWithStressTest) { + using Map = NCMap<int, int>; + using VT = Map::value_type; // const std::pair<int, int> + + auto Combine = [](const VT *A, const VT *B) -> std::pair<int, int> { + int Key = A ? A->first : B->first; + int AV = A ? A->second : 0; + int BV = B ? B->second : 0; + return {Key, 1 + 3 * AV + 5 * BV}; + }; + + std::mt19937 Rng(7); + std::uniform_int_distribution<int> SizeDist(0, 120); + std::uniform_int_distribution<int> KeyDist(0, 200); + std::uniform_int_distribution<int> ValDist(0, 1000); + + for (bool KeepUnmatched : {true, false}) { + for (int Trial = 0; Trial < 300; ++Trial) { + Map::Factory F; + Map A = F.getEmptyMap(), B = F.getEmptyMap(); + std::map<int, int> RefA, RefB; + + int NA = SizeDist(Rng), NB = SizeDist(Rng); + for (int I = 0; I < NA; ++I) { + int K = KeyDist(Rng), V = ValDist(Rng); + A = F.add(A, K, V); + RefA[K] = V; + } + for (int I = 0; I < NB; ++I) { + int K = KeyDist(Rng), V = ValDist(Rng); + B = F.add(B, K, V); + RefB[K] = V; + } + + Map U = F.mergeWith(A, B, Combine, KeepUnmatched); + if (const auto *Root = U.getRootWithoutRetain()) + Root->validateTree(); + + // Reference result: matched keys always combined; unmatched keys are + // kept as-is when KeepUnmatched, else passed through Combine. + std::map<int, int> Exp; + for (auto &[K, AV] : RefA) { + auto It = RefB.find(K); + if (It != RefB.end()) + Exp[K] = 1 + 3 * AV + 5 * It->second; + else + Exp[K] = KeepUnmatched ? AV : 1 + 3 * AV; + } + for (auto &[K, BV] : RefB) + if (!RefA.count(K)) + Exp[K] = KeepUnmatched ? BV : 1 + 5 * BV; + + std::map<int, int> Got; + for (Map::iterator It = U.begin(), E = U.end(); It != E; ++It) + Got[It.getKey()] = It.getData(); + EXPECT_EQ(Exp, Got) << "KeepUnmatched=" << KeepUnmatched; + + // Inputs unchanged. + std::map<int, int> GotA, GotB; + for (Map::iterator It = A.begin(), E = A.end(); It != E; ++It) + GotA[It.getKey()] = It.getData(); + for (Map::iterator It = B.begin(), E = B.end(); It != E; ++It) + GotB[It.getKey()] = It.getData(); + EXPECT_EQ(RefA, GotA); + EXPECT_EQ(RefB, GotB); + } + } + } } diff --git a/llvm/unittests/ADT/ImmutableSetTest.cpp b/llvm/unittests/ADT/ImmutableSetTest.cpp index ee65ac7a94573..8c2eb355cc63f 100644 --- a/llvm/unittests/ADT/ImmutableSetTest.cpp +++ b/llvm/unittests/ADT/ImmutableSetTest.cpp @@ -387,4 +387,58 @@ TEST_F(ImmutableSetTest, IteratorEquality) { EXPECT_TRUE(A == TreeIter()); } } + +// Exercises Factory::unionSets against a std::set oracle over many random +// input pairs, for both the canonicalizing (per-element fallback) and +// non-canonicalizing (bulk tree-tree) code paths. +template <bool Canonicalize> static void unionSetsStress() { + using Set = ImmutableSet<int, ImutContainerInfo<int>, Canonicalize>; + std::mt19937 Rng(42); + std::uniform_int_distribution<int> SizeDist(0, 120); + std::uniform_int_distribution<int> ValDist(0, 200); + + for (int Trial = 0; Trial < 300; ++Trial) { + typename Set::Factory F; + Set A = F.getEmptySet(); + Set B = F.getEmptySet(); + std::set<int> RefA, RefB; + + int NA = SizeDist(Rng), NB = SizeDist(Rng); + for (int I = 0; I < NA; ++I) { + int V = ValDist(Rng); + A = F.add(A, V); + RefA.insert(V); + } + for (int I = 0; I < NB; ++I) { + int V = ValDist(Rng); + B = F.add(B, V); + RefB.insert(V); + } + + Set U = F.unionSets(A, B); + + // The result is a structurally valid AVL tree. + if (const auto *Root = U.getRootWithoutRetain()) + Root->validateTree(); + + // Its contents are exactly the union of the inputs. + std::set<int> RefU = RefA; + RefU.insert(RefB.begin(), RefB.end()); + std::set<int> Got(U.begin(), U.end()); + EXPECT_EQ(RefU, Got); + + // The inputs are unchanged (persistence). + EXPECT_EQ(RefA, std::set<int>(A.begin(), A.end())); + EXPECT_EQ(RefB, std::set<int>(B.begin(), B.end())); + + // Unioning with a subset / self is the identity. + EXPECT_TRUE(F.unionSets(U, A) == U); + EXPECT_TRUE(F.unionSets(A, A) == A); + } +} + +TEST_F(ImmutableSetTest, UnionSetsStressTest) { + unionSetsStress</*Canonicalize=*/false>(); + unionSetsStress</*Canonicalize=*/true>(); +} } // namespace _______________________________________________ cfe-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
