[ARM] Fix crash in chained BFI combine due to incorrectly RAUW'ing a node.

[aemerson]

Summary

For a bfi chain like:

a = bfi input, x, y
b = bfi a, x', y'

The previous code was RAUW'ing a with x, mutating the second b bfi, and when
SelectionDAG's CSE code ended up deleting it unexpectedly, bad things happend.
There's no need to RAUW in this case because we can just return our newly
created replacement BFI node. It also looked incorrect because it didn't account
for other users of the a bfi.

rdar://79095399

[efriedma]

Change makes sense.

llvm/test/CodeGen/ARM/bfi-chain-cse-crash.ll

39

Can you simplify the testcase so it doesn't have any references to "undef"? That tends to make testcases a lot more fragile.

[aemerson]

Rework test.

[efriedma]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14049

Hang on, I'm not sure this is right. At least, it's not obviously equivalent. FindBFIToCombineWith can look through multiple BFI operations. Not sure if we have proper test coverage for this.

[aemerson]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14049

Right, looks like FindBFIToCombineWith tries to skip unrelated BFIs, and the one it can combine with is CombineBFI. I still think this change is correct, but there's not enough test coverage here.

Having looked further, I can't actually force it to generate chains of more than 2 BFIs at all, since this DAG combine ends up running in between each attempt, combining them before 3 chained BFIs can exist at all.

[efriedma]

llvm/lib/Target/ARM/ARMISelLowering.cpp

14049

If you want to change the code so it can't skip unrelated BFIs, that's probably fine.

[aemerson]

Remove BFI walking code.

[efriedma]

LGTM

llvm/lib/Target/ARM/ARMISelLowering.cpp

 // Neither A nor B must be zero.
 static bool BitsProperlyConcatenate(const APInt &A, const APInt &B) {
   unsigned LastActiveBitInA =  A.countTrailingZeros();
   unsigned FirstActiveBitInB = B.getBitWidth() - B.countLeadingZeros() - 1;
   return LastActiveBitInA - 1 == FirstActiveBitInB;
 }
 
 static SDValue FindBFIToCombineWith(SDNode *N) {
-  // We have a BFI in N. Follow a possible chain of BFIs and find a BFI it can combine with,
-  // if one exists.
+  // We have a BFI in N. Find a BFI it can combine with, if one exists.
   APInt ToMask, FromMask;
   SDValue From = ParseBFI(N, ToMask, FromMask);
   SDValue To = N->getOperand(0);
 
-  // Now check for a compatible BFI to merge with. We can pass through BFIs that
-  // aren't compatible, but not if they set the same bit in their destination as
-  // we do (or that of any BFI we're going to combine with).
   SDValue V = To;
-  APInt CombinedToMask = ToMask;
-  while (V.getOpcode() == ARMISD::BFI) {
-    APInt NewToMask, NewFromMask;
-    SDValue NewFrom = ParseBFI(V.getNode(), NewToMask, NewFromMask);
-    if (NewFrom != From) {
-      // This BFI has a different base. Keep going.
-      CombinedToMask |= NewToMask;
-      V = V.getOperand(0);
-      continue;
-    }
+  if (V.getOpcode() != ARMISD::BFI)
+    return SDValue();
 
-    // Do the written bits conflict with any we've seen so far?
-    if ((NewToMask & CombinedToMask).getBoolValue())
-      // Conflicting bits - bail out because going further is unsafe.
-      return SDValue();
+  APInt NewToMask, NewFromMask;
+  SDValue NewFrom = ParseBFI(V.getNode(), NewToMask, NewFromMask);
+  if (NewFrom != From)
 
-    // Are the new bits contiguous when combined with the old bits?
-    if (BitsProperlyConcatenate(ToMask, NewToMask) &&
-        BitsProperlyConcatenate(FromMask, NewFromMask))
-      return V;
-    if (BitsProperlyConcatenate(NewToMask, ToMask) &&
-        BitsProperlyConcatenate(NewFromMask, FromMask))
-      return V;
+  // Do the written bits conflict with any we've seen so far?
+  if ((NewToMask & ToMask).getBoolValue())
+    // Conflicting bits.
+    return SDValue();
 
-    // We've seen a write to some bits, so track it.
-    CombinedToMask |= NewToMask;
-    // Keep going...
-    V = V.getOperand(0);
-  }
+  // Are the new bits contiguous when combined with the old bits?
+  if (BitsProperlyConcatenate(ToMask, NewToMask) &&
+      BitsProperlyConcatenate(FromMask, NewFromMask))
+    return V;
+  if (BitsProperlyConcatenate(NewToMask, ToMask) &&
+      BitsProperlyConcatenate(NewFromMask, FromMask))
+    return V;
 
   return SDValue();
 }
 
 static SDValue PerformBFICombine(SDNode *N,
                                  TargetLowering::DAGCombinerInfo &DCI) {
   SDValue N1 = N->getOperand(1);
   if (N1.getOpcode() == ISD::AND) {

                static_cast<unsigned>(std::numeric_limits<unsigned>::digits) &&
            "undefined behavior");
     unsigned Mask = (1u << Width) - 1;
     unsigned Mask2 = N11C->getZExtValue();
     if ((Mask & (~Mask2)) == 0)
       return DCI.DAG.getNode(ARMISD::BFI, SDLoc(N), N->getValueType(0),
                              N->getOperand(0), N1.getOperand(0),
                              N->getOperand(2));
-  } else if (N->getOperand(0).getOpcode() == ARMISD::BFI) {
-    // We have a BFI of a BFI. Walk up the BFI chain to see how long it goes.
-    // Keep track of any consecutive bits set that all come from the same base
-    // value. We can combine these together into a single BFI.
-    SDValue CombineBFI = FindBFIToCombineWith(N);
-    if (CombineBFI == SDValue())
-      return SDValue();
+    return SDValue();
+  }
+  // Look for another BFI to combine with.
+  SDValue CombineBFI = FindBFIToCombineWith(N);
+  if (CombineBFI == SDValue())
+    return SDValue();
 
