Skip to content

[msan][NFCI] Generalize handleVectorPmaddIntrinsic() #166282

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 5, 2025
Merged

[msan][NFCI] Generalize handleVectorPmaddIntrinsic() #166282

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
989fd15
[msan][NFCI] Generalize handleVectorPmaddIntrinsic()
thurstond Nov 4, 2025
19864e6
Drive-by comment fix
thurstond Nov 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 49 additions & 28 deletions llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3903,7 +3903,12 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
// adding/"accumulating" %s. "Accumulation" stores the result in one
// of the source registers, but this accumulate vs. add distinction
// is lost when dealing with LLVM intrinsics.)
//
// ZeroPurifies means that multiplying a known-zero with an uninitialized
// value results in an initialized value. This is applicable for integer
// multiplication, but not floating-point (counter-example: NaN).
void handleVectorPmaddIntrinsic(IntrinsicInst &I, unsigned ReductionFactor,
bool ZeroPurifies,
unsigned EltSizeInBits = 0) {
IRBuilder<> IRB(&I);

Expand Down Expand Up @@ -3945,7 +3950,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
assert(AccumulatorType == ReturnType);
}

FixedVectorType *ImplicitReturnType = ReturnType;
FixedVectorType *ImplicitReturnType =
cast<FixedVectorType>(getShadowTy(ReturnType));
Comment on lines +3953 to +3954
Copy link

Copilot AI Nov 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The assignment of ImplicitReturnType is overwritten at line 3957 when EltSizeInBits is non-zero. This initialization should be moved inside the else block (line 3968) to avoid unnecessary computation when it will be replaced.

Copilot uses AI. Check for mistakes.
// Step 1: instrument multiplication of corresponding vector elements
if (EltSizeInBits) {
ImplicitReturnType = cast<FixedVectorType>(
Expand All @@ -3964,30 +3970,40 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
ReturnType->getNumElements() * ReductionFactor);
}

// Multiplying an *initialized* zero by an uninitialized element results in
// an initialized zero element.
//
// This is analogous to bitwise AND, where "AND" of 0 and a poisoned value
// results in an unpoisoned value. We can therefore adapt the visitAnd()
// instrumentation:
// OutShadow = (SaNonZero & SbNonZero)
// | (VaNonZero & SbNonZero)
// | (SaNonZero & VbNonZero)
// where non-zero is checked on a per-element basis (not per bit).
Value *SZero = Constant::getNullValue(Va->getType());
Value *VZero = Constant::getNullValue(Sa->getType());
Value *SaNonZero = IRB.CreateICmpNE(Sa, SZero);
Value *SbNonZero = IRB.CreateICmpNE(Sb, SZero);
Value *VaNonZero = IRB.CreateICmpNE(Va, VZero);
Value *VbNonZero = IRB.CreateICmpNE(Vb, VZero);

Value *SaAndSbNonZero = IRB.CreateAnd(SaNonZero, SbNonZero);
Value *VaAndSbNonZero = IRB.CreateAnd(VaNonZero, SbNonZero);
Value *SaAndVbNonZero = IRB.CreateAnd(SaNonZero, VbNonZero);

// Each element of the vector is represented by a single bit (poisoned or
// not) e.g., <8 x i1>.
Value *And = IRB.CreateOr({SaAndSbNonZero, VaAndSbNonZero, SaAndVbNonZero});
Value *SaNonZero = IRB.CreateIsNotNull(Sa);
Value *SbNonZero = IRB.CreateIsNotNull(Sb);
Value *And;
if (ZeroPurifies) {
// Multiplying an *initialized* zero by an uninitialized element results
// in an initialized zero element.
//
// This is analogous to bitwise AND, where "AND" of 0 and a poisoned value
// results in an unpoisoned value. We can therefore adapt the visitAnd()
// instrumentation:
// OutShadow = (SaNonZero & SbNonZero)
// | (VaNonZero & SbNonZero)
// | (SaNonZero & VbNonZero)
// where non-zero is checked on a per-element basis (not per bit).
Value *VaInt = Va;
Value *VbInt = Vb;
if (!Va->getType()->isIntegerTy()) {
Copy link

Copilot AI Nov 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The condition !Va->getType()->isIntegerTy() checks if the type is not integer, but Va is a vector type according to the context. This should check if the vector's element type is not integer using !Va->getType()->getScalarType()->isIntegerTy() instead.

Suggested change
if (!Va->getType()->isIntegerTy()) {
if (!Va->getType()->getScalarType()->isIntegerTy()) {

Copilot uses AI. Check for mistakes.
VaInt = CreateAppToShadowCast(IRB, Va);
VbInt = CreateAppToShadowCast(IRB, Vb);
}

Value *VaNonZero = IRB.CreateIsNotNull(VaInt);
Value *VbNonZero = IRB.CreateIsNotNull(VbInt);

Value *SaAndSbNonZero = IRB.CreateAnd(SaNonZero, SbNonZero);
Value *VaAndSbNonZero = IRB.CreateAnd(VaNonZero, SbNonZero);
Value *SaAndVbNonZero = IRB.CreateAnd(SaNonZero, VbNonZero);

And = IRB.CreateOr({SaAndSbNonZero, VaAndSbNonZero, SaAndVbNonZero});
} else {
And = IRB.CreateOr({SaNonZero, SbNonZero});
}

// Extend <8 x i1> to <8 x i16>.
// (The real pmadd intrinsic would have computed intermediate values of
Expand Down Expand Up @@ -5752,17 +5768,20 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
case Intrinsic::x86_ssse3_pmadd_ub_sw_128:
case Intrinsic::x86_avx2_pmadd_ub_sw:
case Intrinsic::x86_avx512_pmaddubs_w_512:
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/2);
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/2,
/*ZeroPurifies=*/true);
break;

// <1 x i64> @llvm.x86.ssse3.pmadd.ub.sw(<1 x i64>, <1 x i64>)
case Intrinsic::x86_ssse3_pmadd_ub_sw:
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/2, /*EltSize=*/8);
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/2,
/*ZeroPurifies=*/true, /*EltSizeInBits=*/8);
break;

// <1 x i64> @llvm.x86.mmx.pmadd.wd(<1 x i64>, <1 x i64>)
case Intrinsic::x86_mmx_pmadd_wd:
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/2, /*EltSize=*/16);
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/2,
/*ZeroPurifies=*/true, /*EltSizeInBits=*/16);
break;

// AVX Vector Neural Network Instructions: bytes
Expand Down Expand Up @@ -5848,7 +5867,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
case Intrinsic::x86_avx2_vpdpbuuds_128:
case Intrinsic::x86_avx2_vpdpbuuds_256:
case Intrinsic::x86_avx10_vpdpbuuds_512:
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/4, /*EltSize=*/8);
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/4,
/*ZeroPurifies=*/true, /*EltSizeInBits=*/8);
break;

// AVX Vector Neural Network Instructions: words
Expand Down Expand Up @@ -5901,7 +5921,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
case Intrinsic::x86_avx512_vpdpwssds_128:
case Intrinsic::x86_avx512_vpdpwssds_256:
case Intrinsic::x86_avx512_vpdpwssds_512:
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/2, /*EltSize=*/16);
handleVectorPmaddIntrinsic(I, /*ReductionFactor=*/2,
/*ZeroPurifies=*/true, /*EltSizeInBits=*/16);
break;

// TODO: Dot Product of BF16 Pairs Accumulated Into Packed Single
Expand Down