[AVX-512] Add AVX-512 vector shift intrinsics to memory santitizer.

Just needed to add the intrinsics to the exist switch. The code is generic enough to support the wider vectors with no changes.

llvm-svn: 286980

2 files changed, 82 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 6dba8d16516..fafb0fcbd01 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -2362,6 +2362,29 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     case llvm::Intrinsic::x86_sse_cvttps2pi:
       handleVectorConvertIntrinsic(I, 2);
       break;
+
+    case llvm::Intrinsic::x86_avx512_psll_w_512:
+    case llvm::Intrinsic::x86_avx512_psll_d_512:
+    case llvm::Intrinsic::x86_avx512_psll_q_512:
+    case llvm::Intrinsic::x86_avx512_pslli_w_512:
+    case llvm::Intrinsic::x86_avx512_pslli_d_512:
+    case llvm::Intrinsic::x86_avx512_pslli_q_512:
+    case llvm::Intrinsic::x86_avx512_psrl_w_512:
+    case llvm::Intrinsic::x86_avx512_psrl_d_512:
+    case llvm::Intrinsic::x86_avx512_psrl_q_512:
+    case llvm::Intrinsic::x86_avx512_psra_w_512:
+    case llvm::Intrinsic::x86_avx512_psra_d_512:
+    case llvm::Intrinsic::x86_avx512_psra_q_512:
+    case llvm::Intrinsic::x86_avx512_psrli_w_512:
+    case llvm::Intrinsic::x86_avx512_psrli_d_512:
+    case llvm::Intrinsic::x86_avx512_psrli_q_512:
+    case llvm::Intrinsic::x86_avx512_psrai_w_512:
+    case llvm::Intrinsic::x86_avx512_psrai_d_512:
+    case llvm::Intrinsic::x86_avx512_psrai_q_512:
+    case llvm::Intrinsic::x86_avx512_psra_q_256:
+    case llvm::Intrinsic::x86_avx512_psra_q_128:
+    case llvm::Intrinsic::x86_avx512_psrai_q_256:
+    case llvm::Intrinsic::x86_avx512_psrai_q_128:
     case llvm::Intrinsic::x86_avx2_psll_w:
     case llvm::Intrinsic::x86_avx2_psll_d:
     case llvm::Intrinsic::x86_avx2_psll_q:
@@ -2414,14 +2437,22 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       break;
     case llvm::Intrinsic::x86_avx2_psllv_d:
     case llvm::Intrinsic::x86_avx2_psllv_d_256:
+    case llvm::Intrinsic::x86_avx512_psllv_d_512:
     case llvm::Intrinsic::x86_avx2_psllv_q:
     case llvm::Intrinsic::x86_avx2_psllv_q_256:
+    case llvm::Intrinsic::x86_avx512_psllv_q_512:
     case llvm::Intrinsic::x86_avx2_psrlv_d:
     case llvm::Intrinsic::x86_avx2_psrlv_d_256:
+    case llvm::Intrinsic::x86_avx512_psrlv_d_512:
     case llvm::Intrinsic::x86_avx2_psrlv_q:
     case llvm::Intrinsic::x86_avx2_psrlv_q_256:
+    case llvm::Intrinsic::x86_avx512_psrlv_q_512:
     case llvm::Intrinsic::x86_avx2_psrav_d:
     case llvm::Intrinsic::x86_avx2_psrav_d_256:
+    case llvm::Intrinsic::x86_avx512_psrav_d_512:
+    case llvm::Intrinsic::x86_avx512_psrav_q_128:
+    case llvm::Intrinsic::x86_avx512_psrav_q_256:
+    case llvm::Intrinsic::x86_avx512_psrav_q_512:
       handleVectorShiftIntrinsic(I, /* Variable */ true);
       break;
 
diff --git a/llvm/test/Instrumentation/MemorySanitizer/vector_shift.ll b/llvm/test/Instrumentation/MemorySanitizer/vector_shift.ll
index ef437dd79a0..978bad3b697 100644
--- a/llvm/test/Instrumentation/MemorySanitizer/vector_shift.ll
+++ b/llvm/test/Instrumentation/MemorySanitizer/vector_shift.ll
@@ -6,10 +6,13 @@ target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f3
 target triple = "x86_64-unknown-linux-gnu"
 
 declare x86_mmx @llvm.x86.mmx.psll.d(x86_mmx, x86_mmx)
