JIT: Optimize const ShiftRightLogical for byte values on XArch

[MihaZupan]

Contributes to #82564

(First time touching the JIT so I don't really know what I'm doing)

Vector128<byte> v0;

// Expands patterns like
v0 >>> 4;
// to
(v0.AsInt32() >>> 4).AsByte() & Vector128.Create((byte)15);

// and
v0 >>> nonConstAmount;
// to
int maskedShiftAmount = nonConstAmount & 7;
Vector128<int> shiftVector = Sse2.ConvertScalarToVector128Int32(maskedShiftAmount);
Vector128<byte> shiftedInput = Sse2.ShiftRightLogical(v0.AsInt32(), shiftVector).AsByte();
return shiftedInput & Vector128.Create((byte)(255 >> maskedShiftAmount));

Codegen for

[Benchmark]
[Arguments(42)]
public Vector128<byte> SimpleShift(byte input) =>
    Vector128.Create(input) >>> 4;

Before

; VectorTest.SimpleShift(Byte)
       sub       rsp,48
       vzeroupper
       movzx     eax,r8b
       vpbroadcastb xmm0,eax
       vmovaps   [rsp+30],xmm0
       mov       rax,[rsp+30]
       mov       [rsp+20],rax
       movzx     eax,byte ptr [rsp+20]
       shr       eax,4
       movzx     eax,al
       mov       [rsp+28],al
       movzx     eax,byte ptr [rsp+21]
       shr       eax,4
       movzx     eax,al
       mov       [rsp+29],al
       movzx     eax,byte ptr [rsp+22]
       shr       eax,4
       movzx     eax,al
       mov       [rsp+2A],al
       movzx     eax,byte ptr [rsp+23]
       shr       eax,4
       movzx     eax,al
       mov       [rsp+2B],al
       movzx     eax,byte ptr [rsp+24]
       shr       eax,4
       movzx     eax,al
       mov       [rsp+2C],al
       movzx     eax,byte ptr [rsp+25]
       shr       eax,4
       movzx     eax,al
       mov       [rsp+2D],al
       movzx     eax,byte ptr [rsp+26]
       shr       eax,4
       movzx     eax,al
       mov       [rsp+2E],al
       movzx     eax,byte ptr [rsp+27]
       shr       eax,4
       movzx     eax,al
       mov       [rsp+2F],al
       mov       rax,[rsp+28]
       mov       rcx,[rsp+38]
       mov       [rsp+10],rcx
       movzx     ecx,byte ptr [rsp+10]
       shr       ecx,4
       movzx     ecx,cl
       mov       [rsp+18],cl
       movzx     ecx,byte ptr [rsp+11]
       shr       ecx,4
       movzx     ecx,cl
       mov       [rsp+19],cl
       movzx     ecx,byte ptr [rsp+12]
       shr       ecx,4
       movzx     ecx,cl
       mov       [rsp+1A],cl
       movzx     ecx,byte ptr [rsp+13]
       shr       ecx,4
       movzx     ecx,cl
       mov       [rsp+1B],cl
       movzx     ecx,byte ptr [rsp+14]
       shr       ecx,4
       movzx     ecx,cl
       mov       [rsp+1C],cl
       movzx     ecx,byte ptr [rsp+15]
       shr       ecx,4
       movzx     ecx,cl
       mov       [rsp+1D],cl
       movzx     ecx,byte ptr [rsp+16]
       shr       ecx,4
       movzx     ecx,cl
       mov       [rsp+1E],cl
       movzx     ecx,byte ptr [rsp+17]
       shr       ecx,4
       movzx     ecx,cl
       mov       [rsp+1F],cl
       mov       rcx,[rsp+18]
       mov       [rsp],rax
       mov       [rsp+8],rcx
       vmovaps   xmm0,[rsp]
       vmovups   [rdx],xmm0
       mov       rax,rdx
       add       rsp,48
       ret
; Total bytes of code 319

After

; VectorTest.SimpleShift(Byte)
       vzeroupper
       movzx     eax,r8b
       vpbroadcastb xmm0,eax
       vpsrld    xmm0,xmm0,4
       vpand     xmm0,xmm0,[7FFC6F58A2E0]
       vmovups   [rdx],xmm0
       mov       rax,rdx
       ret
; Total bytes of code 34

Codegen for

[Benchmark]
[Arguments(3)]
public Vector128<byte> SimpleShift(int shiftCount) =>
    Vector128.Create((byte)42) >>> shiftCount;

