[RISCV] Lower cmpxchg with success

2 files changed, 131 insertions, 1 deletions

diff --git a/backend/lib/Target/RISCV/RISCVCodeGen.cpp b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
index 506c1e2..9cf7b53 100644
--- a/backend/lib/Target/RISCV/RISCVCodeGen.cpp
+++ b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
@@ -391,6 +391,9 @@ buildBinaryBefore(MachineIRBuilder &Builder, MachineInstr &Anchor,
 
 [[nodiscard]] Expected<void> eraseInstruction(MachineInstr &Instruction);
 
+[[nodiscard]] Expected<void> rewriteUnary(MachineInstr &Instruction, GenericOpcode Generic,
+                                          Register Destination, Register Source);
+
 struct NarrowAnyExtBinaryPlan final {
   MachineInstr *Binary = nullptr;
   GenericOpcode Generic = GenericOpcode::Invalid;
@@ -608,6 +611,77 @@ narrowAnyExtBinaryPlan(const MachineInstr &Extension, const MachineFunction &Fun
   return {};
 }
 
+[[nodiscard]] Expected<void>
+lowerAtomicCompareExchangeWithSuccess(MachineInstr &Instruction, MachineFunction &Function,
+                                      const TargetDescription &Target) {
+  const MachineOperand *OldValueOperand = Instruction.operandAt(0);
+  const MachineOperand *SuccessOperand = Instruction.operandAt(1);
+  const MachineOperand *AddressOperand = Instruction.operandAt(2);
+  const MachineOperand *CompareOperand = Instruction.operandAt(3);
+  const MachineOperand *NewValueOperand = Instruction.operandAt(4);
+  const Register *OldValue =
+      OldValueOperand == nullptr ? nullptr : OldValueOperand->getIf<Register>();
+  const Register *Success = SuccessOperand == nullptr ? nullptr : SuccessOperand->getIf<Register>();
+  const Register *Address = AddressOperand == nullptr ? nullptr : AddressOperand->getIf<Register>();
+  const Register *Compare = CompareOperand == nullptr ? nullptr : CompareOperand->getIf<Register>();
+  const Register *NewValue =
+      NewValueOperand == nullptr ? nullptr : NewValueOperand->getIf<Register>();
+  const std::optional<LLT> ValueType = resultType(Function, Instruction);
+  if (OldValue == nullptr || Success == nullptr || Address == nullptr || Compare == nullptr ||
+      NewValue == nullptr || !ValueType || Instruction.memOperands().size() != 1)
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "RISC-V compare-exchange-with-success has invalid operands");
+
+  auto AtomicResult = Function.registerInfo().createVirtualRegister(*ValueType);
+  if (!AtomicResult)
+    return std::unexpected(AtomicResult.error());
+  MachineIRBuilder Builder(Function, Target);
+  auto Point = Builder.setInsertPointBefore(Instruction);
+  if (!Point)
+    return std::unexpected(Point.error());
+  auto Atomic = Builder.buildInstr(opcode(GenericOpcode::G_ATOMIC_CMPXCHG));
+  if (!Atomic)
+    return std::unexpected(Atomic.error());
+  if (auto Added = Atomic->addDef(*AtomicResult); !Added)
+    return std::unexpected(Added.error());
+  if (auto Added = Atomic->addUse(*Address); !Added)
+    return std::unexpected(Added.error());
+  if (auto Added = Atomic->addUse(*Compare); !Added)
+    return std::unexpected(Added.error());
+  if (auto Added = Atomic->addUse(*NewValue); !Added)
+    return std::unexpected(Added.error());
+  if (auto Added = Atomic->addMemOperand(*Instruction.memOperands().front()); !Added)
+    return std::unexpected(Added.error());
+  auto BuiltAtomic = Atomic->build();
+  if (!BuiltAtomic)
+    return std::unexpected(BuiltAtomic.error());
+
+  Point = Builder.setInsertPointBefore(Instruction);
+  if (!Point)
+    return std::unexpected(Point.error());
+  auto CompareResult = Builder.buildInstr(opcode(GenericOpcode::G_ICMP));
+  if (!CompareResult)
+    return std::unexpected(CompareResult.error());
+  if (auto Added = CompareResult->addDef(*Success); !Added)
+    return std::unexpected(Added.error());
+  if (auto Added = CompareResult->addIntegerPredicate(IntegerPredicate::Equal); !Added)
+    return std::unexpected(Added.error());
+  if (auto Added = CompareResult->addUse(*AtomicResult); !Added)
+    return std::unexpected(Added.error());
+  if (auto Added = CompareResult->addUse(*Compare); !Added)
+    return std::unexpected(Added.error());
+  auto BuiltCompare = CompareResult->build();
+  if (!BuiltCompare)
+    return std::unexpected(BuiltCompare.error());
+
+  while (!Instruction.memOperands().empty()) {
+    auto Erased = Instruction.eraseMemOperand(0);
+    if (!Erased)
+      return std::unexpected(Erased.error());
+  }
+  return rewriteUnary(Instruction, GenericOpcode::COPY, *OldValue, *AtomicResult);
+}
+
 [[nodiscard]] Expected<void> foldNarrowAnyExtBinary(MachineInstr &Extension,
                                                      MachineFunction &Function,
                                                      const TargetDescription &Target) {
@@ -3049,6 +3123,23 @@ public:
       }
       return Type->scalarBitWidth() == 64 ? LegalizationAction::Legal : LegalizationAction::Custom;
     }
