[PATCH] D46575: AMDGPU: Handle partial shift reduction for variable shifts

[Matt Arsenault via Phabricator via llvm-commits]

arsenm created this revision.
arsenm added reviewers: rampitec, kzhuravl, Farhana, cfang.
Herald added subscribers: t-tye, tpr, dstuttard, yaxunl, nhaehnle, wdng.

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


https://reviews.llvm.org/D46575

Files:
  lib/Target/AMDGPU/AMDGPUISelLowering.cpp
  test/CodeGen/AMDGPU/partial-shift-shrink.ll


Index: test/CodeGen/AMDGPU/partial-shift-shrink.ll
===================================================================
--- test/CodeGen/AMDGPU/partial-shift-shrink.ll
+++ test/CodeGen/AMDGPU/partial-shift-shrink.ll
@@ -100,3 +100,40 @@
   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
+}
Index: lib/Target/AMDGPU/AMDGPUISelLowering.cpp
===================================================================
--- lib/Target/AMDGPU/AMDGPUISelLowering.cpp
+++ lib/Target/AMDGPU/AMDGPUISelLowering.cpp
@@ -3153,22 +3153,30 @@
         (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;
-
-          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);
+      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 MidVT = VT.isVector() ?
+          EVT::getVectorVT(*DAG.getContext(), MVT::i32,
+                           VT.getVectorNumElements()) : MVT::i32;
+
+        EVT AmtVT = Amt.getValueType();
+        EVT ShiftTy = getShiftAmountTy(MidVT, DAG.getDataLayout());
+        SDValue Trunc = DAG.getNode(ISD::TRUNCATE, SL, MidVT,
+                                    Src.getOperand(0));
+        DCI.AddToWorklist(Trunc.getNode());
+
+        if (AmtVT != ShiftTy) {
+          Amt = DAG.getZExtOrTrunc(Amt, SL, ShiftTy);
+          DCI.AddToWorklist(Amt.getNode());
         }
+
+        SDValue ShrunkShift = DAG.getNode(Src.getOpcode(), SL, MidVT,
+                                          Trunc, Amt);
+        return DAG.getNode(ISD::TRUNCATE, SL, VT, ShrunkShift);
       }
     }
   }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D46575.145666.patch
Type: text/x-patch
Size: 3814 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20180508/45795d5a/attachment.bin>