Before

; VectorTest.SimpleShift(Int32)
       sub       rsp,58
       vzeroupper
       vmovups   xmm0,[7FFC74E11BB0]
       vmovaps   [rsp+40],xmm0
       mov       rax,[rsp+40]
       mov       [rsp+20],rax
       xor       eax,eax
       mov       ecx,r8d
       and       ecx,7
M00_L00:
       lea       r9,[rsp+20]
       movsxd    r10,eax
       movzx     r9d,byte ptr [r9+r10]
       shrx      r9d,r9d,ecx
       movzx     r9d,r9b
       lea       r11,[rsp+28]
       mov       [r11+r10],r9b
       inc       eax
       cmp       eax,8
       jl        short M00_L00
       vmovaps   [rsp+30],xmm0
       mov       rax,[rsp+28]
       mov       rcx,[rsp+38]
       mov       [rsp+10],rcx
       xor       ecx,ecx
       and       r8d,7
M00_L01:
       lea       r9,[rsp+10]
       movsxd    r10,ecx
       movzx     r9d,byte ptr [r9+r10]
       shrx      r9d,r9d,r8d
       movzx     r9d,r9b
       lea       r11,[rsp+18]
       mov       [r11+r10],r9b
       inc       ecx
       cmp       ecx,8
       jl        short M00_L01
       mov       rcx,[rsp+18]
       mov       [rsp],rax
       mov       [rsp+8],rcx
       vmovaps   xmm0,[rsp]
       vmovups   [rdx],xmm0
       mov       rax,rdx
       add       rsp,58
       ret
; Total bytes of code 173

After

; VectorTest.SimpleShift(Int32)
       vzeroupper
       and       r8d,7
       vmovd     xmm0,r8d
       vmovups   xmm1,[7FFCEB69A330]
       vpsrld    xmm0,xmm1,xmm0
       mov       eax,0FF
       shrx      eax,eax,r8d
       vpbroadcastb xmm1,eax
       vpand     xmm0,xmm0,xmm1
       vmovups   [rdx],xmm0
       mov       rax,rdx
       ret
; Total bytes of code 52

[ghost]

Tagging subscribers to this area: @JulieLeeMSFT, @jakobbotsch
See info in area-owners.md if you want to be subscribed.

Issue Details

---

Contributes to #82564

(First time touching the JIT so I don't really know what I'm doing)

Vector128<byte> v0;

// Expands patterns like
v0 >>> 4;
// to
(v0.AsInt32() >>> 4).AsByte() & Vector128.Create((byte)15);

Author:	MihaZupan
Assignees:	MihaZupan
Labels:	area-CodeGen-coreclr
Milestone:	8.0.0

[tannergooding]

We shouldn't need this limitation. For shift in particular, working with constant or non-constant data is the same since we have overloads that take byte shiftAmount and overloads that take Vector128<T> shiftAmount with the latter simply being Vector128.CreateScalarUnsafe(shiftAmount) of the non-constant byte.

[tannergooding]

That handling of constant vs non-constant is general and simply part of the existing gtNewSimdBinOp handling, so you shouldn't need to do anything additional really.

We'll already have appropriately masked off the shift amount to be 0-7 and converted it to a Vector128 if necessary; so you should just need to generate the shift followed by the and to mask off the upper n-bits of each element.

It basically just becomes that if its a constant, you can generate a new icon directly and otherwise the mask is gtNewSimdBroadcastNode of 0xFF >> shiftAmount

[MihaZupan]

I updated it now to handle non-const shift amounts as well, but it feels a bit hacky with the temporary nonConstantByteShiftCountOp - please let me know if I'm on the right track.

Right now the non-const path is something like

static Vector128<byte> Shift(Vector128<byte> input, int shiftAmount)
{
    int maskedShiftAmount = shiftAmount & 7;
    Vector128<int> shiftVector = Sse2.ConvertScalarToVector128Int32(maskedShiftAmount);
    Vector128<byte> shiftedInput = Sse2.ShiftRightLogical(input.AsInt32(), shiftVector).AsByte();
    return shiftedInput & Vector128.Create((byte)(255 >> maskedShiftAmount));
}

[MihaZupan]

(don't mind the referenced issue spam - I'm just testing stuff and this PR is convenient)

[tannergooding]

CC. @dotnet/jit-contrib, this needs a secondary review before merging.

[MihaZupan]

Test failures were #88582