[x86][AVX-VNNI] Fix VPDPBUSD Argument Types (#155194)

Fixed intrinsic VPDPBUSD[,S]_128/256/512's argument types to match with the ISA.

Fixes part of #97271

1 files changed, 89 insertions, 13 deletions

diff --git a/llvm/lib/IR/AutoUpgrade.cpp b/llvm/lib/IR/AutoUpgrade.cpp
index 8034b3f..8d8120a 100644
--- a/llvm/lib/IR/AutoUpgrade.cpp
+++ b/llvm/lib/IR/AutoUpgrade.cpp
@@ -107,6 +107,24 @@ static bool upgradeX86MaskedFPCompare(Function *F, Intrinsic::ID IID,
   return true;
 }
 
+// Upgrade the declaration of multiply and add bytes intrinsics whose input
+// arguments' types have changed from vectors of i32 to vectors of i8
+static bool upgradeX86MultiplyAddBytes(Function *F, Intrinsic::ID IID,
+                                       Function *&NewFn) {
+  // check if input argument type is a vector of i8
+  Type *Arg1Type = F->getFunctionType()->getParamType(1);
+  Type *Arg2Type = F->getFunctionType()->getParamType(2);
+  if (Arg1Type->isVectorTy() &&
+      cast<VectorType>(Arg1Type)->getElementType()->isIntegerTy(8) &&
+      Arg2Type->isVectorTy() &&
+      cast<VectorType>(Arg2Type)->getElementType()->isIntegerTy(8))
+    return false;
+
+  rename(F);
+  NewFn = Intrinsic::getOrInsertDeclaration(F->getParent(), IID);
+  return true;
+}
+
 static bool upgradeX86BF16Intrinsic(Function *F, Intrinsic::ID IID,
                                     Function *&NewFn) {
   if (F->getReturnType()->getScalarType()->isBFloatTy())
@@ -546,19 +564,34 @@ static bool upgradeX86IntrinsicFunction(Function *F, StringRef Name,
   if (ID != Intrinsic::not_intrinsic)
     return upgradeX86IntrinsicsWith8BitMask(F, ID, NewFn);
 
-  if (Name.consume_front("avx512.mask.cmp.")) {
-    // Added in 7.0
-    ID = StringSwitch<Intrinsic::ID>(Name)
-             .Case("pd.128", Intrinsic::x86_avx512_mask_cmp_pd_128)
-             .Case("pd.256", Intrinsic::x86_avx512_mask_cmp_pd_256)
-             .Case("pd.512", Intrinsic::x86_avx512_mask_cmp_pd_512)
-             .Case("ps.128", Intrinsic::x86_avx512_mask_cmp_ps_128)
-             .Case("ps.256", Intrinsic::x86_avx512_mask_cmp_ps_256)
-             .Case("ps.512", Intrinsic::x86_avx512_mask_cmp_ps_512)
-             .Default(Intrinsic::not_intrinsic);
-    if (ID != Intrinsic::not_intrinsic)
-      return upgradeX86MaskedFPCompare(F, ID, NewFn);
-    return false; // No other 'x86.avx523.mask.cmp.*'.
+  if (Name.consume_front("avx512.")) {
+    if (Name.consume_front("mask.cmp.")) {
+      // Added in 7.0
+      ID = StringSwitch<Intrinsic::ID>(Name)
+               .Case("pd.128", Intrinsic::x86_avx512_mask_cmp_pd_128)
+               .Case("pd.256", Intrinsic::x86_avx512_mask_cmp_pd_256)
+               .Case("pd.512", Intrinsic::x86_avx512_mask_cmp_pd_512)
+               .Case("ps.128", Intrinsic::x86_avx512_mask_cmp_ps_128)
+               .Case("ps.256", Intrinsic::x86_avx512_mask_cmp_ps_256)
+               .Case("ps.512", Intrinsic::x86_avx512_mask_cmp_ps_512)
+               .Default(Intrinsic::not_intrinsic);
+      if (ID != Intrinsic::not_intrinsic)
+        return upgradeX86MaskedFPCompare(F, ID, NewFn);
+    } else if (Name.starts_with("vpdpbusd.") ||
+               Name.starts_with("vpdpbusds.")) {
+      // Added in 21.1
+      ID = StringSwitch<Intrinsic::ID>(Name)
+               .Case("vpdpbusd.128", Intrinsic::x86_avx512_vpdpbusd_128)
+               .Case("vpdpbusd.256", Intrinsic::x86_avx512_vpdpbusd_256)
+               .Case("vpdpbusd.512", Intrinsic::x86_avx512_vpdpbusd_512)
+               .Case("vpdpbusds.128", Intrinsic::x86_avx512_vpdpbusds_128)
+               .Case("vpdpbusds.256", Intrinsic::x86_avx512_vpdpbusds_256)
+               .Case("vpdpbusds.512", Intrinsic::x86_avx512_vpdpbusds_512)
+               .Default(Intrinsic::not_intrinsic);
+      if (ID != Intrinsic::not_intrinsic)
+        return upgradeX86MultiplyAddBytes(F, ID, NewFn);
+    }
+    return false; // No other 'x86.avx512.*'.
   }
 
   if (Name.consume_front("avx512bf16.")) {
@@ -4149,6 +4182,32 @@ static Value *upgradeX86IntrinsicCall(StringRef Name, CallBase *CI, Function *F,
 
     Value *Args[] = {CI->getArgOperand(0), CI->getArgOperand(1),
                      CI->getArgOperand(2)};
+
+    // Input arguments types were incorrectly set to vectors of i32 before but
+    // they should be vectors of i8. Insert bit cast when encountering the old
+    // types
+    if (Args[1]->getType()->isVectorTy() &&
+        cast<VectorType>(Args[1]->getType())
+            ->getElementType()
+            ->isIntegerTy(32) &&
+        Args[2]->getType()->isVectorTy() &&
+        cast<VectorType>(Args[2]->getType())
+            ->getElementType()
+            ->isIntegerTy(32)) {
+      Type *NewArgType = nullptr;
+      if (VecWidth == 128)
+        NewArgType = VectorType::get(Builder.getInt8Ty(), 16, false);
+      else if (VecWidth == 256)
+        NewArgType = VectorType::get(Builder.getInt8Ty(), 32, false);
+      else if (VecWidth == 512)
+        NewArgType = VectorType::get(Builder.getInt8Ty(), 64, false);
+      else
+        llvm_unreachable("Unexpected vector bit width");
+
+      Args[1] = Builder.CreateBitCast(Args[1], NewArgType);
+      Args[2] = Builder.CreateBitCast(Args[2], NewArgType);
+    }
+
     Rep = Builder.CreateIntrinsic(IID, Args);
     Value *PassThru = ZeroMask ? ConstantAggregateZero::get(CI->getType())
                                : CI->getArgOperand(0);
@@ -5156,6 +5215,23 @@ void llvm::UpgradeIntrinsicCall(CallBase *CI, Function *NewFn) {
     CI->eraseFromParent();
     return;
   }
+
+  case Intrinsic::x86_avx512_vpdpbusd_128:
+  case Intrinsic::x86_avx512_vpdpbusd_256:
+  case Intrinsic::x86_avx512_vpdpbusd_512:
+  case Intrinsic::x86_avx512_vpdpbusds_128:
+  case Intrinsic::x86_avx512_vpdpbusds_256:
+  case Intrinsic::x86_avx512_vpdpbusds_512: {
+    unsigned NumElts = CI->getType()->getPrimitiveSizeInBits() / 8;
+    Value *Args[] = {CI->getArgOperand(0), CI->getArgOperand(1),
+                     CI->getArgOperand(2)};
+    Type *NewArgType = VectorType::get(Builder.getInt8Ty(), NumElts, false);
+    Args[1] = Builder.CreateBitCast(Args[1], NewArgType);
+    Args[2] = Builder.CreateBitCast(Args[2], NewArgType);
+
+    NewCall = Builder.CreateCall(NewFn, Args);
+    break;
+  }
   }
   assert(NewCall && "Should have either set this variable or returned through "
                     "the default case");