[RISCV] Use AMOSWAP for 32 and 64-bit atomic stores

[paulkirth]

Summary

These are expected to be more performant and comply with the
recommendations in Hans Boehm's proposal for changes to table A.6 of the
unprivileged spec and the RISC-V psABI. Further details can be found in
the psABI proposal:
https://github.com/riscv-non-isa/riscv-elf-psabi-doc/pull/378.

The proposal suggests that as an optimization, the memory_order_seq_cst
store mapping may be replaced by the following for 32- and 64-bit
operands:

amoswap.rl{w|d};

Depends on D149486

[jrtc27]

I thought this was being abandoned because it has worse performance. Certainly there are codegen regressions here that clearly indicate it being a bad idea.

[paulkirth]

In D149488#4306677, @jrtc27 wrote:

I thought this was being abandoned because it has worse performance. Certainly there are codegen regressions here that clearly indicate it being a bad idea.

The proposed change to the ISA, psABI, and the discussion on tech-unprivileged seem to indicate that this may still be profitable. Since this is only a potential optimization, I think it's fine to wait on this "optimization", or forgo it altogether.

And yes, for forced atomics, it indeed looks like there is a regression in codegen.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

15806

@jrtc27 Maybe we should forgo the optimization when forced atomics are used? As you've pointed out there seem to be some unintended interactions in that configuration.

[paulkirth]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

15806

oh, this is the wrong place anyway, so ignore my suggestion

[paulkirth]

Disable AMOSWAP optimization with forced atomics

[jrtc27]

llvm/test/CodeGen/RISCV/atomic-load-store.ll

1427

Why are we writing a killed value to a0?

[paulkirth]

llvm/test/CodeGen/RISCV/atomic-load-store.ll

1427

After Atomic Expand the store is rewritten from:

store atomic i32 %b, ptr %a seq_cst, align 4

to

%1 = atomicrmw xchg ptr %a, i32 %b release, align 4

The langref says the semantics for xchg here should be :*ptr = val, so I think that's all correct. https://llvm.org/docs/LangRef.html#id229

From there, the rest of the codegen is consistent w/ the contents of https://github.com/llvm/llvm-project/blob/f762798599171efca03964f4371cc8104d2392f9/llvm/test/CodeGen/RISCV/atomic-rmw.ll#L9844

If that is wrong, then I guess we need to adjust the tablegen for AMOSWAP?

[paulkirth]

Rebase

[asb]

Do you know if GCC landed this change as well?

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

10097

Nit: Capitalise 'returns' and add full stop at end of sentence.

[paulkirth]

In D149488#4315076, @asb wrote:

Do you know if GCC landed this change as well?

Based on this test, I'd guess that they haven't. https://github.com/gcc-mirror/gcc/blob/6cff5f3da7f263bb11cf48e6011931422b62b364/gcc/testsuite/gcc.target/riscv/amo-table-a-6-store-compat-3.c#L8

[paulkirth]

Rebase and adpat code to use EnableSeqCstFence.

I should be able to test this on some existing hardware against the existing version and the trailing fence

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

10098

I think LLVM coding guidelines are to use static for functions rather than anonymous namespace. https://llvm.org/docs/CodingStandards.html#anonymous-namespaces

[paulkirth]

Make helper static instead of using anonymous namespace.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

10098

huh, I could have sworn we preferred the anonymous namespace, but obviously not.

[paulkirth]

Rebase and adopt feature/march flag

[paulkirth]

Rebase

llvm/lib/Target/RISCV/RISCVISelLowering.h

public:
 
   bool convertSetCCLogicToBitwiseLogic(EVT VT) const override {
     return VT.isScalarInteger();
   }
   bool convertSelectOfConstantsToMath(EVT VT) const override { return true; }
 
   bool preferZeroCompareBranch() const override { return true; }
 
-  bool shouldInsertFencesForAtomic(const Instruction *I) const override {
-    return isa<LoadInst>(I) || isa<StoreInst>(I);
-  }
+  bool shouldInsertFencesForAtomic(const Instruction *I) const override;
+
   Instruction *emitLeadingFence(IRBuilderBase &Builder, Instruction *Inst,
                                 AtomicOrdering Ord) const override;
   Instruction *emitTrailingFence(IRBuilderBase &Builder, Instruction *Inst,
                                  AtomicOrdering Ord) const override;
 
   bool isFMAFasterThanFMulAndFAdd(const MachineFunction &MF,
                                   EVT VT) const override;
 
public:
   bool shouldConsiderGEPOffsetSplit() const override { return true; }
 
   bool decomposeMulByConstant(LLVMContext &Context, EVT VT,
                               SDValue C) const override;
 
   bool isMulAddWithConstProfitable(SDValue AddNode,
                                    SDValue ConstNode) const override;
 