+    case GenericOpcode::G_ATOMIC_CMPXCHG_WITH_SUCCESS: {
+      const std::optional<LLT> Success = operandType(Function, Instruction, 1);
+      const std::optional<LLT> Address = operandType(Function, Instruction, 2);
+      const std::optional<LLT> Compare = operandType(Function, Instruction, 3);
+      const std::optional<LLT> NewValue = operandType(Function, Instruction, 4);
+      return Type && Success && Address && Compare && NewValue && isI1Type(*Success) &&
+                     *Compare == *Type && *NewValue == *Type &&
+                     (Type->isInteger() || Type->isAnyScalar()) &&
+                     (Type->scalarBitWidth() == 8 || Type->scalarBitWidth() == 16 ||
+                      Type->scalarBitWidth() == 32 || Type->scalarBitWidth() == 64) &&
+                     Address->isPointer() && Address->addressSpace() == 0 &&
+                     Address->scalarBitWidth() == 64 &&
+                     isSupportedAtomicContainerAccess(Instruction, AtomicOperation::CompareExchange,
+                                                      (Type->scalarBitWidth() + 7U) / 8U)
+                 ? LegalizationAction::Custom
+                 : LegalizationAction::Unsupported;
+    }
     case GenericOpcode::G_FENCE: {
       const MachineOperand *OrderingOperand = Instruction.operandAt(0);
       const MachineOperand *ScopeOperand = Instruction.operandAt(1);
@@ -3213,6 +3304,9 @@ public:
         Instruction.opcode() == opcode(GenericOpcode::G_ATOMIC_CMPXCHG))
       return lowerAtomicCompareExchange(Instruction, Function, Target);
     if (Action == LegalizationAction::Custom &&
+        Instruction.opcode() == opcode(GenericOpcode::G_ATOMIC_CMPXCHG_WITH_SUCCESS))
+      return lowerAtomicCompareExchangeWithSuccess(Instruction, Function, Target);
+    if (Action == LegalizationAction::Custom &&
         (Instruction.opcode() == opcode(GenericOpcode::G_ANYEXT) ||
          Instruction.opcode() == opcode(GenericOpcode::G_SEXT) ||
          Instruction.opcode() == opcode(GenericOpcode::G_ZEXT))) {
diff --git a/backend/unittests/Target/RISCVCodeGenTest.cpp b/backend/unittests/Target/RISCVCodeGenTest.cpp
index 17987f1..9c44b28 100644
--- a/backend/unittests/Target/RISCVCodeGenTest.cpp
+++ b/backend/unittests/Target/RISCVCodeGenTest.cpp
@@ -831,12 +831,14 @@ void testAtomicCompareExchangeLegalization(TestContext &Context) {
   if (!Target || !Function)
     return;
   auto Block = (*Function)->createBasicBlock();
+  auto I1 = rcc::LLT::makeAnyScalar(1);
   auto I8 = rcc::LLT::makeAnyScalar(8);
   auto I64 = rcc::LLT::makeAnyScalar(64);
   auto Pointer = rcc::LLT::makePointer(0, 64);
   auto Size = rcc::MemorySize::fixed(1);
   auto Align = rcc::Alignment::create(1);
-  if (!Block || !(*Function)->appendBlock(**Block) || !I8 || !I64 || !Pointer || !Size || !Align)
+  if (!Block || !(*Function)->appendBlock(**Block) || !I1 || !I8 || !I64 || !Pointer || !Size ||
+      !Align)
     return;
   auto Address = (*Function)->registerInfo().createVirtualRegister(*Pointer);
   auto Compare = (*Function)->registerInfo().createVirtualRegister(*I8);
@@ -944,6 +946,40 @@ void testAtomicCompareExchangeLegalization(TestContext &Context) {
                      TruncationSource != nullptr && *TruncationSource == *WideFallbackResult &&
                      Target->validateInstruction(*Truncation, (*Function)->registerInfo()),
                  "general compare-exchange widening must truncate from the new i64 result");
+
+  auto OldValue = (*Function)->registerInfo().createVirtualRegister(*I64);
+  auto Success = (*Function)->registerInfo().createVirtualRegister(*I1);
+  if (!OldValue || !Success || !Builder.setInsertPoint(**Block))
+    return;
+  auto WithSuccess =
+      Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_ATOMIC_CMPXCHG_WITH_SUCCESS));
+  if (!WithSuccess || !WithSuccess->addDef(*OldValue) || !WithSuccess->addDef(*Success) ||
+      !WithSuccess->addUse(*Address) || !WithSuccess->addUse(*Compare) ||
+      !WithSuccess->addUse(*NewValue) || !WithSuccess->addMemOperand(**Memory))
+    return;
+  auto BuiltWithSuccess = WithSuccess->build();
+  if (!BuiltWithSuccess)
+    return;
+  auto WithSuccessAction =
+      rcc::riscv::legalizerInfo().classify(**BuiltWithSuccess, **Function, *Target);
+  auto WithSuccessLowered = WithSuccessAction
+                                ? rcc::riscv::legalizerInfo().lower(**BuiltWithSuccess, **Function,
+                                                                    *Target, *WithSuccessAction)
+                                : rcc::Expected<void>(std::unexpected(WithSuccessAction.error()));
+  rcc::MachineInstr *SuccessCompare = (*BuiltWithSuccess)->previous();
+  rcc::MachineInstr *LoweredAtomic =
+      SuccessCompare == nullptr ? nullptr : SuccessCompare->previous();
+  Context.expect(
+      WithSuccessAction && *WithSuccessAction == rcc::LegalizationAction::Custom &&
+          WithSuccessLowered && LoweredAtomic != nullptr && SuccessCompare != nullptr &&
+          LoweredAtomic->opcode() == rcc::opcode(rcc::GenericOpcode::G_ATOMIC_CMPXCHG) &&
+          SuccessCompare->opcode() == rcc::opcode(rcc::GenericOpcode::G_ICMP) &&
+          (*BuiltWithSuccess)->opcode() == rcc::opcode(rcc::GenericOpcode::COPY) &&
+          (*BuiltWithSuccess)->memOperands().empty() &&
+          Target->validateInstruction(*LoweredAtomic, (*Function)->registerInfo()) &&
+          Target->validateInstruction(*SuccessCompare, (*Function)->registerInfo()) &&
+          Target->validateInstruction(**BuiltWithSuccess, (*Function)->registerInfo()),
+      "compare-exchange-with-success must lower to atomic, equality test, and old-value copy");
 }
 
 void testAtomicSelection(TestContext &Context) {