[msan] Instrument x86 BMI intrinsics.

Summary:
They simply shuffle bits. MSan needs to do the same with shadow bits,
after making sure that the shuffle mask is fully initialized.

Reviewers: pcc, vitalybuka

Subscribers: hiraditya, #sanitizers, llvm-commits

Tags: #sanitizers, #llvm

Differential Revision: https://reviews.llvm.org/D58858

llvm-svn: 355348

3 files changed, 325 insertions, 1 deletions

diff --git a/compiler-rt/lib/msan/tests/msan_test.cc b/compiler-rt/lib/msan/tests/msan_test.cc
index a3a98301c64..1d14604b600 100644
--- a/compiler-rt/lib/msan/tests/msan_test.cc
+++ b/compiler-rt/lib/msan/tests/msan_test.cc
@@ -169,7 +169,7 @@ static bool TrackingOrigins() {
 
 #define EXPECT_POISONED(x) ExpectPoisoned(x)
 
-template<typename T>
+template <typename T>
 void ExpectPoisoned(const T& t) {
   EXPECT_NE(-1, __msan_test_shadow((void*)&t, sizeof(t)));
 }
@@ -4644,3 +4644,149 @@ TEST(MemorySanitizer, MallocUsableSizeTest) {
   delete int_ptr;
 }
 #endif  // SANITIZER_TEST_HAS_MALLOC_USABLE_SIZE
+
+static bool HaveBmi() {
+#ifdef __x86_64__
+  U4 a = 0, b = 0, c = 0, d = 0;
+  asm("cpuid\n\t" : "=a"(a), "=D"(b), "=c"(c), "=d"(d) : "a"(7));
+  const U4 kBmi12Mask = (1U<<3) | (1U<<8);
+  return b | kBmi12Mask;
+#else
+  return false;
+#endif
+}
+
+__attribute__((target("bmi,bmi2")))
+static void TestBZHI() {
+  EXPECT_NOT_POISONED(
+      __builtin_ia32_bzhi_si(Poisoned<U4>(0xABCDABCD, 0xFF000000), 24));
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_si(Poisoned<U4>(0xABCDABCD, 0xFF800000), 24));
+  // Second operand saturates.
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_si(Poisoned<U4>(0xABCDABCD, 0x80000000), 240));
+  // Any poison in the second operand poisons output.
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_si(0xABCDABCD, Poisoned<U4>(1, 1)));
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_si(0xABCDABCD, Poisoned<U4>(1, 0x80000000)));
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_si(0xABCDABCD, Poisoned<U4>(1, 0xFFFFFFFF)));
+
+  EXPECT_NOT_POISONED(
+      __builtin_ia32_bzhi_di(Poisoned<U8>(0xABCDABCDABCDABCD, 0xFF00000000000000ULL), 56));
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_di(Poisoned<U8>(0xABCDABCDABCDABCD, 0xFF80000000000000ULL), 56));
+  // Second operand saturates.
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_di(Poisoned<U8>(0xABCDABCDABCDABCD, 0x8000000000000000ULL), 240));
+  // Any poison in the second operand poisons output.
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_di(0xABCDABCDABCDABCD, Poisoned<U8>(1, 1)));
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_di(0xABCDABCDABCDABCD, Poisoned<U8>(1, 0x8000000000000000ULL)));
+  EXPECT_POISONED(
+      __builtin_ia32_bzhi_di(0xABCDABCDABCDABCD, Poisoned<U8>(1, 0xFFFFFFFF00000000ULL)));
+}
+
+inline U4 bextr_imm(U4 start, U4 len) {
+  start &= 0xFF;
+  len &= 0xFF;
+  return (len << 8) | start;
+}
+
+__attribute__((target("bmi,bmi2")))
+static void TestBEXTR() {
+  EXPECT_POISONED(
+      __builtin_ia32_bextr_u32(Poisoned<U4>(0xABCDABCD, 0xFF), bextr_imm(0, 8)));
+  EXPECT_POISONED(
+      __builtin_ia32_bextr_u32(Poisoned<U4>(0xABCDABCD, 0xFF), bextr_imm(7, 8)));
+  EXPECT_NOT_POISONED(
+      __builtin_ia32_bextr_u32(Poisoned<U4>(0xABCDABCD, 0xFF), bextr_imm(8, 8)));
+  EXPECT_NOT_POISONED(
+      __builtin_ia32_bextr_u32(Poisoned<U4>(0xABCDABCD, 0xFF), bextr_imm(8, 800)));
+  EXPECT_POISONED(
+      __builtin_ia32_bextr_u32(Poisoned<U4>(0xABCDABCD, 0xFF), bextr_imm(7, 800)));
+  EXPECT_NOT_POISONED(
+      __builtin_ia32_bextr_u32(Poisoned<U4>(0xABCDABCD, 0xFF), bextr_imm(5, 0)));
+
+  EXPECT_POISONED(
+      __builtin_ia32_bextr_u32(0xABCDABCD, Poisoned<U4>(bextr_imm(7, 800), 1)));
+  EXPECT_POISONED(__builtin_ia32_bextr_u32(
+      0xABCDABCD, Poisoned<U4>(bextr_imm(7, 800), 0x80000000)));
+
+  EXPECT_POISONED(
+      __builtin_ia32_bextr_u64(Poisoned<U8>(0xABCDABCD, 0xFF), bextr_imm(0, 8)));
+  EXPECT_POISONED(
+      __builtin_ia32_bextr_u64(Poisoned<U8>(0xABCDABCD, 0xFF), bextr_imm(7, 8)));
+  EXPECT_NOT_POISONED(
+      __builtin_ia32_bextr_u64(Poisoned<U8>(0xABCDABCD, 0xFF), bextr_imm(8, 8)));
+  EXPECT_NOT_POISONED(
+      __builtin_ia32_bextr_u64(Poisoned<U8>(0xABCDABCD, 0xFF), bextr_imm(8, 800)));
+  EXPECT_POISONED(
+      __builtin_ia32_bextr_u64(Poisoned<U8>(0xABCDABCD, 0xFF), bextr_imm(7, 800)));
+  EXPECT_NOT_POISONED(
+      __builtin_ia32_bextr_u64(Poisoned<U8>(0xABCDABCD, 0xFF), bextr_imm(5, 0)));
+
+  // Poison in the top half.
+  EXPECT_NOT_POISONED(__builtin_ia32_bextr_u64(
+      Poisoned<U8>(0xABCDABCD, 0xFF0000000000), bextr_imm(32, 8)));
+  EXPECT_POISONED(__builtin_ia32_bextr_u64(
+      Poisoned<U8>(0xABCDABCD, 0xFF0000000000), bextr_imm(32, 9)));
+
+  EXPECT_POISONED(
+      __builtin_ia32_bextr_u64(0xABCDABCD, Poisoned<U8>(bextr_imm(7, 800), 1)));
+  EXPECT_POISONED(__builtin_ia32_bextr_u64(
+      0xABCDABCD, Poisoned<U8>(bextr_imm(7, 800), 0x80000000)));
+}
+
+__attribute__((target("bmi,bmi2")))
+static void TestPDEP() {
+  U4 x = Poisoned<U4>(0, 0xFF00);
+  EXPECT_NOT_POISONED(__builtin_ia32_pdep_si(x, 0xFF));
+  EXPECT_POISONED(__builtin_ia32_pdep_si(x, 0x1FF));
+  EXPECT_NOT_POISONED(__builtin_ia32_pdep_si(x, 0xFF00));
+  EXPECT_POISONED(__builtin_ia32_pdep_si(x, 0x1FF00));
+
+  EXPECT_NOT_POISONED(__builtin_ia32_pdep_si(x, 0x1FF00) & 0xFF);
+  EXPECT_POISONED(__builtin_ia32_pdep_si(0, Poisoned<U4>(0xF, 1)));
+
+  U8 y = Poisoned<U8>(0, 0xFF00);
+  EXPECT_NOT_POISONED(__builtin_ia32_pdep_di(y, 0xFF));
+  EXPECT_POISONED(__builtin_ia32_pdep_di(y, 0x1FF));
+  EXPECT_NOT_POISONED(__builtin_ia32_pdep_di(y, 0xFF0000000000));
+  EXPECT_POISONED(__builtin_ia32_pdep_di(y, 0x1FF000000000000));
+
+  EXPECT_NOT_POISONED(__builtin_ia32_pdep_di(y, 0x1FF00) & 0xFF);
+  EXPECT_POISONED(__builtin_ia32_pdep_di(0, Poisoned<U4>(0xF, 1)));
+}
+
+__attribute__((target("bmi,bmi2")))
+static void TestPEXT() {
+  U4 x = Poisoned<U4>(0, 0xFF00);
+  EXPECT_NOT_POISONED(__builtin_ia32_pext_si(x, 0xFF));
+  EXPECT_POISONED(__builtin_ia32_pext_si(x, 0x1FF));
+  EXPECT_POISONED(__builtin_ia32_pext_si(x, 0x100));
+  EXPECT_POISONED(__builtin_ia32_pext_si(x, 0x1000));
+  EXPECT_NOT_POISONED(__builtin_ia32_pext_si(x, 0x10000));
+
+  EXPECT_POISONED(__builtin_ia32_pext_si(0xFF00, Poisoned<U4>(0xFF, 1)));
+
+  U8 y = Poisoned<U8>(0, 0xFF0000000000);
+  EXPECT_NOT_POISONED(__builtin_ia32_pext_di(y, 0xFF00000000));
+  EXPECT_POISONED(__builtin_ia32_pext_di(y, 0x1FF00000000));
+  EXPECT_POISONED(__builtin_ia32_pext_di(y, 0x10000000000));
+  EXPECT_POISONED(__builtin_ia32_pext_di(y, 0x100000000000));
+  EXPECT_NOT_POISONED(__builtin_ia32_pext_di(y, 0x1000000000000));
+
+  EXPECT_POISONED(__builtin_ia32_pext_di(0xFF00, Poisoned<U8>(0xFF, 1)));
+}
+
+TEST(MemorySanitizer, Bmi) {
+  if (HaveBmi()) {
+    TestBZHI();
+    TestBEXTR();
+    TestPDEP();
+    TestPEXT();
+  }
+}
diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 4b217e947f1..bfa1d0f9f7f 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -2928,6 +2928,26 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     return true;
   }
 
