[CHERI-RISC-V] Support capability-based integer atomic RMW ops

This allows expanding all integer atomic RMW operations using capability
pointers inline in hybrid mode.

llvm/lib/Target/RISCV/RISCVExpandAtomicPseudoInsts.cpp

@@ -235,6 +235,32 @@ bool RISCVExpandAtomicPseudo::expandMI(MachineBasicBlock &MBB,
   case RISCV::PseudoCheriCmpXchgCapExplicitCap:
     return expandAtomicCmpXchgCap(MBB, MBBI, true, NextMBBI,
                                   /*ExplicitAddrMode=*/true);
+
+    // To avoid the combinatorics explosion of cases here, we use a macro to
+    // handle the explicit addressing mode atomics.
+#define EXPLICIT_ADDRMODE_CASE(name, op, fn, size)                             \
+  case RISCV::PseudoCheriAtomicExplicitCap##name##size:                        \
+    return fn(MBB, MBBI, AtomicRMWInst::op, false, size, true, NextMBBI, true);
+#define HANDLE_EXPLICIT_ADDRMODE_RMW_PSEUDOS(size)                             \
+  EXPLICIT_ADDRMODE_CASE(Swap, Xchg, expandAtomicBinOp, size)                  \
+  EXPLICIT_ADDRMODE_CASE(LoadAdd, Add, expandAtomicBinOp, size)                \
+  EXPLICIT_ADDRMODE_CASE(LoadSub, Sub, expandAtomicBinOp, size)                \
+  EXPLICIT_ADDRMODE_CASE(LoadAnd, And, expandAtomicBinOp, size)                \
+  EXPLICIT_ADDRMODE_CASE(LoadOr, Or, expandAtomicBinOp, size)                  \
+  EXPLICIT_ADDRMODE_CASE(LoadXor, Xor, expandAtomicBinOp, size)                \
+  EXPLICIT_ADDRMODE_CASE(LoadNand, Nand, expandAtomicBinOp, size)              \
+  EXPLICIT_ADDRMODE_CASE(LoadMax, Max, expandAtomicMinMaxOp, size)             \
+  EXPLICIT_ADDRMODE_CASE(LoadMin, Min, expandAtomicMinMaxOp, size)             \
+  EXPLICIT_ADDRMODE_CASE(LoadUMax, UMax, expandAtomicMinMaxOp, size)           \
+  EXPLICIT_ADDRMODE_CASE(LoadUMin, UMin, expandAtomicMinMaxOp, size)
+
+    HANDLE_EXPLICIT_ADDRMODE_RMW_PSEUDOS(8)
+    HANDLE_EXPLICIT_ADDRMODE_RMW_PSEUDOS(16)
+    HANDLE_EXPLICIT_ADDRMODE_RMW_PSEUDOS(32)
+    HANDLE_EXPLICIT_ADDRMODE_RMW_PSEUDOS(64)
+
+#undef HANDLE_EXPLICIT_ADDRMODE_RMW_PSEUDOS
+#undef EXPLICIT_ADDRMODE_CASE
   }
 
   return false;

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -9261,9 +9261,14 @@ RISCVTargetLowering::shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const {
   if (AI->isFloatingPointOperation())
     return AtomicExpansionKind::CmpXChg;
 
-  unsigned Size = AI->getType()->getPrimitiveSizeInBits();
-  if ((Size == 8 || Size == 16) &&
-      !RISCVABI::isCheriPureCapABI(Subtarget.getTargetABI()))
+  // For capability pointers we must always use the smallest possible type
+  // instead of promoting to i32/i64 to ensure we don't trigger bounds errors.
+  const DataLayout &DL = AI->getModule()->getDataLayout();
+  if (DL.isFatPointer(AI->getPointerOperand()->getType()))
+    return AtomicExpansionKind::None;
+
+  uint64_t Size = DL.getTypeSizeInBits(AI->getType()).getFixedSize();
+  if (Size == 8 || Size == 16)
     return AtomicExpansionKind::MaskedIntrinsic;
   return AtomicExpansionKind::None;
 }
