[x86] try to widen 'shl' as part of LEA formation

[spatel]

Summary

The test file has pairs of tests that are logically equivalent:
https://rise4fun.com/Alive/2zQ

%t4 = and i8 %t1, 8
%t5 = zext i8 %t4 to i16
%sh = shl i16 %t5, 2
%t6 = add i16 %sh, %t0
=>
%t4 = and i8 %t1, 8
%sh2 = shl i8 %t4, 2
%z5 = zext i8 %sh2 to i16
%t6 = add i16 %z5, %t0

...so if we can fold the shift op into LEA in the 1st pattern, then we should be able to do the same in the 2nd pattern (unnecessary 'movzbl' is a separate bug I think).

We don't want to do this any sooner though because that would conflict with generic transforms that try to narrow the width of the shift.

[lebedev.ri]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1920–1921

This comment reads weird.
It is talking about the case like: https://rise4fun.com/Alive/DQd
Maybe something closer to

// The new shift must be NUW, that is, it must be safe to swap zext and shl around.

1935

What about the insertDAGNode() calls? We simply expect that they will be dropped afterwards?

[craig.topper]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1935

Should we restrict this to legal scale amounts? Can we just update the addressing mode and stop instead of calling matchAddressRecursively?

[spatel]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1920–1921

Agreed - that's not very clear, but let's make sure we're seeing the same thing.
The pattern difference I was trying to highlight looks like this:
https://rise4fun.com/Alive/K0YF

So the *old* (narrow) shift must be 'nuw' (it only shifts out zeros).

1935

Yes, limiting to legal scale factors is probably better. Not sure about stopping the match here, but I'll try some experiments.

When I originally drafted this patch (a few days ago), this was allowing combining a narrow leading add with a wide trailing add by continuing the match. Something recent seems to have changed that though, so we invert the order of an add+shl.

[spatel]

Patch updated:

1. Try to make code comment clearer.
2. Set scale/index parts of the AddressMode directly rather than recursing.
3. Remove dead node and explicitly RAUW.
4. Added some shift constant variation to the tests and a negative test.

#3 is unclear to me...some cases within here appear to do part of this node updating explicitly while others seem to leave it to the callers (or other passes?). My understanding of what's expected at this level isn't good, so what I have there right now might be the worst of all worlds. :)

[craig.topper]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1918

Don't we need to check that IndexReg hasn't already been assigned?

1943

The ZERO_EXTEND hasn't been deleted so its still still using Shl. I don't think Shl is dead.

[spatel]

Patch updated:

1. Added check to make sure that IndexReg has not already been set.
2. Fixed node insertion/deletion/rauw to match existing code patterns in foldMask*() helpers. Dead node removal should be handled by PostprocessISelDAG(), so there's no need to do that explicitly.

[craig.topper]

I believe the dead nodes will still be eligible for instruction selection and they will be pruned as dead by the loop in SelectionDAGISel::DoInstructionSelection.

LGTM

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

case ISD::AND: {
 
     // 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;
   }
+  case ISD::ZERO_EXTEND: {
+    // Try to widen a zexted shift left to the same size as its use, so we can
+    // match the shift as a scale factor.
+    if (N.getOperand(0).getOpcode() != ISD::SHL || !N.getOperand(0).hasOneUse())
+      break;
+
+    // Give up if the shift is not a valid scale factor [1,2,3].
+    SDValue Shl = N.getOperand(0);
+    auto *ShAmtC = dyn_cast<ConstantSDNode>(Shl.getOperand(1));
+    if (!ShAmtC || ShAmtC->getZExtValue() > 3 || AM.Scale != 1)
+      break;
+
+    // The narrow shift must only shift out zero bits (it must be 'nuw').
+    // That makes it safe to widen to the destination type.
+    APInt HighZeros = APInt::getHighBitsSet(Shl.getValueSizeInBits(),
+                                            ShAmtC->getZExtValue());
+    if (!CurDAG->MaskedValueIsZero(Shl.getOperand(0), HighZeros))
+      break;
+
+    // zext (shl nuw i8 %x, C) to i32 --> shl (zext i8 %x to i32), (zext C)
+    MVT VT = N.getSimpleValueType();
+    SDLoc DL(N);
+    SDValue Zext = CurDAG->getNode(ISD::ZERO_EXTEND, DL, VT, Shl.getOperand(0));
+    SDValue NewShl = CurDAG->getNode(ISD::SHL, DL, VT, Zext, Shl.getOperand(1));
+
+    // Convert the shift to scale factor.
+    AM.Scale = 1 << ShAmtC->getZExtValue();
+    AM.IndexReg = Zext;
+
+    // DAG maintenance for new/old nodes:
+    // 1. The original shift had one use (the zext), so it is dead now.
+    // 2. The new shift replaces the zext that we started with.
+    insertDAGNode(*CurDAG, Shl, Zext);
+    insertDAGNode(*CurDAG, N, NewShl);
+    CurDAG->RemoveDeadNode(Shl.getNode());
+    CurDAG->ReplaceAllUsesWith(N, NewShl);
+    return false;
+  }
   }
 
   return matchAddressBase(N, AM);
 }
 
 /// Helper for MatchAddress. Add the specified node to the
 /// specified addressing mode without any further recursion.
 bool X86DAGToDAGISel::matchAddressBase(SDValue N, X86ISelAddressMode &AM) {

llvm/test/CodeGen/X86/lea-dagdag.ll

; CHECK-NEXT:    retq
   %sh = shl i16 %t5, 2
   %t6 = add i16 %sh, %t0
   ret i16 %t6
 }
 
 define i16 @and_i8_shl_zext_add_i16(i16 %t0, i8 %t1) {
 ; CHECK-LABEL: and_i8_shl_zext_add_i16:
 ; CHECK:       # %bb.0:
+; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
 ; CHECK-NEXT:    andb $8, %sil
-; CHECK-NEXT:    shlb $2, %sil
 ; CHECK-NEXT:    movzbl %sil, %eax
-; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    leal (%rdi,%rax,4), %eax
 ; CHECK-NEXT:    # kill: def $ax killed $ax killed $eax
 ; CHECK-NEXT:    retq
   %t4 = and i8 %t1, 8
   %sh = shl i8 %t4, 2
   %t5 = zext i8 %sh to i16
   %t6 = add i16 %t5, %t0
   ret i16 %t6
 }
