[MSAN] Add handleCountZeroes for ctlz and cttz.

[kda]

Summary

This addresses a bug where vector versions of ctlz are creating false positive reports.

Depends on D136369

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2974

isn't this implementation strictly better than default? It handles the second operand correctly.

2976

Is there a problem with scalable vectors? You could just use Type::isVectorTy() otherwise.

2994

use CreateIsNull

[kda]

updated for input from eugenis.

[eugenis]

LGTM, but I'd merge the two changes. I don't really see the point of submitting a test for the broken behavior first.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2982

This line seems unnecessary.

[kda]

In D136523#3880730, @eugenis wrote:

LGTM, but I'd merge the two changes. I don't really see the point of submitting a test for the broken behavior first.

The test does not demonstrate "broken" behavior but rather shows what is generated before my change.
I will keep them separate as I would like to show how my code changes what is generated.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2974

Agreed. I added this comment (conditional) prior to adding handling of second parameter.

2976

dropped (see previous comment).

2982

Oh. Again, I rearranged the code and forgot about this dangling item. I originally had the OR outside the conditional, thus I needed a BoolZeroPoison that was passive.

Removed.

[kda]

removed unused line

[vitalybuka]

In D136523#3880730, @eugenis wrote:

LGTM, but I'd merge the two changes. I don't really see the point of submitting a test for the broken behavior first.

Actually I rely like splitting stuff like this
the effect of non-NFC part is very visible

[vitalybuka]

LGTM with nits

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2975

Operand -> Src to match the name in LangRef

2978

CreateIsNull like you do below is less code

2983

cast<> and it can be only the Constant
then you can remove assert() as cast<> already has one

2983

constOperant -> Poison (or something similar to the name in LangRef)
Note: vars are capitalized

[kda]

improvements based on vitalybuka feedback

[kda]

rebase.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

void handleBswap(IntrinsicInst &I) {
     Value *Op = I.getArgOperand(0);
     Type *OpType = Op->getType();
     Function *BswapFunc = Intrinsic::getDeclaration(
         F.getParent(), Intrinsic::bswap, makeArrayRef(&OpType, 1));
     setShadow(&I, IRB.CreateCall(BswapFunc, getShadow(Op)));
     setOrigin(&I, getOrigin(Op));
   }
 
+  void handleCountZeroes(IntrinsicInst &I) {
+    IRBuilder<> IRB(&I);
+    Value *Src = I.getArgOperand(0);
+
+    // Set the Output shadow based on input Shadow
+    Value *BoolShadow = IRB.CreateIsNotNull(getShadow(Src), "_mscz_bs");
+
+    // If zero poison is requested, mix in with the shadow
+    Constant *IsZeroPoison = cast<Constant>(I.getOperand(1));
+    if (!IsZeroPoison->isZeroValue()) {
+      Value *BoolZeroPoison = IRB.CreateIsNull(Src, "_mscz_bzp");
+      BoolShadow = IRB.CreateOr(BoolShadow, BoolZeroPoison, "_mscz_bs");
+    }
+
+    Value *OutputShadow =
+        IRB.CreateSExt(BoolShadow, getShadowTy(Src), "_mscz_os");
+
+    setShadow(&I, OutputShadow);
+    setOriginForNaryOp(I);
+  }
+
   // Instrument vector convert intrinsic.
   //
   // This function instruments intrinsics like cvtsi2ss:
   // %Out = int_xxx_cvtyyy(%ConvertOp)
   // or
   // %Out = int_xxx_cvtyyy(%CopyOp, %ConvertOp)
   // Intrinsic converts \p NumUsedElements elements of \p ConvertOp to the same
   // number \p Out elements, and (if has 2 arguments) copies the rest of the
   // elements from \p CopyOp.
case Intrinsic::lifetime_start:
       break;
     case Intrinsic::launder_invariant_group:
     case Intrinsic::strip_invariant_group:
       handleInvariantGroup(I);
       break;
     case Intrinsic::bswap:
       handleBswap(I);
       break;
+    case Intrinsic::ctlz:
+    case Intrinsic::cttz:
+      handleCountZeroes(I);
+      break;
     case Intrinsic::masked_compressstore:
       handleMaskedCompressStore(I);
       break;
     case Intrinsic::masked_expandload:
       handleMaskedExpandLoad(I);
       break;
     case Intrinsic::masked_gather:
       handleMaskedGather(I);

