[Reassociate] Convert shl by constant into multiply during tree linearization.

[mcrosier]

Summary

If we're linearizing a multiply expression, convert any left shift by a constant into a multiply, so they can be reassociated. We already do a similar optimization where we convert a negation into a multiply by -1.

See the test case for an example of what this hits.

Also, this addresses PR30256, which is a 4.0 release blocker.

Chad

[mcrosier]

BTW, I'll be on vacation starting this afternoon and won't be returning until the 20th. In my absence, I'll have someone on my team follow up on this patch since it's a release blocker and we need to get a fix in for PR30256.

[mssimpso]

Chad,

This change looks good to me, but I'm not familiar enough with reassociation to approve. Maybe @majnemer can? Thanks!

[hans]

What's the status here?

[mssimpso]

The patch is still awaiting approval. I'm hoping we can get more eyes on this because, although it looks good to me, I don't think I know this code well enough to approve. Has anyone else had a chance to look?

[mcrosier]

Ping. This is a fix for a release blocker (PR30256).

Unfortunately, Reassociation doesn't have a code owner, but this should be fairly simple to understand.

[mcrosier]

Add a few more reviewers.

[efriedma]

I'm not following why this fixes the crash. As far as I can tell, on trunk, Reassociate special-cases shl instructions in exactly one place: the check at the beginning of ReassociatePass::OptimizeInst. In that case, we replace all the uses of the shl instruction, so it shouldn't matter. (I'd like an answer to this to make sure you aren't papering over a real problem.)

lib/Transforms/Scalar/Reassociate.cpp

371

Why don't we don't erase the dead Shl here?

[chandlerc]

Some longer term comments here, but all of this seems reasonably in-line with how we currently do things in reassociate, but raises some long-standing issues in a potentially more concerning area.

I also think Eli's comment is important.

lib/Transforms/Scalar/Reassociate.cpp

597–598

I'm OK doing this temporarily if we need to, but we should have a FIXME to revisit this.

• Does this create an endless cycle between reassociate and instcombine? That seems bad
• Can we just synthesize this by teaching the reassociation itself to magically treat shifts of this form as-if they were multiplies and transforming them as necessary on the fly?
• Alternatively, can we add some strong assurance that if we don't reassociate, we transform everything back to a shl afterward? Maybe test cases for interesting expression trees where we put everything into multiply form and then back?

Essentially, all of this is that we somehow need this to not leak out of reassociate in the cases where we don't change anything (or we need to change our canonicalization).

[mcrosier]

In D29777#682654, @efriedma wrote:

I'm not following why this fixes the crash. As far as I can tell, on trunk, Reassociate special-cases shl instructions in exactly one place: the check at the beginning of ReassociatePass::OptimizeInst. In that case, we replace all the uses of the shl instruction, so it shouldn't matter.

The convert only occurs if the Shl is feeding a Mul or an Add. In my test case, the Shl is feeding a negation. D30223 is an alternative solution that has the same result.

(I'd like an answer to this to make sure you aren't papering over a real problem.)

This fixes PR30256 by avoiding the particular case that is causing the assert. However, I think the underlying issue is that if a factor is a negative constant, we consider the positive value as a factor as well (because we can percolate the negate out). Thus, I think this patch does kinda papers over the issue, but I can also argue it's a reasonable optimization..

I think I've just come up with a patch that can more directly address the underlying issue.. Let me post that before we continue the review here.

lib/Transforms/Scalar/Reassociate.cpp

371

I believe this is to avoid an iterator invalidation bug (but I could be mistaken).

In the current code, the Shl is inserted into the RedoInsts list after the call to ConvertShiftToMul. This will ensure the Shl is removed as dead code. I should do the same for my code (or alternatively pass in the RedoInsts into ConvertShiftToMul and insert the Shl into RedoInsts directly in CovertShiftToMul).

[mcrosier]

In D29777#682749, @chandlerc wrote:

Some longer term comments here, but all of this seems reasonably in-line with how we currently do things in reassociate, but raises some long-standing issues in a potentially more concerning area.

After thinking a bit and discussing this with Eli, this patch (and D30223) really are just papering over the issue. See Eli's comments in D30228 for a good description of what's going on; D30228 also directly address the problem.

