Skip to content

JIT: Accelerate Vector.Dot for all base types #111853

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 7 commits into from
Mar 11, 2025
Merged

JIT: Accelerate Vector.Dot for all base types #111853

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
3562aea
accelerate Vector.Dot for all base types
saucecontrol Jan 27, 2025
41df6d6
add assert
saucecontrol Jan 27, 2025
f4ccc5d
remove sse4.1 long multiply fallback
saucecontrol Jan 27, 2025
e151a21
Merge remote-tracking branch 'upstream/main' into vdot
saucecontrol Feb 4, 2025
95e163d
Merge remote-tracking branch 'upstream/main' into vdot
saucecontrol Feb 18, 2025
02de0e9
Merge remote-tracking branch 'upstream/main' into vdot
saucecontrol Mar 10, 2025
7dea308
remove unnecessary AVX2 requirement for Sum
saucecontrol Mar 10, 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
62 changes: 31 additions & 31 deletions src/coreclr/jit/gentree.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21401,45 +21401,45 @@ GenTree* Compiler::gtNewSimdBinOpNode(
}
else if (varTypeIsLong(simdBaseType))
{
assert((simdSize == 16) || (simdSize == 32) || (simdSize == 64));
// This fallback path will be used only if the vpmullq instruction is not available.
// The implementation is a simple decomposition using pmuludq, which multiplies
// two uint32s and returns a uint64 result.
//
// aLo * bLo + ((aLo * bHi + aHi * bLo) << 32)

assert(((simdSize == 16) && compOpportunisticallyDependsOn(InstructionSet_SSE41)) ||
((simdSize == 32) && compOpportunisticallyDependsOn(InstructionSet_AVX2)));
assert(!canUseEvexEncodingDebugOnly());
assert((simdSize == 16) || compIsaSupportedDebugOnly(InstructionSet_AVX2));

// Make op1 and op2 multi-use:
GenTree* op1Dup = fgMakeMultiUse(&op1);
GenTree* op2Dup = fgMakeMultiUse(&op2);
NamedIntrinsic muludq = (simdSize == 16) ? NI_SSE2_Multiply : NI_AVX2_Multiply;

GenTree* op1Dup1 = fgMakeMultiUse(&op1);
GenTree* op1Dup2 = gtCloneExpr(op1Dup1);
GenTree* op2Dup1 = fgMakeMultiUse(&op2);
GenTree* op2Dup2 = gtCloneExpr(op2Dup1);

const bool is256 = simdSize == 32;
// Vector128<ulong> low = Sse2.Multiply(a.AsUInt32(), b.AsUInt32());
GenTree* low = gtNewSimdHWIntrinsicNode(type, op1, op2, muludq, CORINFO_TYPE_ULONG, simdSize);

// Vector256<ulong> tmp0 = Avx2.Multiply(left, right);
GenTreeHWIntrinsic* tmp0 =
gtNewSimdHWIntrinsicNode(type, op1, op2, is256 ? NI_AVX2_Multiply : NI_SSE2_Multiply,
CORINFO_TYPE_ULONG, simdSize);
// Vector128<ulong> mid = (b >>> 32).AsUInt64();
GenTree* mid = gtNewSimdBinOpNode(GT_RSZ, type, op2Dup1, gtNewIconNode(32), simdBaseJitType, simdSize);

// Vector256<uint> tmp1 = Avx2.Shuffle(right.AsUInt32(), ZWXY);
GenTree* shuffleMask = gtNewIconNode(SHUFFLE_ZWXY, TYP_INT);
GenTreeHWIntrinsic* tmp1 =
gtNewSimdHWIntrinsicNode(type, op2Dup, shuffleMask, is256 ? NI_AVX2_Shuffle : NI_SSE2_Shuffle,
CORINFO_TYPE_UINT, simdSize);
// mid = Sse2.Multiply(mid.AsUInt32(), a.AsUInt32());
mid = gtNewSimdHWIntrinsicNode(type, mid, op1Dup1, muludq, CORINFO_TYPE_ULONG, simdSize);

// Vector256<uint> tmp2 = Avx2.MultiplyLow(left.AsUInt32(), tmp1);
GenTree* tmp2 = gtNewSimdBinOpNode(GT_MUL, type, op1Dup, tmp1, CORINFO_TYPE_UINT, simdSize);
// Vector128<ulong> tmp = (a >>> 32).AsUInt64();
GenTree* tmp = gtNewSimdBinOpNode(GT_RSZ, type, op1Dup2, gtNewIconNode(32), simdBaseJitType, simdSize);

// Vector256<int> tmp3 = Avx2.HorizontalAdd(tmp2.AsInt32(), Vector256<int>.Zero);
GenTreeHWIntrinsic* tmp3 =
gtNewSimdHWIntrinsicNode(type, tmp2, gtNewZeroConNode(type),
is256 ? NI_AVX2_HorizontalAdd : NI_SSSE3_HorizontalAdd, CORINFO_TYPE_UINT,
simdSize);
// tmp = Sse2.Multiply(tmp.AsUInt32(), b.AsUInt32());
tmp = gtNewSimdHWIntrinsicNode(type, tmp, op2Dup2, muludq, CORINFO_TYPE_ULONG, simdSize);

// Vector256<int> tmp4 = Avx2.Shuffle(tmp3, YWXW);
shuffleMask = gtNewIconNode(SHUFFLE_YWXW, TYP_INT);
GenTreeHWIntrinsic* tmp4 =
gtNewSimdHWIntrinsicNode(type, tmp3, shuffleMask, is256 ? NI_AVX2_Shuffle : NI_SSE2_Shuffle,
CORINFO_TYPE_UINT, simdSize);
// mid += tmp;
mid = gtNewSimdBinOpNode(GT_ADD, type, mid, tmp, simdBaseJitType, simdSize);

// result = tmp0 + tmp4;
return gtNewSimdBinOpNode(GT_ADD, type, tmp0, tmp4, simdBaseJitType, simdSize);
// mid <<= 32;
mid = gtNewSimdBinOpNode(GT_LSH, type, mid, gtNewIconNode(32), simdBaseJitType, simdSize);

// return low + mid;
return gtNewSimdBinOpNode(GT_ADD, type, low, mid, simdBaseJitType, simdSize);
}
#elif defined(TARGET_ARM64)
if (varTypeIsLong(simdBaseType))
Expand Down Expand Up @@ -26070,7 +26070,7 @@ GenTree* Compiler::gtNewSimdSumNode(var_types type, GenTree* op1, CorInfoType si

if (simdSize == 32)
{
assert(compIsaSupportedDebugOnly(InstructionSet_AVX2));
assert(IsBaselineVector256IsaSupportedDebugOnly());
GenTree* op1Dup = fgMakeMultiUse(&op1);

op1 = gtNewSimdGetLowerNode(TYP_SIMD16, op1, simdBaseJitType, simdSize);
Expand Down
1 change: 1 addition & 0 deletions src/coreclr/jit/hwintrinsiclistxarch.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -302,6 +302,7 @@ HARDWARE_INTRINSIC(Vector512, Create,
HARDWARE_INTRINSIC(Vector512, CreateScalar, 64, -1, {INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid}, HW_Category_Helper, HW_Flag_SpecialImport|HW_Flag_NoCodeGen)
HARDWARE_INTRINSIC(Vector512, CreateScalarUnsafe, 64, 1, {INS_movd, INS_movd, INS_movd, INS_movd, INS_movd, INS_movd, INS_movd, INS_movd, INS_movss, INS_movsd_simd}, HW_Category_SIMDScalar, HW_Flag_SpecialImport|HW_Flag_SpecialCodeGen)
HARDWARE_INTRINSIC(Vector512, CreateSequence, 64, 2, {INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid}, HW_Category_Helper, HW_Flag_InvalidNodeId)
HARDWARE_INTRINSIC(Vector512, Dot, 64, 2, {INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid}, HW_Category_Helper, HW_Flag_SpecialImport|HW_Flag_NoCodeGen|HW_Flag_BaseTypeFromFirstArg)
HARDWARE_INTRINSIC(Vector512, Equals, 64, 2, {INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid}, HW_Category_Helper, HW_Flag_InvalidNodeId|HW_Flag_BaseTypeFromFirstArg)
HARDWARE_INTRINSIC(Vector512, EqualsAny, 64, 2, {INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid}, HW_Category_Helper, HW_Flag_InvalidNodeId|HW_Flag_BaseTypeFromFirstArg)
HARDWARE_INTRINSIC(Vector512, ExtractMostSignificantBits, 64, 1, {INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid, INS_invalid}, HW_Category_Helper, HW_Flag_InvalidNodeId|HW_Flag_BaseTypeFromFirstArg)
Expand Down
96 changes: 25 additions & 71 deletions src/coreclr/jit/hwintrinsicxarch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2450,40 +2450,40 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic intrinsic,

case NI_Vector128_Dot:
case NI_Vector256_Dot:
case NI_Vector512_Dot:
{
assert(sig->numArgs == 2);
var_types simdType = getSIMDTypeForSize(simdSize);

if (varTypeIsByte(simdBaseType) || varTypeIsLong(simdBaseType))
if ((simdSize == 32) && !varTypeIsFloating(simdBaseType) &&
!compOpportunisticallyDependsOn(InstructionSet_AVX2))
{
// TODO-XARCH-CQ: We could support dot product for 8-bit and
// 64-bit integers if we support multiplication for the same
// We can't deal with TYP_SIMD32 for integral types if the compiler doesn't support AVX2
break;
}

if (simdSize == 32)
{
if (!varTypeIsFloating(simdBaseType) && !compOpportunisticallyDependsOn(InstructionSet_AVX2))
{
// We can't deal with TYP_SIMD32 for integral types if the compiler doesn't support AVX2
break;
}
}
else if ((simdBaseType == TYP_INT) || (simdBaseType == TYP_UINT))
#if defined(TARGET_X86)
if (varTypeIsLong(simdBaseType) && !compOpportunisticallyDependsOn(InstructionSet_SSE41))
{
if (!compOpportunisticallyDependsOn(InstructionSet_SSE41))
{
// TODO-XARCH-CQ: We can support 32-bit integers if we updating multiplication
// to be lowered rather than imported as the relevant operations.
break;
}
// We need SSE41 to handle long, use software fallback
break;
}
#endif // TARGET_X86

op2 = impSIMDPopStack();
op1 = impSIMDPopStack();

if ((simdSize == 64) || varTypeIsByte(simdBaseType) || varTypeIsLong(simdBaseType) ||
(varTypeIsInt(simdBaseType) && !compOpportunisticallyDependsOn(InstructionSet_SSE41)))
{
// The lowering for Dot doesn't handle these cases, so import as Sum(left * right)
retNode = gtNewSimdBinOpNode(GT_MUL, simdType, op1, op2, simdBaseJitType, simdSize);
retNode = gtNewSimdSumNode(retType, retNode, simdBaseJitType, simdSize);
break;
}

retNode = gtNewSimdDotProdNode(simdType, op1, op2, simdBaseJitType, simdSize);
retNode = gtNewSimdGetElementNode(retType, retNode, gtNewIconNode(0), simdBaseJitType, simdSize);
retNode = gtNewSimdToScalarNode(retType, retNode, simdBaseJitType, simdSize);
break;
}

Expand Down Expand Up @@ -3349,28 +3349,14 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic intrinsic,
break;
}

#if defined(TARGET_X86)
if (varTypeIsLong(simdBaseType))
{
if (TARGET_POINTER_SIZE == 4)
{
// TODO-XARCH-CQ: 32bit support
break;
}

if ((simdSize == 32) && compOpportunisticallyDependsOn(InstructionSet_AVX2))
{
// Emulate NI_AVX512DQ_VL_MultiplyLow with AVX2 for SIMD32
}
else if ((simdSize == 16) && compOpportunisticallyDependsOn(InstructionSet_SSE41))
{
// Emulate NI_AVX512DQ_VL_MultiplyLow with SSE41 for SIMD16
}
else if (simdSize != 64)
{
// Software fallback
break;
}
// TODO-XARCH-CQ: We can't handle long here, only because one of the args might
// be scalar, and gtNewSimdCreateBroadcastNode doesn't handle long on x86.
break;
}
#endif // TARGET_X86

CORINFO_ARG_LIST_HANDLE arg1 = sig->args;
CORINFO_ARG_LIST_HANDLE arg2 = info.compCompHnd->getArgNext(arg1);
Expand Down Expand Up @@ -3405,29 +3391,6 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic intrinsic,
break;
}

if (varTypeIsLong(simdBaseType))
{
if (TARGET_POINTER_SIZE == 4)
{
// TODO-XARCH-CQ: 32bit support
break;
}

if ((simdSize == 32) && compOpportunisticallyDependsOn(InstructionSet_AVX2))
{
// Emulate NI_AVX512DQ_VL_MultiplyLow with AVX2 for SIMD32
}
else if ((simdSize == 16) && compOpportunisticallyDependsOn(InstructionSet_SSE41))
{
// Emulate NI_AVX512DQ_VL_MultiplyLow with SSE41 for SIMD16
}
else if (simdSize != 64)
{
// Software fallback
break;
}
}

op3 = impSIMDPopStack();
op2 = impSIMDPopStack();
op1 = impSIMDPopStack();
Expand Down Expand Up @@ -3818,17 +3781,8 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic intrinsic,
{
assert(sig->numArgs == 1);

if ((simdSize == 32) && !compOpportunisticallyDependsOn(InstructionSet_AVX2))
{
// Vector256 requires AVX2
break;
}
else if ((simdSize == 16) && !compOpportunisticallyDependsOn(InstructionSet_SSE2))
{
break;
}
#if defined(TARGET_X86)
else if (varTypeIsLong(simdBaseType) && !compOpportunisticallyDependsOn(InstructionSet_SSE41))
if (varTypeIsLong(simdBaseType) && !compOpportunisticallyDependsOn(InstructionSet_SSE41))
{
// We need SSE41 to handle long, use software fallback
break;
Expand Down
Loading