From 53a500edf1d6527b4519f462901667d2ab233a04 Mon Sep 17 00:00:00 2001
From: Egor Bogatov <egorbo@gmail.com>
Date: Mon, 5 Aug 2024 18:09:40 +0200
Subject: [PATCH] Don't use ZR as target in LSE atomics (#105854)

* Don't use ZR as target in LSE atomics

* Update codegenarm64.cpp

* Update lsraarm64.cpp

* Update lsraarm64.cpp
---
 src/coreclr/jit/codegenarm64.cpp | 13 +++++++------
 src/coreclr/jit/lsraarm64.cpp    | 11 +++++++++--
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/coreclr/jit/codegenarm64.cpp b/src/coreclr/jit/codegenarm64.cpp
index 2ac66ba5ea7..ac28ebd30a1 100644
--- a/src/coreclr/jit/codegenarm64.cpp
+++ b/src/coreclr/jit/codegenarm64.cpp
@@ -3878,19 +3878,21 @@ void CodeGen::genLockedInstructions(GenTreeOp* treeNode)
     {
         assert(!data->isContainedIntOrIImmed());
 
+        // These instructions change semantics when targetReg is ZR (the memory ordering becomes weaker).
+        // See atomicBarrierDroppedOnZero in LLVM
+        assert((targetReg != REG_NA) && (targetReg != REG_ZR));
+
         switch (treeNode->gtOper)
         {
             case GT_XORR:
-                GetEmitter()->emitIns_R_R_R(INS_ldsetal, dataSize, dataReg, (targetReg == REG_NA) ? REG_ZR : targetReg,
-                                            addrReg);
+                GetEmitter()->emitIns_R_R_R(INS_ldsetal, dataSize, dataReg, targetReg, addrReg);
                 break;
             case GT_XAND:
             {
                 // Grab a temp reg to perform `MVN` for dataReg first.
                 regNumber tempReg = internalRegisters.GetSingle(treeNode);
                 GetEmitter()->emitIns_R_R(INS_mvn, dataSize, tempReg, dataReg);
-                GetEmitter()->emitIns_R_R_R(INS_ldclral, dataSize, tempReg, (targetReg == REG_NA) ? REG_ZR : targetReg,
-                                            addrReg);
+                GetEmitter()->emitIns_R_R_R(INS_ldclral, dataSize, tempReg, targetReg, addrReg);
                 break;
             }
             case GT_XCHG:
@@ -3908,8 +3910,7 @@ void CodeGen::genLockedInstructions(GenTreeOp* treeNode)
                 break;
             }
             case GT_XADD:
-                GetEmitter()->emitIns_R_R_R(INS_ldaddal, dataSize, dataReg, (targetReg == REG_NA) ? REG_ZR : targetReg,
-                                            addrReg);
+                GetEmitter()->emitIns_R_R_R(INS_ldaddal, dataSize, dataReg, targetReg, addrReg);
                 break;
             default:
                 assert(!"Unexpected treeNode->gtOper");
diff --git a/src/coreclr/jit/lsraarm64.cpp b/src/coreclr/jit/lsraarm64.cpp
index 2352fa5ec02..3906f913527 100644
--- a/src/coreclr/jit/lsraarm64.cpp
+++ b/src/coreclr/jit/lsraarm64.cpp
@@ -1086,10 +1086,17 @@ int LinearScan::BuildNode(GenTree* tree)
                     }
                     setInternalRegsDelayFree = true;
                 }
+                buildInternalRegisterUses();
+                if (dstCount == 1)
+                {
+                    BuildDef(tree);
+                }
             }
-            buildInternalRegisterUses();
-            if (dstCount == 1)
+            else
             {
+                // We always need the target reg for LSE, even if
+                // return value is unused, see genLockedInstructions
+                buildInternalRegisterUses();
                 BuildDef(tree);
             }
         }