[AArch64] Fix useful bits detection for BFM instructions

[sbaranga]

Summary

When computing useful bits for a BFM instruction, we need
to take into consideration the case where both operands
of the BFM are equal and provide data that we need to track.

Not doing this can cause us to miss useful bits.

Fixes PR31138 (https://llvm.org/bugs/show_bug.cgi?id=31138)

[srhines]

This is looking great. We already verified that the patch resolves the original bug that was reported.

test/CodeGen/AArch64/arm64-bitfield-extract.ll

544

Should all these values be renamed to something more clear? For example, this is probably better as %shr47. Several other long value names below could similarly be improved.

[sbaranga]

Update variable names names in the test and make it more readable overall.

[srhines]

Thanks for cleaning up the names in the test. I would prefer someone more familiar with this backend to approve it, but everything looks great now.

[jmolloy]

LGTM too. Thanks!

[sbaranga]

Thanks, both! r288253

-Silviu

lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

 
 static void getUsefulBitsFromBFM(SDValue Op, SDValue Orig, APInt &UsefulBits,
                                  unsigned Depth) {
   uint64_t Imm =
       cast<const ConstantSDNode>(Op.getOperand(2).getNode())->getZExtValue();
   uint64_t MSB =
       cast<const ConstantSDNode>(Op.getOperand(3).getNode())->getZExtValue();
 
-  if (Op.getOperand(1) == Orig)
-    return getUsefulBitsFromBitfieldMoveOpd(Op, UsefulBits, Imm, MSB, Depth);
-
   APInt OpUsefulBits(UsefulBits);
   OpUsefulBits = 1;
 
+  APInt ResultUsefulBits(UsefulBits.getBitWidth(), 0);
+  ResultUsefulBits.flipAllBits();
+  APInt Mask(UsefulBits.getBitWidth(), 0);
+
+  getUsefulBits(Op, ResultUsefulBits, Depth + 1);
+
   if (MSB >= Imm) {
-    OpUsefulBits = OpUsefulBits.shl(MSB - Imm + 1);
+    // The instruction is a BFXIL.
+    uint64_t Width = MSB - Imm + 1;
+    uint64_t LSB = Imm;
+
+    OpUsefulBits = OpUsefulBits.shl(Width);
     --OpUsefulBits;
-    UsefulBits &= ~OpUsefulBits;
-    getUsefulBits(Op, UsefulBits, Depth + 1);
+
+    if (Op.getOperand(1) == Orig) {
+      // Copy the low bits from the result to bits starting from LSB.
+      Mask = ResultUsefulBits & OpUsefulBits;
+      Mask = Mask.shl(LSB);
+    }
+
+    if (Op.getOperand(0) == Orig)
+      // Bits starting from LSB in the input contribute to the result.
+      Mask |= (ResultUsefulBits & ~OpUsefulBits);
   } else {
-    OpUsefulBits = OpUsefulBits.shl(MSB + 1);
+    // The instruction is a BFI.
+    uint64_t Width = MSB + 1;
+    uint64_t LSB = UsefulBits.getBitWidth() - Imm;
+
+    OpUsefulBits = OpUsefulBits.shl(Width);
     --OpUsefulBits;
-    UsefulBits = ~(OpUsefulBits.shl(OpUsefulBits.getBitWidth() - Imm));
-    getUsefulBits(Op, UsefulBits, Depth + 1);
+    OpUsefulBits = OpUsefulBits.shl(LSB);
+
+    if (Op.getOperand(1) == Orig) {
+      // Copy the bits from the result to the zero bits.
+      Mask = ResultUsefulBits & OpUsefulBits;
+      Mask = Mask.lshr(LSB);
     }
+
+    if (Op.getOperand(0) == Orig)
+      Mask |= (ResultUsefulBits & ~OpUsefulBits);
+  }
+
+  UsefulBits &= Mask;
 }
 
 static void getUsefulBitsForUse(SDNode *UserNode, APInt &UsefulBits,
                                 SDValue Orig, unsigned Depth) {
 
   // Users of this node should have already been instruction selected
   // FIXME: Can we turn that into an assert?
   if (!UserNode->isMachineOpcode())

test/CodeGen/AArch64/arm64-bitfield-extract.ll

; CHECK: {{bfm|bfi|bfxil}}
 
   %shl16 = shl i32 %insertion, 8
   %or18 = or i32 %shl16, %insertion
   %conv19 = trunc i32 %or18 to i16
 ; CHECK: bfi {{w[0-9]+}}, {{w[0-9]+}}, #8, #7
 
   ret i16 %conv19
 }
+
+; The following test excercises the case where we have a BFI
+; instruction with the same input in both operands. We need to
+; track the useful bits through both operands.
+; CHECK-LABEL: sameOperandBFI
+; CHECK: lsr
+; CHECK: and
+; CHECK: bfi
+; CHECK: bfi
+define void @sameOperandBFI(i64 %src, i64 %src2, i16 *%ptr) {
+entry:
+  %shr47 = lshr i64 %src, 47
+  %src2.trunc = trunc i64 %src2 to i32
+  br i1 undef, label %end, label %if.else
+
+if.else:
+  %and3 = and i32 %src2.trunc, 3
+  %shl2 = shl nuw nsw i64 %shr47, 2
+  %shl2.trunc = trunc i64 %shl2 to i32
+  %and12 = and i32 %shl2.trunc, 12
+  %BFISource = or i32 %and3, %and12         ; ...00000ABCD
+  %BFIRHS = shl nuw nsw i32 %BFISource, 4   ; ...0ABCD0000
+  %BFI = or i32 %BFIRHS, %BFISource         ; ...0ABCDABCD
+  %BFItrunc = trunc i32 %BFI to i16
+  store i16 %BFItrunc, i16* %ptr, align 4
+  br label %end
+
+end:
+  ret void
+}