-    // We've found a BFI.
-    APInt ToMask1, FromMask1;
-    SDValue From1 = ParseBFI(N, ToMask1, FromMask1);
 
-    APInt ToMask2, FromMask2;
-    SDValue From2 = ParseBFI(CombineBFI.getNode(), ToMask2, FromMask2);
-    assert(From1 == From2);
-    (void)From2;
 
-    // First, unlink CombineBFI.
-    DCI.DAG.ReplaceAllUsesWith(CombineBFI, CombineBFI.getOperand(0));
-    // Then create a new BFI, combining the two together.
-    APInt NewFromMask = FromMask1 | FromMask2;
-    APInt NewToMask = ToMask1 | ToMask2;
+  // Create a new BFI, combining the two together.
+  APInt NewFromMask = FromMask1 | FromMask2;
+  APInt NewToMask = ToMask1 | ToMask2;
 
-    EVT VT = N->getValueType(0);
-    SDLoc dl(N);
 
-    if (NewFromMask[0] == 0)
-      From1 = DCI.DAG.getNode(
         ISD::SRL, dl, VT, From1,
         DCI.DAG.getConstant(NewFromMask.countTrailingZeros(), dl, VT));
-    return DCI.DAG.getNode(ARMISD::BFI, dl, VT, N->getOperand(0), From1,
-                           DCI.DAG.getConstant(~NewToMask, dl, VT));
-  }
-  return SDValue();
+  return DCI.DAG.getNode(ARMISD::BFI, dl, VT, CombineBFI.getOperand(0), From1,
 }
 
 /// PerformVMOVRRDCombine - Target-specific dag combine xforms for
 /// ARMISD::VMOVRRD.
 static SDValue PerformVMOVRRDCombine(SDNode *N,
                                      TargetLowering::DAGCombinerInfo &DCI,
                                      const ARMSubtarget *Subtarget) {
   // vmovrrd(vmovdrr x, y) -> x,y
   SDValue InDouble = N->getOperand(0);
   if (InDouble.getOpcode() == ARMISD::VMOVDRR && Subtarget->hasFP64())
     return DCI.CombineTo(N, InDouble.getOperand(0), InDouble.getOperand(1));
 
   // vmovrrd(load f64) -> (load i32), (load i32)
   SDNode *InNode = InDouble.getNode();
   if (ISD::isNormalLoad(InNode) && InNode->hasOneUse() &&
       InNode->getValueType(0) == MVT::f64 &&
       InNode->getOperand(1).getOpcode() == ISD::FrameIndex &&
       !cast<LoadSDNode>(InNode)->isVolatile()) {
     // TODO: Should this be done for non-FrameIndex operands?
     LoadSDNode *LD = cast<LoadSDNode>(InNode);
 

llvm/test/CodeGen/ARM/bfi-chain-cse-crash.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=thumbv7s | FileCheck %s
+target datalayout = "e-m:o-p:32:32-f64:32:64-v64:32:64-v128:32:128-a:0:32-n32-S32"
+target triple = "thumbv7s-apple-ios3.1.3"
+
+define void @bfi_chain_cse_crash(i8* %0, i8 *%ptr) {
+; CHECK-LABEL: bfi_chain_cse_crash:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    ldrb r2, [r0]
+; CHECK-NEXT:    and r3, r2, #1
+; CHECK-NEXT:    lsr.w r12, r2, #3
+; CHECK-NEXT:    bfi r3, r12, #3, #1
+; CHECK-NEXT:    strb r3, [r0]
+; CHECK-NEXT:    and r0, r2, #4
+; CHECK-NEXT:    bfi r0, r12, #3, #1
+; CHECK-NEXT:    strb r0, [r1]
+; CHECK-NEXT:    bx lr
+entry:
+  %1 = load i8, i8* %0, align 1
+  %2 = and i8 %1, 1
+  %3 = select i1 false, i8 %2, i8 0
+  %4 = and i8 %1, 4
+  %5 = icmp eq i8 %4, 0
+  %6 = zext i8 %3 to i32
+  %7 = or i32 %6, 4
+  %8 = trunc i32 %7 to i8
+  %9 = select i1 %5, i8 %3, i8 %8
+  %10 = and i8 %1, 8
+  %11 = icmp eq i8 %10, 0
+  %12 = zext i8 %2 to i32
+  %13 = or i32 %12, 8
+  %14 = trunc i32 %13 to i8
+  %15 = zext i8 %9 to i32
+  %16 = or i32 %15, 8
+  %17 = trunc i32 %16 to i8
+  %18 = select i1 %11, i8 %2, i8 %14
+  %19 = select i1 %11, i8 %9, i8 %17
+  store i8 %18, i8* %0, align 1
+  store i8 %19, i8* %ptr, align 1
+  ret void
+}