[Reassociate] Keep flags for more unchanged operations

[dmgreen]

Summary

Reassociation destroys nsw/nuw flags from BinOps that are changed. If the expression in at the end of a tree that was altered, but didn't change itself the flags do not need to be removed though. For example, if %a, %b and %c are reassociated in

%x = add nsw i32 %a, %c
%y = add nsw i32 %x, %b
%z = add nsw i32 %y, %d

The value of %y and so add %y %d remains the same, and %z needn't drop the nsw flags.
https://alive2.llvm.org/ce/z/_juAiV

[nikic]

LGTM

llvm/lib/Transforms/Scalar/Reassociate.cpp

832–833

As a side note, I don't think we need to do this bit either for unchanged expressions.

llvm/lib/Transforms/Scalar/Reassociate.cpp

void ReassociatePass::RewriteExprTree(BinaryOperator *I,
   /// of it, or it may momentarily become reassociable during rewriting (below)
   /// due it being removed as an operand of one of its uses.  Ensure that misuse
   /// of leaf nodes as inner nodes cannot occur by remembering all of the future
   /// leaves and refusing to reuse any of them as inner nodes.
   SmallPtrSet<Value*, 8> NotRewritable;
   for (unsigned i = 0, e = Ops.size(); i != e; ++i)
     NotRewritable.insert(Ops[i].Op);
 
-  // ExpressionChanged - Non-null if the rewritten expression differs from the
-  // original in some non-trivial way, requiring the clearing of optional flags.
-  // Flags are cleared from the operator in ExpressionChanged up to I inclusive.
-  BinaryOperator *ExpressionChanged = nullptr;
+  // ExpressionChangedStart - Non-null if the rewritten expression differs from
+  // the original in some non-trivial way, requiring the clearing of optional
+  // flags. Flags are cleared from the operator in ExpressionChangedStart up to
+  // ExpressionChangedEnd inclusive.
+  BinaryOperator *ExpressionChangedStart = nullptr,
+                 *ExpressionChangedEnd = nullptr;
   for (unsigned i = 0; ; ++i) {
     // The last operation (which comes earliest in the IR) is special as both
     // operands will come from Ops, rather than just one with the other being
     // a subexpression.
     if (i+2 == Ops.size()) {
       Value *NewLHS = Ops[i].Op;
       Value *NewRHS = Ops[i+1].Op;
       Value *OldLHS = Op->getOperand(0);
if (i+2 == Ops.size()) {
       if (NewRHS != OldRHS) {
         BinaryOperator *BO = isReassociableOp(OldRHS, Opcode);
         if (BO && !NotRewritable.count(BO))
           NodesToRewrite.push_back(BO);
         Op->setOperand(1, NewRHS);
       }
       LLVM_DEBUG(dbgs() << "TO: " << *Op << '\n');
 
-      ExpressionChanged = Op;
+      ExpressionChangedStart = Op;
+      if (!ExpressionChangedEnd)
+        ExpressionChangedEnd = Op;
       MadeChange = true;
       ++NumChanged;
 
       break;
     }
 
     // Not the last operation.  The left-hand side will be a sub-expression
     // while the right-hand side will be the current element of Ops.
     Value *NewRHS = Ops[i].Op;
     if (NewRHS != Op->getOperand(1)) {
       LLVM_DEBUG(dbgs() << "RA: " << *Op << '\n');
       if (NewRHS == Op->getOperand(0)) {
         // The new right-hand side was already present as the left operand.  If
         // we are lucky then swapping the operands will sort out both of them.
         Op->swapOperands();
       } else {
         // Overwrite with the new right-hand side.
         BinaryOperator *BO = isReassociableOp(Op->getOperand(1), Opcode);
         if (BO && !NotRewritable.count(BO))
           NodesToRewrite.push_back(BO);
         Op->setOperand(1, NewRHS);
-        ExpressionChanged = Op;
+        ExpressionChangedStart = Op;
+        if (!ExpressionChangedEnd)
+          ExpressionChangedEnd = Op;
       }
       LLVM_DEBUG(dbgs() << "TO: " << *Op << '\n');
       MadeChange = true;
       ++NumChanged;
     }
 
     // Now deal with the left-hand side.  If this is already an operation node
     // from the original expression then just rewrite the rest of the expression
if (NodesToRewrite.empty()) {
         NewOp->setFastMathFlags(I->getFastMathFlags());
     } else {
       NewOp = NodesToRewrite.pop_back_val();
     }
 
     LLVM_DEBUG(dbgs() << "RA: " << *Op << '\n');
     Op->setOperand(0, NewOp);
     LLVM_DEBUG(dbgs() << "TO: " << *Op << '\n');
-    ExpressionChanged = Op;
+    ExpressionChangedStart = Op;
+    if (!ExpressionChangedEnd)
+      ExpressionChangedEnd = Op;
     MadeChange = true;
     ++NumChanged;
     Op = NewOp;
   }
 
   // If the expression changed non-trivially then clear out all subclass data
   // starting from the operator specified in ExpressionChanged, and compactify
   // the operators to just before the expression root to guarantee that the
   // expression tree is dominated by all of Ops.
-  if (ExpressionChanged)
+  if (ExpressionChangedStart) {
+    bool ClearFlags = true;
     do {
       // Preserve FastMathFlags.
+      if (ClearFlags) {
         if (isa<FPMathOperator>(I)) {
           FastMathFlags Flags = I->getFastMathFlags();
-        ExpressionChanged->clearSubclassOptionalData();
-        ExpressionChanged->setFastMathFlags(Flags);
+          ExpressionChangedStart->clearSubclassOptionalData();
+          ExpressionChangedStart->setFastMathFlags(Flags);
         } else
-        ExpressionChanged->clearSubclassOptionalData();
+          ExpressionChangedStart->clearSubclassOptionalData();
+      }
 
-      if (ExpressionChanged == I)
+      if (ExpressionChangedStart == ExpressionChangedEnd)
+        ClearFlags = false;
+      if (ExpressionChangedStart == I)
         break;
 
       // Discard any debug info related to the expressions that has changed (we
-      // can leave debug infor related to the root, since the result of the
-      // expression tree should be the same even after reassociation).
-      replaceDbgUsesWithUndef(ExpressionChanged);
-
-      ExpressionChanged->moveBefore(I);
-      ExpressionChanged = cast<BinaryOperator>(*ExpressionChanged->user_begin());
+      // can leave debug info related to the root and any operation that didn't
+      // change, since the result of the expression tree should be the same
+      // even after reassociation).
+      if (ClearFlags)
+        replaceDbgUsesWithUndef(ExpressionChangedStart);
+
+      ExpressionChangedStart->moveBefore(I);
+      ExpressionChangedStart =
+          cast<BinaryOperator>(*ExpressionChangedStart->user_begin());
     } while (true);
+  }
 
   // Throw away any left over nodes from the original expression.
   for (unsigned i = 0, e = NodesToRewrite.size(); i != e; ++i)
     RedoInsts.insert(NodesToRewrite[i]);
 }
 
 /// Insert instructions before the instruction pointed to by BI,
 /// that computes the negative version of the value specified.  The negative

llvm/test/Transforms/Reassociate/cse-pairs.ll

;
   store i32 %11, ptr @num3, align 4
   ret i32 undef
 }
 
 define signext i32 @twoPairsAllOpInPairs(i32 signext %0, i32 signext %1, i32 signext %2, i32 signext %3) {
 ; CHECK-LABEL: @twoPairsAllOpInPairs(
 ; CHECK-NEXT:    [[TMP5:%.*]] = add i32 [[TMP2:%.*]], [[TMP1:%.*]]
 ; CHECK-NEXT:    [[TMP6:%.*]] = add i32 [[TMP5]], [[TMP0:%.*]]
-; CHECK-NEXT:    [[TMP7:%.*]] = add i32 [[TMP6]], [[TMP3:%.*]]
+; CHECK-NEXT:    [[TMP7:%.*]] = add nsw i32 [[TMP6]], [[TMP3:%.*]]
 ; CHECK-NEXT:    store i32 [[TMP7]], ptr @num1, align 4
 ; CHECK-NEXT:    store i32 [[TMP5]], ptr @num2, align 4
 ; CHECK-NEXT:    [[TMP8:%.*]] = add nsw i32 [[TMP3]], [[TMP0]]
 ; CHECK-NEXT:    store i32 [[TMP8]], ptr @num3, align 4
 ; CHECK-NEXT:    ret i32 undef
 ;
   %5 = add nsw i32 %0, %1
   %6 = add nsw i32 %5, %2
   %7 = add nsw i32 %6, %3
   store i32 %7, ptr @num1, align 4
   %8 = add nsw i32 %1, %2
   store i32 %8, ptr @num2, align 4
   %9 = add nsw i32 %0, %3
   store i32 %9, ptr @num3, align 4
   ret i32 undef
 }
 
 define signext i32 @threePairsAllOpInPairs(i32 signext %0, i32 signext %1, i32 signext %2, i32 signext %3, i32 signext %4, i32 signext %5) {
 ; CHECK-LABEL: @threePairsAllOpInPairs(
 ; CHECK-NEXT:    [[TMP7:%.*]] = add i32 [[TMP3:%.*]], [[TMP2:%.*]]
 ; CHECK-NEXT:    [[TMP8:%.*]] = add i32 [[TMP7]], [[TMP0:%.*]]
 ; CHECK-NEXT:    [[TMP9:%.*]] = add i32 [[TMP8]], [[TMP1:%.*]]
-; CHECK-NEXT:    [[TMP10:%.*]] = add i32 [[TMP9]], [[TMP4:%.*]]
-; CHECK-NEXT:    [[TMP11:%.*]] = add i32 [[TMP10]], [[TMP5:%.*]]
+; CHECK-NEXT:    [[TMP10:%.*]] = add nsw i32 [[TMP9]], [[TMP4:%.*]]
+; CHECK-NEXT:    [[TMP11:%.*]] = add nsw i32 [[TMP10]], [[TMP5:%.*]]
 ; CHECK-NEXT:    store i32 [[TMP11]], ptr @num1, align 4
 ; CHECK-NEXT:    [[TMP12:%.*]] = add nsw i32 [[TMP5]], [[TMP0]]
 ; CHECK-NEXT:    store i32 [[TMP12]], ptr @num2, align 4
 ; CHECK-NEXT:    [[TMP13:%.*]] = add nsw i32 [[TMP4]], [[TMP1]]
 ; CHECK-NEXT:    store i32 [[TMP13]], ptr @num3, align 4
 ; CHECK-NEXT:    store i32 [[TMP7]], ptr @num4, align 4
 ; CHECK-NEXT:    ret i32 undef
 ;

llvm/test/Transforms/Reassociate/local-cse.ll

 ; CSE-LABEL: define void @chain_spanning_several_blocks
 ; CSE-SAME: (i64 [[INV1:%.*]], i64 [[INV2:%.*]], i64 [[INV3:%.*]], i64 [[INV4:%.*]], i64 [[INV5:%.*]]) {
 ; CSE-NEXT:  bb1:
 ; CSE-NEXT:    br label [[BB2:%.*]]
 ; CSE:       bb2:
 ; CSE-NEXT:    [[VAL_BB2:%.*]] = call i64 @get_val()
 ; CSE-NEXT:    [[CHAIN_A0:%.*]] = add i64 [[VAL_BB2]], [[INV1]]
 ; CSE-NEXT:    [[CHAIN_A1:%.*]] = add i64 [[CHAIN_A0]], [[INV2]]
-; CSE-NEXT:    [[CHAIN_A2:%.*]] = add i64 [[CHAIN_A1]], [[INV4]]
+; CSE-NEXT:    [[CHAIN_A2:%.*]] = add nuw nsw i64 [[CHAIN_A1]], [[INV4]]
 ; CSE-NEXT:    [[CHAIN_B2:%.*]] = add nuw nsw i64 [[CHAIN_A1]], [[INV5]]
 ; CSE-NEXT:    [[CHAIN_C1:%.*]] = add i64 [[CHAIN_A0]], [[INV3]]
 ; CSE-NEXT:    call void @keep_alive(i64 [[CHAIN_A2]])
 ; CSE-NEXT:    call void @keep_alive(i64 [[CHAIN_B2]])
 ; CSE-NEXT:    call void @keep_alive(i64 [[CHAIN_C1]])
 ; CSE-NEXT:    ret void
 ;
 bb1:

 ; CSE-NEXT:    br label [[BB2:%.*]]
 ; CSE:       bb2:
 ; CSE-NEXT:    [[INV3_BB2:%.*]] = call i64 @get_val()
 ; CSE-NEXT:    [[INV4_BB2:%.*]] = call i64 @get_val()
 ; CSE-NEXT:    [[INV5_BB2:%.*]] = call i64 @get_val()
 ; CSE-NEXT:    [[VAL_BB2:%.*]] = call i64 @get_val()
 ; CSE-NEXT:    [[CHAIN_A0:%.*]] = add i64 [[VAL_BB2]], [[INV1_BB1]]
 ; CSE-NEXT:    [[CHAIN_A1:%.*]] = add i64 [[CHAIN_A0]], [[INV2_BB0]]
-; CSE-NEXT:    [[CHAIN_A2:%.*]] = add i64 [[CHAIN_A1]], [[INV4_BB2]]
+; CSE-NEXT:    [[CHAIN_A2:%.*]] = add nuw nsw i64 [[CHAIN_A1]], [[INV4_BB2]]
 ; CSE-NEXT:    [[CHAIN_B2:%.*]] = add nuw nsw i64 [[CHAIN_A1]], [[INV5_BB2]]
 ; CSE-NEXT:    [[CHAIN_C1:%.*]] = add i64 [[CHAIN_A0]], [[INV3_BB2]]
 ; CSE-NEXT:    call void @keep_alive(i64 [[CHAIN_A2]])
 ; CSE-NEXT:    call void @keep_alive(i64 [[CHAIN_B2]])
 ; CSE-NEXT:    call void @keep_alive(i64 [[CHAIN_C1]])
 ; CSE-NEXT:    ret void
 ;
 bb0: