Add Avx512 support to ProbabilisticMap #96931

MihaZupan · 2024-01-13T05:05:31Z

public class ProbMap
{
    private static readonly SearchValues<char> s_values = SearchValues.Create("abcABC123ű");
    private string _text = "ő";

    [Params(16, 32, 64, 128, 256, 10_000)]
    public int Length;

    [GlobalSetup]
    public void Setup() => _text += new string('\n', Length - 1);

    [Benchmark]
    public int IndexOfAny() => _text.AsSpan().IndexOfAny(s_values);
}

Method	Toolchain	Length	Mean	Error	Ratio
IndexOfAny	main	16	8.471 ns	0.0203 ns	1.00
IndexOfAny	pr	16	7.466 ns	0.0391 ns	0.88

IndexOfAny	main	32	9.284 ns	0.0215 ns	1.00
IndexOfAny	pr	32	7.313 ns	0.0242 ns	0.79

IndexOfAny	main	64	12.980 ns	0.0082 ns	1.00
IndexOfAny	pr	64	8.339 ns	0.0619 ns	0.64

IndexOfAny	main	128	21.795 ns	0.0089 ns	1.00
IndexOfAny	pr	128	11.330 ns	0.0359 ns	0.52

IndexOfAny	main	256	38.733 ns	0.0193 ns	1.00
IndexOfAny	pr	256	18.775 ns	0.0245 ns	0.48

IndexOfAny	main	10000	1,347.139 ns	0.7767 ns	1.00
IndexOfAny	pr	10000	625.296 ns	0.3626 ns	0.46

Besides just doubling the vector width, Avx512Vbmi.VL gives us access to better permute instructions, allowing us to replace

Vector256<byte> bitMaskLower = Avx2.Shuffle(charMapLower, index);
Vector256<byte> bitMaskUpper = Avx2.Shuffle(charMapUpper, index - Vector256.Create((byte)16));
Vector256<byte> mask = Vector256.GreaterThan(index, Vector256.Create((byte)15));
Vector256<byte> bitMask = Vector256.ConditionalSelect(mask, bitMaskUpper, bitMaskLower);

with

Vector256<byte> bitMask = Avx512Vbmi.VL.PermuteVar32x8(charMap, index);

AVX2 vs AVX512 main loop

// AVX2 loop to process 32 characters   // AVX512 loop to process 64 characters
M02_L00:                                M02_L00:
       vmovups   ymm4,[r14]                    vmovups   zmm1,[rbx]
       vmovups   ymm5,[r14+20]                 vmovups   zmm2,[rbx+40]
       vmovups   ymm6,[7FC7B99432C0]           vmovups   zmm3,[7F3FF2744E40]
       vpand     ymm7,ymm4,ymm6                vpandd    zmm4,zmm1,zmm3
       vpand     ymm6,ymm5,ymm6                vpandd    zmm3,zmm2,zmm3
       vpackuswb ymm6,ymm7,ymm6                vpackuswb zmm3,zmm4,zmm3
       vpsrlw    ymm4,ymm4,8                   vpsrlw    zmm1,zmm1,8
       vpsrlw    ymm5,ymm5,8                   vpsrlw    zmm2,zmm2,8
       vpackuswb ymm4,ymm4,ymm5                vpackuswb zmm1,zmm1,zmm2
       vpsrld    ymm5,ymm6,5                   vmovups   zmm2,[7F3FF2744E80]
       vmovups   ymm7,[7FC7B99432E0]           vpandd    zmm4,zmm3,zmm2
       vpand     ymm5,ymm5,ymm7                vpermb    zmm4,zmm4,zmm0
       vmovups   ymm8,[7FC7B9943300]           vpsrld    zmm3,zmm3,5
       vpshufb   ymm5,ymm8,ymm5                vmovups   zmm5,[7F3FF2744EC0]
       vmovups   ymm8,[7FC7B9943320]           vpandd    zmm3,zmm3,zmm5
       vpand     ymm6,ymm6,ymm8                vmovups   zmm6,[7F3FF2744F00]
       vmovups   ymm9,[7FC7B9943340]           vpshufb   zmm3,zmm6,zmm3
       vpsubb    ymm10,ymm6,ymm9               vpandd    zmm3,zmm4,zmm3
       vpcmpgtb  ymm10,ymm10,[7FC7B9943360]    vptestnmb k1,zmm3,zmm3
       vmovups   ymm11,[7FC7B9943380]          vpandd    zmm2,zmm1,zmm2
       vpsubb    ymm12,ymm6,ymm11              vpermb    zmm2,zmm2,zmm0
       vpshufb   ymm12,ymm3,ymm12              vpsrld    zmm1,zmm1,5
       vpshufb   ymm6,ymm2,ymm6                vpandd    zmm1,zmm1,zmm5
       vpblendvb ymm6,ymm6,ymm12,ymm10         vpshufb   zmm1,zmm6,zmm1
       vpand     ymm5,ymm6,ymm5                vpandd    zmm1,zmm2,zmm1
       vxorps    ymm6,ymm6,ymm6                vptestnmb k2,zmm1,zmm1
       vpcmpeqb  ymm5,ymm5,ymm6                korq      k1,k1,k2
       vpsrld    ymm6,ymm4,5                   vpmovm2b  zmm1,k1
       vpand     ymm6,ymm6,ymm7                vpternlogd zmm2,zmm2,zmm2,0FF
       vmovups   ymm7,[7FC7B9943300]           vpxord    zmm1,zmm1,zmm2
       vpshufb   ymm6,ymm7,ymm6                vptestmb  k1,zmm1,zmm1
       vpand     ymm4,ymm4,ymm8                kortestq  k1,k1
       vpsubb    ymm7,ymm4,ymm9                jne       short M02_L03
       vpcmpgtb  ymm7,ymm7,[7FC7B9943360]
       vpsubb    ymm8,ymm4,ymm11
       vpshufb   ymm8,ymm3,ymm8
       vpshufb   ymm4,ymm2,ymm4
       vpblendvb ymm4,ymm4,ymm8,ymm7
       vpand     ymm4,ymm4,ymm6
       vxorps    ymm6,ymm6,ymm6
       vpcmpeqb  ymm4,ymm4,ymm6
       vpor      ymm4,ymm5,ymm4
       vpcmpeqd  ymm5,ymm5,ymm5
       vpxor     ymm4,ymm4,ymm5
       vptest    ymm4,ymm4
       jne       near ptr M02_L10
M02_L01:                                M02_L01:
       add       r14,40                        add       rbx,80
       cmp       r14,r13                       cmp       rbx,r13
       vmovups   ymm2,[rbp-70]                 vmovups   zmm0,[rbp-70]
       vmovups   ymm3,[rbp-90]                 jbe       near ptr M02_L00
       jbe       near ptr M02_L00

ghost · 2024-01-13T05:05:37Z

Tagging subscribers to this area: @dotnet/area-system-memory
See info in area-owners.md if you want to be subscribed.

Issue Details

public class ProbMap
{
    private static readonly SearchValues<char> s_values = SearchValues.Create("abcABC123ű");
    private string _text = "ő";

    [Params(16, 32, 64, 128, 256, 10_000)]
    public int Length;

    [GlobalSetup]
    public void Setup() => _text += new string('\n', Length - 1);

    [Benchmark]
    public int IndexOfAny() => _text.AsSpan().IndexOfAny(s_values);
}

Method	Toolchain	Length	Mean	Error	Ratio
IndexOfAny	main	16	8.471 ns	0.0203 ns	1.00
IndexOfAny	pr	16	7.466 ns	0.0391 ns	0.88

IndexOfAny	main	32	9.284 ns	0.0215 ns	1.00
IndexOfAny	pr	32	7.313 ns	0.0242 ns	0.79

IndexOfAny	main	64	12.980 ns	0.0082 ns	1.00
IndexOfAny	pr	64	8.339 ns	0.0619 ns	0.64

IndexOfAny	main	128	21.795 ns	0.0089 ns	1.00
IndexOfAny	pr	128	11.330 ns	0.0359 ns	0.52

