Unmasking EGPRs in register allocator #114867

khushal1996 · 2025-04-21T20:08:30Z

Public issue: #114559

Description:
Currently in JIT, we have enabled extended general-purpose registers (EGPRs) in most of the instructions via REX2 or extended EVEX prefix. But previously, out of the testing stability, we have masked out EGPRs for most of the gentree nodes when building them in register allocator (LinearScan::BuildNode). This PR is to revert the masking and let EGPR ACTUALLY available to those nodes.

In detail, we used to have masks in those nodes, and this PR will remove all of them.

BuildHWIntrinsic
BuildIntrinsic
BuildBlockStore
BuildCall
BuildShiftRotate
GT_INDEX_ADDR
GT_CKFINITE
GT_RETURNTRAP
GT_SWITCH_TABLE
GT_JMPTABLE
BuildMul
BuildIndir
Buildcast
genFnProlog

Results:

SuperPMI asmdiff:
latest main + APX on v.s. latest main + changes + APX on

Note: We tend to demonstrate the improvement of our changes, so we need to have both base and diff with APX on. We have added an icount (Instruction Count) column in superpmi results to measure how well our changes are working. The following columns are added in the superpmi run below ->
Base Instruction Count - Total Instruction in Base
Diff Instruction Count - delta from total instructions in Base(percentage change considering only methods which had a diff)(percentage change considering only the methods which had icount change)

dotnet-policy-service · 2025-04-21T20:09:20Z

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

GT_RETURNTRAP GT_SWITCH_TABLE GT_JMPTABLE genFnProlog

Copilot

Pull Request Overview

This PR reverts the masking of extended GPRs (EGPRs) in the register allocator to allow their full usage in JIT code generation.

Removed calls to BuildApxIncompatibleGPRMask in favor of unmasked register definitions when extended registers are available.
Updated multiple LSRA routines and code emission paths to use the new useEvex/canUseApxRegs driven logic.
Introduced the BuildApxIncompatibleGPRMaskIfNeeded helper to streamline candidates selection when EVEX support is available.

Reviewed Changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
src/coreclr/jit/lsraxarch.cpp	Removed masking calls in node building and introduced useEvex checks.
src/coreclr/jit/lsrabuild.cpp	Updated BuildBinaryUses and related routines to conditionally use APX masks.
src/coreclr/jit/lsra.h	Added inline declaration for BuildApxIncompatibleGPRMaskIfNeeded.
src/coreclr/jit/instrsxarch.h	Added Encoding_REX2 to imul instructions for extended registers.
src/coreclr/jit/emitxarch.cpp	Adjusted movbe emission to select the APX variant based on register usage.
src/coreclr/jit/codegenxarch.cpp	Updated code emission paths for movbe and internal register extraction.
src/coreclr/jit/codegencommon.cpp	Modified function prologue to conditionally exclude high registers.

khushal1996 · 2025-05-02T22:04:24Z

@dotnet/intel for review

BruceForstall · 2025-05-02T22:46:12Z

cc @kunalspathak @dotnet/jit-contrib

khushal1996 · 2025-05-07T22:58:22Z

@kunalspathak for reviews

kunalspathak · 2025-05-08T15:28:46Z

src/coreclr/jit/instrsxarch.h

-INSTMUL(imul_SI,        "imul",             IUM_RD, BAD_CODE,     0x003068,     BAD_CODE,                                INS_TT_NONE,    Writes_OF      | Undefined_SF  | Undefined_ZF  | Undefined_AF  | Undefined_PF  | Writes_CF     | INS_FLAGS_Has_Sbit | INS_Flags_Has_NF)
-INSTMUL(imul_DI,        "imul",             IUM_RD, BAD_CODE,     0x003868,     BAD_CODE,                                INS_TT_NONE,    Writes_OF      | Undefined_SF  | Undefined_ZF  | Undefined_AF  | Undefined_PF  | Writes_CF     | INS_FLAGS_Has_Sbit | INS_Flags_Has_NF)
+INSTMUL(imul_AX,        "imul",             IUM_RD, BAD_CODE,     0x000068,     BAD_CODE,                                INS_TT_NONE,    Writes_OF      | Undefined_SF  | Undefined_ZF  | Undefined_AF  | Undefined_PF  | Writes_CF     | INS_FLAGS_Has_Sbit | INS_Flags_Has_NF | Encoding_REX2)
+INSTMUL(imul_CX,        "imul",             IUM_RD, BAD_CODE,     0x000868,     BAD_CODE,                                INS_TT_NONE,    Writes_OF      | Undefined_SF  | Undefined_ZF  | Undefined_AF  | Undefined_PF  | Writes_CF     | INS_FLAGS_Has_Sbit | INS_Flags_Has_NF | Encoding_REX2)


What changed in this section? Is it mostly spaces?

kunalspathak · 2025-05-08T15:31:18Z

src/coreclr/jit/lsrabuild.cpp

@@ -3841,19 +3841,24 @@ int LinearScan::BuildBinaryUses(GenTreeOp* node, SingleTypeRegSet candidates)
    GenTree* op1 = node->gtGetOp1();
    GenTree* op2 = node->gtGetOp2IfPresent();
 #ifdef TARGET_XARCH
