[ConstraintElimination] Remove dead variables when dropping constraints.

This patch extends ConstraintElimination to also remove dead variables
when removing a constraint. When a constraint is removed because it is
out of scope, all new variables added for this constraint can also be
removed.

This keeps the total size of the systems much smaller, because it
reduces the number of variables drastically.

It also fixes a bug where variables where removed incorrectly.

Fixes https://github.com/llvm/llvm-project/issues/54228

llvm/include/llvm/Analysis/ConstraintSystem.h

@@ -53,6 +53,11 @@ public:
   }
 
   bool addVariableRowFill(ArrayRef<int64_t> R) {
+    // If all variable coefficients are 0, the constraint does not provide any
+    // usable information.
+    if (all_of(makeArrayRef(R).drop_front(1), [](int64_t C) { return C == 0; }))
+      return false;
+
     for (auto &CR : Constraints) {
       while (CR.size() != R.size())
         CR.push_back(0);
@@ -75,6 +80,12 @@ public:
   bool isConditionImplied(SmallVector<int64_t, 8> R) const;
 
   void popLastConstraint() { Constraints.pop_back(); }
+  void popLastNVariables(unsigned N) {
+    for (auto &C : Constraints) {
+      for (unsigned i = 0; i < N; i++)
+        C.pop_back();
+    }
+  }
 
   /// Returns the number of rows in the constraint system.
   unsigned size() const { return Constraints.size(); }

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

@@ -74,6 +74,9 @@ public:
   }
 
   void popLastConstraint(bool Signed) { getCS(Signed).popLastConstraint(); }
+  void popLastNVariables(bool Signed, unsigned N) {
+    getCS(Signed).popLastNVariables(N);
+  }
 };
 
 /// Struct to express a pre-condition of the form %Op0 Pred %Op1.
@@ -371,11 +374,14 @@ struct StackEntry {
   Instruction *Condition;
   bool IsNot;
   bool IsSigned = false;
+  /// Variables that can be removed from the system once the stack entry gets
+  /// removed.
+  SmallVector<Value *, 2> ValuesToRelease;
 
-  StackEntry(unsigned NumIn, unsigned NumOut, Instruction *Condition,
-             bool IsNot, bool IsSigned)
+  StackEntry(unsigned NumIn, unsigned NumOut, CmpInst *Condition, bool IsNot,
+             bool IsSigned, SmallVector<Value *, 2> ValuesToRelease)
       : NumIn(NumIn), NumOut(NumOut), Condition(Condition), IsNot(IsNot),
-        IsSigned(IsSigned) {}
+        IsSigned(IsSigned), ValuesToRelease(ValuesToRelease) {}
 };
 } // namespace
 
@@ -512,8 +518,13 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT) {
         break;
       LLVM_DEBUG(dbgs() << "Removing " << *E.Condition << " " << E.IsNot
                         << "\n");
-      DFSInStack.pop_back();
       Info.popLastConstraint(E.IsSigned);
+      // Remove variables in the system that went out of scope.
+      auto &Mapping = Info.getValue2Index(E.IsSigned);
+      for (Value *V : E.ValuesToRelease)
+        Mapping.erase(V);
+      Info.popLastNVariables(E.IsSigned, E.ValuesToRelease.size());
+      DFSInStack.pop_back();
     }
 
     LLVM_DEBUG({
@@ -603,10 +614,6 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT) {
     if (!R.isValid(Info))
       continue;
 
-    for (auto &KV : NewIndices)
-      Info.getValue2Index(CmpInst::isSigned(CB.Condition->getPredicate()))
-          .insert(KV);
-
     LLVM_DEBUG(dbgs() << "Adding " << *CB.Condition << " " << CB.Not << "\n");
     bool Added = false;
     assert(CmpInst::isSigned(CB.Condition->getPredicate()) == R.IsSigned &&
@@ -620,8 +627,11 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT) {
     // If R has been added to the system, queue it for removal once it goes
     // out-of-scope.
     if (Added) {
-      for (auto &KV : NewIndices)
+      SmallVector<Value *, 2> ValuesToRelease;
+      for (auto &KV : NewIndices) {
         Info.getValue2Index(R.IsSigned).insert(KV);
+        ValuesToRelease.push_back(KV.first);
+      }
 
       LLVM_DEBUG({
         dbgs() << "  constraint: ";
@@ -629,7 +639,7 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT) {
       });
 
       DFSInStack.emplace_back(CB.NumIn, CB.NumOut, CB.Condition, CB.Not,
-                              R.IsSigned);
+                              R.IsSigned, ValuesToRelease);
 
       if (R.IsEq) {
         // Also add the inverted constraint for equality constraints.
@@ -638,7 +648,7 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT) {
         CSToUse.addVariableRowFill(R.Coefficients);
 
         DFSInStack.emplace_back(CB.NumIn, CB.NumOut, CB.Condition, CB.Not,
-                                R.IsSigned);
+                                R.IsSigned, SmallVector<Value *, 2>());
       }
     }
   }

llvm/test/Transforms/ConstraintElimination/pr54228-variable-name-order.ll

@@ -24,7 +24,7 @@ define i1 @test_pr54228(i32 %a, i32 %b, i1 %i.0, i1 %i.1) {
 ; CHECK-NEXT:    br i1 [[C_3]], label [[EXIT:%.*]], label [[LOOP_HEADER]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    [[C_4:%.*]] = icmp eq i32 [[A]], 0
-; CHECK-NEXT:    ret i1 true
+; CHECK-NEXT:    ret i1 [[C_4]]
 ;
 entry:
   br i1 %i.0, label %ph.1, label %loop.header