[RISCV] Preserve narrow cmpxchg results when widening

2 files changed, 40 insertions, 2 deletions

diff --git a/backend/lib/Target/RISCV/RISCVCodeGen.cpp b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
index 44f2e7e..506c1e2 100644
--- a/backend/lib/Target/RISCV/RISCVCodeGen.cpp
+++ b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
@@ -546,6 +546,7 @@ narrowAnyExtBinaryPlan(const MachineInstr &Extension, const MachineFunction &Fun
       OriginalType->scalarBitWidth() >= 64)
     return makeUnexpected(ErrorCode::VerificationError,
                           "RISC-V narrow compare-exchange has invalid operands");
+  const Register OriginalDestination = *Destination;
 
   auto WideType = LLT::makeAnyScalar(64);
   if (!WideType)
@@ -561,7 +562,7 @@ narrowAnyExtBinaryPlan(const MachineInstr &Extension, const MachineFunction &Fun
     return std::unexpected(WideNewValue.error());
 
   Register WideDestination;
-  MachineInstr *FoldedExtension = singleRegisterUse(Function, *Destination);
+  MachineInstr *FoldedExtension = singleRegisterUse(Function, OriginalDestination);
   if (FoldedExtension != nullptr && FoldedExtension->opcode() == opcode(GenericOpcode::G_ANYEXT) &&
       operandType(Function, *FoldedExtension, 0)
           .transform([](LLT Type) { return Type.scalarBitWidth() == 64; })
@@ -597,7 +598,7 @@ narrowAnyExtBinaryPlan(const MachineInstr &Extension, const MachineFunction &Fun
   auto Truncation = Builder.buildInstr(opcode(GenericOpcode::G_TRUNC));
   if (!Truncation)
     return std::unexpected(Truncation.error());
-  if (auto Added = Truncation->addDef(*Destination); !Added)
+  if (auto Added = Truncation->addDef(OriginalDestination); !Added)
     return std::unexpected(Added.error());
   if (auto Added = Truncation->addUse(WideDestination); !Added)
     return std::unexpected(Added.error());
diff --git a/backend/unittests/Target/RISCVCodeGenTest.cpp b/backend/unittests/Target/RISCVCodeGenTest.cpp
index 3a08f92..17987f1 100644
--- a/backend/unittests/Target/RISCVCodeGenTest.cpp
+++ b/backend/unittests/Target/RISCVCodeGenTest.cpp
@@ -907,6 +907,43 @@ void testAtomicCompareExchangeLegalization(TestContext &Context) {
                      Target->validateInstruction(**BuiltAtomic, (*Function)->registerInfo()) &&
                      LegalAction && *LegalAction == rcc::LegalizationAction::Legal,
                  "compare-exchange widening must preserve the narrow MMO in a legal i64 container");
+
+  auto FallbackCompare = (*Function)->registerInfo().createVirtualRegister(*I8);
+  auto FallbackNewValue = (*Function)->registerInfo().createVirtualRegister(*I8);
+  auto FallbackResult = (*Function)->registerInfo().createVirtualRegister(*I8);
+  if (!FallbackCompare || !FallbackNewValue || !FallbackResult || !Builder.setInsertPoint(**Block))
+    return;
+  for (rcc::Register Definition : {*FallbackCompare, *FallbackNewValue}) {
+    auto Implicit = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_IMPLICIT_DEF));
+    if (!Implicit || !Implicit->addDef(Definition) || !Implicit->build())
+      return;
+  }
+  auto Fallback = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_ATOMIC_CMPXCHG));
+  if (!Fallback || !Fallback->addDef(*FallbackResult) || !Fallback->addUse(*Address) ||
+      !Fallback->addUse(*FallbackCompare) || !Fallback->addUse(*FallbackNewValue) ||
+      !Fallback->addMemOperand(**Memory))
+    return;
+  auto BuiltFallback = Fallback->build();
+  if (!BuiltFallback)
+    return;
+  auto FallbackAction = rcc::riscv::legalizerInfo().classify(**BuiltFallback, **Function, *Target);
+  auto FallbackLowered =
+      FallbackAction
+          ? rcc::riscv::legalizerInfo().lower(**BuiltFallback, **Function, *Target, *FallbackAction)
+          : rcc::Expected<void>(std::unexpected(FallbackAction.error()));
+  rcc::MachineInstr *Truncation = (*BuiltFallback)->next();
+  const rcc::Register *WideFallbackResult = (*BuiltFallback)->operandAt(0)->getIf<rcc::Register>();
+  const rcc::Register *TruncatedResult =
+      Truncation == nullptr ? nullptr : Truncation->operandAt(0)->getIf<rcc::Register>();
+  const rcc::Register *TruncationSource =
+      Truncation == nullptr ? nullptr : Truncation->operandAt(1)->getIf<rcc::Register>();
+  Context.expect(FallbackLowered && WideFallbackResult != nullptr &&
+                     *WideFallbackResult != *FallbackResult && Truncation != nullptr &&
+                     Truncation->opcode() == rcc::opcode(rcc::GenericOpcode::G_TRUNC) &&
+                     TruncatedResult != nullptr && *TruncatedResult == *FallbackResult &&
+                     TruncationSource != nullptr && *TruncationSource == *WideFallbackResult &&
+                     Target->validateInstruction(*Truncation, (*Function)->registerInfo()),
+                 "general compare-exchange widening must truncate from the new i64 result");
 }
 
 void testAtomicSelection(TestContext &Context) {