JIT: Rematerialize spilled FP/SIMD/mask constants instead of stack spill/reload

src/coreclr/jit/codegen.h

@@ -1121,6 +1121,7 @@ class CodeGen final : public CodeGenInterface
         unsigned varNum, var_types type, GenTreeLclVar* lclNode, regNumber regNum, bool reSpill, bool isLastUse);
     void      genUnspillRegIfNeeded(GenTree* tree);
     void      genUnspillRegIfNeeded(GenTree* tree, unsigned multiRegIndex);
+    bool      isRematerializableConstant(GenTree* tree);
     regNumber genConsumeReg(GenTree* tree);
     regNumber genConsumeReg(GenTree* tree, unsigned multiRegIndex);
     void      genCopyRegIfNeeded(GenTree* tree, regNumber needReg);

src/coreclr/jit/codegenlinear.cpp

@@ -1404,23 +1404,142 @@ void CodeGen::genUnspillRegIfNeeded(GenTree* tree)
         }
         else
         {
-            // Here we may have a GT_RELOAD.
-            // The spill temp allocated for it is associated with the original tree that defined the
-            // register that it was spilled from.
-            // So we use 'unspillTree' to recover that spill temp.
-            TempDsc* t        = regSet.rsUnspillInPlace(unspillTree, unspillTree->GetRegNum());
-            emitAttr emitType = emitActualTypeSize(unspillTree->TypeGet());
-            // Reload into the register specified by 'tree' which may be a GT_RELOAD.
             regNumber dstReg = tree->GetRegNum();
-            GetEmitter()->emitIns_R_S(ins_Load(unspillTree->gtType), emitType, dstReg, t->tdTempNum(), 0);
-            regSet.tmpRlsTemp(t);
 
-            unspillTree->gtFlags &= ~GTF_SPILLED;
-            gcInfo.gcMarkRegPtrVal(dstReg, unspillTree->TypeGet());
+#if defined(TARGET_XARCH) || defined(TARGET_ARM64)
+            // A floating-point, SIMD, or mask constant is not spilled to the stack (see
+            // genProduceReg); rematerialize it directly into the reload target register rather
+            // than loading it from a spill temp that was never created.
+            if (isRematerializableConstant(unspillTree) && ((unspillTree->gtFlags & GTF_NOREG_AT_USE) == 0))
+            {
+#if defined(TARGET_XARCH)
+                // The GT_CNS_VEC and GT_CNS_MSK GenTree* overloads of genSetRegToConst materialize
+                // into the node's own register, so the explicit-register overloads are used to honor
+                // 'dstReg' (which may be a GT_RELOAD target register that differs from the original
+                // definition's register).
+                switch (unspillTree->OperGet())
+                {
+                    case GT_CNS_DBL:
+                        genSetRegToConst(dstReg, unspillTree->TypeGet(), unspillTree);
+                        break;
+                    case GT_CNS_VEC:
+                        genSetRegToConst(dstReg, unspillTree->TypeGet(), &unspillTree->AsVecCon()->gtSimdVal);
+                        break;
+                    case GT_CNS_MSK:
+                        genSetRegToConst(dstReg, unspillTree->TypeGet(), &unspillTree->AsMskCon()->gtSimdMaskVal);
+                        break;
+                    default:
+                        unreached();
+                }
+#else  // TARGET_ARM64
+                // The arm64 genSetRegToConst honors the explicit target register for every operator.
+                genSetRegToConst(dstReg, unspillTree->TypeGet(), unspillTree);
+#endif // TARGET_XARCH
+
+                unspillTree->gtFlags &= ~GTF_SPILLED;
+            }
+            else
+#endif // TARGET_XARCH || TARGET_ARM64
+            {
+                // Here we may have a GT_RELOAD.
+                // The spill temp allocated for it is associated with the original tree that defined the
+                // register that it was spilled from.
+                // So we use 'unspillTree' to recover that spill temp.
+                TempDsc* t        = regSet.rsUnspillInPlace(unspillTree, unspillTree->GetRegNum());
+                emitAttr emitType = emitActualTypeSize(unspillTree->TypeGet());
+                // Reload into the register specified by 'tree' which may be a GT_RELOAD.
+                GetEmitter()->emitIns_R_S(ins_Load(unspillTree->gtType), emitType, dstReg, t->tdTempNum(), 0);
+                regSet.tmpRlsTemp(t);
+
+                unspillTree->gtFlags &= ~GTF_SPILLED;
+                gcInfo.gcMarkRegPtrVal(dstReg, unspillTree->TypeGet());
+            }
         }
     }
 }
 
