[DAGCombiner] Fold (zext (and/or/xor (shl/shr (load x), cst), cst))

[Carrot]

Summary

In our real world application, we found the following optimization is missed in DAGCombiner

(zext (and/or/xor (shl/shr (load x), cst), cst)) -> (and/or/xor (shl/shr (zextload x), (zext cst)), (zext cst))

If the user of original zext is an add, it may enable further lea optimization on x86.

This patch add a new function CombineZExtLogicopShiftLoad to do this optimization.

[samparker]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

7463 ↗

(On Diff #138092)

Assert that N is a zext?

test/CodeGen/X86/zext-logicop-shift-load.ll

1 ↗

(On Diff #138092)

You're missing test cases for SRL and OR, as well as non constant operand to the logic op. You could also test indexed and sext loads.

[samparker]

Thanks, LGTM.

[mstorsjo]

This change caused widespread breakage when compiling a number of different projects, see PR36873. I'll revert this in the meantime.

[Carrot]

Reopened to fix pr36873.

[Carrot]

The previous revision caused PR36873. The problem is for some targets, the VT of shift constant may be larger than the other operand. So for the extended shift, we can't simply use the result VT as the constant VT, instead we should use the max of result VT and the original constant VT.

[samparker]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

7527 ↗

(On Diff #139641)

I would of thought that the actual problem was that previously you didn't use the VT from the shift at all, so I'm not convinced that you need to perform this check and select the max. In my mind, ShiftVT should be fine.

[Carrot]

Since both original ShiftVT and shift value are used, we can simply use the original shift constant node.

[Carrot]

ping

[samparker]

Sorry for the delay, LGTM. Thanks.

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     bool isSetCCEquivalent(SDValue N, SDValue &LHS, SDValue &RHS,
                            SDValue &CC) const;
     bool isOneUseSetCC(SDValue N) const;
 
     SDValue SimplifyNodeWithTwoResults(SDNode *N, unsigned LoOp,
                                          unsigned HiOp);
     SDValue CombineConsecutiveLoads(SDNode *N, EVT VT);
     SDValue CombineExtLoad(SDNode *N);
+    SDValue CombineZExtLogicopShiftLoad(SDNode *N);
     SDValue combineRepeatedFPDivisors(SDNode *N);
     SDValue combineInsertEltToShuffle(SDNode *N, unsigned InsIndex);
     SDValue ConstantFoldBITCASTofBUILD_VECTOR(SDNode *, EVT);
     SDValue BuildSDIV(SDNode *N);
     SDValue BuildSDIVPow2(SDNode *N);
     SDValue BuildUDIV(SDNode *N);
     SDValue BuildLogBase2(SDValue Op, const SDLoc &DL);
     SDValue BuildReciprocalEstimate(SDValue Op, SDNodeFlags Flags);
SDValue DAGCombiner::CombineExtLoad(SDNode *N) {
   SDValue Trunc =
       DAG.getNode(ISD::TRUNCATE, SDLoc(N0), N0.getValueType(), NewValue);
   ExtendSetCCUses(SetCCs, N0, NewValue, DL,
                   (ISD::NodeType)N->getOpcode());
   CombineTo(N0.getNode(), Trunc, NewChain);
   return SDValue(N, 0); // Return N so it doesn't get rechecked!
 }
 
+// fold (zext (and/or/xor (shl/shr (load x), cst), cst)) ->
+//      (and/or/xor (shl/shr (zextload x), (zext cst)), (zext cst))
+SDValue DAGCombiner::CombineZExtLogicopShiftLoad(SDNode *N) {
+  assert(N->getOpcode() == ISD::ZERO_EXTEND);
+  EVT VT = N->getValueType(0);
+
+  // and/or/xor
+  SDValue N0 = N->getOperand(0);
+  if (!(N0.getOpcode() == ISD::AND || N0.getOpcode() == ISD::OR ||
+        N0.getOpcode() == ISD::XOR) ||
+      N0.getOperand(1).getOpcode() != ISD::Constant ||
+      (LegalOperations && !TLI.isOperationLegal(N0.getOpcode(), VT)))
+    return SDValue();
+
+  // shl/shr
+  SDValue N1 = N0->getOperand(0);
+  if (!(N1.getOpcode() == ISD::SHL || N1.getOpcode() == ISD::SRL) ||
+      N1.getOperand(1).getOpcode() != ISD::Constant ||
+      (LegalOperations && !TLI.isOperationLegal(N1.getOpcode(), VT)))
+    return SDValue();
+
+  // load
+  if (!isa<LoadSDNode>(N1.getOperand(0)))
+    return SDValue();
+  LoadSDNode *Load = cast<LoadSDNode>(N1.getOperand(0));
+  EVT MemVT = Load->getMemoryVT();
+  if (!TLI.isLoadExtLegal(ISD::ZEXTLOAD, VT, MemVT) ||
+      Load->getExtensionType() == ISD::SEXTLOAD || Load->isIndexed())
+    return SDValue();
+
+
+  // If the shift op is SHL, the logic op must be AND, otherwise the result
+  // will be wrong.
+  if (N1.getOpcode() == ISD::SHL && N0.getOpcode() != ISD::AND)
+    return SDValue();
+
+  if (!N0.hasOneUse() || !N1.hasOneUse())
+    return SDValue();
+
+  SmallVector<SDNode*, 4> SetCCs;
+  if (!ExtendUsesToFormExtLoad(VT, N1.getNode(), N1.getOperand(0),
+                               ISD::ZERO_EXTEND, SetCCs, TLI))
+    return SDValue();
+
+  // Actually do the transformation.
+  SDValue ExtLoad = DAG.getExtLoad(ISD::ZEXTLOAD, SDLoc(Load), VT,
+                                   Load->getChain(), Load->getBasePtr(),
+                                   Load->getMemoryVT(), Load->getMemOperand());
+
+  SDLoc DL1(N1);
+  SDValue Shift = DAG.getNode(N1.getOpcode(), DL1, VT, ExtLoad,
+                              N1.getOperand(1));
+
+  APInt Mask = cast<ConstantSDNode>(N0.getOperand(1))->getAPIntValue();
+  Mask = Mask.zext(VT.getSizeInBits());
+  SDLoc DL0(N0);
+  SDValue And = DAG.getNode(N0.getOpcode(), DL0, VT, Shift,
+                            DAG.getConstant(Mask, DL0, VT));
+
+  ExtendSetCCUses(SetCCs, N1.getOperand(0), ExtLoad, SDLoc(Load),
+                  ISD::ZERO_EXTEND);
+  CombineTo(N, And);
+  if (SDValue(Load, 0).hasOneUse()) {
+    DAG.ReplaceAllUsesOfValueWith(SDValue(Load, 1), ExtLoad.getValue(1));
+  } else {
+    SDValue Trunc = DAG.getNode(ISD::TRUNCATE, SDLoc(Load),
+                                Load->getValueType(0), ExtLoad);
+    CombineTo(Load, Trunc, ExtLoad.getValue(1));
+  }
+  return SDValue(N,0); // Return N so it doesn't get rechecked!
+}
+
 /// If we're narrowing or widening the result of a vector select and the final
 /// size is the same size as a setcc (compare) feeding the select, then try to
 /// apply the cast operation to the select's operands because matching vector
 /// sizes for a select condition and other operands should be more efficient.
 SDValue DAGCombiner::matchVSelectOpSizesWithSetCC(SDNode *Cast) {
   unsigned CastOpcode = Cast->getOpcode();
   assert((CastOpcode == ISD::SIGN_EXTEND || CastOpcode == ISD::ZERO_EXTEND ||
           CastOpcode == ISD::TRUNCATE || CastOpcode == ISD::FP_EXTEND ||
if (TLI.isLoadExtLegal(ISD::ZEXTLOAD, VT, MemVT) &&
                                       LN00->getValueType(0), ExtLoad);
           CombineTo(LN00, Trunc, ExtLoad.getValue(1));
         }
         return SDValue(N,0); // Return N so it doesn't get rechecked!
       }
     }
   }
 
+  // fold (zext (and/or/xor (shl/shr (load x), cst), cst)) ->
+  //      (and/or/xor (shl/shr (zextload x), (zext cst)), (zext cst))
+  if (SDValue ZExtLoad = CombineZExtLogicopShiftLoad(N))
+    return ZExtLoad;
+
   // fold (zext (zextload x)) -> (zext (truncate (zextload x)))
   // fold (zext ( extload x)) -> (zext (truncate (zextload x)))
   if ((ISD::isZEXTLoad(N0.getNode()) || ISD::isEXTLoad(N0.getNode())) &&
       ISD::isUNINDEXEDLoad(N0.getNode()) && N0.hasOneUse()) {
     LoadSDNode *LN0 = cast<LoadSDNode>(N0);
     EVT MemVT = LN0->getMemoryVT();
     if ((!LegalOperations && !LN0->isVolatile()) ||
         TLI.isLoadExtLegal(ISD::ZEXTLOAD, VT, MemVT)) {

llvm/trunk/test/CodeGen/AArch64/zext-logic-shift-load.ll

+; RUN: llc -mtriple=aarch64-linux-gnu < %s -o - | FileCheck %s
+
+define i32 @test1(i8* %p) {
+; CHECK:       ldrb
+; CHECK-NEXT:  ubfx
+; CHECK-NEXT:  ret
+
+  %1 = load i8, i8* %p
+  %2 = lshr i8 %1, 1
+  %3 = and i8 %2, 1
+  %4 = zext i8 %3 to i32
+  ret i32 %4
+}
+

llvm/trunk/test/CodeGen/ARM/zext-logic-shift-load.ll

+; RUN: llc -mtriple=armv7-linux-gnu < %s -o - | FileCheck %s
+
+define void @test1(i8* %p, i16* %q) {
+; CHECK:       ldrb
+; CHECK-NEXT:  mov
+; CHECK-NEXT:  and
+; CHECK-NEXT:  strh
+; CHECK-NEXT:  bx
+
+  %1 = load i8, i8* %p
+  %2 = shl i8 %1, 2
+  %3 = and i8 %2, 12
+  %4 = zext i8 %3 to i16
+  store i16 %4, i16* %q
+  ret void
+}
+

llvm/trunk/test/CodeGen/X86/zext-logicop-shift-load.ll

+; RUN: llc < %s -mtriple=x86_64-unknown-unknown | FileCheck %s
+
+
+define i64 @test1(i8* %data) {
+; CHECK-LABEL: test1:
+; CHECK:       movzbl
+; CHECK-NEXT:  shlq
+; CHECK-NEXT:  andl
+; CHECK-NEXT:  retq
+entry:
+  %bf.load = load i8, i8* %data, align 4
+  %bf.clear = shl i8 %bf.load, 2
+  %0 = and i8 %bf.clear, 60
+  %mul = zext i8 %0 to i64
+  ret i64 %mul
+}
+
+define i8* @test2(i8* %data) {
+; CHECK-LABEL: test2:
+; CHECK:       movzbl
+; CHECK-NEXT:  andl
+; CHECK-NEXT:  leaq
+; CHECK-NEXT:  retq
+entry:
+  %bf.load = load i8, i8* %data, align 4
+  %bf.clear = shl i8 %bf.load, 2
+  %0 = and i8 %bf.clear, 60
+  %mul = zext i8 %0 to i64
+  %add.ptr = getelementptr inbounds i8, i8* %data, i64 %mul
+  ret i8* %add.ptr
+}
+
+; If the shift op is SHL, the logic op can only be AND.
+define i64 @test3(i8* %data) {
+; CHECK-LABEL: test3:
+; CHECK:       movb
+; CHECK-NEXT:  shlb
+; CHECK-NEXT:  xorb
+; CHECK-NEXT:  movzbl
+; CHECK-NEXT:  retq
+entry:
+  %bf.load = load i8, i8* %data, align 4
+  %bf.clear = shl i8 %bf.load, 2
+  %0 = xor i8 %bf.clear, 60
+  %mul = zext i8 %0 to i64
+  ret i64 %mul
+}
+
+define i64 @test4(i8* %data) {
+; CHECK-LABEL: test4:
+; CHECK:       movzbl
+; CHECK-NEXT:  shrq
+; CHECK-NEXT:  andl
+; CHECK-NEXT:  retq
+entry:
+  %bf.load = load i8, i8* %data, align 4
+  %bf.clear = lshr i8 %bf.load, 2
+  %0 = and i8 %bf.clear, 60
+  %1 = zext i8 %0 to i64
+  ret i64 %1
+}
+
+define i64 @test5(i8* %data) {
+; CHECK-LABEL: test5:
+; CHECK:       movzbl
+; CHECK-NEXT:  shrq
+; CHECK-NEXT:  xorq
+; CHECK-NEXT:  retq
+entry:
+  %bf.load = load i8, i8* %data, align 4
+  %bf.clear = lshr i8 %bf.load, 2
+  %0 = xor i8 %bf.clear, 60
+  %1 = zext i8 %0 to i64
+  ret i64 %1
+}
+
+define i64 @test6(i8* %data) {
+; CHECK-LABEL: test6:
+; CHECK:       movzbl
+; CHECK-NEXT:  shrq
+; CHECK-NEXT:  orq
+; CHECK-NEXT:  retq
+entry:
+  %bf.load = load i8, i8* %data, align 4
+  %bf.clear = lshr i8 %bf.load, 2
+  %0 = or i8 %bf.clear, 60
+  %1 = zext i8 %0 to i64
+  ret i64 %1
+}
+
+; Don't do the folding if the other operand isn't a constant.
+define i64 @test7(i8* %data, i8 %logop) {
+; CHECK-LABEL: test7:
+; CHECK:       movb
+; CHECK-NEXT:  shrb
+; CHECK-NEXT:  orb
+; CHECK-NEXT:  movzbl
+; CHECK-NEXT:  retq
+entry:
+  %bf.load = load i8, i8* %data, align 4
+  %bf.clear = lshr i8 %bf.load, 2
+  %0 = or i8 %bf.clear, %logop
+  %1 = zext i8 %0 to i64
+  ret i64 %1
+}
+
+; Load is folded with sext.
+define i64 @test8(i8* %data) {
+; CHECK-LABEL: test8:
+; CHECK:       movsbl
+; CHECK-NEXT:  movzwl
+; CHECK-NEXT:  shrl
+; CHECK-NEXT:  orl
+entry:
+  %bf.load = load i8, i8* %data, align 4
+  %ext = sext i8 %bf.load to i16
+  %bf.clear = lshr i16 %ext, 2
+  %0 = or i16 %bf.clear, 60
+  %1 = zext i16 %0 to i64
+  ret i64 %1
+}
+