[clang] [Headers][X86] VisitCallExpr constexpr immediate shifts (#154293) (PR #155542)

[Simon Pilgrim via cfe-commits]

================
@@ -11644,64 +11679,181 @@ bool VectorExprEvaluator::VisitCallExpr(const CallExpr *E) {
 
     for (unsigned EltNum = 0; EltNum < SourceLen; ++EltNum) {
       APSInt LHS = SourceLHS.getVectorElt(EltNum).getInt();
-      APSInt RHS = SourceRHS.getVectorElt(EltNum).getInt();
-      switch (E->getBuiltinCallee()) {
-      case Builtin::BI__builtin_elementwise_add_sat:
-        ResultElements.push_back(APValue(
-            APSInt(LHS.isSigned() ? LHS.sadd_sat(RHS) : LHS.uadd_sat(RHS),
-                   DestUnsigned)));
-        break;
-      case Builtin::BI__builtin_elementwise_sub_sat:
-        ResultElements.push_back(APValue(
-            APSInt(LHS.isSigned() ? LHS.ssub_sat(RHS) : LHS.usub_sat(RHS),
-                   DestUnsigned)));
-        break;
-      case clang::X86::BI__builtin_ia32_pmulhuw128:
-      case clang::X86::BI__builtin_ia32_pmulhuw256:
-      case clang::X86::BI__builtin_ia32_pmulhuw512:
-        ResultElements.push_back(APValue(APSInt(llvm::APIntOps::mulhu(LHS, RHS),
-                                                /*isUnsigned=*/true)));
-        break;
-      case clang::X86::BI__builtin_ia32_pmulhw128:
-      case clang::X86::BI__builtin_ia32_pmulhw256:
-      case clang::X86::BI__builtin_ia32_pmulhw512:
-        ResultElements.push_back(APValue(APSInt(llvm::APIntOps::mulhs(LHS, RHS),
-                                                /*isUnsigned=*/false)));
-        break;
-      case clang::X86::BI__builtin_ia32_psllv2di:
-      case clang::X86::BI__builtin_ia32_psllv4di:
-      case clang::X86::BI__builtin_ia32_psllv4si:
-      case clang::X86::BI__builtin_ia32_psllv8si:
-        if (RHS.uge(RHS.getBitWidth())) {
-          ResultElements.push_back(
-              APValue(APSInt(APInt::getZero(RHS.getBitWidth()), DestUnsigned)));
+
+      if (SourceRHS.isInt()) {
+        uint64_t LaneWidth = 0;
+        bool IsLeftShift = false;
+        bool IsRightShift = false;
+        bool IsArithmeticRightShift = false;
+
+        switch (E->getBuiltinCallee()) {
+        case clang::X86::BI__builtin_ia32_psllwi128:
+        case clang::X86::BI__builtin_ia32_psllwi256:
+        case clang::X86::BI__builtin_ia32_psllwi512:
+          IsLeftShift = true;
+          LaneWidth = 16;
+          break;
+        case clang::X86::BI__builtin_ia32_pslldi128:
+        case clang::X86::BI__builtin_ia32_pslldi256:
+        case clang::X86::BI__builtin_ia32_pslldi512:
+          IsLeftShift = true;
+          LaneWidth = 32;
+          break;
+        case clang::X86::BI__builtin_ia32_psllqi128:
+        case clang::X86::BI__builtin_ia32_psllqi256:
+        case clang::X86::BI__builtin_ia32_psllqi512:
+          IsLeftShift = true;
+          LaneWidth = 64;
           break;
+
+        case clang::X86::BI__builtin_ia32_psrlwi128:
+        case clang::X86::BI__builtin_ia32_psrlwi256:
+        case clang::X86::BI__builtin_ia32_psrlwi512:
+          IsRightShift = true;
+          LaneWidth = 16;
+          break;
+        case clang::X86::BI__builtin_ia32_psrldi128:
+        case clang::X86::BI__builtin_ia32_psrldi256:
+        case clang::X86::BI__builtin_ia32_psrldi512:
+          IsRightShift = true;
+          LaneWidth = 32;
+          break;
+        case clang::X86::BI__builtin_ia32_psrlqi128:
+        case clang::X86::BI__builtin_ia32_psrlqi256:
+        case clang::X86::BI__builtin_ia32_psrlqi512:
+          IsRightShift = true;
+          LaneWidth = 64;
+          break;
+
+        case clang::X86::BI__builtin_ia32_psrawi128:
+        case clang::X86::BI__builtin_ia32_psrawi256:
+        case clang::X86::BI__builtin_ia32_psrawi512:
+          IsArithmeticRightShift = true;
+          LaneWidth = 16;
+          break;
+        case clang::X86::BI__builtin_ia32_psradi128:
+        case clang::X86::BI__builtin_ia32_psradi256:
+        case clang::X86::BI__builtin_ia32_psradi512:
+          IsArithmeticRightShift = true;
+          LaneWidth = 32;
+          break;
+        case clang::X86::BI__builtin_ia32_psraqi128:
+        case clang::X86::BI__builtin_ia32_psraqi256:
+        case clang::X86::BI__builtin_ia32_psraqi512:
+          IsArithmeticRightShift = true;
+          LaneWidth = 64;
+          break;
+
+        default:
+          llvm_unreachable("Unexpected builtin callee");
         }
-        ResultElements.push_back(
-            APValue(APSInt(LHS.shl(RHS.getZExtValue()), DestUnsigned)));
-        break;
-      case clang::X86::BI__builtin_ia32_psrav4si:
-      case clang::X86::BI__builtin_ia32_psrav8si:
-        if (RHS.uge(RHS.getBitWidth())) {
+
+        const APSInt RHS = SourceRHS.getInt();
+        const auto ShiftAmount = RHS.getZExtValue();
+        APInt ResultOut;
+        if (IsArithmeticRightShift) {
+          ResultOut = LHS.ashr(std::min(ShiftAmount, LaneWidth));
+        } else if (ShiftAmount >= LaneWidth) {
+          ResultOut = APInt(LaneWidth, 0);
+        } else if (IsLeftShift) {
+          ResultOut = LHS.shl(ShiftAmount);
+        } else if (IsRightShift) {
+          ResultOut = LHS.lshr(ShiftAmount);
+        } else {
+          llvm_unreachable("Invalid shift type");
+        }
+        ResultElements.push_back(APValue(APSInt(
+            std::move(ResultOut),
+            /*isUnsigned=*/DestEltTy->isUnsignedIntegerOrEnumerationType())));
+      } else {
+        APSInt RHS = SourceRHS.getVectorElt(EltNum).getInt();
----------------
RKSimon wrote:

You can reduce the diff and avoid the else case if you continue at the end of the SourceRHS.isInt() block

https://github.com/llvm/llvm-project/pull/155542