llvm/test/Instrumentation/MemorySanitizer/count-zeroes.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt %s -S -passes=msan 2>&1 | FileCheck %s
 
 target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 declare i64 @llvm.ctlz.i64(i64, i1) nounwind readnone
 define i64 @test_ctlz_i64_zeropoison(i64 %v) #0 {
 ; CHECK-LABEL: @test_ctlz_i64_zeropoison(
 ; CHECK-NEXT:    [[TMP1:%.*]] = load i64, ptr @__msan_param_tls, align 8
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i64 [[TMP1]], 0
-; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP2:%.*]], label [[TMP3:%.*]], !prof [[PROF0:![0-9]+]]
-; CHECK:       2:
-; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR4:[0-9]+]]
-; CHECK-NEXT:    unreachable
-; CHECK:       3:
-; CHECK-NEXT:    [[RES:%.*]] = call i64 @llvm.ctlz.i64(i64 [[V:%.*]], i1 true)
-; CHECK-NEXT:    store i64 0, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    [[_MSCZ_BS:%.*]] = icmp ne i64 [[TMP1]], 0
+; CHECK-NEXT:    [[_MSCZ_BZP:%.*]] = icmp eq i64 [[V:%.*]], 0
+; CHECK-NEXT:    [[_MSCZ_BS1:%.*]] = or i1 [[_MSCZ_BS]], [[_MSCZ_BZP]]
+; CHECK-NEXT:    [[_MSCZ_OS:%.*]] = sext i1 [[_MSCZ_BS1]] to i64
+; CHECK-NEXT:    [[RES:%.*]] = call i64 @llvm.ctlz.i64(i64 [[V]], i1 true)
+; CHECK-NEXT:    store i64 [[_MSCZ_OS]], ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret i64 [[RES]]
 ;
   %res = call i64 @llvm.ctlz.i64(i64 %v, i1 true) ; <<i64>> [#uses=1]
   ret i64 %res
 }
 define i64 @test_ctlz_i64_nozeropoison(i64 %v) #0 {
 ; CHECK-LABEL: @test_ctlz_i64_nozeropoison(
 ; CHECK-NEXT:    [[TMP1:%.*]] = load i64, ptr @__msan_param_tls, align 8
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i64 [[TMP1]], 0
-; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP2:%.*]], label [[TMP3:%.*]], !prof [[PROF0]]
-; CHECK:       2:
-; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR4]]
-; CHECK-NEXT:    unreachable
-; CHECK:       3:
+; CHECK-NEXT:    [[_MSCZ_BS:%.*]] = icmp ne i64 [[TMP1]], 0
+; CHECK-NEXT:    [[_MSCZ_OS:%.*]] = sext i1 [[_MSCZ_BS]] to i64
 ; CHECK-NEXT:    [[RES:%.*]] = call i64 @llvm.ctlz.i64(i64 [[V:%.*]], i1 false)
