[AMDGPU][DAG] Fix trunc/shift combine condition

[Pierre-vh]

Summary

The condition needs to be different for right-shifts, else we may lose information in some cases.

[foad]

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

3248–3255

The whole if condition can be replaced with: Known.getMaxValue().ule(MaxCstSize)

3253–3254

For left shifts you can do it so long as the new shift amount is still valid, so ShiftAmt < 32 (aka <= 31).

So this can be: ... ? 31 : 32 - Size

[Pierre-vh]

Comments

[Pierre-vh]

Update comments, remove dead variable

[foad]

LGTM, thanks!

[arsenm]

Can you check the equivalent combine in globalisel? I think it's just handling the left shifts but should be generalized

[Pierre-vh]

In D136059#3868436, @arsenm wrote:

Can you check the equivalent combine in globalisel? I think it's just handling the left shifts but should be generalized

Sure,, I'll check it in a different patch. Can this one land in the meantime?

[Pierre-vh]

In D136059#3868436, @arsenm wrote:

Can you check the equivalent combine in globalisel? I think it's just handling the left shifts but should be generalized

I took a look at it and it seems like it's a different combine. This one explicitly tries to reduce 64 bits shift to 32 bits, but matchCombineTruncOfShl/applyCombineTruncOfShl simply moves the trunc into the operand of the shift:

// Fold trunc (shl x, K) -> shl (trunc x), K 
//    => K < VT.getScalarSizeInBits()

I don't think this can work on right shifts unless we add another trunc in front of the shl again.

Maybe another (AMDGPU-specific since I think it just benefits us?) combine would be better for 64 bits right shifts reduction? I wrote it in D136319

[arsenm]

In D136059#3870167, @Pierre-vh wrote:

In D136059#3868436, @arsenm wrote:

Can you check the equivalent combine in globalisel? I think it's just handling the left shifts but should be generalized

Sure,, I'll check it in a different patch. Can this one land in the meantime?

Yes, not relevant to this one

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

SDValue AMDGPUTargetLowering::performTruncateCombine(
   if (VT.getScalarSizeInBits() < 32) {
     EVT SrcVT = Src.getValueType();
     if (SrcVT.getScalarSizeInBits() > 32 &&
         (Src.getOpcode() == ISD::SRL ||
          Src.getOpcode() == ISD::SRA ||
          Src.getOpcode() == ISD::SHL)) {
       SDValue Amt = Src.getOperand(1);
       KnownBits Known = DAG.computeKnownBits(Amt);
-      unsigned Size = VT.getScalarSizeInBits();
-      if ((Known.isConstant() && Known.getConstant().ule(Size)) ||
-          (Known.countMaxActiveBits() <= Log2_32(Size))) {
+
+      // - For left shifts, do the transform as long as the shift
+      //   amount is still legal for i32, so when ShiftAmt < 32 (<= 31)
+      // - For right shift, do it if ShiftAmt <= (32 - Size) to avoid
+      //   losing information stored in the high bits when truncating.
+      const unsigned MaxCstSize =
+          (Src.getOpcode() == ISD::SHL) ? 31 : (32 - VT.getScalarSizeInBits());
+      if (Known.getMaxValue().ule(MaxCstSize)) {
         EVT MidVT = VT.isVector() ?
           EVT::getVectorVT(*DAG.getContext(), MVT::i32,
                            VT.getVectorNumElements()) : MVT::i32;
 
         EVT NewShiftVT = getShiftAmountTy(MidVT, DAG.getDataLayout());
         SDValue Trunc = DAG.getNode(ISD::TRUNCATE, SL, MidVT,
                                     Src.getOperand(0));
         DCI.AddToWorklist(Trunc.getNode());

llvm/test/CodeGen/AMDGPU/partial-shift-shrink.ll

; GCN-NEXT:    s_setpc_b64 s[30:31]
   %trunc = trunc i64 %shift to i16
   ret i16 %trunc
 }
 
 define i16 @trunc_srl_i64_var_mask16_to_i16(i64 %x, i64 %amt) {
 ; GCN-LABEL: trunc_srl_i64_var_mask16_to_i16:
 ; GCN:       ; %bb.0:
 ; GCN-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
-; GCN-NEXT:    v_and_b32_e32 v2, 16, v2
-; GCN-NEXT:    v_lshrrev_b64 v[0:1], v2, v[0:1]
+; GCN-NEXT:    v_and_b32_e32 v1, 16, v2
+; GCN-NEXT:    v_lshrrev_b32_e32 v0, v1, v0
 ; GCN-NEXT:    s_setpc_b64 s[30:31]
   %amt.masked = and i64 %amt, 16
   %shift = lshr i64 %x, %amt.masked
   %trunc = trunc i64 %shift to i16
   ret i16 %trunc
 }
 
 define i16 @trunc_srl_i64_var_mask31_to_i16(i64 %x, i64 %amt) {
 ; GCN-LABEL: trunc_srl_i64_var_mask31_to_i16:
 ; GCN:       ; %bb.0:
 ; GCN-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GCN-NEXT:    v_and_b32_e32 v2, 31, v2
 ; GCN-NEXT:    v_lshrrev_b64 v[0:1], v2, v[0:1]
 ; GCN-NEXT:    s_setpc_b64 s[30:31]
   %amt.masked = and i64 %amt, 31
   %shift = lshr i64 %x, %amt.masked
   %trunc = trunc i64 %shift to i16
   ret i16 %trunc
 }
+
+define i32 @trunc_srl_i64_25_to_i26(i64 %x) {
+; GCN-LABEL: trunc_srl_i64_25_to_i26:
+; GCN:       ; %bb.0:
+; GCN-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GCN-NEXT:    v_and_b32_e32 v0, 0xa000000, v0
+; GCN-NEXT:    v_alignbit_b32 v0, 0, v0, 25
+; GCN-NEXT:    v_add_u32_e32 v0, 55, v0
+; GCN-NEXT:    s_setpc_b64 s[30:31]
+  %value.knownbits2 = and i64 %x, 167772160 ; 0xA000000
+  %shift = lshr i64 %value.knownbits2, 25
+  %trunc = trunc i64 %shift to i26
+  %add = add i26 %trunc, 55
+  %ext = zext i26 %add to i32
+  ret i32 %ext
+}