Remove AND mask of bottom 6 bits when result is used for SystemZ shift/rotate

[colpell]

Summary

On SystemZ, shift and rotate instructions only use the bottom 6 bits of the shift/rotate amount. Therefore, if the amount is ANDed with an immediate mask that has all of the bottom 6 bits set, we can remove the AND operation entirely.

[uweigand]

I'm not sure that adding so many extra patterns in the .td file is the best way to implement this. Maybe it would be better to implement the transformation of the mask in PerformDAGCombine? This is done earlier, and might in theory expose more combination opportunities ...

[colpell]

Good point, thanks Ulrich. As suggested, I've moved the implementation to a more concise one in PerformDAGCombine.

[uweigand]

Looks good in general. However, it seems you missed the recent refactoring to the PerformDAGCombine routine (rev. 274191), so your patch won't apply. Please update to current mainline (your code should now go into a subroutine like combineShift).

As further enhancement I'm wondering if it might be useful to tweak the ANDed constant even if the AND cannot be optimized away completely. Not sure if that really makes a difference to real-world code though.

lib/Target/SystemZ/SystemZISelLowering.cpp

5064 ↗

(On Diff #62425)

In principle it ought to be possible to do that optimization even if there are multiple uses of the AND result (the AND would still stay there, just not used for the shift).

[colpell]

Update to latest trunk and remove hasOneUse check.

[uweigand]

In D21854#474231, @colpell wrote:

Update to latest trunk and remove hasOneUse check.

Ah, now it seems the CombineTo use is incorrect (this will replace *all* uses of the AND with its input value, which is incorrect if there is indeed another use). So I think we have to generate a new shift/rotate output node instead of using CombineTo here.

[colpell]

If the AND value has more than one use, generate a new shift/rotate node with the AND's first operand rather than removing the AND completely. Also add tests for this situation.

[uweigand]

LGTM, thanks!

lib/Target/SystemZ/SystemZInstrInfo.td

def : Pat<(sra (shl (i64 (anyext (i32 (z_select_ccmask 1, 0, imm32zx4:$valid,
                (i32 63)),
           (Select64 (LGHI -1), (LGHI 0), imm32zx4:$valid, imm32zx4:$cc)>;
 
 // Avoid generating 2 XOR instructions. (xor (and x, y), y) is
 // equivalent to (and (xor x, -1), y)
 def : Pat<(and (xor GR64:$x, (i64 -1)), GR64:$y),
                           (XGR GR64:$y, (NGR GR64:$y, GR64:$x))>;
 
-// Shift/rotate instructions only use the last 6 bits of the second operand
-// register, so we can safely use NILL (16 fewer bits than NILF) to only AND the
-// last 16 bits.
 // Complexity is added so that we match this before we match NILF on the AND
 // operation alone.
 let AddedComplexity = 4 in {
+  // Shift/rotate instructions only use the last 6 bits of the second operand
+  // register, so we can safely use NILL (16 fewer bits than NILF) to only AND the
+  // last 16 bits.
   def : Pat<(shl GR32:$val, (and GR32:$shift, uimm32:$imm)),
             (SLL GR32:$val, (NILL GR32:$shift, uimm32:$imm), 0)>;
 
   def : Pat<(sra GR32:$val, (and GR32:$shift, uimm32:$imm)),
             (SRA GR32:$val, (NILL GR32:$shift, uimm32:$imm), 0)>;
 
   def : Pat<(srl GR32:$val, (and GR32:$shift, uimm32:$imm)),
             (SRL GR32:$val, (NILL GR32:$shift, uimm32:$imm), 0)>;
let AddedComplexity = 4 in {
 
   def : Pat<(rotl GR32:$val, (and GR32:$shift, uimm32:$imm)),
             (RLL GR32:$val, (NILL GR32:$shift, uimm32:$imm), 0)>;
 
   def : Pat<(rotl GR64:$val, (and GR32:$shift, uimm32:$imm)),
             (RLLG GR64:$val, (NILL GR32:$shift, uimm32:$imm), 0)>;
 }
 
+let AddedComplexity = 5 in {
+  // If an AND operation doesn't change the last 6 bits before the result is
+  // used for a shift/rotate, we can remove it entirely.
+  def : Pat<(shl GR32:$val, (and GR32:$shift, imm32bottom6)),
+            (SLL GR32:$val, GR32:$shift, 0)>;
+
+  def : Pat<(sra GR32:$val, (and GR32:$shift, imm32bottom6)),
+            (SRA GR32:$val, GR32:$shift, 0)>;
+
+  def : Pat<(srl GR32:$val, (and GR32:$shift, imm32bottom6)),
+            (SRL GR32:$val, GR32:$shift, 0)>;
+
+  def : Pat<(shl GR64:$val, (and GR32:$shift, imm32bottom6)),
+            (SLLG GR64:$val, GR32:$shift, 0)>;
+
+  def : Pat<(sra GR64:$val, (and GR32:$shift, imm32bottom6)),
+            (SRAG GR64:$val, GR32:$shift, 0)>;
+
+  def : Pat<(srl GR64:$val, (and GR32:$shift, imm32bottom6)),
+            (SRLG GR64:$val, GR32:$shift, 0)>;
+
+  def : Pat<(rotl GR32:$val, (and GR32:$shift, imm32bottom6)),
+            (RLL GR32:$val, GR32:$shift, 0)>;
+
+  def : Pat<(rotl GR64:$val, (and GR32:$shift, imm32bottom6)),
+            (RLLG GR64:$val, GR32:$shift, 0)>;
+}
+
 // Peepholes for turning scalar operations into block operations.
 defm : BlockLoadStore<anyextloadi8, i32, MVCSequence, NCSequence, OCSequence,
                       XCSequence, 1>;
 defm : BlockLoadStore<anyextloadi16, i32, MVCSequence, NCSequence, OCSequence,
                       XCSequence, 2>;
 defm : BlockLoadStore<load, i32, MVCSequence, NCSequence, OCSequence,
                       XCSequence, 4>;
 defm : BlockLoadStore<anyextloadi8, i64, MVCSequence, NCSequence,
                       OCSequence, XCSequence, 1>;
 defm : BlockLoadStore<anyextloadi16, i64, MVCSequence, NCSequence, OCSequence,
                       XCSequence, 2>;
 defm : BlockLoadStore<anyextloadi32, i64, MVCSequence, NCSequence, OCSequence,
                       XCSequence, 4>;
 defm : BlockLoadStore<load, i64, MVCSequence, NCSequence, OCSequence,
                       XCSequence, 8>;

lib/Target/SystemZ/SystemZOperands.td

 }], SIMM16, "S16Imm">;
 
 def imm32zx16 : Immediate<i32, [{
   return isUInt<16>(N->getZExtValue());
 }], UIMM16, "U16Imm">;
 
 def imm32sx16trunc : Immediate<i32, [{}], SIMM16, "S16Imm">;
 
+// Bottom 6 bits are all set
+def imm32bottom6 : Immediate<i32, [{
+  return (N->getZExtValue() & 0x3f) == 0x3f;
+}], NOOP_SDNodeXForm, "U32Imm">;
+
 // Full 32-bit immediates.  we need both signed and unsigned versions
 // because the assembler is picky.  E.g. AFI requires signed operands
 // while NILF requires unsigned ones.
 def simm32 : Immediate<i32, [{}], SIMM32, "S32Imm">;
 def uimm32 : Immediate<i32, [{}], UIMM32, "U32Imm">;
 
 def imm32 : ImmLeaf<i32, [{}]>;
 

test/CodeGen/SystemZ/rot-01.ll

; CHECK: rll %r2, %r2, 0(%r3)
   %rotl = or i32 %parta, %partb
 
   ret i32 %rotl
 }
 
 ; Test 64-bit rotate.
 define i64 @f2(i64 %val, i64 %amt) {
 ; CHECK-LABEL: f2:
-; CHECK: nill %r3, 63
+; CHECK: nill %r3, 31
 ; CHECK: rllg %r2, %r2, 0(%r3)
-  %mod = urem i64 %amt, 64
+  %mod = urem i64 %amt, 32
 
   %inv = sub i64 64, %mod
   %parta = shl i64 %val, %mod
   %partb = lshr i64 %val, %inv
 
   %rotl = or i64 %parta, %partb
 
   ret i64 %rotl
 }

test/CodeGen/SystemZ/rot-02.ll

+; Test removal of AND operations that don't affect last 6 bits of rotate amount
+; operand.
+;
+; RUN: llc < %s -mtriple=s390x-linux-gnu | FileCheck %s
+
+; Test that AND is not removed when some lower 6 bits are not set.
+define i32 @f1(i32 %val, i32 %amt) {
+; CHECK-LABEL: f1:
+; CHECK: nil{{[lf]}} %r3, 31
+; CHECK: rll %r2, %r2, 0(%r3)
+  %and = and i32 %amt, 31
+
+  %inv = sub i32 32, %and
+  %parta = shl i32 %val, %and
+  %partb = lshr i32 %val, %inv
+
+  %rotl = or i32 %parta, %partb
+
+  ret i32 %rotl
+}
+
+; Test removal of AND mask with only bottom 6 bits set.
+define i32 @f2(i32 %val, i32 %amt) {
+; CHECK-LABEL: f2:
+; CHECK-NOT: nil{{[lf]}} %r3, 63
+; CHECK: rll %r2, %r2, 0(%r3)
+  %and = and i32 %amt, 63
+
+  %inv = sub i32 32, %and
+  %parta = shl i32 %val, %and
+  %partb = lshr i32 %val, %inv
+
+  %rotl = or i32 %parta, %partb
+
+  ret i32 %rotl
+}
+
+; Test removal of AND mask including but not limited to bottom 6 bits.
+define i32 @f3(i32 %val, i32 %amt) {
+; CHECK-LABEL: f3:
+; CHECK-NOT: nil{{[lf]}} %r3, 255
+; CHECK: rll %r2, %r2, 0(%r3)
+  %and = and i32 %amt, 255
+
+  %inv = sub i32 32, %and
+  %parta = shl i32 %val, %and
+  %partb = lshr i32 %val, %inv
+
+  %rotl = or i32 %parta, %partb
+
+  ret i32 %rotl
+}
+
+; Test removal of AND mask from RLLG.
+define i64 @f4(i64 %val, i64 %amt) {
+; CHECK-LABEL: f4:
+; CHECK-NOT: nil{{[lf]}} %r3, 63
+; CHECK: rllg %r2, %r2, 0(%r3)
+  %and = and i64 %amt, 63
+
+  %inv = sub i64 64, %and
+  %parta = shl i64 %val, %and
+  %partb = lshr i64 %val, %inv
+
+  %rotl = or i64 %parta, %partb
+
+  ret i64 %rotl
+}

test/CodeGen/SystemZ/shift-12.ll

+; Test removal of AND operations that don't affect last 6 bits of shift amount
+; operand.
+;
+; RUN: llc < %s -mtriple=s390x-linux-gnu | FileCheck %s
+
+; Test that AND is not removed when some lower 6 bits are not set.
+define i32 @f1(i32 %a, i32 %sh) {
+; CHECK-LABEL: f1:
+; CHECK: nil{{[lf]}} %r3, 31
+; CHECK: sll %r2, 0(%r3)
+  %and = and i32 %sh, 31
+  %shift = shl i32 %a, %and
+  ret i32 %shift
+}
+
+; Test removal of AND mask with only bottom 6 bits set.
+define i32 @f2(i32 %a, i32 %sh) {
+; CHECK-LABEL: f2:
+; CHECK-NOT: nil{{[lf]}} %r3, 63
+; CHECK: sll %r2, 0(%r3)
+  %and = and i32 %sh, 63
+  %shift = shl i32 %a, %and
+  ret i32 %shift
+}
+
+; Test removal of AND mask including but not limited to bottom 6 bits.
+define i32 @f3(i32 %a, i32 %sh) {
+; CHECK-LABEL: f3:
+; CHECK-NOT: nil{{[lf]}} %r3, 255
+; CHECK: sll %r2, 0(%r3)
+  %and = and i32 %sh, 255
+  %shift = shl i32 %a, %and
+  ret i32 %shift
+}
+
+; Test removal of AND mask from SRA.
+define i32 @f4(i32 %a, i32 %sh) {
+; CHECK-LABEL: f4:
+; CHECK-NOT: nil{{[lf]}} %r3, 63
+; CHECK: sra %r2, 0(%r3)
+  %and = and i32 %sh, 63
+  %shift = ashr i32 %a, %and
+  ret i32 %shift
+}
+
+; Test removal of AND mask from SRL.
+define i32 @f5(i32 %a, i32 %sh) {
+; CHECK-LABEL: f5:
+; CHECK-NOT: nil{{[lf]}} %r3, 63
+; CHECK: srl %r2, 0(%r3)
+  %and = and i32 %sh, 63
+  %shift = lshr i32 %a, %and
+  ret i32 %shift
+}
+
+; Test removal of AND mask from SLLG.
+define i64 @f6(i64 %a, i64 %sh) {
+; CHECK-LABEL: f6:
+; CHECK-NOT: nil{{[lf]}} %r3, 63
+; CHECK: sllg %r2, %r2, 0(%r3)
+  %and = and i64 %sh, 63
+  %shift = shl i64 %a, %and
+  ret i64 %shift
+}
+
+; Test removal of AND mask from SRAG.
+define i64 @f7(i64 %a, i64 %sh) {
+; CHECK-LABEL: f7:
+; CHECK-NOT: nil{{[lf]}} %r3, 63
+; CHECK: srag %r2, %r2, 0(%r3)
+  %and = and i64 %sh, 63
+  %shift = ashr i64 %a, %and
+  ret i64 %shift
+}
+
+; Test removal of AND mask from SRLG.
+define i64 @f8(i64 %a, i64 %sh) {
+; CHECK-LABEL: f8:
+; CHECK-NOT: nil{{[lf]}} %r3, 63
+; CHECK: srlg %r2, %r2, 0(%r3)
+  %and = and i64 %sh, 63
+  %shift = lshr i64 %a, %and
+  ret i64 %shift
+}