1 files changed, 22 insertions, 15 deletions
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp b/llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp
index 32dc2a7afce..8d0ab496938 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp
@@ -3154,22 +3154,29 @@ SDValue AMDGPUTargetLowering::performTruncateCombine(
         (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 NewShiftVT = getShiftAmountTy(MidVT, DAG.getDataLayout());
+        SDValue Trunc = DAG.getNode(ISD::TRUNCATE, SL, MidVT,
+                                    Src.getOperand(0));
+        DCI.AddToWorklist(Trunc.getNode());
+
+        if (Amt.getValueType() != NewShiftVT) {
+          Amt = DAG.getZExtOrTrunc(Amt, SL, NewShiftVT);
+          DCI.AddToWorklist(Amt.getNode());
         }
+
+        SDValue ShrunkShift = DAG.getNode(Src.getOpcode(), SL, MidVT,
+                                          Trunc, Amt);
+        return DAG.getNode(ISD::TRUNCATE, SL, VT, ShrunkShift);
       }
     }
   }