; CHECK-NEXT:    retq
   %sh = shl i32 %t5, 3
   %t6 = add i32 %sh, %t0
   ret i32 %t6
 }
 
 define i32 @and_i8_shl_zext_add_i32(i32 %t0, i8 %t1) {
 ; CHECK-LABEL: and_i8_shl_zext_add_i32:
 ; CHECK:       # %bb.0:
+; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
 ; CHECK-NEXT:    andb $8, %sil
-; CHECK-NEXT:    shlb $3, %sil
 ; CHECK-NEXT:    movzbl %sil, %eax
-; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    leal (%rdi,%rax,8), %eax
 ; CHECK-NEXT:    retq
   %t4 = and i8 %t1, 8
   %sh = shl i8 %t4, 3
   %t5 = zext i8 %sh to i32
   %t6 = add i32 %t5, %t0
   ret i32 %t6
 }
 
; CHECK-NEXT:    retq
   %t6 = add i64 %sh, %t0
   ret i64 %t6
 }
 
 define i64 @and_i8_shl_zext_add_i64(i64 %t0, i8 %t1) {
 ; CHECK-LABEL: and_i8_shl_zext_add_i64:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    andb $8, %sil
-; CHECK-NEXT:    addb %sil, %sil
 ; CHECK-NEXT:    movzbl %sil, %eax
-; CHECK-NEXT:    addq %rdi, %rax
+; CHECK-NEXT:    leaq (%rdi,%rax,2), %rax
 ; CHECK-NEXT:    retq
   %t4 = and i8 %t1, 8
   %sh = shl i8 %t4, 1
   %t5 = zext i8 %sh to i64
   %t6 = add i64 %t5, %t0
   ret i64 %t6
 }
 
; CHECK-NEXT:    retq
   ret i64 %t6
 }
 
 define i64 @and_i32_shl_zext_add_i64(i64 %t0, i32 %t1) {
 ; CHECK-LABEL: and_i32_shl_zext_add_i64:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    # kill: def $esi killed $esi def $rsi
 ; CHECK-NEXT:    andl $8, %esi
-; CHECK-NEXT:    leal (,%rsi,8), %eax
-; CHECK-NEXT:    addq %rdi, %rax
+; CHECK-NEXT:    leaq (%rdi,%rsi,8), %rax
 ; CHECK-NEXT:    retq
   %t4 = and i32 %t1, 8
   %sh = shl i32 %t4, 3
   %t5 = zext i32 %sh to i64
   %t6 = add i64 %t5, %t0
   ret i64 %t6
 }
 
+; Negative test - shift can't be converted to scale factor.
+
 define i64 @and_i32_zext_shl_add_i64_overshift(i64 %t0, i32 %t1) {
 ; CHECK-LABEL: and_i32_zext_shl_add_i64_overshift:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    # kill: def $esi killed $esi def $rsi
 ; CHECK-NEXT:    andl $8, %esi
 ; CHECK-NEXT:    shlq $4, %rsi
 ; CHECK-NEXT:    leaq (%rsi,%rdi), %rax
 ; CHECK-NEXT:    retq