+    const bool canUseApxRegs = compiler->canUseApxEncoding() && compiler->canUseEvexEncoding();


canUseApxEncoding

can you cache this value in compiler object or LinearScan?

kunalspathak · 2025-05-08T15:31:51Z

src/coreclr/jit/lsrabuild.cpp

-        if (op1->isContainedIndir() &&
-            ((varTypeUsesFloatReg(node) || node->OperGet() == GT_BSWAP || node->OperGet() == GT_BSWAP16)) &&
-            candidates == RBM_NONE)
+        if (op1->isContainedIndir() && candidates == RBM_NONE && !canUseApxRegs)


Suggested change

if (op1->isContainedIndir() && candidates == RBM_NONE && !canUseApxRegs)

if (op1->isContainedIndir() && (candidates == RBM_NONE) && !canUseApxRegs)

kunalspathak · 2025-05-08T15:32:16Z

src/coreclr/jit/lsrabuild.cpp

@@ -3879,7 +3882,7 @@ int LinearScan::BuildBinaryUses(GenTreeOp* node, SingleTypeRegSet candidates)
    {

 #ifdef TARGET_XARCH
-        if (op2->isContainedIndir() && varTypeUsesFloatReg(op1) && candidates == RBM_NONE)
+        if (op2->isContainedIndir() && candidates == RBM_NONE && !canUseApxRegs)


Suggested change

if (op2->isContainedIndir() && candidates == RBM_NONE && !canUseApxRegs)

if (op2->isContainedIndir() && (candidates == RBM_NONE) && !canUseApxRegs)

It seems here, we should have if ((op2->isContainedIndir() && (candidates == RBM_NONE)) || !canUseApxRegs)...basically even if op2 is not contained indir and candidates != RBM_NONE, we should still use lowGprRegs if we cannot use APX registers.

kunalspathak · 2025-05-08T15:35:52Z

src/coreclr/jit/lsrabuild.cpp

@@ -3879,7 +3882,7 @@ int LinearScan::BuildBinaryUses(GenTreeOp* node, SingleTypeRegSet candidates)
    {

 #ifdef TARGET_XARCH
-        if (op2->isContainedIndir() && varTypeUsesFloatReg(op1) && candidates == RBM_NONE)


likewise here...can you double check if && !canUseApxRegs is the right check?

kunalspathak · 2025-05-08T15:37:25Z

src/coreclr/jit/lsrabuild.cpp

@@ -3982,14 +3984,6 @@ void LinearScan::BuildStoreLocDef(GenTreeLclVarCommon* storeLoc,
    defCandidates = allRegs(type);
 #endif // TARGET_X86

-#ifdef TARGET_AMD64
-    if (op1->isContained() && op1->OperIs(GT_BITCAST) && varTypeUsesIntReg(varDsc->GetRegisterType(storeLoc)))


so we don't care if APX regs is available here and below for GT_BITCAST too?

kunalspathak · 2025-05-08T15:41:35Z

src/coreclr/jit/lsraxarch.cpp

@@ -1422,6 +1409,7 @@ int LinearScan::BuildBlockStore(GenTreeBlk* blkNode)
    SingleTypeRegSet dstAddrRegMask = RBM_NONE;
    SingleTypeRegSet srcRegMask     = RBM_NONE;
    SingleTypeRegSet sizeRegMask    = RBM_NONE;
+    bool             useEvex        = compiler->canUseEvexEncoding();


see if you can cache this value in LinearScan object. It is used at multiple places.

kunalspathak · 2025-05-08T15:42:36Z

src/coreclr/jit/lsraxarch.cpp

@@ -1542,7 +1529,8 @@ int LinearScan::BuildBlockStore(GenTreeBlk* blkNode)
                    // or if are but the remainder is a power of 2 and less than the
                    // size of a register

-                    SingleTypeRegSet regMask = BuildApxIncompatibleGPRMask(blkNode, availableIntRegs, true);
+                    // SingleTypeRegSet regMask = BuildApxIncompatibleGPRMask(blkNode, availableIntRegs, true);
+                    SingleTypeRegSet regMask = availableIntRegs;


at some of these places, don't we need useEvex check, the way it is done at other places?

kunalspathak · 2025-05-08T15:42:45Z

src/coreclr/jit/lsraxarch.cpp

@@ -1542,7 +1529,8 @@ int LinearScan::BuildBlockStore(GenTreeBlk* blkNode)
                    // or if are but the remainder is a power of 2 and less than the
                    // size of a register

-                    SingleTypeRegSet regMask = BuildApxIncompatibleGPRMask(blkNode, availableIntRegs, true);
+                    // SingleTypeRegSet regMask = BuildApxIncompatibleGPRMask(blkNode, availableIntRegs, true);


delete the comment

kunalspathak · 2025-05-08T15:45:33Z

src/coreclr/jit/lsraxarch.cpp

-            BuildUse(srcAddrOrFill, BuildApxIncompatibleGPRMask(srcAddrOrFill, srcRegMask));
+            if (useEvex)
+            {
+                BuildUse(srcAddrOrFill, srcRegMask);


useEvex check is scattered through out the code. Is it possible to create a helper method or update BuildApxIncompatibleGPRMask that takes the original mask. It will check if useEvex == true, return that mask as it is, otherwise do BuildApxIncompatibleGPRMask() operation on it to get right mask.

Just noticed you already have BuildApxIncompatibleGPRMaskIfNeeded. Can we reuse it at other places as well?

kunalspathak · 2025-05-08T15:46:50Z

src/coreclr/jit/lsraxarch.cpp

-                GenTree* user = nullptr;
-
-                if (LIR::AsRange(blockSequence[curBBSeqNum]).TryGetUse(intrinsicTree, &use))
+                // TODO-XArch-APX: some of the permute intrinsics are APX-EVEX compatible, we need to separate and


do we need a GH issue for this?

kunalspathak · 2025-05-08T15:49:07Z

src/coreclr/jit/emitinl.h

+    assert(INS_imul_21 - INS_imul_AX == REG_R21);
+    assert(INS_imul_22 - INS_imul_AX == REG_R22);
+    assert(INS_imul_23 - INS_imul_AX == REG_R23);
+    // TODO-XArch-APX: The asserts below need the register definition from R24~R31.


we already have it right?

kunalspathak · 2025-05-08T15:51:16Z

SuperPMI asmdiff:

Can you also share the tpdiff numbers for similar configuration?

Also I was hoping the asmdiff on CI to be zero diff, but still seeing differences. Do we know why?

ghost added the area-CodeGen-coreclr CLR JIT compiler in src/coreclr/src/jit and related components such as SuperPMI label Apr 21, 2025

khushal1996 changed the title ~~unmasking EGPRs in register allocator~~ Unmasking EGPRs in register allocator Apr 21, 2025

dotnet-policy-service bot added the community-contribution Indicates that the PR has been added by a community member label Apr 21, 2025

khushal1996 force-pushed the kcm-egpr-mask branch from ca02c0a to 2efec54 Compare April 21, 2025 21:15

Ruihan-Yin and others added 13 commits May 1, 2025 18:33

unmask HWIntrinsic node.

439d511

IMUL, GT_INDEX_ADDR

92ee763

GT_CKFINITE

908bcdc

GT_RETURNTRAP GT_SWITCH_TABLE GT_JMPTABLE genFnProlog

Complete BuildHWIntrinsics

2d41d38

Let genFnProlog have EGPR when APX and EVEX are both available.

23bf4cf

Indir

7ad0fed

clean-up in lsrabuild

73f9e42

Enable egprs in cast nodes

b8dad4b

unmask BuildIntrinsic

81d76db

unmask BuildShiftRotate

7c77deb

Unmask BlockStore

c50c26c

run formatting

c765707

Use regMask when looking for srcReg candidate in BlkStore

52a6234

khushal1996 force-pushed the kcm-egpr-mask branch from a7ac128 to 52a6234 Compare May 2, 2025 01:46

Merge branch 'main' into kcm-egpr-mask

e4d5618

khushal1996 marked this pull request as ready for review May 2, 2025 09:59

Copilot AI review requested due to automatic review settings May 2, 2025 09:59

Copilot AI reviewed May 2, 2025

View reviewed changes

kunalspathak self-requested a review May 2, 2025 23:13

kunalspathak reviewed May 8, 2025

View reviewed changes

Merge branch 'main' into kcm-egpr-mask

aa9b236

build-analysis bot mentioned this pull request May 8, 2025

Checkout failure: "Git fetch failed with exit code 128" dotnet/arcade#9009

Open

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Unmasking EGPRs in register allocator #114867

Unmasking EGPRs in register allocator #114867

khushal1996 commented Apr 21, 2025 •

edited

Loading

dotnet-policy-service bot commented Apr 21, 2025

Copilot AI left a comment

khushal1996 commented May 2, 2025

BruceForstall commented May 2, 2025

khushal1996 commented May 7, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak May 8, 2025 •

edited

Loading

kunalspathak May 8, 2025

kunalspathak May 8, 2025

kunalspathak commented May 8, 2025

	if (op1->isContainedIndir() && candidates == RBM_NONE && !canUseApxRegs)
	if (op1->isContainedIndir() && (candidates == RBM_NONE) && !canUseApxRegs)

	if (op2->isContainedIndir() && candidates == RBM_NONE && !canUseApxRegs)
	if (op2->isContainedIndir() && (candidates == RBM_NONE) && !canUseApxRegs)

Unmasking EGPRs in register allocator #114867

Are you sure you want to change the base?

Unmasking EGPRs in register allocator #114867

Conversation

khushal1996 commented Apr 21, 2025 • edited Loading

dotnet-policy-service bot commented Apr 21, 2025

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

khushal1996 commented May 2, 2025

BruceForstall commented May 2, 2025

khushal1996 commented May 7, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kunalspathak May 8, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kunalspathak commented May 8, 2025

khushal1996 commented Apr 21, 2025 •

edited

Loading

kunalspathak May 8, 2025 •

edited

Loading