+  TargetLoweringBase::AtomicExpansionKind
+  shouldExpandAtomicStoreInIR(StoreInst *SI) const override;
+
   TargetLowering::AtomicExpansionKind
   shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const override;
   Value *emitMaskedAtomicRMWIntrinsic(IRBuilderBase &Builder, AtomicRMWInst *AI,
                                       Value *AlignedAddr, Value *Incr,
                                       Value *Mask, Value *ShiftAmt,
                                       AtomicOrdering Ord) const override;
   TargetLowering::AtomicExpansionKind
   shouldExpandAtomicCmpXchgInIR(AtomicCmpXchgInst *CI) const override;

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

public:
   SDValue const &getCOp() const { return operator[](CommonPos); }
 
   // Returns differ operand.
   SDValue const &getDOp() const { return operator[](DifferPos); }
 
   // Returns condition code of comparison operation.
   ISD::CondCode getCondCode() const { return CCode; }
 };
+
+// returns true for seq_cst stores of 32/64bit
+bool canAmoSwapStoreInst(const StoreInst *SI, const RISCVSubtarget &Subtarget) {
+  unsigned Size = SI->getValueOperand()->getType()->getPrimitiveSizeInBits();
+  return EnableSeqCstFence && !Subtarget.hasForcedAtomics() &&
+         SI->getOrdering() == AtomicOrdering::SequentiallyConsistent &&
+         (Size == 32 || Size == 64);
+}
+
 } // namespace
 
 // Verifies conditions to apply an optimization.
 // Returns Reference comparison code and three operands A, B, C.
 // Conditions for optimization:
 //   One operand of the compasions has to be common.
 //   This operand is written to C.
 //   Two others operands are differend. They are written to A and B.

 Instruction *RISCVTargetLowering::emitTrailingFence(IRBuilderBase &Builder,
                                                     Instruction *Inst,
                                                     AtomicOrdering Ord) const {
   if (Subtarget.hasStdExtZtso())
     return nullptr;
 
   if (isa<LoadInst>(Inst) && isAcquireOrStronger(Ord))
     return Builder.CreateFence(AtomicOrdering::Acquire);
+
   if (EnableSeqCstFence && isa<StoreInst>(Inst) &&
-      Ord == AtomicOrdering::SequentiallyConsistent)
+      Ord == AtomicOrdering::SequentiallyConsistent) {
+    // We only use the trailing fence for s{b,h}.
+    // Table A.6 prescribes AMOSWAP for s{w,d}
+    //   fence rw, w
+    //   s{b,h}
+    //   fence rw, rw
+    auto *SI = cast<StoreInst>(Inst);
+    if (!canAmoSwapStoreInst(SI, Subtarget))
       return Builder.CreateFence(AtomicOrdering::SequentiallyConsistent);
+  }
   return nullptr;
 }
 
