AMDGPU: Handle partial shift reduction for variable shifts

[arsenm]

Summary

If the variable shift amount has known bits, we can still reduce the shift.

[rampitec]

LGTM

[arsenm]

r331917

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

SDValue AMDGPUTargetLowering::performTruncateCombine(
   // i16 (trunc (srl i64:x, K)), K <= 16 ->
   //     i16 (trunc (srl (i32 (trunc x), K)))
   if (VT.getScalarSizeInBits() < 32) {
     EVT SrcVT = Src.getValueType();
     if (SrcVT.getScalarSizeInBits() > 32 &&
         (Src.getOpcode() == ISD::SRL ||
          Src.getOpcode() == ISD::SRA ||
          Src.getOpcode() == ISD::SHL)) {
-      if (auto ShiftAmount = isConstOrConstSplat(Src.getOperand(1))) {
-        if (ShiftAmount->getZExtValue() <= VT.getScalarSizeInBits()) {
-          EVT MidVT = VT.isVector() ?
-            EVT::getVectorVT(*DAG.getContext(), MVT::i32,
-                             VT.getVectorNumElements()) : MVT::i32;
+      SDValue Amt = Src.getOperand(1);
+      KnownBits Known;
+      DAG.computeKnownBits(Amt, Known);
+      unsigned Size = VT.getScalarSizeInBits();
+      if ((Known.isConstant() && Known.getConstant().ule(Size)) ||
+          (Known.getBitWidth() - Known.countMinLeadingZeros() <= Log2_32(Size))) {
 
-          EVT ShiftTy = getShiftAmountTy(MidVT, DAG.getDataLayout());
-          SDValue NewShiftAmt = DAG.getConstant(ShiftAmount->getZExtValue(),
-                                                SL, ShiftTy);
-          SDValue Trunc = DAG.getNode(ISD::TRUNCATE, SL, MidVT,
-                                      Src.getOperand(0));
-          DCI.AddToWorklist(Trunc.getNode());
-          SDValue ShrunkShift = DAG.getNode(Src.getOpcode(), SL, MidVT,
-                                            Trunc, NewShiftAmt);
-          return DAG.getNode(ISD::TRUNCATE, SL, VT, ShrunkShift);
-        }
-      }
-    }
-  }
+        EVT AmtVT = Amt.getValueType();
+
+        if (AmtVT != ShiftTy) {
+          Amt = DAG.getZExtOrTrunc(Amt, SL, ShiftTy);
+          DCI.AddToWorklist(Amt.getNode());
+        }
+
+                                          Trunc, Amt);
 
   return SDValue();
 }
 
 // We need to specifically handle i64 mul here to avoid unnecessary conversion
 // instructions. If we only match on the legalized i64 mul expansion,
 // SimplifyDemandedBits will be unable to remove them because there will be
 // multiple uses due to the separate mul + mulh[su].

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

 ; GCN: global_store_short v{{\[[0-9]+:[0-9]+\]}}, [[V_RESULT]]
 define amdgpu_kernel void @s_trunc_srl_i64_16_to_i16(i64 %x) {
   %shift = lshr i64 %x, 16
   %trunc = trunc i64 %shift to i16
   %add = or i16 %trunc, 4
   store i16 %add, i16 addrspace(1)* undef
   ret void
 }
+
+; GCN-LABEL: {{^}}trunc_srl_i64_var_mask15_to_i16:
+; GCN: s_waitcnt
+; GCN-NEXT: v_and_b32_e32 v1, 15, v2
+; GCN-NEXT: v_lshrrev_b32_e32 v0, v1, v0
+; GCN-NEXT: s_setpc_b64
+define i16 @trunc_srl_i64_var_mask15_to_i16(i64 %x, i64 %amt) {
+  %amt.masked = and i64 %amt, 15
+  %shift = lshr i64 %x, %amt.masked
+  %trunc = trunc i64 %shift to i16
+  ret i16 %trunc
+}
+
+
+; GCN-LABEL: {{^}}trunc_srl_i64_var_mask16_to_i16:
+; GCN: s_waitcnt
+; GCN-NEXT: v_and_b32_e32 v2, 16, v2
+; GCN-NEXT: v_lshrrev_b64 v[0:1], v2, v[0:1]
+; GCN-NEXT: s_setpc_b64
+define i16 @trunc_srl_i64_var_mask16_to_i16(i64 %x, i64 %amt) {
+  %amt.masked = and i64 %amt, 16
+  %shift = lshr i64 %x, %amt.masked
+  %trunc = trunc i64 %shift to i16
+  ret i16 %trunc
+}
+
+; GCN-LABEL: {{^}}trunc_srl_i64_var_mask31_to_i16:
+; GCN: 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]
+define i16 @trunc_srl_i64_var_mask31_to_i16(i64 %x, i64 %amt) {
+  %amt.masked = and i64 %amt, 31
+  %shift = lshr i64 %x, %amt.masked
+  %trunc = trunc i64 %shift to i16
+  ret i16 %trunc
+}