diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp --- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp @@ -2970,6 +2970,27 @@ 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(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: @@ -3651,6 +3672,10 @@ case Intrinsic::bswap: handleBswap(I); break; + case Intrinsic::ctlz: + case Intrinsic::cttz: + handleCountZeroes(I); + break; case Intrinsic::masked_compressstore: handleMaskedCompressStore(I); break; diff --git a/llvm/test/Instrumentation/MemorySanitizer/count-zeroes.ll b/llvm/test/Instrumentation/MemorySanitizer/count-zeroes.ll --- a/llvm/test/Instrumentation/MemorySanitizer/count-zeroes.ll +++ b/llvm/test/Instrumentation/MemorySanitizer/count-zeroes.ll @@ -9,14 +9,12 @@ ; 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) ; <> [#uses=1] @@ -26,14 +24,10 @@ ; 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) ; <> [#uses=1] @@ -45,15 +39,12 @@ ; 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] @@ -63,15 +54,10 @@ ; 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] @@ -83,14 +69,12 @@ ; 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) ; <> [#uses=1] @@ -100,14 +84,10 @@ ; 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) ; <> [#uses=1] @@ -119,15 +99,12 @@ ; 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] @@ -137,15 +114,10 @@ ; 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]