+declare <16 x i32> @llvm.x86.avx512.psllv.d.512(<16 x i32>, <16 x i32>)
 declare <8 x i32> @llvm.x86.avx2.psllv.d.256(<8 x i32>, <8 x i32>)
 declare <4 x i32> @llvm.x86.avx2.psllv.d(<4 x i32>, <4 x i32>)
 declare <8 x i16> @llvm.x86.sse2.psrl.w(<8 x i16>, <8 x i16>)
+declare <32 x i16> @llvm.x86.avx512.psrl.w.512(<32 x i16>, <8 x i16>)
 declare <8 x i16> @llvm.x86.sse2.pslli.w(<8 x i16>, i32)
+declare <32 x i16> @llvm.x86.avx512.pslli.w.512(<32 x i16>, i32)
 
 define i64 @test_mmx(i64 %x.coerce, i64 %y.coerce) sanitize_memory {
 entry:
@@ -49,6 +52,22 @@ entry:
 ; CHECK: ret <8 x i16>
 
 
+define <32 x i16> @test_avx512_scalar(<32 x i16> %x, i32 %y) sanitize_memory {
+entry:
+  %0 = tail call <32 x i16> @llvm.x86.avx512.pslli.w.512(<32 x i16> %x, i32 %y)
+  ret <32 x i16> %0
+}
+
+; CHECK-LABEL: @test_avx512_scalar
+; CHECK: = icmp ne i32 {{.*}}, 0
+; CHECK: = sext i1 {{.*}} to i512
+; CHECK: = bitcast i512 {{.*}} to <32 x i16>
+; CHECK: = call <32 x i16> @llvm.x86.avx512.pslli.w.512(
+; CHECK: = or <32 x i16>
+; CHECK: call <32 x i16> @llvm.x86.avx512.pslli.w.512(
+; CHECK: ret <32 x i16>
+
+
 define <8 x i16> @test_sse2(<8 x i16> %x, <8 x i16> %y) sanitize_memory {
 entry:
   %0 = tail call <8 x i16> @llvm.x86.sse2.psrl.w(<8 x i16> %x, <8 x i16> %y)
@@ -67,6 +86,24 @@ entry:
 ; CHECK: ret <8 x i16>
 
 
+define <32 x i16> @test_avx512(<32 x i16> %x, <8 x i16> %y) sanitize_memory {
+entry:
+  %0 = tail call <32 x i16> @llvm.x86.avx512.psrl.w.512(<32 x i16> %x, <8 x i16> %y)
+  ret <32 x i16> %0
+}
+
+; CHECK-LABEL: @test_avx512
+; CHECK: = bitcast <8 x i16> {{.*}} to i128
+; CHECK: = trunc i128 {{.*}} to i64
+; CHECK: = icmp ne i64 {{.*}}, 0
+; CHECK: = sext i1 {{.*}} to i512
+; CHECK: = bitcast i512 {{.*}} to <32 x i16>
+; CHECK: = call <32 x i16> @llvm.x86.avx512.psrl.w.512(
+; CHECK: = or <32 x i16>
+; CHECK: call <32 x i16> @llvm.x86.avx512.psrl.w.512(
+; CHECK: ret <32 x i16>
+
+
 ; Test variable shift (i.e. vector by vector).
 
 define <4 x i32> @test_avx2(<4 x i32> %x, <4 x i32> %y) sanitize_memory {
@@ -96,3 +133,17 @@ entry:
 ; CHECK: = or <8 x i32>
 ; CHECK: = tail call <8 x i32> @llvm.x86.avx2.psllv.d.256(
 ; CHECK: ret <8 x i32>
+
+define <16 x i32> @test_avx512_512(<16 x i32> %x, <16 x i32> %y) sanitize_memory {
+entry:
+  %0 = tail call <16 x i32> @llvm.x86.avx512.psllv.d.512(<16 x i32> %x, <16 x i32> %y)
+  ret <16 x i32> %0
+}
+
+; CHECK-LABEL: @test_avx512_512
+; CHECK: = icmp ne <16 x i32> {{.*}}, zeroinitializer
+; CHECK: = sext <16 x i1> {{.*}} to <16 x i32>
+; CHECK: = call <16 x i32> @llvm.x86.avx512.psllv.d.512(
+; CHECK: = or <16 x i32>
+; CHECK: = tail call <16 x i32> @llvm.x86.avx512.psllv.d.512(
+; CHECK: ret <16 x i32>