[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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@355348 91177308-0d34-0410-b5e6-96231b3b80d8

diff --git a/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 4b217e947f1f28e01f978d070a1c6b85a2b745f5..bfa1d0f9f7fc3c039dcbfe57c51da3e1ab74f2aa 100644 (file)
--- a/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/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/test/Instrumentation/MemorySanitizer/bmi.ll b/test/Instrumentation/MemorySanitizer/bmi.ll
new file mode 100644 (file)
index 0000000..27f0097
--- /dev/null
+++ b/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