diff --git a/llvm/include/llvm/ADT/SparseMultiSet.h b/llvm/include/llvm/ADT/SparseMultiSet.h
index 0aa7edbcea673..5e4e1703373a9 100644
--- a/llvm/include/llvm/ADT/SparseMultiSet.h
+++ b/llvm/include/llvm/ADT/SparseMultiSet.h
@@ -21,9 +21,9 @@
 #ifndef LLVM_ADT_SPARSEMULTISET_H
 #define LLVM_ADT_SPARSEMULTISET_H
 
+#include "llvm/ADT/STLForwardCompat.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/SparseSet.h"
-#include "llvm/ADT/identity.h"
 #include <cassert>
 #include <cstdint>
 #include <cstdlib>
@@ -77,11 +77,12 @@ namespace llvm {
 /// intuitive and fast removal.
 ///
 /// @tparam ValueT      The type of objects in the set.
+/// @tparam KeyT        The type of the key that identifies objects in the set.
 /// @tparam KeyFunctorT A functor that computes an unsigned index from KeyT.
 /// @tparam SparseT     An unsigned integer type. See above.
 ///
-template <typename ValueT, typename KeyFunctorT = identity<unsigned>,
-          typename SparseT = uint8_t>
+template <typename ValueT, typename KeyT = unsigned,
+          typename KeyFunctorT = identity_cxx20, typename SparseT = uint8_t>
 class SparseMultiSet {
   static_assert(std::is_unsigned_v<SparseT>,
                 "SparseT must be an unsigned integer type");
@@ -112,7 +113,6 @@ class SparseMultiSet {
     bool isValid() const { return Prev != INVALID; }
   };
 
-  using KeyT = typename KeyFunctorT::argument_type;
   using DenseT = SmallVector<SMSNode, 8>;
   DenseT Dense;
   SparseT *Sparse = nullptr;
diff --git a/llvm/include/llvm/CodeGen/ScheduleDAGInstrs.h b/llvm/include/llvm/CodeGen/ScheduleDAGInstrs.h
index ba5b2da64fd80..4eacbdca163d9 100644
--- a/llvm/include/llvm/CodeGen/ScheduleDAGInstrs.h
+++ b/llvm/include/llvm/CodeGen/ScheduleDAGInstrs.h
@@ -90,15 +90,16 @@ namespace llvm {
   /// allocated once for the pass. It can be cleared in constant time and reused
   /// without any frees.
   using RegUnit2SUnitsMap =
-      SparseMultiSet<PhysRegSUOper, identity<unsigned>, uint16_t>;
+      SparseMultiSet<PhysRegSUOper, unsigned, identity_cxx20, uint16_t>;
 
   /// Track local uses of virtual registers. These uses are gathered by the DAG
   /// builder and may be consulted by the scheduler to avoid iterating an entire
   /// vreg use list.
-  using VReg2SUnitMultiMap = SparseMultiSet<VReg2SUnit, VirtReg2IndexFunctor>;
+  using VReg2SUnitMultiMap =
+      SparseMultiSet<VReg2SUnit, Register, VirtReg2IndexFunctor>;
 
   using VReg2SUnitOperIdxMultiMap =
-      SparseMultiSet<VReg2SUnitOperIdx, VirtReg2IndexFunctor>;
+      SparseMultiSet<VReg2SUnitOperIdx, Register, VirtReg2IndexFunctor>;
 
   using ValueType = PointerUnion<const Value *, const PseudoSourceValue *>;