[IRSim] Ensure that assignment accurately reduces potential mapping between different candidates

[AndrewLitteken]

Summary

Previous
When we do not make decisions about commutative operands, we can end up in a situation where two values have two potential canonical numbers between two regions. This ensures that an ordering is decided after the initial structure between two regions is determined.

Current:
Previously the outliner only checked that assignment to a value matched what was already known, this patch makes sure that it matches what has already been found, and creates a mapping between the two values where it is a one-to-one mapping.

[paquette]

Is this one of the compile time patches we discussed offline? If so, could you mention that + an overview of the improvements in the commit message?

[AndrewLitteken]

In D139336#3975102, @paquette wrote:

Is this one of the compile time patches we discussed offline? If so, could you mention that + an overview of the improvements in the commit message?

This is a bug fix that needs to be made prior to committing the compile time patches that shows up when you start making that optimization.

[paquette]

Is it possible to write a testcase?

[AndrewLitteken]

I could write a test case that would become a problem if this was not implemented when https://reviews.llvm.org/D139338 is merged on top of it. But I can't write a test case that would fail without this check for the current version of the outliner.

[paquette]

I think this looks good, with a nit.

llvm/lib/Analysis/IRSimilarityIdentifier.cpp

833

would it be possible to pull this out into a lambda/function?

llvm/include/llvm/Analysis/IRSimilarityIdentifier.h

public:
   /// \param A - The first IRInstructionCandidate, operand values, and current
   /// operand mappings to compare.
   /// \param B - The second IRInstructionCandidate, operand values, and current
   /// operand mappings to compare.
   /// \returns true if the IRSimilarityCandidates operands are compatible.
   static bool compareCommutativeOperandMapping(OperandMapping A,
                                                OperandMapping B);
 
+  /// Compare the GVN of the assignment value in corresponding instructions in
+  /// IRSimilarityCandidates \p A and \p B and check that there exists a mapping
+  /// between the values and replaces the mapping with a one-to-one value if
+  /// needed.
+  ///
+  /// \param InstValA - The assignment GVN from the first IRSimilarityCandidate.
+  /// \param InstValB - The assignment GVN from the second
+  /// IRSimilarityCandidate.
+  /// \param [in,out] ValueNumberMappingA - A mapping of value numbers from 
+  /// candidate \p A to candidate \B.
+  /// \param [in,out] ValueNumberMappingB - A mapping of value numbers from 
+  /// candidate \p B to candidate \A.
+  /// \returns true if the IRSimilarityCandidates assignments are compatible.
+  static bool compareAssignmentMapping(
+      const unsigned InstValA, const unsigned &InstValB,
+      DenseMap<unsigned, DenseSet<unsigned>> &ValueNumberMappingA,
+      DenseMap<unsigned, DenseSet<unsigned>> &ValueNumberMappingB);
+
   /// Compare the relative locations in \p A and \p B and check that the
   /// distances match if both locations are contained in the region, and that
   /// the branches both point outside the region if they do not.
   /// Example Region:
   /// \code
   /// entry:
   ///   br i1 %0, label %block_1, label %block_3
   /// block_0:

llvm/lib/Analysis/IRSimilarityIdentifier.cpp

bool IRSimilarityCandidate::compareCommutativeOperandMapping(
   if (!checkNumberingAndReplaceCommutative(B.IRSC.ValueToNumber,
                                            B.ValueNumberMapping, B.OperVals,
                                            ValueNumbersA))
     return false;
 
   return true;
 }
 
