From c1d62627ff93c0fedbf26752a7dcbf0645cc0abc Mon Sep 17 00:00:00 2001
From: Tanner Gooding <tagoo@outlook.com>
Date: Tue, 4 Jun 2024 08:07:28 -0700
Subject: [PATCH] Ensure we call genProduceReg for the embedded mask node
 (#103024)

---
 src/coreclr/jit/hwintrinsiccodegenxarch.cpp | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/src/coreclr/jit/hwintrinsiccodegenxarch.cpp b/src/coreclr/jit/hwintrinsiccodegenxarch.cpp
index 11d746111be5a..4aabaeda47cdd 100644
--- a/src/coreclr/jit/hwintrinsiccodegenxarch.cpp
+++ b/src/coreclr/jit/hwintrinsiccodegenxarch.cpp
@@ -180,6 +180,7 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
     CORINFO_InstructionSet isa         = HWIntrinsicInfo::lookupIsa(intrinsicId);
     HWIntrinsicCategory    category    = HWIntrinsicInfo::lookupCategory(intrinsicId);
     size_t                 numArgs     = node->GetOperandCount();
+    GenTree*               embMaskNode = nullptr;
     GenTree*               embMaskOp   = nullptr;
 
     // We need to validate that other phases of the compiler haven't introduced unsupported intrinsics
@@ -233,6 +234,9 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
                 op2->ClearContained();
                 op2->SetRegNum(targetReg);
 
+                // Track the original mask node so we can call genProduceReg
+                embMaskNode = node;
+
                 // Fixup all the already initialized variables
                 node        = op2->AsHWIntrinsic();
                 intrinsicId = node->GetHWIntrinsicId();
@@ -696,13 +700,27 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
             // Handle an extra operand we need to consume so that
             // embedded masking can work without making the overall
             // logic significantly more complex.
+
+            assert(embMaskNode != nullptr);
             genConsumeReg(embMaskOp);
         }
 
         genProduceReg(node);
+
+        if (embMaskNode != nullptr)
+        {
+            // Similarly to the mask operand, we need to handle the
+            // mask node to ensure everything works correctly, particularly
+            // lifetimes and spilling if required. Doing it this way avoids
+            // needing to duplicate all our existing handling.
+
+            assert(embMaskOp != nullptr);
+            genProduceReg(embMaskNode);
+        }
         return;
     }
 
+    assert(embMaskNode == nullptr);
     assert(embMaskOp == nullptr);
 
     switch (isa)