[Reassociate] Keep NSW/NUW flags on binary ops whenever possible.

[ab]

Summary

The previous - conservative - behavior was to always drop the overflow flags
when reassociating scalar expressions.

This kept address computing code done on int/i32 and sign-extended to i64,
because the nsw/nuw flags are needed to promote. In turn, this made it
impossible to fold the address computation into the addressing mode.

This (simple) patch tries to keep the nsw/nuw flags, but only when they are
consistent in the complete expression tree. My understanding is, this should be
valid, because the poison has to propagate to the root, no matter how the
expression is reassociated.
My reading of the rest of the code tells me the expression tree only consists of
the same operator, and all expressions have only one use.

Now for measurements:

Compile Time	Δ	Previous	Current	σ
MultiSource/Benchmarks/ASC_Sequoia/IRSmk/IRSmk	49.81%	0.1321	0.1979	0.0005
SingleSource/Benchmarks/Misc/himenobmtxpa	7.40%	0.1839	0.1975	0.0009

Execution Time	Δ	Previous	Current	σ
SingleSource/Benchmarks/Linpack/linpack-pc	1.35%	2.3403	2.3719	0.0033
MultiSource/Benchmarks/ASC_Sequoia/IRSmk/IRSmk	-8.58%	4.7386	4.3318	0.0052
SingleSource/Benchmarks/Misc/himenobmtxpa	-5.91%	1.6205	1.5248	0.0178
MultiSource/Benchmarks/tramp3d-v4/tramp3d-v4	-2.39%	0.4302	0.4199	0.0011

I'm not exactly sure why linpack is slower; I'm (lightly) investigating, but can
look into it more seriously if people think it's a big deal?

Thanks,
-Ahmed

lib/Transforms/Scalar/Reassociate.cpp

namespace {
   }
 }
 
 #ifndef NDEBUG
 /// PrintOps - Print out the expression identified in the Ops list.
 ///
 static void PrintOps(Instruction *I, const SmallVectorImpl<ValueEntry> &Ops) {
   Module *M = I->getParent()->getParent()->getParent();
-  dbgs() << Instruction::getOpcodeName(I->getOpcode()) << " "
-       << *Ops[0].Op->getType() << '\t';
+  dbgs() << Instruction::getOpcodeName(I->getOpcode()) << " ";
+  if (OverflowingBinaryOperator *BO = dyn_cast<OverflowingBinaryOperator>(I)) {
+    if (BO->hasNoUnsignedWrap())
+      dbgs() << "nuw ";
+    if (BO->hasNoSignedWrap())
+      dbgs() << "nsw ";
+  }
+  dbgs() << *Ops[0].Op->getType() << '\t';
   for (unsigned i = 0, e = Ops.size(); i != e; ++i) {
     dbgs() << "[ ";
     Ops[i].Op->printAsOperand(dbgs(), false, M);
     dbgs() << ", #" << Ops[i].Rank << "] ";
   }
 }
 #endif
 
public:
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.setPreservesCFG();
     }
   private:
     void BuildRankMap(Function &F);
     unsigned getRank(Value *V);
     void ReassociateExpression(BinaryOperator *I);