+//------------------------------------------------------------------------
+// isRematerializableConstant: Determine whether a spilled constant tree temp can be
+//    rematerialized at its reload point instead of being spilled to and reloaded
+//    from the stack.
+//
+// Arguments:
+//    tree - the node that defines the value
+//
+// Return Value:
+//    True if 'tree' is a floating-point, SIMD, or mask constant that this target can
+//    rematerialize without requiring a scratch register or GC bookkeeping.
+//
+// Notes:
+//    These constants have no GC liveness. On xarch every such constant can be
+//    rematerialized without a scratch register (zero via `xorps`, all-bits-set via
+//    `pcmpeqd`/`vpternlogd`, or a RIP-relative load from the constant's memory home),
+//    so all of them are eligible.
+//
+//    On arm64 only the subset that genSetRegToConst encodes directly into an instruction
+//    is eligible; the remaining values are loaded from the constant pool via `adrp`/`ldr`
+//    which needs a scratch address register that is not reserved at the reload point. The
+//    conditions below must therefore stay in sync with the direct-encoding fast paths in
+//    genSetRegToConst: 0.0/fmov-immediate for CNS_DBL, zero/all-bits-set/movi-broadcast for
+//    CNS_VEC, and every CNS_MSK (encoded via `pfalse`/`ptrue`, which never needs a scratch).
+//    Other targets can be enabled later for their own scratch-free subset.
+//
+bool CodeGen::isRematerializableConstant(GenTree* tree)
+{
+#if defined(TARGET_XARCH)
+    return tree->OperIs(GT_CNS_DBL, GT_CNS_VEC, GT_CNS_MSK);
+#elif defined(TARGET_ARM64)
+    switch (tree->OperGet())
+    {
+        case GT_CNS_DBL:
+        {
+            double constValue = tree->AsDblCon()->DconValue();
+            return (*(int64_t*)&constValue == 0) || emitter::emitIns_valid_imm_for_fmov(constValue);
+        }
+
+#if defined(FEATURE_SIMD)
+        case GT_CNS_VEC:
+        {
+            // genSetRegToConst only encodes SIMD8/12/16 constants without a scratch register; any
+            // wider type falls through to its 'unreached' path, so reject those up front (mirroring
+            // the type switch in genSetRegToConst).
+            if (!tree->TypeIs(TYP_SIMD8, TYP_SIMD12, TYP_SIMD16))
+            {
+                return false;
+            }
+
+            GenTreeVecCon* vecCon = tree->AsVecCon();
+
+            if (vecCon->IsAllBitsSet() || vecCon->IsZero())
+            {
+                return true;
+            }
+
+            const bool is8 = tree->TypeIs(TYP_SIMD8);
+            simd16_t   val = vecCon->gtSimd16Val;
+
+            return (ElementsAreSame(val.i32, is8 ? 2 : 4) &&
+                    emitter::emitIns_valid_imm_for_movi(val.i32[0], EA_4BYTE)) ||
+                   (ElementsAreSame(val.i16, is8 ? 4 : 8) &&
+                    emitter::emitIns_valid_imm_for_movi(val.i16[0], EA_2BYTE)) ||
+                   (ElementsAreSame(val.i8, is8 ? 8 : 16) &&
+                    emitter::emitIns_valid_imm_for_movi(val.i8[0], EA_1BYTE));
+        }
+#endif // FEATURE_SIMD
+
+#if defined(FEATURE_MASKED_HW_INTRINSICS)
+        case GT_CNS_MSK:
+            return true;
+#endif // FEATURE_MASKED_HW_INTRINSICS
+
+        default:
+            return false;
+    }
+#else
+    return false;
+#endif // TARGET_XARCH
+}
+
 //------------------------------------------------------------------------
 // genCopyRegIfNeeded: Copy the given node into the specified register
 //
@@ -2242,6 +2361,20 @@ void CodeGen::genProduceReg(GenTree* tree)
             }
             else
             {
+                // Floating-point, SIMD, and mask constants have no GC liveness and (on the targets
+                // for which isRematerializableConstant is enabled) can be rematerialized cheaply
+                // without a scratch register. Rather than spilling such a value to the stack and
+                // reloading it, skip creating the spill temp and emitting the store;
+                // genUnspillRegIfNeeded will rematerialize it at the reload point. If the value will
+                // instead be consumed directly from the spill temp as a contained memory operand
+                // (GTF_NOREG_AT_USE), it must still be spilled normally.
+                if (isRematerializableConstant(tree) && ((tree->gtFlags & GTF_NOREG_AT_USE) == 0))
+                {
+                    tree->gtFlags |= GTF_SPILLED;
+                    tree->gtFlags &= ~GTF_SPILL;
+                    return;
+                }
+
                 regSet.rsSpillTree(tree->GetRegNum(), tree);
                 gcInfo.gcMarkRegSetNpt(genRegMask(tree->GetRegNum()));
             }