LSRA: Add support to track more than 64 registers

[kunalspathak]

Currently LSRA supports following number of registers for various platforms:

Architecture	# of General purpose	# of Float/Vector	# of Predicate/Mask	Total registers
x86	8	8	8	24
x64	16	32	8	**56
arm	16	32	0	48
arm64	32	32	*16	80

*16 new predicate registers for SVE
** There are 16 new GPR registers that will be added with APX, which will make the total for x64 to 72

Until now, the number of registers for all platforms were at most 64, we used regMaskTP data structure (typedef unsigned __int64) to represent them and pass them around. Throughout the RyuJIT codebase, whenever we have to pass, return or track a pool of registers, we use regMaskTP. Since it is 64 bits long, each bit represents a register. Lower bits represent GPR registers, while higher bits represent float/vector registers. However, with the #93095 to add SVE support for arm64, we need to add 16 predicate/mask registers, totaling the number of registers to be tracked from 64 to 80. They will not fit in regMaskTP anymore and we need an alternate way to represent these registers so that we can track the new 16 predicate registers that need to be added for SVE/arm64 work, but also 16 new GPR registers that will be added to support APX/x64 work.

Here are few options to solve this problem:

1. Convert regMaskTP to struct

The first option that we tried in #96196 was to just convert the regMaskTP to a struct that looks something like this:

typedef struct __regMaskTP
{
  unsigned __int64 low;
  unsigned __int64 high;
} regMaskTP;

To avoid refactoring all the code paths that use regMaskTP, we overloaded various operators for this struct. We found out that it regressed the throughput by around 10% for MinOpts and 7% for FullOpts as seen here.

2. Convert regMaskTP to intrinsic vector128

Next option we explored in #94589 was to use unsigned __int128 (only supported by clang on linux) which under the hood uses vector128. Our assumption was that compiler can optimize the access pattern of regMaskTP_128 and we will see lesser TP impact than option 1. However, when we cross compiled this change on linux/x64, we started seeing lot of seg faults in places where regMaskTP was initialized. The problem, as mentioned here was that __int128 assumed that regMaskTP field is at 16-byte aligned and would generate seg fault, whenever that was not the case. So, we had to give up on this option.

3. Segregate the gpr/float/predicate registers usage

This is WIP that I am working on currently in #98258. I created regMaskGpr, regMaskFloat and regMaskPredicate. Then I went through all the places in the code base and used the relevant types. Places where any register pool is accessed, I created a struct AllRegMask() that looks like this:

struct AllRegMask()
{
  regMaskGpr gprRegs;
  regMaskFloat floatRegs;
  regMaskPredicate predicateRegs;
}

 
In the code, I then pass AllRegMask() around and whenever we have to update (add/remove/track) a register from the mask, I added a check to see if the register in question is GPR/float/predicate and update the relevant field accordingly. Currently, we see the TP impact for this is around 6% regression in Minopts and 2% in FullOpts. With that, this is much better than option 1, but still is not acceptable.

4. Track predicate registers separately

Another option is to just track predicate registers separately and pass them around. There are not many places where we need to track them. The GPR/float registers will continue to get represented as regMaskTP 64-bits, and predicate registers will be tracked separately for platforms that has more than 64 registers (SVE/arm and in future APX/x64). The downside is, in future when number of GPR+float registers go beyond 64 registers, we will have to fall back to option 3. The other drawback of this approach is there are lot of places, more relevant is GenTree*, RefPosition and Interval that has regMaskTP as a field. Adding another field for "predicate registers" will consume more memory in these data structures, so probably a union and a bit to indicate if the regMaskTP indicates gpr+float or predicate might do the trick.

[kunalspathak]

@dotnet/arm64-contrib

[BruceForstall]

@dotnet/jit-contrib @dotnet/avx512-contrib

[tannergooding]

3 has always been my preferred approach here.

I expect, especially long term, it will result in the best performance and the most extensibility for new features/architectures. It does, however, require the most up front work.

I would expect that most places don't actually need AllRegMask() as most instructions only access 1 register set. Those that can access multiple tend to be well known (like movd which explicitly moves data between the SIMD and GPR sets), but even then it's not a case that needs more than 1x 64-bit mask per node, as just like with how its handled in emit (and the underlying insDesc) we intrinsically know the register set (GPR or SIMD) based on the gtType of the node.

I would expect the relatively few places that LSRA needs to look at all register sets would be better handled by passing such an AllRegMask struct byref (as AllRegMask&) or better, seeing if a minimal refactoring can be done to iterate through each mask independently since they rarely ever need to be processed in coordination with eachother.

[BruceForstall]