IndexOfAny	main	256	38.733 ns	0.0193 ns	1.00
IndexOfAny	pr	256	18.775 ns	0.0245 ns	0.48

IndexOfAny	main	10000	1,347.139 ns	0.7767 ns	1.00
IndexOfAny	pr	10000	625.296 ns	0.3626 ns	0.46

Besides just doubling the vector width, Avx512Vbmi.VL gives us access to better permute instructions, allowing us to replace

Vector256<byte> bitMaskLower = Avx2.Shuffle(charMapLower, index);
Vector256<byte> bitMaskUpper = Avx2.Shuffle(charMapUpper, index - Vector256.Create((byte)16));
Vector256<byte> mask = Vector256.GreaterThan(index, Vector256.Create((byte)15));
Vector256<byte> bitMask = Vector256.ConditionalSelect(mask, bitMaskUpper, bitMaskLower);

with

Vector256<byte> bitMask = Avx512Vbmi.VL.PermuteVar32x8(charMap, index);

AVX2 vs AVX512 main loop

// AVX2 loop to process 32 characters   // AVX512 loop to process 64 characters
M02_L00:                                M02_L00:
       vmovups   ymm4,[r14]                    vmovups   zmm1,[rbx]
       vmovups   ymm5,[r14+20]                 vmovups   zmm2,[rbx+40]
       vmovups   ymm6,[7FC7B99432C0]           vmovups   zmm3,[7F3FF2744E40]
       vpand     ymm7,ymm4,ymm6                vpandd    zmm4,zmm1,zmm3
       vpand     ymm6,ymm5,ymm6                vpandd    zmm3,zmm2,zmm3
       vpackuswb ymm6,ymm7,ymm6                vpackuswb zmm3,zmm4,zmm3
       vpsrlw    ymm4,ymm4,8                   vpsrlw    zmm1,zmm1,8
       vpsrlw    ymm5,ymm5,8                   vpsrlw    zmm2,zmm2,8
       vpackuswb ymm4,ymm4,ymm5                vpackuswb zmm1,zmm1,zmm2
       vpsrld    ymm5,ymm6,5                   vmovups   zmm2,[7F3FF2744E80]
       vmovups   ymm7,[7FC7B99432E0]           vpandd    zmm4,zmm3,zmm2
       vpand     ymm5,ymm5,ymm7                vpermb    zmm4,zmm4,zmm0
       vmovups   ymm8,[7FC7B9943300]           vpsrld    zmm3,zmm3,5
       vpshufb   ymm5,ymm8,ymm5                vmovups   zmm5,[7F3FF2744EC0]
       vmovups   ymm8,[7FC7B9943320]           vpandd    zmm3,zmm3,zmm5
       vpand     ymm6,ymm6,ymm8                vmovups   zmm6,[7F3FF2744F00]
       vmovups   ymm9,[7FC7B9943340]           vpshufb   zmm3,zmm6,zmm3
       vpsubb    ymm10,ymm6,ymm9               vpandd    zmm3,zmm4,zmm3
       vpcmpgtb  ymm10,ymm10,[7FC7B9943360]    vptestnmb k1,zmm3,zmm3
       vmovups   ymm11,[7FC7B9943380]          vpandd    zmm2,zmm1,zmm2
       vpsubb    ymm12,ymm6,ymm11              vpermb    zmm2,zmm2,zmm0
       vpshufb   ymm12,ymm3,ymm12              vpsrld    zmm1,zmm1,5
       vpshufb   ymm6,ymm2,ymm6                vpandd    zmm1,zmm1,zmm5
       vpblendvb ymm6,ymm6,ymm12,ymm10         vpshufb   zmm1,zmm6,zmm1
       vpand     ymm5,ymm6,ymm5                vpandd    zmm1,zmm2,zmm1
       vxorps    ymm6,ymm6,ymm6                vptestnmb k2,zmm1,zmm1
       vpcmpeqb  ymm5,ymm5,ymm6                korq      k1,k1,k2
       vpsrld    ymm6,ymm4,5                   vpmovm2b  zmm1,k1
       vpand     ymm6,ymm6,ymm7                vpternlogd zmm2,zmm2,zmm2,0FF
       vmovups   ymm7,[7FC7B9943300]           vpxord    zmm1,zmm1,zmm2
       vpshufb   ymm6,ymm7,ymm6                vptestmb  k1,zmm1,zmm1
       vpand     ymm4,ymm4,ymm8                kortestq  k1,k1
       vpsubb    ymm7,ymm4,ymm9                jne       short M02_L03
       vpcmpgtb  ymm7,ymm7,[7FC7B9943360]
       vpsubb    ymm8,ymm4,ymm11
       vpshufb   ymm8,ymm3,ymm8
       vpshufb   ymm4,ymm2,ymm4
       vpblendvb ymm4,ymm4,ymm8,ymm7
       vpand     ymm4,ymm4,ymm6
       vxorps    ymm6,ymm6,ymm6
       vpcmpeqb  ymm4,ymm4,ymm6
       vpor      ymm4,ymm5,ymm4
       vpcmpeqd  ymm5,ymm5,ymm5
       vpxor     ymm4,ymm4,ymm5
       vptest    ymm4,ymm4
       jne       near ptr M02_L10