-    void RewriteExprTree(BinaryOperator *I, SmallVectorImpl<ValueEntry> &Ops);
+    void RewriteExprTree(BinaryOperator *I, SmallVectorImpl<ValueEntry> &Ops,
+                         unsigned NoWrapFlags);
     Value *OptimizeExpression(BinaryOperator *I,
                               SmallVectorImpl<ValueEntry> &Ops);
     Value *OptimizeAdd(Instruction *I, SmallVectorImpl<ValueEntry> &Ops);
     Value *OptimizeXor(Instruction *I, SmallVectorImpl<ValueEntry> &Ops);
     bool CombineXorOpnd(Instruction *I, XorOpnd *Opnd1, APInt &ConstOpnd,
                         Value *&Res);
     bool CombineXorOpnd(Instruction *I, XorOpnd *Opnd1, XorOpnd *Opnd2,
                         APInt &ConstOpnd, Value *&Res);

 /// which requires passing through a phi node.
 ///
 /// Note that this routine may also mutate binary operators of the wrong type
 /// that have all uses inside the expression (i.e. only used by non-leaf nodes
 /// of the expression) if it can turn them into binary operators of the right
 /// type and thus make the expression bigger.
 
 static bool LinearizeExprTree(BinaryOperator *I,
-                              SmallVectorImpl<RepeatedValue> &Ops) {
+                              SmallVectorImpl<RepeatedValue> &Ops,
+                              unsigned &NoWrapFlags) {
   DEBUG(dbgs() << "LINEARIZE: " << *I << '\n');
   unsigned Bitwidth = I->getType()->getScalarType()->getPrimitiveSizeInBits();
   unsigned Opcode = I->getOpcode();
   assert(I->isAssociative() && I->isCommutative() &&
          "Expected an associative and commutative operation!");
 
   // Visit all operands of the expression, keeping track of their weight (the
   // number of paths from the expression root to the operand, or if you like
   // the number of times that operand occurs in the linearized expression).
   // For example, if I = X + A, where X = A + B, then I, X and B have weight 1
   // while A has weight two.
 
   // Worklist of non-leaf nodes (their operands are in the expression too) along
   // with their weights, representing a certain number of paths to the operator.
   // If an operator occurs in the worklist multiple times then we found multiple
   // ways to get to it.
   SmallVector<std::pair<BinaryOperator*, APInt>, 8> Worklist; // (Op, Weight)
   Worklist.push_back(std::make_pair(I, APInt(Bitwidth, 1)));
   bool MadeChange = false;
 
+  // Keep track of the NSW/NUW flags found in the expression.  If any one of
+  // the operators doesn't have a flag, don't use it in the entire expression.
+  NoWrapFlags = (OverflowingBinaryOperator::NoUnsignedWrap |
+                 OverflowingBinaryOperator::NoSignedWrap);
+
   // Leaves of the expression are values that either aren't the right kind of
   // operation (eg: a constant, or a multiply in an add tree), or are, but have
   // some uses that are not inside the expression.  For example, in I = X + X,
   // X = A + B, the value X has two uses (by I) that are in the expression.  If
   // X has any other uses, for example in a return instruction, then we consider
   // X to be a leaf, and won't analyze it further.  When we first visit a value,
   // if it has more than one use then at first we conservatively consider it to
   // be a leaf.  Later, as the expression is explored, we may discover some more

 
 #ifndef NDEBUG
   SmallPtrSet<Value*, 8> Visited; // For sanity checking the iteration scheme.
 #endif
   while (!Worklist.empty()) {
     std::pair<BinaryOperator*, APInt> P = Worklist.pop_back_val();
     I = P.first; // We examine the operands of this binary operator.
 
+    // If operator I can't or doesn't have one of the NSW/NUW flags, remove it.
+    if (!isa<OverflowingBinaryOperator>(I) || !I->hasNoSignedWrap())
+      NoWrapFlags &= ~OverflowingBinaryOperator::NoSignedWrap;
+    if (!isa<OverflowingBinaryOperator>(I) || !I->hasNoUnsignedWrap())
+      NoWrapFlags &= ~OverflowingBinaryOperator::NoUnsignedWrap;
+
     for (unsigned OpIdx = 0; OpIdx < 2; ++OpIdx) { // Visit operands.
       Value *Op = I->getOperand(OpIdx);
       APInt Weight = P.second; // Number of paths to this operand.
       DEBUG(dbgs() << "OPERAND: " << *Op << " (" << Weight << ")\n");
       assert(!Op->use_empty() && "No uses, so how did we get to it?!");
 
       // If this is a binary operation of the right kind with only one use then
       // add its operands to the expression.
#endif
   }
 
   return MadeChange;
 }
 
 // RewriteExprTree - Now that the operands for this expression tree are
 // linearized and optimized, emit them in-order.
 void Reassociate::RewriteExprTree(BinaryOperator *I,
-                                  SmallVectorImpl<ValueEntry> &Ops) {
+                                  SmallVectorImpl<ValueEntry> &Ops,
+                                  unsigned NoWrapFlags) {
   assert(Ops.size() > 1 && "Single values should be used directly!");
 
   // Since our optimizations should never increase the number of operations, the
   // new expression can usually be written reusing the existing binary operators
   // from the original expression tree, without creating any new instructions,
   // though the rewritten expression may have a completely different topology.
   // We take care to not change anything if the new expression will be the same
   // as the original.  If more than trivial changes (like commuting operands)
void Reassociate::RewriteExprTree(BinaryOperator *I,
   // expression tree is dominated by all of Ops.
   if (ExpressionChanged)
     do {
       // Preserve FastMathFlags.
       if (isa<FPMathOperator>(I)) {
         FastMathFlags Flags = I->getFastMathFlags();
         ExpressionChanged->clearSubclassOptionalData();
         ExpressionChanged->setFastMathFlags(Flags);
-      } else
+      } else {
         ExpressionChanged->clearSubclassOptionalData();
+        if (isa<OverflowingBinaryOperator>(ExpressionChanged)) {
+          ExpressionChanged->setHasNoUnsignedWrap(
+              NoWrapFlags & OverflowingBinaryOperator::NoUnsignedWrap);
+          ExpressionChanged->setHasNoSignedWrap(
+              NoWrapFlags & OverflowingBinaryOperator::NoSignedWrap);
+        }
+      }
 
       if (ExpressionChanged == I)
         break;
       ExpressionChanged->moveBefore(I);
       ExpressionChanged = cast<BinaryOperator>(*ExpressionChanged->user_begin());
     } while (1);
 
   // Throw away any left over nodes from the original expression.

 /// multiplication sequence, and if this sequence contains a multiply by Factor,
 /// remove Factor from the tree and return the new tree.
 Value *Reassociate::RemoveFactorFromExpression(Value *V, Value *Factor) {
   BinaryOperator *BO = isReassociableOp(V, Instruction::Mul, Instruction::FMul);
   if (!BO)
     return nullptr;
 
   SmallVector<RepeatedValue, 8> Tree;
-  MadeChange |= LinearizeExprTree(BO, Tree);
+  unsigned NoWrapFlags = 0;
+  MadeChange |= LinearizeExprTree(BO, Tree, NoWrapFlags);
   SmallVector<ValueEntry, 8> Factors;
   Factors.reserve(Tree.size());
   for (unsigned i = 0, e = Tree.size(); i != e; ++i) {
     RepeatedValue E = Tree[i];
     Factors.append(E.second.getZExtValue(),
                    ValueEntry(getRank(E.first), E.first));
   }
 
if (ConstantInt *FC1 = dyn_cast<ConstantInt>(Factor)) {
           break;
         }
       }
     }
   }
 
   if (!FoundFactor) {
     // Make sure to restore the operands to the expression tree.
-    RewriteExprTree(BO, Factors);
+    RewriteExprTree(BO, Factors, NoWrapFlags);
     return nullptr;
   }
 
   BasicBlock::iterator InsertPt = BO; ++InsertPt;
 
   // If this was just a single multiply, remove the multiply and return the only
   // remaining operand.
   if (Factors.size() == 1) {
     RedoInsts.insert(BO);
     V = Factors[0].Op;
   } else {
-    RewriteExprTree(BO, Factors);
+    RewriteExprTree(BO, Factors, NoWrapFlags);
     V = BO;
   }
 
   if (NeedsNegate)
     V = CreateNeg(V, "neg", InsertPt, BO);
 
   return V;
 }

 
 void Reassociate::ReassociateExpression(BinaryOperator *I) {
   assert(!I->getType()->isVectorTy() &&
          "Reassociation of vector instructions is not supported.");
 
   // First, walk the expression tree, linearizing the tree, collecting the
   // operand information.
   SmallVector<RepeatedValue, 8> Tree;
-  MadeChange |= LinearizeExprTree(I, Tree);
+  unsigned NoWrapFlags = 0;
+  MadeChange |= LinearizeExprTree(I, Tree, NoWrapFlags);
   SmallVector<ValueEntry, 8> Ops;
   Ops.reserve(Tree.size());
   for (unsigned i = 0, e = Tree.size(); i != e; ++i) {
     RepeatedValue E = Tree[i];
     Ops.append(E.second.getZExtValue(),
                ValueEntry(getRank(E.first), E.first));
   }
 
if (Ops.size() == 1) {
     if (Instruction *OI = dyn_cast<Instruction>(Ops[0].Op))
       OI->setDebugLoc(I->getDebugLoc());
     RedoInsts.insert(I);
     return;
   }
 
   // Now that we ordered and optimized the expressions, splat them back into
   // the expression tree, removing any unneeded nodes.
-  RewriteExprTree(I, Ops);
+  RewriteExprTree(I, Ops, NoWrapFlags);
 }
 
 bool Reassociate::runOnFunction(Function &F) {
   if (skipOptnoneFunction(F))
     return false;
 
   // Calculate the rank map for F
   BuildRankMap(F);

test/Transforms/Reassociate/no-op.ll

 ; CHECK-NEXT: %m1 = mul nsw i32 %m0, %b
   call void @use(i32 %a1)
 ; CHECK-NEXT: call void @use
   call void @use(i32 %m1)
   ret void
 }
 
 define void @test2(i32 %a, i32 %b, i32 %c, i32 %d) {
-; The initial add doesn't change so should not lose the nsw flag.
+; The initial add doesn't change.
 ; CHECK-LABEL: @test2(
   %a0 = add nsw i32 %b, %a
 ; CHECK-NEXT: %a0 = add nsw i32 %b, %a
   %a1 = add nsw i32 %a0, %d
-; CHECK-NEXT: %a1 = add i32 %a0, %c
+; CHECK-NEXT: %a1 = add nsw i32 %a0, %c
   %a2 = add nsw i32 %a1, %c
-; CHECK-NEXT: %a2 = add i32 %a1, %d
+; CHECK-NEXT: %a2 = add nsw i32 %a1, %d
   call void @use(i32 %a2)
 ; CHECK-NEXT: call void @use
   ret void
 }

test/Transforms/Reassociate/nsw-nuw.ll

+; RUN: opt %s -reassociate -S | FileCheck %s
+
+;; Check that we keep overflow flags on binary operators whenever possible.
+
+define i32 @test_all_nsw(i32 %reg109, i32 %reg1111) {
+  %reg115 = add nsw i32 %reg109, -30
+  %reg116 = add nsw i32 %reg115, %reg1111
+  %reg117 = add nsw i32 %reg116, 30
+  ret i32 %reg117
+
+; CHECK-LABEL: @test_all_nsw
+; CHECK-NEXT: %reg117 = add nsw i32 %reg1111, %reg109
+; CHECK-NEXT: ret i32 %reg117
+}
+
+define i32 @test_all_nuw(i32 %reg109, i32 %reg1111) {
+  %reg115 = add nuw i32 %reg109, -30
+  %reg116 = add nuw i32 %reg115, %reg1111
+  %reg117 = add nuw i32 %reg116, 30
+  ret i32 %reg117
+
+; CHECK-LABEL: @test_all_nuw
+; CHECK-NEXT: %reg117 = add nuw i32 %reg1111, %reg109
+; CHECK-NEXT: ret i32 %reg117
+}
+
+define i32 @test_all_both(i32 %reg109, i32 %reg1111) {
+  %reg115 = add nuw nsw i32 %reg109, -30
+  %reg116 = add nuw nsw i32 %reg115, %reg1111
+  %reg117 = add nuw nsw i32 %reg116, 30
+  ret i32 %reg117
+
+; CHECK-LABEL: @test_all_both
+; CHECK-NEXT: %reg117 = add nuw nsw i32 %reg1111, %reg109
+; CHECK-NEXT: ret i32 %reg117
+}
+
+;; The rest of the tests check that the flags are dropped.
+;; We only keep flags when reassociating an expression consisting of a single
+;; operation, into one consisting of the same operation.
+
+define i32 @test_allbutone_nsw(i32 %reg109, i32 %reg1111) {
+  %reg115 = add i32 %reg109, -30
+  %reg116 = add nsw i32 %reg115, %reg1111
+  %reg117 = add nsw i32 %reg116, 30
+  ret i32 %reg117
+
+; CHECK-LABEL: @test_allbutone_nsw
+; CHECK-NEXT: %reg117 = add i32 %reg1111, %reg109
+; CHECK-NEXT: ret i32 %reg117
+}
+
+define i32 @test_allbutone_nsw_2(i32 %reg109, i32 %reg1111) {
+  %reg115 = add nsw i32 %reg109, -30
+  %reg116 = add nsw i32 %reg115, %reg1111
+  %reg117 = add i32 %reg116, 30
+  ret i32 %reg117
+
+; CHECK-LABEL: @test_allbutone_nsw
+; CHECK-NEXT: %reg117 = add i32 %reg1111, %reg109
+; CHECK-NEXT: ret i32 %reg117
+}
+
+define i32 @test_morph_adds(i32 %X) {
+  %Y = add nuw i32 %X ,%X
+  %Z = add nuw i32 %Y, %X
+  ret i32 %Z
+; CHECK-LABEL: @test_morph_adds
+; CHECK-NEXT: mul i32 %X, 3
+; CHECK-NEXT: ret i32
+}
+
+define i32 @test_mixed_flags(i32 %reg109, i32 %reg1111) {
+  %reg115 = add nuw i32 %reg109, -30
+  %reg116 = add nsw i32 %reg115, %reg1111
+  %reg117 = add nuw i32 %reg116, 30
+  ret i32 %reg117
+
+; CHECK-LABEL: @test_mixed_flags
+; CHECK-NEXT: %reg117 = add i32 %reg1111, %reg109
+; CHECK-NEXT: ret i32 %reg117
+}
+
+; This should be one add and two multiplies.
+define i32 @test_mixed_ops(i32 %A, i32 %B, i32 %C) {
+  ; A*A*B + A*C*A
+  %aa = mul nsw i32 %A, %A
+  %aab = mul nsw i32 %aa, %B
+  %ac = mul nsw i32 %A, %C
+  %aac = mul nsw i32 %ac, %A
+  %r = add nsw i32 %aab, %aac
+  ret i32 %r
+
+; CHECK-LABEL: @test_mixed_ops
+; CHECK-NEXT: add i32 %C, %B
+; CHECK-NEXT: mul i32
+; CHECK-NEXT: mul i32
+; CHECK-NEXT: ret i32
+}