[X86] Teach foldMaskedShiftToScaledMask to look throw an any_extend from i32 to i64 between the and & shl

[craig.topper]

Summary

foldMaskedShiftToScaledMask tries to reorder and & shl to enable the shl to fold into an LEA. But if there is an any_extend between them it doesn't work.

This patch modifies the code to look through any_extend from i32 to i64 when the and mask only uses bits that weren't from the extended part.

This will prevent a regression from D60358 caused by 64-bit SHL being narrowed to 32-bits when their upper bits aren't demanded.

[craig.topper]

Upload the right patch. The other was partial

[craig.topper]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1877 ↗

(On Diff #194569)

This is a little hard to see here, but 3 of the functions look for SRL and one looks for SHL. But we used to check for either SHL or SRL earlier. Now that we also need to consider ANY_EXTEND. I've removed the early check for SHL/SHR. Covered the 3 functions that use SRL with a check for that. Each of the functions all check ISD::SRL internally, but expect the input to be passed as X. This could probably be cleaned up a little though one of the functions does have 2 callers.

The SHL function is now unqualified and will check for ANY_EXTEND/SHL and create X internally.

[spatel]

LGTM - thanks!

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1389 ↗

(On Diff #194569)

We have -> If we have

1391 ↗

(On Diff #194569)

behind -> beside?

llvm/test/CodeGen/X86/fold-and-shift-x86_64.ll

79 ↗

(On Diff #194569)

Remove stale comment

llvm/trunk/lib/Target/X86/X86ISelDAGToDAG.cpp

static bool foldMaskAndShiftToExtract(SelectionDAG &DAG, SDValue N,
   AM.Scale = (1 << ScaleLog);
   return false;
 }
 
 // Transforms "(X << C1) & C2" to "(X & (C2>>C1)) << C1" if safe and if this
 // allows us to fold the shift into this addressing mode. Returns false if the
 // transform succeeded.
 static bool foldMaskedShiftToScaledMask(SelectionDAG &DAG, SDValue N,
-                                        SDValue Shift, SDValue X,
                                         X86ISelAddressMode &AM) {
+  SDValue Shift = N.getOperand(0);
+
+  // Use a signed mask so that shifting right will insert sign bits. These
+  // bits will be removed when we shift the result left so it doesn't matter
+  // what we use. This might allow a smaller immediate encoding.
+  int64_t Mask = cast<ConstantSDNode>(N->getOperand(1))->getSExtValue();
+
+  // If we have an any_extend feeding the AND, look through it to see if there
+  // is a shift behind it. But only if the AND doesn't use the extended bits.
+  // FIXME: Generalize this to other ANY_EXTEND than i32 to i64?
+  bool FoundAnyExtend = false;
+  if (Shift.getOpcode() == ISD::ANY_EXTEND && Shift.hasOneUse() &&
+      Shift.getOperand(0).getSimpleValueType() == MVT::i32 &&
+      isUInt<32>(Mask)) {
+    FoundAnyExtend = true;
+    Shift = Shift.getOperand(0);
+  }
+
   if (Shift.getOpcode() != ISD::SHL ||
       !isa<ConstantSDNode>(Shift.getOperand(1)))
     return true;
 
+  SDValue X = Shift.getOperand(0);
+
   // Not likely to be profitable if either the AND or SHIFT node has more
   // than one use (unless all uses are for address computation). Besides,
   // isel mechanism requires their node ids to be reused.
   if (!N.hasOneUse() || !Shift.hasOneUse())
     return true;
 
   // Verify that the shift amount is something we can fold.
   unsigned ShiftAmt = Shift.getConstantOperandVal(1);
   if (ShiftAmt != 1 && ShiftAmt != 2 && ShiftAmt != 3)
     return true;
 
-  // Use a signed mask so that shifting right will insert sign bits. These
-  // bits will be removed when we shift the result left so it doesn't matter
-  // what we use. This might allow a smaller immediate encoding.
-  int64_t Mask = cast<ConstantSDNode>(N->getOperand(1))->getSExtValue();
-
   MVT VT = N.getSimpleValueType();
   SDLoc DL(N);
+  if (FoundAnyExtend) {
+    SDValue NewX = DAG.getNode(ISD::ANY_EXTEND, DL, VT, X);
+    insertDAGNode(DAG, N, NewX);
+    X = NewX;
+  }
+
   SDValue NewMask = DAG.getConstant(Mask >> ShiftAmt, DL, VT);
   SDValue NewAnd = DAG.getNode(ISD::AND, DL, VT, X, NewMask);
   SDValue NewShift = DAG.getNode(ISD::SHL, DL, VT, NewAnd, Shift.getOperand(1));
 
   // Insert the new nodes into the topological ordering. We must do this in
   // a valid topological ordering as nothing is going to go back and re-sort
   // these nodes. We continually insert before 'N' in sequence as this is
   // essentially a pre-flattened and pre-sorted sequence of nodes. There is no
case ISD::AND: {
     // Scale must not be used already.
     if (AM.IndexReg.getNode() != nullptr || AM.Scale != 1) break;
 
     // We only handle up to 64-bit values here as those are what matter for
     // addressing mode optimizations.
     assert(N.getSimpleValueType().getSizeInBits() <= 64 &&
            "Unexpected value size!");
 
+    if (!isa<ConstantSDNode>(N.getOperand(1)))
+      break;
+
+    if (N.getOperand(0).getOpcode() == ISD::SRL) {
       SDValue Shift = N.getOperand(0);
-    if (Shift.getOpcode() != ISD::SRL && Shift.getOpcode() != ISD::SHL) break;
       SDValue X = Shift.getOperand(0);
 
-    if (!isa<ConstantSDNode>(N.getOperand(1)))
-      break;
       uint64_t Mask = N.getConstantOperandVal(1);
 
       // Try to fold the mask and shift into an extract and scale.
       if (!foldMaskAndShiftToExtract(*CurDAG, N, Mask, Shift, X, AM))
         return false;
 
       // Try to fold the mask and shift directly into the scale.
       if (!foldMaskAndShiftToScale(*CurDAG, N, Mask, Shift, X, AM))
         return false;
 
-    // Try to swap the mask and shift to place shifts which can be done as
-    // a scale on the outside of the mask.
-    if (!foldMaskedShiftToScaledMask(*CurDAG, N, Shift, X, AM))
-      return false;
-
       // Try to fold the mask and shift into BEXTR and scale.
       if (!foldMaskedShiftToBEXTR(*CurDAG, N, Mask, Shift, X, AM, *Subtarget))
         return false;
+    }
+
+    // Try to swap the mask and shift to place shifts which can be done as
+    // a scale on the outside of the mask.
+    if (!foldMaskedShiftToScaledMask(*CurDAG, N, AM))
+      return false;
 
     break;
   }
   }
 
   return matchAddressBase(N, AM);
 }
 

llvm/trunk/test/CodeGen/X86/fold-and-shift-x86_64.ll

 entry:
   %tmp2 = shl i64 %i, 2
   %tmp4 = and i64 %tmp2, 17178869176
   %tmp7 = getelementptr i8, i8* %X, i64 %tmp4
   %tmp9 = load i8, i8* %tmp7
   ret i8 %tmp9
 }
 
-; FIXME should be able to fold shift into address.
 define i8 @t6(i8* %X, i32 %i) {
 ; CHECK-LABEL: t6:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    # kill: def $esi killed $esi def $rsi
-; CHECK-NEXT:    shll $2, %esi
-; CHECK-NEXT:    andl $60, %esi
-; CHECK-NEXT:    movb (%rdi,%rsi), %al
+; CHECK-NEXT:    andl $15, %esi
+; CHECK-NEXT:    movb (%rdi,%rsi,4), %al
 ; CHECK-NEXT:    retq
 entry:
   %tmp2 = shl i32 %i, 2
   %tmp3 = zext i32 %tmp2 to i64
   %tmp4 = and i64 %tmp3, 60
   %tmp7 = getelementptr i8, i8* %X, i64 %tmp4
   %tmp9 = load i8, i8* %tmp7
   ret i8 %tmp9
 }