-; CHECK-NEXT:    store i64 0, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    store i64 [[_MSCZ_OS]], ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret i64 [[RES]]
 ;
   %res = call i64 @llvm.ctlz.i64(i64 %v, i1 false) ; <<i64>> [#uses=1]
   ret i64 %res
 }
 
 declare <2 x i64> @llvm.ctlz.v2i64(<2 x i64>, i1) nounwind readnone
 define <2 x i64> @test_ctlz_v2i64_zeropoison(<2 x i64> %v) #0 {
 ; CHECK-LABEL: @test_ctlz_v2i64_zeropoison(
 ; CHECK-NEXT:    [[TMP1:%.*]] = load <2 x i64>, ptr @__msan_param_tls, align 8
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    [[TMP2:%.*]] = bitcast <2 x i64> [[TMP1]] to i128
-; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i128 [[TMP2]], 0
-; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP3:%.*]], label [[TMP4:%.*]], !prof [[PROF0]]
-; CHECK:       3:
-; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR4]]
-; CHECK-NEXT:    unreachable
-; CHECK:       4:
-; CHECK-NEXT:    [[RES:%.*]] = call <2 x i64> @llvm.ctlz.v2i64(<2 x i64> [[V:%.*]], i1 true)
-; CHECK-NEXT:    store <2 x i64> zeroinitializer, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    [[_MSCZ_BS:%.*]] = icmp ne <2 x i64> [[TMP1]], zeroinitializer
+; CHECK-NEXT:    [[_MSCZ_BZP:%.*]] = icmp eq <2 x i64> [[V:%.*]], zeroinitializer
+; CHECK-NEXT:    [[_MSCZ_BS1:%.*]] = or <2 x i1> [[_MSCZ_BS]], [[_MSCZ_BZP]]
+; CHECK-NEXT:    [[_MSCZ_OS:%.*]] = sext <2 x i1> [[_MSCZ_BS1]] to <2 x i64>
+; CHECK-NEXT:    [[RES:%.*]] = call <2 x i64> @llvm.ctlz.v2i64(<2 x i64> [[V]], i1 true)
+; CHECK-NEXT:    store <2 x i64> [[_MSCZ_OS]], ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret <2 x i64> [[RES]]
 ;
   %res = call <2 x i64> @llvm.ctlz.v2i64(<2 x i64> %v, i1 true) ; <<2 x i64>> [#uses=1]
   ret <2 x i64> %res
 }
 define <2 x i64> @test_ctlz_v2i64_nozeropoison(<2 x i64> %v) #0 {
 ; CHECK-LABEL: @test_ctlz_v2i64_nozeropoison(
 ; CHECK-NEXT:    [[TMP1:%.*]] = load <2 x i64>, ptr @__msan_param_tls, align 8
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    [[TMP2:%.*]] = bitcast <2 x i64> [[TMP1]] to i128
-; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i128 [[TMP2]], 0
-; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP3:%.*]], label [[TMP4:%.*]], !prof [[PROF0]]
-; CHECK:       3:
-; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR4]]
-; CHECK-NEXT:    unreachable
-; CHECK:       4:
+; CHECK-NEXT:    [[_MSCZ_BS:%.*]] = icmp ne <2 x i64> [[TMP1]], zeroinitializer
+; CHECK-NEXT:    [[_MSCZ_OS:%.*]] = sext <2 x i1> [[_MSCZ_BS]] to <2 x i64>
 ; CHECK-NEXT:    [[RES:%.*]] = call <2 x i64> @llvm.ctlz.v2i64(<2 x i64> [[V:%.*]], i1 false)
-; CHECK-NEXT:    store <2 x i64> zeroinitializer, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    store <2 x i64> [[_MSCZ_OS]], ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret <2 x i64> [[RES]]
 ;
   %res = call <2 x i64> @llvm.ctlz.v2i64(<2 x i64> %v, i1 false) ; <<2 x i64>> [#uses=1]
   ret <2 x i64> %res
 }
 
 declare i64 @llvm.cttz.i64(i64, i1) nounwind readnone
 define i64 @test_cttz_i64_zeropoison(i64 %v) #0 {
 ; CHECK-LABEL: @test_cttz_i64_zeropoison(
 ; CHECK-NEXT:    [[TMP1:%.*]] = load i64, ptr @__msan_param_tls, align 8
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i64 [[TMP1]], 0
-; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP2:%.*]], label [[TMP3:%.*]], !prof [[PROF0]]
-; CHECK:       2:
-; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR4]]
-; CHECK-NEXT:    unreachable
-; CHECK:       3:
-; CHECK-NEXT:    [[RES:%.*]] = call i64 @llvm.cttz.i64(i64 [[V:%.*]], i1 true)
-; CHECK-NEXT:    store i64 0, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    [[_MSCZ_BS:%.*]] = icmp ne i64 [[TMP1]], 0
+; CHECK-NEXT:    [[_MSCZ_BZP:%.*]] = icmp eq i64 [[V:%.*]], 0
+; CHECK-NEXT:    [[_MSCZ_BS1:%.*]] = or i1 [[_MSCZ_BS]], [[_MSCZ_BZP]]
+; CHECK-NEXT:    [[_MSCZ_OS:%.*]] = sext i1 [[_MSCZ_BS1]] to i64
+; CHECK-NEXT:    [[RES:%.*]] = call i64 @llvm.cttz.i64(i64 [[V]], i1 true)
+; CHECK-NEXT:    store i64 [[_MSCZ_OS]], ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret i64 [[RES]]
 ;
   %res = call i64 @llvm.cttz.i64(i64 %v, i1 true) ; <<i64>> [#uses=1]
   ret i64 %res
 }
 define i64 @test_cttz_i64_nozeropoison(i64 %v) #0 {
 ; CHECK-LABEL: @test_cttz_i64_nozeropoison(
 ; CHECK-NEXT:    [[TMP1:%.*]] = load i64, ptr @__msan_param_tls, align 8
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i64 [[TMP1]], 0
-; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP2:%.*]], label [[TMP3:%.*]], !prof [[PROF0]]
-; CHECK:       2:
-; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR4]]
-; CHECK-NEXT:    unreachable
-; CHECK:       3:
+; CHECK-NEXT:    [[_MSCZ_BS:%.*]] = icmp ne i64 [[TMP1]], 0
+; CHECK-NEXT:    [[_MSCZ_OS:%.*]] = sext i1 [[_MSCZ_BS]] to i64
 ; CHECK-NEXT:    [[RES:%.*]] = call i64 @llvm.cttz.i64(i64 [[V:%.*]], i1 false)