[mcrosier]

lib/Transforms/Scalar/Reassociate.cpp

597–598

Yes, reassociation relies on instcombine to undo the canonicalizations that enable reassocation performed in this pass. While I'm likely going to abandon this patch in favor of D30228, I wanted to acknowledge that I agree with Chandler's comments. Specifically, we should improve reassociate to only perform these type of transforms when we know reassociation will be successful. Of course that is much easier said than done.

[mcrosier]

Abandon in favor of rL296003.

lib/Transforms/Scalar/Reassociate.cpp

assert(LHS.getZExtValue() < Threshold && RHS.getZExtValue() < Threshold &&
            "Weights not reduced!");
     unsigned Total = LHS.getZExtValue() + RHS.getZExtValue();
     while (Total >= Threshold)
       Total -= CM;
     LHS = Total;
   }
 }
 
+/// If this is a shift of a reassociable multiply or is used by one, change
+/// this into a multiply by a constant to assist with further reassociation.
+static BinaryOperator *ConvertShiftToMul(Instruction *Shl) {
+  Constant *MulCst = ConstantInt::get(Shl->getType(), 1);
+  MulCst = ConstantExpr::getShl(MulCst, cast<Constant>(Shl->getOperand(1)));
+
+  BinaryOperator *Mul =
+    BinaryOperator::CreateMul(Shl->getOperand(0), MulCst, "", Shl);
+  Shl->setOperand(0, UndefValue::get(Shl->getType())); // Drop use of op.
+  Mul->takeName(Shl);
+
+  // Everyone now refers to the mul instruction.
+  Shl->replaceAllUsesWith(Mul);
+  Mul->setDebugLoc(Shl->getDebugLoc());
+
+  // We can safely preserve the nuw flag in all cases.  It's also safe to turn a
+  // nuw nsw shl into a nuw nsw mul.  However, nsw in isolation requires special
+  // handling.
+  bool NSW = cast<BinaryOperator>(Shl)->hasNoSignedWrap();
+  bool NUW = cast<BinaryOperator>(Shl)->hasNoUnsignedWrap();
+  if (NSW && NUW)
+    Mul->setHasNoSignedWrap(true);
+  Mul->setHasNoUnsignedWrap(NUW);
+  return Mul;
+}
+
 typedef std::pair<Value*, APInt> RepeatedValue;
 
 /// Given an associative binary expression, return the leaf
 /// nodes in Ops along with their weights (how many times the leaf occurs).  The
 /// original expression is the same as
 ///   (Ops[0].first op Ops[0].first op ... Ops[0].first)  <- Ops[0].second times
 /// op
 ///   (Ops[1].first op Ops[1].first op ... Ops[1].first)  <- Ops[1].second times
#endif
       // can usefully morph it into an expression of the right kind.
       assert((!isa<Instruction>(Op) ||
               cast<Instruction>(Op)->getOpcode() != Opcode
               || (isa<FPMathOperator>(Op) &&
                   !cast<Instruction>(Op)->hasUnsafeAlgebra())) &&
              "Should have been handled above!");
       assert(Op->hasOneUse() && "Has uses outside the expression tree!");
 
+      if (BinaryOperator *BO = dyn_cast<BinaryOperator>(Op)) {
         // If this is a multiply expression, turn any internal negations into
         // multiplies by -1 so they can be reassociated.
-      if (BinaryOperator *BO = dyn_cast<BinaryOperator>(Op))
         if ((Opcode == Instruction::Mul && BinaryOperator::isNeg(BO)) ||
             (Opcode == Instruction::FMul && BinaryOperator::isFNeg(BO))) {
           DEBUG(dbgs() << "MORPH LEAF: " << *Op << " (" << Weight << ") TO ");
           BO = LowerNegateToMultiply(BO);
           DEBUG(dbgs() << *BO << '\n');
           Worklist.push_back(std::make_pair(BO, Weight));
           Changed = true;
           continue;
         }
-
+        // If this is a multiply expression, convert any left shift by a
+        // constant into a multiply so they can be reassociated.
+        if (Opcode == Instruction::Mul && BO->getOpcode() == Instruction::Shl &&
+            isa<ConstantInt>(BO->getOperand(1))) {
+          DEBUG(dbgs() << "MORPH LEAF: " << *Op << " (" << Weight << ") TO ");
+          BO = ConvertShiftToMul(BO);
+          DEBUG(dbgs() << *BO << '\n');
+          Worklist.push_back(std::make_pair(BO, Weight));
+          Changed = true;
+          continue;
+        }
+      }
       // Failed to morph into an expression of the right type.  This really is
       // a leaf.
       DEBUG(dbgs() << "ADD LEAF: " << *Op << " (" << Weight << ")\n");
       assert(!isReassociableOp(Op, Opcode) && "Value was morphed?");
       LeafOrder.push_back(Op);
       Leaves[Op] = Weight;
     }
   }