@@ -9405,11 +9410,14 @@ bool RISCVTargetLowering::supportsAtomicOperation(const DataLayout &DL,
                                                   Type *ValueTy,
                                                   Type *PointerTy,
                                                   Align Alignment) const {
-  // FIXME: we current have to expand RMW to libcalls since we are missing
-  //  the SelectionDAG nodes+expansions to use the explicit addressing mode.
-  if (DL.isFatPointer(PointerTy) && isa<AtomicRMWInst>(AI) &&
-      !RISCVABI::isCheriPureCapABI(Subtarget.getTargetABI()))
-    return false;
+  if (auto *RMWI = dyn_cast<AtomicRMWInst>(AI)) {
+    // FIXME: We do not support AtomicRMW for a capability arguments (other than
+    //  XCHG in the purecap ABI).
+    if (DL.isFatPointer(ValueTy) && DL.isFatPointer(PointerTy) &&
+        (RMWI->getOperation() != AtomicRMWInst::Xchg ||
+         !RISCVABI::isCheriPureCapABI(Subtarget.getTargetABI())))
+      return false;
+  }
   return TargetLowering::supportsAtomicOperation(DL, AI, ValueTy, PointerTy,
                                                  Alignment);
 }

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

@@ -832,6 +832,11 @@ bool RISCVInstrInfo::isBranchOffsetInRange(unsigned BranchOp,
   }
 }
 
+// clang-format off
+#define FOR_ALL_INT_SIZES(Pseudo)                                              \
+  Pseudo##8: case Pseudo##16: case Pseudo##32: case Pseudo##64
+// clang-format on
+
 unsigned RISCVInstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
   unsigned Opcode = MI.getOpcode();
 
@@ -908,11 +913,15 @@ unsigned RISCVInstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
   case RISCV::PseudoCheriAtomicLoadXor16:
   case RISCV::PseudoCheriAtomicLoadSub8:
   case RISCV::PseudoCheriAtomicLoadSub16:
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapSwap):
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadAdd):
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadAnd):
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadOr):
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadXor):
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadSub):
     return 16;
-  case RISCV::PseudoCheriAtomicLoadNand8:
-  case RISCV::PseudoCheriAtomicLoadNand16:
-  case RISCV::PseudoCheriAtomicLoadNand32:
-  case RISCV::PseudoCheriAtomicLoadNand64:
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicLoadNand):
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadNand):
     return 20;
   case RISCV::PseudoCheriAtomicLoadMax8:
   case RISCV::PseudoCheriAtomicLoadMax16:
@@ -922,11 +931,12 @@ unsigned RISCVInstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
   case RISCV::PseudoCheriAtomicLoadUMax16:
   case RISCV::PseudoCheriAtomicLoadUMin8:
   case RISCV::PseudoCheriAtomicLoadUMin16:
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadMax):
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadMin):
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadUMax):
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriAtomicExplicitCapLoadUMin):
     return 24;
-  case RISCV::PseudoCheriCmpXchg8:
-  case RISCV::PseudoCheriCmpXchg16:
-  case RISCV::PseudoCheriCmpXchg32:
-  case RISCV::PseudoCheriCmpXchg64:
+  case FOR_ALL_INT_SIZES(RISCV::PseudoCheriCmpXchg):
     return 16;
   case RISCV::PseudoCheriCmpXchg8ExplicitCap:
   case RISCV::PseudoCheriCmpXchg16ExplicitCap:
@@ -971,6 +981,7 @@ unsigned RISCVInstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
   }
   }
 }