+  // Instrument BMI / BMI2 intrinsics.
+  // All of these intrinsics are Z = I(X, Y)
+  // where the types of all operands and the result match, and are either i32 or i64.
+  // The following instrumentation happens to work for all of them:
+  //   Sz = I(Sx, Y) | (sext (Sy != 0))
+  void handleBmiIntrinsic(IntrinsicInst &I) {
+    IRBuilder<> IRB(&I);
+    Type *ShadowTy = getShadowTy(&I);
+
+    // If any bit of the mask operand is poisoned, then the whole thing is.
+    Value *SMask = getShadow(&I, 1);
+    SMask = IRB.CreateSExt(IRB.CreateICmpNE(SMask, getCleanShadow(ShadowTy)),
+                           ShadowTy);
+    // Apply the same intrinsic to the shadow of the first operand.
+    Value *S = IRB.CreateCall(I.getCalledFunction(),
+                              {getShadow(&I, 0), I.getOperand(1)});
+    S = IRB.CreateOr(SMask, S);
+    setShadow(&I, S);
+    setOriginForNaryOp(I);
+  }
 
   void visitIntrinsicInst(IntrinsicInst &I) {
     switch (I.getIntrinsicID()) {
@@ -3144,6 +3164,17 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       handleVectorComparePackedIntrinsic(I);
       break;
 
+    case Intrinsic::x86_bmi_bextr_32:
+    case Intrinsic::x86_bmi_bextr_64:
+    case Intrinsic::x86_bmi_bzhi_32:
+    case Intrinsic::x86_bmi_bzhi_64:
+    case Intrinsic::x86_bmi_pdep_32:
+    case Intrinsic::x86_bmi_pdep_64:
+    case Intrinsic::x86_bmi_pext_32:
+    case Intrinsic::x86_bmi_pext_64:
+      handleBmiIntrinsic(I);
+      break;
+
     case Intrinsic::is_constant:
       // The result of llvm.is.constant() is always defined.
       setShadow(&I, getCleanShadow(&I));
diff --git a/llvm/test/Instrumentation/MemorySanitizer/bmi.ll b/llvm/test/Instrumentation/MemorySanitizer/bmi.ll
new file mode 100644
index 00000000000..27f0097cd7d
--- /dev/null
+++ b/llvm/test/Instrumentation/MemorySanitizer/bmi.ll
@@ -0,0 +1,147 @@
+; RUN: opt < %s -msan-check-access-address=0 -S -passes=msan 2>&1 | FileCheck  \
+; RUN: %s
+; RUN: opt < %s -msan -msan-check-access-address=0 -S | FileCheck %s
+; REQUIRES: x86-registered-target
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare i32 @llvm.x86.bmi.bzhi.32(i32, i32)
+declare i32 @llvm.x86.bmi.bextr.32(i32, i32)
+declare i32 @llvm.x86.bmi.pdep.32(i32, i32)
+declare i32 @llvm.x86.bmi.pext.32(i32, i32)
+
+declare i64 @llvm.x86.bmi.bzhi.64(i64, i64)
+declare i64 @llvm.x86.bmi.bextr.64(i64, i64)
+declare i64 @llvm.x86.bmi.pdep.64(i64, i64)
+declare i64 @llvm.x86.bmi.pext.64(i64, i64)
+
+define i32 @Test_bzhi_32(i32 %a, i32 %b) sanitize_memory {
+entry:
+  %c = tail call i32 @llvm.x86.bmi.bzhi.32(i32 %a, i32 %b)
+  ret i32 %c
+}
+
+; CHECK-LABEL: @Test_bzhi_32(
+; CHECK-DAG: %[[SA:.*]] = load i32, {{.*}}@__msan_param_tls to i32*)
+; CHECK-DAG: %[[SB:.*]] = load i32, {{.*}}@__msan_param_tls to i64), i64 8)
+; CHECK-DAG: %[[SB0:.*]] = icmp ne i32 %[[SB]], 0
+; CHECK-DAG: %[[SB1:.*]] = sext i1 %[[SB0]] to i32
+; CHECK-DAG: %[[X:.*]] = call i32 @llvm.x86.bmi.bzhi.32(i32 %[[SA]], i32 %b)
+; CHECK-DAG: %[[S:.*]] = or i32 %[[SB1]], %[[X]]
+; CHECK-DAG: store i32 %[[S]], {{.*}}@__msan_retval_tls
+; CHECK: ret i32
+
+define i64 @Test_bzhi_64(i64 %a, i64 %b) sanitize_memory {
+entry:
+  %c = tail call i64 @llvm.x86.bmi.bzhi.64(i64 %a, i64 %b)
+  ret i64 %c
+}
+
+; CHECK-LABEL: @Test_bzhi_64(
+; CHECK-DAG: %[[SA:.*]] = load i64, {{.*}}@__msan_param_tls, i32 0, i32 0
+; CHECK-DAG: %[[SB:.*]] = load i64, {{.*}}@__msan_param_tls to i64), i64 8)
+; CHECK-DAG: %[[SB0:.*]] = icmp ne i64 %[[SB]], 0
+; CHECK-DAG: %[[SB1:.*]] = sext i1 %[[SB0]] to i64
+; CHECK-DAG: %[[X:.*]] = call i64 @llvm.x86.bmi.bzhi.64(i64 %[[SA]], i64 %b)
+; CHECK-DAG: %[[S:.*]] = or i64 %[[SB1]], %[[X]]
+; CHECK-DAG: store i64 %[[S]], {{.*}}@__msan_retval_tls
+; CHECK: ret i64
+
+
+define i32 @Test_bextr_32(i32 %a, i32 %b) sanitize_memory {
+entry:
+  %c = tail call i32 @llvm.x86.bmi.bextr.32(i32 %a, i32 %b)
+  ret i32 %c
+}
+
+; CHECK-LABEL: @Test_bextr_32(
+; CHECK-DAG: %[[SA:.*]] = load i32, {{.*}}@__msan_param_tls to i32*)
+; CHECK-DAG: %[[SB:.*]] = load i32, {{.*}}@__msan_param_tls to i64), i64 8)
+; CHECK-DAG: %[[SB0:.*]] = icmp ne i32 %[[SB]], 0
+; CHECK-DAG: %[[SB1:.*]] = sext i1 %[[SB0]] to i32
+; CHECK-DAG: %[[X:.*]] = call i32 @llvm.x86.bmi.bextr.32(i32 %[[SA]], i32 %b)
+; CHECK-DAG: %[[S:.*]] = or i32 %[[SB1]], %[[X]]
+; CHECK-DAG: store i32 %[[S]], {{.*}}@__msan_retval_tls
+; CHECK: ret i32
+
+define i64 @Test_bextr_64(i64 %a, i64 %b) sanitize_memory {
+entry:
+  %c = tail call i64 @llvm.x86.bmi.bextr.64(i64 %a, i64 %b)
+  ret i64 %c
+}
+
+; CHECK-LABEL: @Test_bextr_64(
+; CHECK-DAG: %[[SA:.*]] = load i64, {{.*}}@__msan_param_tls, i32 0, i32 0
+; CHECK-DAG: %[[SB:.*]] = load i64, {{.*}}@__msan_param_tls to i64), i64 8)
+; CHECK-DAG: %[[SB0:.*]] = icmp ne i64 %[[SB]], 0
+; CHECK-DAG: %[[SB1:.*]] = sext i1 %[[SB0]] to i64
+; CHECK-DAG: %[[X:.*]] = call i64 @llvm.x86.bmi.bextr.64(i64 %[[SA]], i64 %b)
+; CHECK-DAG: %[[S:.*]] = or i64 %[[SB1]], %[[X]]
+; CHECK-DAG: store i64 %[[S]], {{.*}}@__msan_retval_tls
+; CHECK: ret i64
+
+
+define i32 @Test_pdep_32(i32 %a, i32 %b) sanitize_memory {
+entry:
+  %c = tail call i32 @llvm.x86.bmi.pdep.32(i32 %a, i32 %b)
+  ret i32 %c
+}
+
+; CHECK-LABEL: @Test_pdep_32(
+; CHECK-DAG: %[[SA:.*]] = load i32, {{.*}}@__msan_param_tls to i32*)
+; CHECK-DAG: %[[SB:.*]] = load i32, {{.*}}@__msan_param_tls to i64), i64 8)
+; CHECK-DAG: %[[SB0:.*]] = icmp ne i32 %[[SB]], 0
+; CHECK-DAG: %[[SB1:.*]] = sext i1 %[[SB0]] to i32
+; CHECK-DAG: %[[X:.*]] = call i32 @llvm.x86.bmi.pdep.32(i32 %[[SA]], i32 %b)
+; CHECK-DAG: %[[S:.*]] = or i32 %[[SB1]], %[[X]]
+; CHECK-DAG: store i32 %[[S]], {{.*}}@__msan_retval_tls
+; CHECK: ret i32
+
+define i64 @Test_pdep_64(i64 %a, i64 %b) sanitize_memory {
+entry:
+  %c = tail call i64 @llvm.x86.bmi.pdep.64(i64 %a, i64 %b)
+  ret i64 %c
+}
+
+; CHECK-LABEL: @Test_pdep_64(
+; CHECK-DAG: %[[SA:.*]] = load i64, {{.*}}@__msan_param_tls, i32 0, i32 0
+; CHECK-DAG: %[[SB:.*]] = load i64, {{.*}}@__msan_param_tls to i64), i64 8)
+; CHECK-DAG: %[[SB0:.*]] = icmp ne i64 %[[SB]], 0
+; CHECK-DAG: %[[SB1:.*]] = sext i1 %[[SB0]] to i64
+; CHECK-DAG: %[[X:.*]] = call i64 @llvm.x86.bmi.pdep.64(i64 %[[SA]], i64 %b)
+; CHECK-DAG: %[[S:.*]] = or i64 %[[SB1]], %[[X]]
+; CHECK-DAG: store i64 %[[S]], {{.*}}@__msan_retval_tls
+; CHECK: ret i64
+
+define i32 @Test_pext_32(i32 %a, i32 %b) sanitize_memory {
+entry:
+  %c = tail call i32 @llvm.x86.bmi.pext.32(i32 %a, i32 %b)
+  ret i32 %c
+}
+
+; CHECK-LABEL: @Test_pext_32(
+; CHECK-DAG: %[[SA:.*]] = load i32, {{.*}}@__msan_param_tls to i32*)
+; CHECK-DAG: %[[SB:.*]] = load i32, {{.*}}@__msan_param_tls to i64), i64 8)
+; CHECK-DAG: %[[SB0:.*]] = icmp ne i32 %[[SB]], 0
+; CHECK-DAG: %[[SB1:.*]] = sext i1 %[[SB0]] to i32
+; CHECK-DAG: %[[X:.*]] = call i32 @llvm.x86.bmi.pext.32(i32 %[[SA]], i32 %b)
+; CHECK-DAG: %[[S:.*]] = or i32 %[[SB1]], %[[X]]
+; CHECK-DAG: store i32 %[[S]], {{.*}}@__msan_retval_tls
+; CHECK: ret i32
+
+define i64 @Test_pext_64(i64 %a, i64 %b) sanitize_memory {
+entry:
+  %c = tail call i64 @llvm.x86.bmi.pext.64(i64 %a, i64 %b)
+  ret i64 %c
+}
+
+; CHECK-LABEL: @Test_pext_64(
+; CHECK-DAG: %[[SA:.*]] = load i64, {{.*}}@__msan_param_tls, i32 0, i32 0
+; CHECK-DAG: %[[SB:.*]] = load i64, {{.*}}@__msan_param_tls to i64), i64 8)
+; CHECK-DAG: %[[SB0:.*]] = icmp ne i64 %[[SB]], 0
+; CHECK-DAG: %[[SB1:.*]] = sext i1 %[[SB0]] to i64
+; CHECK-DAG: %[[X:.*]] = call i64 @llvm.x86.bmi.pext.64(i64 %[[SA]], i64 %b)
+; CHECK-DAG: %[[S:.*]] = or i64 %[[SB1]], %[[X]]
+; CHECK-DAG: store i64 %[[S]], {{.*}}@__msan_retval_tls
+; CHECK: ret i64