-; CHECK-NEXT:    store i64 0, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    store i64 [[_MSCZ_OS]], ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret i64 [[RES]]
 ;
   %res = call i64 @llvm.cttz.i64(i64 %v, i1 false) ; <<i64>> [#uses=1]
   ret i64 %res
 }
 
 declare <2 x i64> @llvm.cttz.v2i64(<2 x i64>, i1) nounwind readnone
 define <2 x i64> @test_cttz_v2i64_zeropoison(<2 x i64> %v) #0 {
 ; CHECK-LABEL: @test_cttz_v2i64_zeropoison(
 ; CHECK-NEXT:    [[TMP1:%.*]] = load <2 x i64>, ptr @__msan_param_tls, align 8
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    [[TMP2:%.*]] = bitcast <2 x i64> [[TMP1]] to i128
-; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i128 [[TMP2]], 0
-; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP3:%.*]], label [[TMP4:%.*]], !prof [[PROF0]]
-; CHECK:       3:
-; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR4]]
-; CHECK-NEXT:    unreachable
-; CHECK:       4:
-; CHECK-NEXT:    [[RES:%.*]] = call <2 x i64> @llvm.cttz.v2i64(<2 x i64> [[V:%.*]], i1 true)
-; CHECK-NEXT:    store <2 x i64> zeroinitializer, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    [[_MSCZ_BS:%.*]] = icmp ne <2 x i64> [[TMP1]], zeroinitializer
+; CHECK-NEXT:    [[_MSCZ_BZP:%.*]] = icmp eq <2 x i64> [[V:%.*]], zeroinitializer
+; CHECK-NEXT:    [[_MSCZ_BS1:%.*]] = or <2 x i1> [[_MSCZ_BS]], [[_MSCZ_BZP]]
+; CHECK-NEXT:    [[_MSCZ_OS:%.*]] = sext <2 x i1> [[_MSCZ_BS1]] to <2 x i64>
+; CHECK-NEXT:    [[RES:%.*]] = call <2 x i64> @llvm.cttz.v2i64(<2 x i64> [[V]], i1 true)
+; CHECK-NEXT:    store <2 x i64> [[_MSCZ_OS]], ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret <2 x i64> [[RES]]
 ;
   %res = call <2 x i64> @llvm.cttz.v2i64(<2 x i64> %v, i1 true) ; <<2 x i64>> [#uses=1]
   ret <2 x i64> %res
 }
 define <2 x i64> @test_cttz_v2i64_nozeropoison(<2 x i64> %v) #0 {
 ; CHECK-LABEL: @test_cttz_v2i64_nozeropoison(
 ; CHECK-NEXT:    [[TMP1:%.*]] = load <2 x i64>, ptr @__msan_param_tls, align 8
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    [[TMP2:%.*]] = bitcast <2 x i64> [[TMP1]] to i128
-; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i128 [[TMP2]], 0
-; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP3:%.*]], label [[TMP4:%.*]], !prof [[PROF0]]
-; CHECK:       3:
-; CHECK-NEXT:    call void @__msan_warning_noreturn() #[[ATTR4]]
-; CHECK-NEXT:    unreachable
-; CHECK:       4:
+; CHECK-NEXT:    [[_MSCZ_BS:%.*]] = icmp ne <2 x i64> [[TMP1]], zeroinitializer
+; CHECK-NEXT:    [[_MSCZ_OS:%.*]] = sext <2 x i1> [[_MSCZ_BS]] to <2 x i64>
 ; CHECK-NEXT:    [[RES:%.*]] = call <2 x i64> @llvm.cttz.v2i64(<2 x i64> [[V:%.*]], i1 false)
-; CHECK-NEXT:    store <2 x i64> zeroinitializer, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    store <2 x i64> [[_MSCZ_OS]], ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret <2 x i64> [[RES]]
 ;
   %res = call <2 x i64> @llvm.cttz.v2i64(<2 x i64> %v, i1 false) ; <<2 x i64>> [#uses=1]
   ret <2 x i64> %res
 }
 
 
 attributes #0 = { sanitize_memory }