BreakUpSubtract(Instruction *Sub, SetVector<AssertingVH<Instruction>> &ToRedo) {
   // Everyone now refers to the add instruction.
   Sub->replaceAllUsesWith(New);
   New->setDebugLoc(Sub->getDebugLoc());
 
   DEBUG(dbgs() << "Negated: " << *New << '\n');
   return New;
 }
 
-/// If this is a shift of a reassociable multiply or is used by one, change
-/// this into a multiply by a constant to assist with further reassociation.
-static BinaryOperator *ConvertShiftToMul(Instruction *Shl) {
-  Constant *MulCst = ConstantInt::get(Shl->getType(), 1);
-  MulCst = ConstantExpr::getShl(MulCst, cast<Constant>(Shl->getOperand(1)));
-
-  BinaryOperator *Mul =
-    BinaryOperator::CreateMul(Shl->getOperand(0), MulCst, "", Shl);
-  Shl->setOperand(0, UndefValue::get(Shl->getType())); // Drop use of op.
-  Mul->takeName(Shl);
-
-  // Everyone now refers to the mul instruction.
-  Shl->replaceAllUsesWith(Mul);
-  Mul->setDebugLoc(Shl->getDebugLoc());
-
-  // We can safely preserve the nuw flag in all cases.  It's also safe to turn a
-  // nuw nsw shl into a nuw nsw mul.  However, nsw in isolation requires special
-  // handling.
-  bool NSW = cast<BinaryOperator>(Shl)->hasNoSignedWrap();
-  bool NUW = cast<BinaryOperator>(Shl)->hasNoUnsignedWrap();
-  if (NSW && NUW)
-    Mul->setHasNoSignedWrap(true);
-  Mul->setHasNoUnsignedWrap(NUW);
-  return Mul;
-}
-
 /// Scan backwards and forwards among values with the same rank as element i
 /// to see if X exists.  If X does not exist, return i.  This is useful when
 /// scanning for 'x' when we see '-x' because they both get the same rank.
 static unsigned FindInOperandList(SmallVectorImpl<ValueEntry> &Ops, unsigned i,
                                   Value *X) {
   unsigned XRank = Ops[i].Rank;
   unsigned e = Ops.size();
   for (unsigned j = i+1; j != e && Ops[j].Rank == XRank; ++j) {

test/Transforms/Reassociate/basictest.ll

define i32 @test15(i32 %X1, i32 %X2, i32 %X3) {
   %A = icmp ne i32 %X1, 0
   %B = icmp slt i32 %X2, %X3
   %C = and i1 %A, %B
   %D = select i1 %C, i32 %X1, i32 0
   ret i32 %D
 ; CHECK-LABEL: @test15
 ; CHECK: and i1 %A, %B
 }
+
+
+; -(a << 1)*2 + b --> -4a + b
+define i32 @test16(i1 %cmp, i32 %a, i32 %b) {
+entry:
+  %shl = shl i32 %a, 1
+  %shl.neg = sub i32 0, %shl
+  %add1 = mul i32 %shl.neg, 2
+  %add2 = add i32 %add1, %b
+  ret i32 %add2
+
+; CHECK-LABEL: @test16
+; CHECK: %[[MUL:.*]] = mul i32 %a, -4
+; CHECK: %[[ADD:.*]] = add i64 %[[MUL]], %b
+; CHECK: ret i32 %[[ADD]]
+}