[MapVector] Introduce {keys,values} iterator, use in VPlan

[artagnon]

Similar to DenseMap::{keys,values}, introduce MapVector::{keys,values}, and use it to simplify some code in VPlan.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-adt

Author: Ramkumar Ramachandra (artagnon)

Changes

Similar to DenseMap::{keys,values}, introduce MapVector::{keys,values}, and use it to simplify some code in VPlan.

---

Full diff: https://github.com/llvm/llvm-project/pull/169675.diff

4 Files Affected:

• (modified) llvm/include/llvm/ADT/MapVector.h (+20)
• (modified) llvm/lib/Transforms/Vectorize/VPlan.h (+7-19)
• (modified) llvm/lib/Transforms/Vectorize/VPlanValue.h (-5)
• (modified) llvm/unittests/ADT/MapVectorTest.cpp (+18)

diff --git a/llvm/include/llvm/ADT/MapVector.h b/llvm/include/llvm/ADT/MapVector.h index 80bcb7e0b7ba4..b1e543e282a03 100644 --- a/llvm/include/llvm/ADT/MapVector.h +++ b/llvm/include/llvm/ADT/MapVector.h @@ -99,6 +99,26 @@ class MapVector { return try_emplace_impl(Key).first->second; } + [[nodiscard]] inline auto keys() { + return map_range( + Vector, [](const std::pair<KeyT, ValueT> &KV) { return KV.first; }); + } + + [[nodiscard]] inline auto values() { + return map_range( + Vector, [](const std::pair<KeyT, ValueT> &KV) { return KV.second; }); + } + + [[nodiscard]] inline auto keys() const { + return map_range( + Vector, [](const std::pair<KeyT, ValueT> &KV) { return KV.first; }); + } + + [[nodiscard]] inline auto values() const { + return map_range( + Vector, [](const std::pair<KeyT, ValueT> &KV) { return KV.second; }); + } + // Returns a copy of the value. Only allowed if ValueT is copyable. [[nodiscard]] ValueT lookup(const KeyT &Key) const { static_assert(std::is_copy_constructible_v<ValueT>, diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h index 0c7d9c0193a03..b22d2083adf0b 100644 --- a/llvm/lib/Transforms/Vectorize/VPlan.h +++ b/llvm/lib/Transforms/Vectorize/VPlan.h @@ -25,7 +25,7 @@ #define LLVM_TRANSFORMS_VECTORIZE_VPLAN_H #include "VPlanValue.h" -#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/MapVector.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/Twine.h" @@ -4336,13 +4336,9 @@ class VPlan { /// Represents the loop-invariant VF * UF of the vector loop region. VPValue VFxUF; - /// Holds a mapping between Values and their corresponding VPValue inside - /// VPlan. - Value2VPValueTy Value2VPValue; - - /// Contains all the external definitions created for this VPlan. External - /// definitions are VPValues that hold a pointer to their underlying IR. - SmallVector<VPValue *, 16> VPLiveIns; + /// Contains all the external definitions created for this VPlan, as a mapping + /// from IR Values to VPValues. + SmallMapVector<Value *, VPValue *, 16> LiveIns; /// Blocks allocated and owned by the VPlan. They will be deleted once the /// VPlan is destroyed. @@ -4539,10 +4535,9 @@ class VPlan { /// yet) for \p V. VPValue *getOrAddLiveIn(Value *V) { assert(V && "Trying to get or add the VPValue of a null Value"); - auto [It, Inserted] = Value2VPValue.try_emplace(V); + auto [It, Inserted] = LiveIns.try_emplace(V); if (Inserted) { VPValue *VPV = new VPValue(V); - VPLiveIns.push_back(VPV); assert(VPV->isLiveIn() && "VPV must be a live-in."); It->second = VPV; } @@ -4574,17 +4569,10 @@ class VPlan { } /// Return the live-in VPValue for \p V, if there is one or nullptr otherwise. - VPValue *getLiveIn(Value *V) const { return Value2VPValue.lookup(V); } + VPValue *getLiveIn(Value *V) const { return LiveIns.lookup(V); } /// Return the list of live-in VPValues available in the VPlan. - ArrayRef<VPValue *> getLiveIns() const { - assert(all_of(Value2VPValue, - [this](const auto &P) { - return is_contained(VPLiveIns, P.second); - }) && - "all VPValues in Value2VPValue must also be in VPLiveIns"); - return VPLiveIns; - } + auto getLiveIns() const { return LiveIns.values(); } #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) /// Print the live-ins of this VPlan to \p O. diff --git a/llvm/lib/Transforms/Vectorize/VPlanValue.h b/llvm/lib/Transforms/Vectorize/VPlanValue.h index 63eacd3d75721..6d7c8312afa7b 100644 --- a/llvm/lib/Transforms/Vectorize/VPlanValue.h +++ b/llvm/lib/Transforms/Vectorize/VPlanValue.h @@ -20,10 +20,8 @@ #ifndef LLVM_TRANSFORMS_VECTORIZE_VPLAN_VALUE_H #define LLVM_TRANSFORMS_VECTORIZE_VPLAN_VALUE_H -#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringMap.h" #include "llvm/ADT/TinyPtrVector.h" #include "llvm/ADT/iterator_range.h" #include "llvm/Support/Compiler.h" @@ -196,9 +194,6 @@ class LLVM_ABI_FOR_TEST VPValue { } }; -typedef DenseMap<Value *, VPValue *> Value2VPValueTy; -typedef DenseMap<VPValue *, Value *> VPValue2ValueTy; - LLVM_ABI_FOR_TEST raw_ostream &operator<<(raw_ostream &OS, const VPRecipeBase &R); diff --git a/llvm/unittests/ADT/MapVectorTest.cpp b/llvm/unittests/ADT/MapVectorTest.cpp index e0589445e3271..b11d4603b90b7 100644 --- a/llvm/unittests/ADT/MapVectorTest.cpp +++ b/llvm/unittests/ADT/MapVectorTest.cpp @@ -307,6 +307,24 @@ TEST(MapVectorTest, AtTest) { EXPECT_EQ(ConstMV.at(1), 12); } +TEST(MapVectorTest, KeysValuesIterator) { + MapVector<int, int> MV; + + MV.insert(std::make_pair(1, 11)); + MV.insert(std::make_pair(2, 12)); + MV.insert(std::make_pair(3, 13)); + MV.insert(std::make_pair(4, 14)); + MV.insert(std::make_pair(5, 15)); + MV.insert(std::make_pair(6, 16)); + + EXPECT_THAT(MV.keys(), testing::ElementsAre(1, 2, 3, 4, 5, 6)); + EXPECT_THAT(MV.values(), testing::ElementsAre(11, 12, 13, 14, 15, 16)); + + const MapVector<int, int> &ConstMV = MV; + EXPECT_THAT(ConstMV.keys(), testing::ElementsAre(1, 2, 3, 4, 5, 6)); + EXPECT_THAT(ConstMV.values(), testing::ElementsAre(11, 12, 13, 14, 15, 16)); +} + template <class IntType> struct MapVectorMappedTypeTest : ::testing::Test { using int_type = IntType; }; 

[huntergr-arm]

LGTM

[fhahn]

keys() appears unused (other than tests)? Are there other places that could use it? If not, better only introduce once there are users?

[fhahn]

Suggested change

	[[nodiscard]] inline auto keys() { return make_first_range(Vector); }
	[[nodiscard]] inline auto values() { return make_second_range(Vector); }
	[[nodiscard]] inline auto keys() const { return make_first_range(Vector); }
	[[nodiscard]] inline auto values() const { return make_second_range(Vector); }
	[[nodiscard]] inline auto keys() { return make_first_range(Vector); }
	[[nodiscard]] inline auto keys() const { return make_first_range(Vector); }
	[[nodiscard]] inline auto values() { return make_second_range(Vector); }
	[[nodiscard]] inline auto values() const { return make_second_range(Vector); }

Keep values(), keys() together?

[kuhar]

Can you land ADT changes (w/ unit tests) separately from the use in vplan? This way potential reverts will be less disruptive