+bool RISCVTargetLowering::shouldInsertFencesForAtomic(
+    const Instruction *I) const {
+  if (isa<LoadInst>(I))
+    return true;
+  if (auto *SI = dyn_cast<StoreInst>(I))
+    return !canAmoSwapStoreInst(SI, Subtarget);
+  return false;
+}
+
+TargetLoweringBase::AtomicExpansionKind
+RISCVTargetLowering::shouldExpandAtomicStoreInIR(StoreInst *SI) const {
+  if (canAmoSwapStoreInst(SI, Subtarget)) {
+    SI->setOrdering(AtomicOrdering::Release);
+    return AtomicExpansionKind::Expand;
+  }
+  return AtomicExpansionKind::None;
+}
+
 TargetLowering::AtomicExpansionKind
 RISCVTargetLowering::shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const {
   // atomicrmw {fadd,fsub} must be expanded to use compare-exchange, as floating
   // point operations can't be used in an lr/sc sequence without breaking the
   // forward-progress guarantee.
   if (AI->isFloatingPointOperation() ||
       AI->getOperation() == AtomicRMWInst::UIncWrap ||
       AI->getOperation() == AtomicRMWInst::UDecWrap)

llvm/test/CodeGen/RISCV/atomic-load-store.ll

 ; RV32IA-WMO-NEXT:    sw a1, 0(a0)
 ; RV32IA-WMO-NEXT:    ret
 ;
 ; RV32IA-TSO-LABEL: atomic_store_i32_seq_cst:
 ; RV32IA-TSO:       # %bb.0:
 ; RV32IA-TSO-NEXT:    sw a1, 0(a0)
 ; RV32IA-TSO-NEXT:    ret
 ;
 ; RV64I-LABEL: atomic_store_i32_seq_cst:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    addi sp, sp, -16
 ; RV64I-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
 ; RV64I-NEXT:    li a2, 5
 ; RV64I-NEXT:    call __atomic_store_4@plt
 ; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IA-WMO-LABEL: atomic_store_i32_seq_cst:
 ; RV64IA-WMO:       # %bb.0:
 ; RV64IA-WMO-NEXT:    fence rw, w
 ; RV64IA-WMO-NEXT:    sw a1, 0(a0)
 ; RV64IA-WMO-NEXT:    ret
 ;
 ; RV64IA-TSO-LABEL: atomic_store_i32_seq_cst:
 ; RV64IA-TSO:       # %bb.0:
 ; RV64IA-TSO-NEXT:    sw a1, 0(a0)
 ; RV64IA-TSO-NEXT:    ret
 ;
 ; RV32IA-WMO-TRAILING-FENCE-LABEL: atomic_store_i32_seq_cst:
 ; RV32IA-WMO-TRAILING-FENCE:       # %bb.0:
-; RV32IA-WMO-TRAILING-FENCE-NEXT:    fence rw, w
-; RV32IA-WMO-TRAILING-FENCE-NEXT:    sw a1, 0(a0)
-; RV32IA-WMO-TRAILING-FENCE-NEXT:    fence rw, rw
+; RV32IA-WMO-TRAILING-FENCE-NEXT:    amoswap.w.rl a0, a1, (a0)
 ; RV32IA-WMO-TRAILING-FENCE-NEXT:    ret
 ;
 ; RV32IA-TSO-TRAILING-FENCE-LABEL: atomic_store_i32_seq_cst:
 ; RV32IA-TSO-TRAILING-FENCE:       # %bb.0:
-; RV32IA-TSO-TRAILING-FENCE-NEXT:    sw a1, 0(a0)
+; RV32IA-TSO-TRAILING-FENCE-NEXT:    amoswap.w.rl a0, a1, (a0)
 ; RV32IA-TSO-TRAILING-FENCE-NEXT:    ret
 ;
 ; RV64IA-WMO-TRAILING-FENCE-LABEL: atomic_store_i32_seq_cst:
 ; RV64IA-WMO-TRAILING-FENCE:       # %bb.0:
-; RV64IA-WMO-TRAILING-FENCE-NEXT:    fence rw, w
-; RV64IA-WMO-TRAILING-FENCE-NEXT:    sw a1, 0(a0)
-; RV64IA-WMO-TRAILING-FENCE-NEXT:    fence rw, rw
+; RV64IA-WMO-TRAILING-FENCE-NEXT:    amoswap.w.rl a0, a1, (a0)
 ; RV64IA-WMO-TRAILING-FENCE-NEXT:    ret
 ;
 ; RV64IA-TSO-TRAILING-FENCE-LABEL: atomic_store_i32_seq_cst:
 ; RV64IA-TSO-TRAILING-FENCE:       # %bb.0:
-; RV64IA-TSO-TRAILING-FENCE-NEXT:    sw a1, 0(a0)
+; RV64IA-TSO-TRAILING-FENCE-NEXT:    amoswap.w.rl a0, a1, (a0)
 ; RV64IA-TSO-TRAILING-FENCE-NEXT:    ret
   store atomic i32 %b, ptr %a seq_cst, align 4
   ret void
 }
 
 define void @atomic_store_i64_unordered(ptr %a, i64 %b) nounwind {
 ; RV32I-LABEL: atomic_store_i64_unordered:
 ; RV32I:       # %bb.0:

 ;
 ; RV64IA-TSO-LABEL: atomic_store_i64_seq_cst:
 ; RV64IA-TSO:       # %bb.0:
 ; RV64IA-TSO-NEXT:    sd a1, 0(a0)
 ; RV64IA-TSO-NEXT:    ret
 ;
 ; RV64IA-WMO-TRAILING-FENCE-LABEL: atomic_store_i64_seq_cst:
 ; RV64IA-WMO-TRAILING-FENCE:       # %bb.0:
-; RV64IA-WMO-TRAILING-FENCE-NEXT:    fence rw, w
-; RV64IA-WMO-TRAILING-FENCE-NEXT:    sd a1, 0(a0)
-; RV64IA-WMO-TRAILING-FENCE-NEXT:    fence rw, rw
+; RV64IA-WMO-TRAILING-FENCE-NEXT:    amoswap.d.rl a0, a1, (a0)
 ; RV64IA-WMO-TRAILING-FENCE-NEXT:    ret
 ;
 ; RV64IA-TSO-TRAILING-FENCE-LABEL: atomic_store_i64_seq_cst:
 ; RV64IA-TSO-TRAILING-FENCE:       # %bb.0:
-; RV64IA-TSO-TRAILING-FENCE-NEXT:    sd a1, 0(a0)
+; RV64IA-TSO-TRAILING-FENCE-NEXT:    amoswap.d.rl a0, a1, (a0)
 ; RV64IA-TSO-TRAILING-FENCE-NEXT:    ret
   store atomic i64 %b, ptr %a seq_cst, align 8
   ret void
 }