M02_L01:                                M02_L01:
       add       r14,40                        add       rbx,80
       cmp       r14,r13                       cmp       rbx,r13
       vmovups   ymm2,[rbp-70]                 vmovups   zmm0,[rbp-70]
       vmovups   ymm3,[rbp-90]                 jbe       near ptr M02_L00
       jbe       near ptr M02_L00

Author:	MihaZupan
Assignees:	MihaZupan
Labels:	`area-System.Memory`
Milestone:	9.0.0

tannergooding · 2024-02-01T17:26:59Z

@MihaZupan is this just pending review or is it blocked on something? Noticed you have it self-assigned so want to make sure.

If it is ready for review, could you click the update branch button so we can get fresh CI?

MihaZupan · 2024-02-01T17:32:24Z

It's ready for review, I updated the branch now

MihaZupan added the area-System.Memory label Jan 13, 2024

MihaZupan added this to the 9.0.0 milestone Jan 13, 2024

MihaZupan requested a review from stephentoub January 13, 2024 05:05

MihaZupan self-assigned this Jan 13, 2024

This was referenced Jan 13, 2024

Segmentation fault in System.Text.RegularExpressions.Tests #93206

Closed

Test failure - System.NullReferenceException in System.Threading.Lock.TryInitializeStatics #94728

Closed

Assert failure: orObject->HasEmptySyncBlockInfo() #96790

Closed

MihaZupan added 3 commits February 1, 2024 18:31

Add Avx512 support to ProbabilisticMap

771179a

Add Vector512 condition

d87fc68

Fix assert

68a5b3f

MihaZupan force-pushed the searchvalues-probmap-avx512 branch from 0d21d17 to 68a5b3f Compare February 1, 2024 17:32

tannergooding approved these changes Feb 1, 2024

View reviewed changes

MihaZupan mentioned this pull request Feb 1, 2024

System.Net.Security.Tests.SslStreamCertificateContextOcspLinuxTests.FetchOcspResponse_FirstInvalidThenValid test failure #97836

Closed

MihaZupan merged commit 900bc5a into dotnet:main Feb 1, 2024
171 of 178 checks passed

github-actions bot locked and limited conversation to collaborators Mar 3, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Avx512 support to ProbabilisticMap #96931

Add Avx512 support to ProbabilisticMap #96931

MihaZupan commented Jan 13, 2024

ghost commented Jan 13, 2024

tannergooding commented Feb 1, 2024

MihaZupan commented Feb 1, 2024

Add Avx512 support to ProbabilisticMap #96931

Add Avx512 support to ProbabilisticMap #96931

Conversation

MihaZupan commented Jan 13, 2024

ghost commented Jan 13, 2024

tannergooding commented Feb 1, 2024

MihaZupan commented Feb 1, 2024