+#undef FOR_ALL_INT_SIZES
 
 bool RISCVInstrInfo::isAsCheapAsAMove(const MachineInstr &MI) const {
   const unsigned Opcode = MI.getOpcode();

llvm/lib/Target/RISCV/RISCVInstrInfoXCheri.td

@@ -1574,6 +1574,15 @@ class PseudoCheriAMO : Pseudo<(outs GPR:$res, GPR:$scratch),
   let mayStore = 1;
   let hasSideEffects = 0;
 }
+class PseudoCheriAMOExplicitAddrMode<RegisterClass PtrTy,
+                                     RegisterClass ValTy=GPR>
+    : Pseudo<(outs ValTy:$res, ValTy:$scratch),
+             (ins PtrTy:$addr, ValTy:$incr, ixlenimm:$ordering), []> {
+  let Constraints = "@earlyclobber $res,@earlyclobber $scratch";
+  let mayLoad = 1;
+  let mayStore = 1;
+  let hasSideEffects = 0;
+}
 
 multiclass PseudoCheriAMOPat<string AtomicOp, Instruction BaseInst> {
   def : Pat<(!cast<PatFrag>(AtomicOp#"_monotonic") GPCR:$addr, GPR:$incr),
@@ -1588,6 +1597,45 @@ multiclass PseudoCheriAMOPat<string AtomicOp, Instruction BaseInst> {
             (BaseInst GPCR:$addr, GPR:$incr, 7)>;
 }
 
+multiclass CheriAMOExplicitCapPseudos<string Size, RegisterClass PtrTy,
+                                      RegisterClass ValTy=GPR> {
+  def NAME#Swap#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadAdd#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadSub#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadAnd#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadOr#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadXor#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadMax#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadMin#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadUMax#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadUMin#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+  def NAME#LoadNand#Size : PseudoCheriAMOExplicitAddrMode<PtrTy, ValTy>;
+}
+
+// The explicit address mode atomics are always relaxed, so we only need
+// _monotonic patterns.
+class CheriAMOExplicitModePat<string PatOp, string InsnOp, string Size,
+                              RegisterClass PtrTy, RegisterClass ValTy>
+    : Pat<(!cast<PatFrag>("atomic_" # PatOp # "_" # Size # "_monotonic")
+            PtrTy:$addr, ValTy:$incr),
+          (!cast<Pseudo>("PseudoCheriAtomicExplicitCap" # InsnOp # Size)
+            PtrTy:$addr, ValTy:$incr, 2)>;
+
+multiclass PseudoCheriAMOExplicitCapPats<string Size, RegisterClass PtrTy,
+                                         RegisterClass ValTy> {
+  def : CheriAMOExplicitModePat<"swap", "Swap", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_add", "LoadAdd", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_sub", "LoadSub", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_and", "LoadAnd", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_or", "LoadOr", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_xor", "LoadXor", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_nand", "LoadNand", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_max", "LoadMax", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_min", "LoadMin", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_umax", "LoadUMax", Size, PtrTy, ValTy>;
+  def : CheriAMOExplicitModePat<"load_umin", "LoadUMin", Size, PtrTy, ValTy>;
+}
+
 class PseudoCheriCmpXchg<RegisterClass ValTy=GPR>
     : Pseudo<(outs ValTy:$res, GPR:$scratch),
              (ins GPCR:$addr, ValTy:$cmpval, ValTy:$newval, ixlenimm:$ordering), []> {
@@ -1654,6 +1702,7 @@ def PseudoCheriAtomicLoadMin8 : PseudoCheriAMO;
 def PseudoCheriAtomicLoadUMax8 : PseudoCheriAMO;
 def PseudoCheriAtomicLoadUMin8 : PseudoCheriAMO;
 def PseudoCheriAtomicLoadNand8 : PseudoCheriAMO;
+defm PseudoCheriAtomicExplicitCap : CheriAMOExplicitCapPseudos<"8", GPCR>;
 def PseudoCheriCmpXchg8 : PseudoCheriCmpXchg;
 def PseudoCheriCmpXchg8ExplicitCap : PseudoCheriCmpXchgExplicitAddrMode<GPCR>;
 
@@ -1668,10 +1717,12 @@ def PseudoCheriAtomicLoadMin16 : PseudoCheriAMO;
 def PseudoCheriAtomicLoadUMax16 : PseudoCheriAMO;
 def PseudoCheriAtomicLoadUMin16 : PseudoCheriAMO;
 def PseudoCheriAtomicLoadNand16 : PseudoCheriAMO;
+defm PseudoCheriAtomicExplicitCap : CheriAMOExplicitCapPseudos<"16", GPCR>;
 def PseudoCheriCmpXchg16 : PseudoCheriCmpXchg;
 def PseudoCheriCmpXchg16ExplicitCap : PseudoCheriCmpXchgExplicitAddrMode<GPCR>;
 
 def PseudoCheriAtomicLoadNand32 : PseudoCheriAMO;
+defm PseudoCheriAtomicExplicitCap : CheriAMOExplicitCapPseudos<"32", GPCR>;
 def PseudoCheriCmpXchg32 : PseudoCheriCmpXchg;
 def PseudoCheriCmpXchg32ExplicitCap : PseudoCheriCmpXchgExplicitAddrMode<GPCR>;
 
@@ -1681,6 +1732,7 @@ def PseudoCheriCmpXchgCapExplicitCap : PseudoCheriCmpXchgExplicitAddrMode<GPCR,
 
 let Predicates = [HasCheri, HasStdExtA, IsRV64] in {
 def PseudoCheriAtomicLoadNand64 : PseudoCheriAMO;
+defm PseudoCheriAtomicExplicitCap : CheriAMOExplicitCapPseudos<"64", GPCR>;
 def PseudoCheriCmpXchg64 : PseudoCheriCmpXchg;
 def PseudoCheriCmpXchg64ExplicitCap : PseudoCheriCmpXchgExplicitAddrMode<GPCR>;
 } // Predicates = [HasCheri, HasStdExtA, IsRV64]
@@ -1915,7 +1967,6 @@ defm : CheriAMOCapPat<"128", "atomic_swap_cap", "CAMOSWAP_C">;
 
 
 /// Patterns for atomics with explicit capability base register
-/// Only atomic_cmp_swap is supported since all other operations are expanded.
 
 let Predicates = [HasCheri, HasStdExtA, NotCapMode] in {
 defm : PseudoCheriCmpXchgPatExplicit<"atomic_cmp_swap_8",
@@ -1926,10 +1977,16 @@ defm : PseudoCheriCmpXchgPatExplicit<"atomic_cmp_swap_32",
                                      PseudoCheriCmpXchg32ExplicitCap>;
 defm : PseudoCheriCmpXchgPatExplicit<"atomic_cmp_swap_cap",
                                      PseudoCheriCmpXchgCapExplicitCap, GPCR>;
+defm : PseudoCheriAMOExplicitCapPats<"8", GPCR, GPR>;
+defm : PseudoCheriAMOExplicitCapPats<"16", GPCR, GPR>;
+defm : PseudoCheriAMOExplicitCapPats<"32", GPCR, GPR>;
+
 } // Predicates = [HasCheri, HasStdExtA, NotCapMode]
-let Predicates = [HasCheri, HasStdExtA, IsRV64, NotCapMode] in
+let Predicates = [HasCheri, HasStdExtA, IsRV64, NotCapMode] in {
 defm : PseudoCheriCmpXchgPatExplicit<"atomic_cmp_swap_64",
                                      PseudoCheriCmpXchg64ExplicitCap>;
+defm : PseudoCheriAMOExplicitCapPats<"64", GPCR, GPR>;
+} // Predicates = [HasCheri, HasStdExtA, IsRV64, NotCapMode]
 
 
 /// 'F' (Single-Precision Floating-Point) extension