+bool IRSimilarityCandidate::compareAssignmentMapping(
+    const unsigned InstValA, const unsigned &InstValB,
+    DenseMap<unsigned, DenseSet<unsigned>> &ValueNumberMappingA,
+    DenseMap<unsigned, DenseSet<unsigned>> &ValueNumberMappingB) {
+  DenseMap<unsigned, DenseSet<unsigned>>::iterator ValueMappingIt;
+  bool WasInserted;
+  std::tie(ValueMappingIt, WasInserted) = ValueNumberMappingA.insert(
+      std::make_pair(InstValA, DenseSet<unsigned>({InstValB})));
+  if (!WasInserted && !ValueMappingIt->second.contains(InstValB))
+    return false;
+  else if (ValueMappingIt->second.size() != 1) {
+    for (unsigned OtherVal : ValueMappingIt->second) {
+      if (OtherVal == InstValB)
+        continue;
+      if (ValueNumberMappingA.find(OtherVal) == ValueNumberMappingA.end())
+        continue;
+      if (!ValueNumberMappingA[OtherVal].contains(InstValA))
+        continue;
+      ValueNumberMappingA[OtherVal].erase(InstValA);
+    }
+    ValueNumberMappingA.erase(ValueMappingIt);
+    std::tie(ValueMappingIt, WasInserted) = ValueNumberMappingA.insert(
+      std::make_pair(InstValA, DenseSet<unsigned>({InstValB})));
+  }
+
+  return true;
+}
+
 bool IRSimilarityCandidate::checkRelativeLocations(RelativeLocMapping A,
                                                    RelativeLocMapping B) {
   // Get the basic blocks the label refers to.
   BasicBlock *ABB = cast<BasicBlock>(A.OperVal);
   BasicBlock *BBB = cast<BasicBlock>(B.OperVal);
 
   // Get the basic blocks contained in each region.
   DenseSet<BasicBlock *> BasicBlockA;
bool IRSimilarityCandidate::compareStructure(
 
   iterator ItA = A.begin();
   iterator ItB = B.begin();
 
   // These ValueNumber Mapping sets create a create a mapping between the values
   // in one candidate to values in the other candidate.  If we create a set with
   // one element, and that same element maps to the original element in the
   // candidate we have a good mapping.
-  DenseMap<unsigned, DenseSet<unsigned>>::iterator ValueMappingIt;
-
 
   // Iterate over the instructions contained in each candidate
   unsigned SectionLength = A.getStartIdx() + A.getLength();
   for (unsigned Loc = A.getStartIdx(); Loc < SectionLength;
        ItA++, ItB++, Loc++) {
     // Make sure the instructions are similar to one another.
     if (!isClose(*ItA, *ItB))
       return false;
 
     Instruction *IA = ItA->Inst;
     Instruction *IB = ItB->Inst;
 
     if (!ItA->Legal || !ItB->Legal)
       return false;
 
     // Get the operand sets for the instructions.
     ArrayRef<Value *> OperValsA = ItA->OperVals;
     ArrayRef<Value *> OperValsB = ItB->OperVals;
 
     unsigned InstValA = A.ValueToNumber.find(IA)->second;
     unsigned InstValB = B.ValueToNumber.find(IB)->second;
 
-    bool WasInserted;
     // Ensure that the mappings for the instructions exists.
-    std::tie(ValueMappingIt, WasInserted) = ValueNumberMappingA.insert(
-        std::make_pair(InstValA, DenseSet<unsigned>({InstValB})));
-    if (!WasInserted && !ValueMappingIt->second.contains(InstValB))
+    if (!compareAssignmentMapping(InstValA, InstValB, ValueNumberMappingA,
+                                  ValueNumberMappingB))
       return false;
     
-    std::tie(ValueMappingIt, WasInserted) = ValueNumberMappingB.insert(
-        std::make_pair(InstValB, DenseSet<unsigned>({InstValA})));
-    if (!WasInserted && !ValueMappingIt->second.contains(InstValA))
+    if (!compareAssignmentMapping(InstValB, InstValA, ValueNumberMappingB,
+                                  ValueNumberMappingA))
       return false;
 
     // We have different paths for commutative instructions and non-commutative
     // instructions since commutative instructions could allow multiple mappings
     // to certain values.
     if (IA->isCommutative() && !isa<FPMathOperator>(IA) &&
         !isa<IntrinsicInst>(IA)) {
       if (!compareCommutativeOperandMapping(

llvm/test/Transforms/IROutliner/outlining-larger-size-commutative.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -p iroutliner,verify -ir-outlining-no-cost < %s | FileCheck %s
+
+; This test checks that commutative instructions where the operands are
+; swapped are outlined as the same function.
+
+; It also checks that non-commutative instructions outlined as different
+; functions when the operands are swapped;
+
+; These are identical functions, except that in the flipped functions,
+; the operands in the adds are commuted.  However, since add instructions
+; are commutative, we should still outline from all four as the same
+; instruction.
+
+define void @function1(i32 %a, i32 %b) {
+; CHECK-LABEL: @function1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[BLOCK_1:%.*]]
+; CHECK:       block_0:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[TMP4:%.*]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = add i32 [[TMP0]], [[TMP0]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp sgt i32 [[TMP0]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[TMP3]], label [[BLOCK_1]], label [[BLOCK_2:%.*]]
+; CHECK:       block_1:
+; CHECK-NEXT:    [[TMP4]] = phi i32 [ [[TMP1]], [[BLOCK_0:%.*]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32 [[B]])
+; CHECK-NEXT:    br label [[BLOCK_0]]
+; CHECK:       block_2:
+; CHECK-NEXT:    [[TMP5:%.*]] = add i32 [[TMP2]], [[TMP2]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %block_1
+
+block_0:
+  %0 = add i32 %a, %b
+  %1 = add i32 %4, 1
+  %2 = add i32 %0, %0
+  %3 = icmp sgt i32 %0, %0
+  br i1 %3, label %block_1, label %block_2
+
+block_1:
+  %4 = phi i32 [ %1, %block_0 ], [ 0, %entry ]
+  %5 = add i32 %b, %b
+  br label %block_0
+
+block_2:
+  %6 = add i32 %2, %2
+  ret void
+}
+
+define void @function2(i32 %a, i32 %b) {
+; CHECK-LABEL: @function2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[BLOCK_1:%.*]]
+; CHECK:       block_0:
+; CHECK-NEXT:    [[TMP0:%.*]] = sub i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 1, [[TMP4:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = add i32 [[TMP0]], [[TMP0]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp sgt i32 [[TMP0]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[TMP3]], label [[BLOCK_1]], label [[BLOCK_2:%.*]]
+; CHECK:       block_1:
+; CHECK-NEXT:    [[TMP4]] = phi i32 [ [[TMP1]], [[BLOCK_0:%.*]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32 [[B]])
+; CHECK-NEXT:    br label [[BLOCK_0]]
+; CHECK:       block_2:
+; CHECK-NEXT:    [[TMP5:%.*]] = sub i32 [[TMP2]], [[TMP2]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %block_1
+
+block_0:
+  %0 = sub i32 %a, %b
+  %1 = add i32 1, %4
+  %2 = add i32 %0, %0
+  %3 = icmp sgt i32 %0, %0
+  br i1 %3, label %block_1, label %block_2
+
+block_1:
+  %4 = phi i32 [ %1, %block_0 ], [ 0, %entry ]
+  %5 = add i32 %b, %b
+  br label %block_0
+
+block_2:
+  %6 = sub i32 %2, %2
+  ret void
+}