[llvm-branch-commits] [flang] [mlir] [mlir][Transforms] Support 1:N mappings in `ConversionValueMapping` (PR #116524)

[Markus Böck via llvm-branch-commits]

================
@@ -63,80 +113,140 @@ static OpBuilder::InsertPoint computeInsertPoint(Value value) {
   return OpBuilder::InsertPoint(insertBlock, insertPt);
 }
 
+/// Helper function that computes an insertion point where the given values are
+/// defined and can be used without a dominance violation.
+static OpBuilder::InsertPoint computeInsertPoint(ArrayRef<Value> vals) {
+  assert(!vals.empty() && "expected at least one value");
+  OpBuilder::InsertPoint pt = computeInsertPoint(vals.front());
+  for (Value v : vals.drop_front())
+    pt = chooseLaterInsertPoint(pt, computeInsertPoint(v));
+  return pt;
+}
+
 //===----------------------------------------------------------------------===//
 // ConversionValueMapping
 //===----------------------------------------------------------------------===//
 
+/// A vector of SSA values, optimized for the most common case of a single
+/// value.
+using ValueVector = SmallVector<Value, 1>;
+
 namespace {
+
+/// Helper class to make it possible to use `ValueVector` as a key in DenseMap.
+struct ValueVectorMapInfo {
+  static ValueVector getEmptyKey() { return ValueVector{}; }
+  static ValueVector getTombstoneKey() { return ValueVector{}; }
+  static ::llvm::hash_code getHashValue(ValueVector val) {
+    return ::llvm::hash_combine_range(val.begin(), val.end());
+  }
+  static bool isEqual(ValueVector LHS, ValueVector RHS) { return LHS == RHS; }
+};
+
 /// This class wraps a IRMapping to provide recursive lookup
 /// functionality, i.e. we will traverse if the mapped value also has a mapping.
 struct ConversionValueMapping {
   /// Return "true" if an SSA value is mapped to the given value. May return
   /// false positives.
   bool isMappedTo(Value value) const { return mappedTo.contains(value); }
 
-  /// Lookup the most recently mapped value with the desired type in the
+  /// Lookup the most recently mapped values with the desired types in the
   /// mapping.
   ///
   /// Special cases:
-  /// - If the desired type is "null", simply return the most recently mapped
-  ///   value.
-  /// - If there is no mapping to the desired type, also return the most
-  ///   recently mapped value.
-  /// - If there is no mapping for the given value at all, return the given
-  ///   value.
-  Value lookupOrDefault(Value from, Type desiredType = nullptr) const;
-
-  /// Lookup a mapped value within the map, or return null if a mapping does not
-  /// exist. If a mapping exists, this follows the same behavior of
-  /// `lookupOrDefault`.
-  Value lookupOrNull(Value from, Type desiredType = nullptr) const;
+  /// - If the desired type range is empty, simply return the most recently
+  ///   mapped values.
+  /// - If there is no mapping to the desired types, also return the most
+  ///   recently mapped values.
+  /// - If there is no mapping for the given values at all, return the given
+  ///   values.
+  ValueVector lookupOrDefault(ValueVector from,
+                              TypeRange desiredTypes = {}) const;
+
+  /// Lookup the given values within the map, or return an empty vector if the
+  /// values are not mapped. If they are mapped, this follows the same behavior
+  /// as `lookupOrDefault`.
+  ValueVector lookupOrNull(const ValueVector &from,
+                           TypeRange desiredTypes = {}) const;
 
   /// Map a value to the one provided.
-  void map(Value oldVal, Value newVal) {
+  void map(const ValueVector &oldVal, const ValueVector &newVal) {
----------------
zero9178 wrote:

There's a few callers that call these with r-values where a move construction would avoid copying entirely.
Could we change these to be universal references and then use `std::forward` for the `operator[]` + assignment at line 184?

I prototyped this a bit roughly here if you'd like to cherry-pick in the changes: https://github.com/llvm/llvm-project/commit/5e857ac6348e7455f1c72fd49eb887fc1406949c

https://github.com/llvm/llvm-project/pull/116524