[SystemZ] Recognize RISBG opportunities involving a truncate

[RolandF]

Summary

Recognize RISBG opportunities where the end result is narrower than the original input - where a truncate separates the shift/and operations.

The motivating case is some code in postgres which looks like:

srlg %r2, %r0, 11
nilh %r2, 255

[uweigand]

Please also update the comment before "struct RxSBGOperands", where it explains:

The output value has BitSize bits, although Input may be
narrower (in which case the upper bits are don't care).

to now also include the possibility that Input is wider than BitSize bits.

Finally, please extend the test cases in test/CodeGen/SystemZ/r?sbg-*.ll to also cover this case (both good cases where a truncation can be eliminated, and bad cases where it cannot).

Thanks for working on improving the SystemZ back end!

lib/Target/SystemZ/SystemZISelDAGToDAG.cpp

748 ↗

(On Diff #61046)

I think we need the same guard against RNSBG here that we have in the AND case.

928 ↗

(On Diff #61046)

Can you explain why you think this is needed? (And if it is indeed needed here, why not below in tryRxSBG?)

[RolandF]

Added tests.
Added check for RNSBG.
Added TRUNCATE profitability check to tryRxSBG.
Added comment to TRUNCATE profitability check.
Updated comment on struct RxSBGOperands.

[RolandF]

lib/Target/SystemZ/SystemZISelDAGToDAG.cpp

928 ↗

(On Diff #61046)

The existing approach is to prefer a simple and/shift/rotate operation over an R*SBG operation if it will do the same job. Since there are both 32-bit and 64-bit ands/shifts/rotates, truncation is typically free. Counting truncation as an operation for the profitability check will result in preferring an R*SBG operation over its simple equivalent. I can add a comment.

And yes, I missed the check in tryRxSBG. It should check for truncate there as well.

[uweigand]

LGTM, thanks!

llvm/trunk/lib/Target/SystemZ/SystemZISelDAGToDAG.cpp

 //
 //   (or (rotl Input, Rotate), ~Mask)
 //
 // for RNSBG and:
 //
 //   (and (rotl Input, Rotate), Mask)
 //
 // otherwise.  The output value has BitSize bits, although Input may be
-// narrower (in which case the upper bits are don't care).
+// narrower (in which case the upper bits are don't care), or wider (in which
+// case the result will be truncated as part of the operation).
 struct RxSBGOperands {
   RxSBGOperands(unsigned Op, SDValue N)
     : Opcode(Op), BitSize(N.getValueType().getSizeInBits()),
       Mask(allOnes(BitSize)), Input(N), Start(64 - BitSize), End(63),
       Rotate(0) {}
 
   unsigned Opcode;
   unsigned BitSize;
if (RxSBG.Rotate != 0)
     Mask = ((Mask << RxSBG.Rotate) | (Mask >> (64 - RxSBG.Rotate)));
   return (Mask & RxSBG.Mask) != 0;
 }
 
 bool SystemZDAGToDAGISel::expandRxSBG(RxSBGOperands &RxSBG) const {
   SDValue N = RxSBG.Input;
   unsigned Opcode = N.getOpcode();
   switch (Opcode) {
+  case ISD::TRUNCATE: {
+    if (RxSBG.Opcode == SystemZ::RNSBG)
+      return false;
+    uint64_t BitSize = N.getValueType().getSizeInBits();
+    uint64_t Mask = allOnes(BitSize);
+    if (!refineRxSBGMask(RxSBG, Mask))
+      return false;
+    RxSBG.Input = N.getOperand(0);
+    return true;
+  }
   case ISD::AND: {
     if (RxSBG.Opcode == SystemZ::RNSBG)
       return false;
 
     auto *MaskNode = dyn_cast<ConstantSDNode>(N.getOperand(1).getNode());
     if (!MaskNode)
       return false;
 

 bool SystemZDAGToDAGISel::tryRISBGZero(SDNode *N) {
   SDLoc DL(N);
   EVT VT = N->getValueType(0);
   if (!VT.isInteger() || VT.getSizeInBits() > 64)
     return false;
   RxSBGOperands RISBG(SystemZ::RISBG, SDValue(N, 0));
   unsigned Count = 0;
   while (expandRxSBG(RISBG))
-    if (RISBG.Input.getOpcode() != ISD::ANY_EXTEND)
+    // The widening or narrowing is expected to be free.
+    // Counting widening or narrowing as a saved operation will result in
+    // preferring an R*SBG over a simple shift/logical instruction.
+    if (RISBG.Input.getOpcode() != ISD::ANY_EXTEND &&
+        RISBG.Input.getOpcode() != ISD::TRUNCATE)
       Count += 1;
   if (Count == 0)
     return false;
   if (Count == 1) {
     // Prefer to use normal shift instructions over RISBG, since they can handle
     // all cases and are sometimes shorter.
     if (N->getOpcode() != ISD::AND)
       return false;
bool SystemZDAGToDAGISel::tryRxSBG(SDNode *N, unsigned Opcode) {
   // and see which goes deepest.
   RxSBGOperands RxSBG[] = {
     RxSBGOperands(Opcode, N->getOperand(0)),
     RxSBGOperands(Opcode, N->getOperand(1))
   };
   unsigned Count[] = { 0, 0 };
   for (unsigned I = 0; I < 2; ++I)
     while (expandRxSBG(RxSBG[I]))
-      if (RxSBG[I].Input.getOpcode() != ISD::ANY_EXTEND)
+      // The widening or narrowing is expected to be free.
+      // Counting widening or narrowing as a saved operation will result in
+      // preferring an R*SBG over a simple shift/logical instruction.
+      if (RxSBG[I].Input.getOpcode() != ISD::ANY_EXTEND &&
+          RxSBG[I].Input.getOpcode() != ISD::TRUNCATE)
         Count[I] += 1;
 
   // Do nothing if neither operand is suitable.
   if (Count[0] == 0 && Count[1] == 0)
     return false;
 
   // Pick the deepest second operand.
   unsigned I = Count[0] > Count[1] ? 0 : 1;

llvm/trunk/test/CodeGen/SystemZ/risbg-01.ll

 ; CHECK: sll %r2, 31
 ; CHECK: sra %r2, 31
 ; CHECK: llgcr %r2, %r2
 ; CHECK: br %r14
   %ext = sext i1 %x to i8
   %ext2 = zext i8 %ext to i64
   ret i64 %ext2
 }
+
+; Check that we get the case where a 64-bit shift is used by a 32-bit and.
+define signext i32 @f43(i64 %x) {
+; CHECK-LABEL: f43:
+; CHECK: risbg [[REG:%r[0-5]]], %r2, 32, 189, 52
+; CHECK: lgfr %r2, [[REG]]
+  %shr3 = lshr i64 %x, 12
+  %shr3.tr = trunc i64 %shr3 to i32
+  %conv = and i32 %shr3.tr, -4
+  ret i32 %conv
+}
+
+; Check that we don't get the case where the 32-bit and mask is not contiguous
+define signext i32 @f44(i64 %x) {
+; CHECK-LABEL: f44:
+; CHECK: srlg [[REG:%r[0-5]]], %r2, 12
+  %shr4 = lshr i64 %x, 12
+  %conv = trunc i64 %shr4 to i32
+  %and = and i32 %conv, 10
+  ret i32 %and
+}

llvm/trunk/test/CodeGen/SystemZ/risbg-02.ll

 ; CHECK-LABEL: f8:
 ; CHECK: risbg %r2, %r3, 8, 63, 56
 ; CHECK: br %r14
   %anda = and i64 %a, -72057594037927936
   %shrb = lshr i64 %b, 8
   %or = or i64 %anda, %shrb
   ret i64 %or
 }
+
+; Check that we can get the case where a 64-bit shift feeds a 32-bit or of
+; ands with complement masks.
+define signext i32 @f9(i64 %x, i32 signext %y) {
+; CHECK-LABEL: f9:
+; CHECK: risbg [[REG:%r[0-5]]], %r2, 48, 63, 16
+; CHECK: lgfr %r2, [[REG]]
+  %shr6 = lshr i64 %x, 48
+  %conv = trunc i64 %shr6 to i32
+  %and1 = and i32 %y, -65536
+  %or = or i32 %conv, %and1
+  ret i32 %or
+}
+
+; Check that we don't get the case where a 64-bit shift feeds a 32-bit or of
+; ands with incompatible masks.
+define signext i32 @f10(i64 %x, i32 signext %y) {
+; CHECK-LABEL: f10:
+; CHECK: nilf %r3, 4278190080
+  %shr6 = lshr i64 %x, 48
+  %conv = trunc i64 %shr6 to i32
+  %and1 = and i32 %y, -16777216
+  %or = or i32 %conv, %and1
+  ret i32 %or
+}