I'm in favor of (4). There is no platform known or expected where we need more than 64 GPR + float/SIMD registers. On our 64-bit platforms (the ones we are focusing on), this means the most important case of GPR/float all fits in one machine register. It's also useful that GPR and float register mask values are non-overlapping, to avoid cases of confusion (and bugs) where "0x1" could be a mask for r0 or v0, say. Mask/predicate registers I expect will continue to be rare, meaning checking for a mask register set equal to zero would be a fast path exit.

[kunalspathak]

I just want to give an update on these efforts and things I have explored so far with regards to option# 3.

Branch-free access

I converted the fields of AllRegsMask to array as seen below and then read/write the field based on the index depending on if the reg or type represents GPR, float or predicate registers.

struct AllRegMask
{
   uint32_t registers[3];
}

This reduced the TP regression by around 3% as seen in #98258 (comment). However, there is still around 3.5% TP regression on MinOpts for arm64 and my goal is bring it down to 0 or signficantly less for arm64. The reason being that arm64 does not have (yet) the predicate registers and so this work should not impact it. There will be some TP regression on x64, but improving that will come only after I improve the TP for arm64. So all my discussion below is tailored around arm64 momentarily.

---

Convert regMaskFloat from 8-bytes to 4-bytes

One of the ambition I had was to make all the regMaskTP i.e. the new regMaskGpr and regMaskFloat to 4-bytes long instead of 8-bytes because each category has 32 registers maximum and should be represented by just 4-bytes. Getting it down to 4-bytes will improve memory significantly because this data structure is a field in some of the common other data structures like GenTree, RefPosition and Interval. Today, the way regMaskTP works is the GPR have numbers 0 (REG_R0) thru 31 (REG_R31) and are mapped to the mask representation value of 0 thru 0x80000000 respectively. After that, the float registers are numbered from 32 (REG_V0) thru 63 (REG_V31) and are mapped to the mask representation of 0x100000000 thru 0x8000000000000000. We can make the size of gpr registers regMaskGpr to 4-bytes without any problems, because they can be represented in the 4-bytes space. However, the only way that we can make float registers fit in 4-bytes space is to re-map the values for float registers such that REG_V0 will have numeric value of 31 but will have mask representation as 0, REG_V1 will be 0x1 and so forth until REG_V31 will map to 0x80000000. With that in place, the code that operates on mask, will not have any idea if it is operating on the gpr or float register under the hood. For e.g. genRegMask() that has around 281 references would start returning value 0x100000 for both REG_R20 (gpr) and REG_V20 (float) and it will be hard to differentiate between these without knowing the actual type we are operating on, which means extra checks, which means TP regression. Additionally, the checks such as mask & REG_INT_CALLEE_TRASH will return true, even if the underlying mask was meant to be for float registers. We will also have to switch float register mask to << 32 once in a while and back to >> 32, which will just add complexity and will be hard to future proof it for other developers. So for now, I will hold off on this effort. I think I can still get the regMaskGpr to 4-bytes so at least at 166 places or so, where we operate on just gpr, we will just pass aronud 4 bytes instead of 8 bytes. However, regMaskOnlyOne will continue to be 8-bytes.

---

AllRegMask to work on single entity

So assuming that float register's mask will still be 8-bytes long and they will be only present in high-32 bites, another prototype I am working is to represent gpr and float registers as a single entity in AllRegMask instead of 2 separate fields or an array of 2 entries. I can possibly make a union of those, so I can access in whichever relevant way is easy to understand in the code base.

struct AllRegsMask
{
    union 
    {
        struct
        {
            uint64_t combined;
#ifdef HAS_PREDICATE_REG
            uint32_t predicate;
#endif
        };
#ifdef HAS_PREDICATE_REG
        uint32_t registers[3];
#else
        uint32_t registers[2];
#endif
        struct
        {
            uint32_t gpr;
            uint32_t fpr;
#ifdef HAS_PREDICATE_REG
            uint32_t pr;
#endif
        };
    };
};

And of course, I have not yet started working on passing AllRegsMask by-ref, which should further bring down the TP regression. I could also inherit a struct AllRegsMaskWithPredicate which add the field for predicate registers, while keeping AllRegsMask just for gpr/float. All the places where it is guaranteed to not make use of predicate registers, we can continue using AllRegsMask.

---

Regarding option 4

Tracking predicate registers is not an easy thing, specially because regMaskTP currently is used in basic data structures GenTree, Interval, etc. We will have to still have a way to differentiate if the underlying mask is predicate or gpr/float. The number of function calls where predicate registers need to be sent will doubled, once for gpr/float and then for predicate and some endless problems to iron out. However, I am trying to blend this option in AllRegsMask that I described above so that we will track it separately only at places where it is used and not at any other places.