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

5115

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/SystemZISelLowering.h

SDValue combineTruncateExtract(const SDLoc &DL, EVT TruncVT, SDValue Op,
                                  DAGCombinerInfo &DCI) const;
   SDValue combineSIGN_EXTEND(SDNode *N, DAGCombinerInfo &DCI) const;
   SDValue combineMERGE(SDNode *N, DAGCombinerInfo &DCI) const;
   SDValue combineSTORE(SDNode *N, DAGCombinerInfo &DCI) const;
   SDValue combineEXTRACT_VECTOR_ELT(SDNode *N, DAGCombinerInfo &DCI) const;
   SDValue combineJOIN_DWORDS(SDNode *N, DAGCombinerInfo &DCI) const;
   SDValue combineFP_ROUND(SDNode *N, DAGCombinerInfo &DCI) const;
   SDValue combineBSWAP(SDNode *N, DAGCombinerInfo &DCI) const;
+  SDValue combineSHIFTROT(SDNode *N, DAGCombinerInfo &DCI) const;
 
   // If the last instruction before MBBI in MBB was some form of COMPARE,
   // try to replace it with a COMPARE AND BRANCH just before MBBI.
   // CCMask and Target are the BRC-like operands for the branch.
   // Return true if the change was made.
   bool convertPrevCompareToBranch(MachineBasicBlock *MBB,
                                   MachineBasicBlock::iterator MBBI,
                                   unsigned CCMask,

lib/Target/SystemZ/SystemZISelLowering.cpp

SystemZTargetLowering::SystemZTargetLowering(const TargetMachine &TM,
   setOperationAction(ISD::VAEND,   MVT::Other, Expand);
 
   // Codes for which we want to perform some z-specific combinations.
   setTargetDAGCombine(ISD::SIGN_EXTEND);
   setTargetDAGCombine(ISD::STORE);
   setTargetDAGCombine(ISD::EXTRACT_VECTOR_ELT);
   setTargetDAGCombine(ISD::FP_ROUND);
   setTargetDAGCombine(ISD::BSWAP);
+  setTargetDAGCombine(ISD::SHL);
+  setTargetDAGCombine(ISD::SRA);
+  setTargetDAGCombine(ISD::SRL);
+  setTargetDAGCombine(ISD::ROTL);
 
   // Handle intrinsics.
   setOperationAction(ISD::INTRINSIC_W_CHAIN, MVT::Other, Custom);
   setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::Other, Custom);
 
   // We want to use MVC in preference to even a single load/store pair.
   MaxStoresPerMemcpy = 0;
   MaxStoresPerMemcpyOptSize = 0;
lowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const {
   SDValue Backchain;
   if (StoreBackchain)
     Backchain = DAG.getLoad(MVT::i64, DL, Chain, OldSP, MachinePointerInfo(),
                             false, false, false, 0);
 
   // Add extra space for alignment if needed.
   if (ExtraAlignSpace)
     NeededSpace = DAG.getNode(ISD::ADD, DL, MVT::i64, NeededSpace,
-                              DAG.getConstant(ExtraAlignSpace, DL, MVT::i64)); 
+                              DAG.getConstant(ExtraAlignSpace, DL, MVT::i64));
 
   // Get the new stack pointer value.
   SDValue NewSP = DAG.getNode(ISD::SUB, DL, MVT::i64, OldSP, NeededSpace);
 
   // Copy the new stack pointer back.
   Chain = DAG.getCopyToReg(Chain, DL, SPReg, NewSP);
 
   // The allocated data lives above the 160 bytes allocated for the standard
if (ISD::isNON_EXTLoad(N->getOperand(0).getNode()) &&
       DCI.CombineTo(Load.getNode(), ResVal, BSLoad.getValue(1));
 
       // Return N so it doesn't get rechecked!
       return SDValue(N, 0);
     }
   return SDValue();
 }
 
+SDValue SystemZTargetLowering::combineSHIFTROT(
+    SDNode *N, DAGCombinerInfo &DCI) const {
+
+  // Shift/rotate instructions only use the last 6 bits of the second operand
+  // register. If the second operand is the result of an AND with an immediate
+  // value that has its last 6 bits set, we can safely remove the AND operation.
+  SDValue N1 = N->getOperand(1);
+  if (N1.getOpcode() == ISD::AND) {
+    auto *AndMask = dyn_cast<ConstantSDNode>(N1.getOperand(1));
+
+    // The AND mask is constant
+    if (AndMask) {
+      auto AmtVal = AndMask->getZExtValue();
+
+      // Bottom 6 bits are set
+      if ((AmtVal & 0x3f) == 0x3f) {
+        SDValue AndOp = N1->getOperand(0);
+
+        // Combine the AND away
+        DCI.CombineTo(N1.getNode(), AndOp);
+
+        // Return N so it isn't rechecked
+        return SDValue(N, 0);
+      }
+    }
+  }
+
+  return SDValue();
+}
+
 SDValue SystemZTargetLowering::PerformDAGCombine(SDNode *N,
                                                  DAGCombinerInfo &DCI) const {
   switch(N->getOpcode()) {
   default: break;
   case ISD::SIGN_EXTEND:        return combineSIGN_EXTEND(N, DCI);
   case SystemZISD::MERGE_HIGH:
   case SystemZISD::MERGE_LOW:   return combineMERGE(N, DCI);
   case ISD::STORE:              return combineSTORE(N, DCI);
   case ISD::EXTRACT_VECTOR_ELT: return combineEXTRACT_VECTOR_ELT(N, DCI);
   case SystemZISD::JOIN_DWORDS: return combineJOIN_DWORDS(N, DCI);
   case ISD::FP_ROUND:           return combineFP_ROUND(N, DCI);
   case ISD::BSWAP:              return combineBSWAP(N, DCI);
+  case ISD::SHL:
+  case ISD::SRA:
+  case ISD::SRL:
+  case ISD::ROTL:               return combineSHIFTROT(N, DCI);
   }
+
   return SDValue();
 }
 
 //===----------------------------------------------------------------------===//
 // Custom insertion
 //===----------------------------------------------------------------------===//
 
 // Create a new basic block after MBB.

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
+}
+
+; Test that AND with two register operands is not affected.
+define i32 @f9(i32 %a, i32 %b, i32 %sh) {
+; CHECK-LABEL: f9:
+; CHECK: nr %r3, %r4
+; CHECK: sll %r2, 0(%r3)
+  %and = and i32 %sh, %b
+  %shift = shl i32 %a, %and
+  ret i32 %shift
+}