[LoongArch] Optimize the atomic store with amswap_db.[w/d]

[gonglingqin]

Summary

When AtomicOrdering is release or stronger, use

amswap_db.[w/d] $zero, $a1, $a0

instead of

dbar 0
st.[w/d] $a0, $a1, 0

Thanks to @xry111 for the suggestion: https://reviews.llvm.org/D128901#3626635

[xry111]

llvm/test/CodeGen/LoongArch/ir-instruction/load-store-atomic.ll

308

The manual says LA32 may lack amswap_db.w (its not in Table 2-1, page 12 in the Chinese version). If it's not a documentation error we need to keep the old behavior for LA32.

[gonglingqin]

llvm/test/CodeGen/LoongArch/ir-instruction/load-store-atomic.ll

308

Thanks, I will change that.

[gonglingqin]

Address @xry111's comments.

[gonglingqin]

A polite ping. Thanks:)

[xen0n]

I'm not sure about the changes (they are a little bit too much and I must admit I'm relatively weak at memory models) but the test case changes are good. I'd leave the LGTM to someone else for now...

[gonglingqin]

rebase

[SixWeining]

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

1106

No need to use else after return.

1107–1109

Size = isa<StoreInst>(I) ? I->getOperand(0)->getType()->getIntegerBitWidth() : 0;

[SixWeining]

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

1107–1113

How about:

if (isa<LoadInst>(I))
  return true;

// On LA64, atomic store operations with IntegerBitWidth of 32 and 64 do not
// require fences beacuse we can use amswap_db.[w/d].
if (isa<StoreInst>(I)) {
  unsigned Size = I->getOperand(0)->getType()->getIntegerBitWidth();
  return (Size == 8 || Size == 16);
}

return false;

[gonglingqin]

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

1106

Thanks.

1107–1109

Thanks.

1107–1113

Thanks, I will change it.

[gonglingqin]

Address @SixWeining's comments.

[gonglingqin]

rebase

[SixWeining]

LGTM.

llvm/lib/Target/LoongArch/LoongArchISelLowering.h

private:
   SDValue lowerConstantPool(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerFP_TO_SINT(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerBITCAST(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerUINT_TO_FP(SDValue Op, SelectionDAG &DAG) const;
 
   bool isFPImmLegal(const APFloat &Imm, EVT VT,
                     bool ForCodeSize) const override;
 
-  bool shouldInsertFencesForAtomic(const Instruction *I) const override {
-    return isa<LoadInst>(I) || isa<StoreInst>(I);
-  }
+  bool shouldInsertFencesForAtomic(const Instruction *I) const override;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_LOONGARCH_LOONGARCHISELLOWERING_H

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

bool LoongArchTargetLowering::isFPImmLegal(const APFloat &Imm, EVT VT,
   assert((VT == MVT::f32 || VT == MVT::f64) && "Unexpected VT");
 
   if (VT == MVT::f32 && !Subtarget.hasBasicF())
     return false;
   if (VT == MVT::f64 && !Subtarget.hasBasicD())
     return false;
   return (Imm.isZero() || Imm.isExactlyValue(+1.0));
 }
+
+bool LoongArchTargetLowering::shouldInsertFencesForAtomic(
+    const Instruction *I) const {
+  if (!Subtarget.is64Bit()) {
+    return isa<LoadInst>(I) || isa<StoreInst>(I);
+  } else {
+    unsigned Size = 0;
+    if (isa<StoreInst>(I))
+      Size = I->getOperand(0)->getType()->getIntegerBitWidth();
+    // On LA64, atomic store operations with IntegerBitWidth of 32 and 64 do not
+    // require fences beacuse we can use amswap_db.[w/d].
+    return isa<LoadInst>(I) || (isa<StoreInst>(I) && (Size == 8 || Size == 16));
+  }
+}

llvm/lib/Target/LoongArch/LoongArchInstrInfo.td

 /// Atomic loads and stores
 
 def : Pat<(atomic_fence timm, timm), (DBAR 0)>;
 
 defm : LdPat<atomic_load_8, LD_B>;
 defm : LdPat<atomic_load_16, LD_H>;
 defm : LdPat<atomic_load_32, LD_W>;
 
+class release_seqcst_store<PatFrag base>
+    : PatFrag<(ops node:$ptr, node:$val), (base node:$ptr, node:$val), [{
+  AtomicOrdering Ordering = cast<AtomicSDNode>(N)->getSuccessOrdering();
+  return isReleaseOrStronger(Ordering);
+}]>;
+
+class unordered_monotonic_store<PatFrag base>
+    : PatFrag<(ops node:$ptr, node:$val), (base node:$ptr, node:$val), [{
+  AtomicOrdering Ordering = cast<AtomicSDNode>(N)->getSuccessOrdering();
+  return !isReleaseOrStronger(Ordering);
+}]>;
+
+def atomic_store_release_seqcst_32 : release_seqcst_store<atomic_store_32>;
+def atomic_store_release_seqcst_64 : release_seqcst_store<atomic_store_64>;
+def atomic_store_unordered_monotonic_32
+    : unordered_monotonic_store<atomic_store_32>;
+def atomic_store_unordered_monotonic_64
+    : unordered_monotonic_store<atomic_store_64>;
+
 defm : StPat<atomic_store_8, ST_B, GPR, GRLenVT>;
 defm : StPat<atomic_store_16, ST_H, GPR, GRLenVT>;
-defm : StPat<atomic_store_32, ST_W, GPR, i32>, Requires<[IsLA32]>;
+defm : StPat<atomic_store_unordered_monotonic_32, ST_W, GPR, i32>,
+             Requires<[IsLA32]>;
+
+def PseudoAtomicStoreW : Pseudo<(outs GPR:$dst), (ins GPR:$rj, GPR:$rk), []>,
+                                 PseudoInstExpansion<(AMSWAP_DB_W R0,
+                                                      GPR:$rk, GPR:$rj)>;
+
+def : Pat<(atomic_store_release_seqcst_32 GPR:$rj, GPR:$rk),
+          (PseudoAtomicStoreW GPR:$rj, GPR:$rk)>;
+
 let Predicates = [IsLA64] in {
+def PseudoAtomicStoreD : Pseudo<(outs GPR:$dst), (ins GPR:$rj, GPR:$rk), []>,
+                                 PseudoInstExpansion<(AMSWAP_DB_D R0,
+                                                      GPR:$rk, GPR:$rj)>;
+
+def : Pat<(atomic_store_release_seqcst_64 GPR:$rj, GPR:$rk),
+          (PseudoAtomicStoreD GPR:$rj, GPR:$rk)>;
+
 defm : LdPat<atomic_load_64, LD_D>;
-defm : StPat<atomic_store_32, ST_W, GPR, i64>;
-defm : StPat<atomic_store_64, ST_D, GPR, i64>;
+defm : StPat<atomic_store_unordered_monotonic_32, ST_W, GPR, i64>;
+defm : StPat<atomic_store_unordered_monotonic_64, ST_D, GPR, i64>;
 } // Predicates = [IsLA64]
 
 /// Other pseudo-instructions
 
 // Pessimistically assume the stack pointer will be clobbered
 let Defs = [R3], Uses = [R3] in {
 def ADJCALLSTACKDOWN : Pseudo<(outs), (ins i32imm:$amt1, i32imm:$amt2),
                               [(callseq_start timm:$amt1, timm:$amt2)]>;

llvm/test/CodeGen/LoongArch/ir-instruction/load-store-atomic.ll

 ; LA32-LABEL: store_release_i32:
 ; LA32:       # %bb.0:
 ; LA32-NEXT:    dbar 0
 ; LA32-NEXT:    st.w $a0, $a1, 0
 ; LA32-NEXT:    jirl $zero, $ra, 0
 ;
 ; LA64-LABEL: store_release_i32:
 ; LA64:       # %bb.0:
-; LA64-NEXT:    dbar 0
-; LA64-NEXT:    st.w $a0, $a1, 0
+; LA64-NEXT:    amswap_db.w $zero, $a1, $a0
 ; LA64-NEXT:    jirl $zero, $ra, 0
   store atomic i32 %v, ptr %ptr release, align 4
   ret void
 }
 
 define void @store_release_i64(ptr %ptr, i64 %v) {
 ; LA32-LABEL: store_release_i64:
 ; LA32:       # %bb.0:
 ; LA32-NEXT:    addi.w $sp, $sp, -16
 ; LA32-NEXT:    .cfi_def_cfa_offset 16
 ; LA32-NEXT:    st.w $ra, $sp, 12 # 4-byte Folded Spill
 ; LA32-NEXT:    .cfi_offset 1, -4
 ; LA32-NEXT:    ori $a3, $zero, 3
 ; LA32-NEXT:    bl __atomic_store_8
 ; LA32-NEXT:    ld.w $ra, $sp, 12 # 4-byte Folded Reload
 ; LA32-NEXT:    addi.w $sp, $sp, 16
 ; LA32-NEXT:    jirl $zero, $ra, 0
 ;
 ; LA64-LABEL: store_release_i64:
 ; LA64:       # %bb.0:
-; LA64-NEXT:    dbar 0
-; LA64-NEXT:    st.d $a0, $a1, 0
+; LA64-NEXT:    amswap_db.d $zero, $a1, $a0
 ; LA64-NEXT:    jirl $zero, $ra, 0
   store atomic i64 %v, ptr %ptr release, align 8
   ret void
 }
+
+define void @store_unordered_i8(ptr %ptr, i8 signext %v) {
+; LA32-LABEL: store_unordered_i8:
+; LA32:       # %bb.0:
+; LA32-NEXT:    st.b $a0, $a1, 0
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_unordered_i8:
+; LA64:       # %bb.0:
+; LA64-NEXT:    st.b $a0, $a1, 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i8 %v, ptr %ptr unordered, align 1
+  ret void
+}
+
+define void @store_unordered_i16(ptr %ptr, i16 signext %v) {
+; LA32-LABEL: store_unordered_i16:
+; LA32:       # %bb.0:
+; LA32-NEXT:    st.h $a0, $a1, 0
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_unordered_i16:
+; LA64:       # %bb.0:
+; LA64-NEXT:    st.h $a0, $a1, 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i16 %v, ptr %ptr unordered, align 2
+  ret void
+}
+
+define void @store_unordered_i32(ptr %ptr, i32 signext %v) {
+; LA32-LABEL: store_unordered_i32:
+; LA32:       # %bb.0:
+; LA32-NEXT:    st.w $a0, $a1, 0
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_unordered_i32:
+; LA64:       # %bb.0:
+; LA64-NEXT:    st.w $a0, $a1, 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i32 %v, ptr %ptr unordered, align 4
+  ret void
+}
+
+define void @store_unordered_i64(ptr %ptr, i64 %v) {
+; LA32-LABEL: store_unordered_i64:
+; LA32:       # %bb.0:
+; LA32-NEXT:    addi.w $sp, $sp, -16
+; LA32-NEXT:    .cfi_def_cfa_offset 16
+; LA32-NEXT:    st.w $ra, $sp, 12 # 4-byte Folded Spill
+; LA32-NEXT:    .cfi_offset 1, -4
+; LA32-NEXT:    move $a3, $zero
+; LA32-NEXT:    bl __atomic_store_8
+; LA32-NEXT:    ld.w $ra, $sp, 12 # 4-byte Folded Reload
+; LA32-NEXT:    addi.w $sp, $sp, 16
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_unordered_i64:
+; LA64:       # %bb.0:
+; LA64-NEXT:    st.d $a0, $a1, 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i64 %v, ptr %ptr unordered, align 8
+  ret void
+}
+
+define void @store_monotonic_i8(ptr %ptr, i8 signext %v) {
+; LA32-LABEL: store_monotonic_i8:
+; LA32:       # %bb.0:
+; LA32-NEXT:    st.b $a0, $a1, 0
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_monotonic_i8:
+; LA64:       # %bb.0:
+; LA64-NEXT:    st.b $a0, $a1, 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i8 %v, ptr %ptr monotonic, align 1
+  ret void
+}
+
+define void @store_monotonic_i16(ptr %ptr, i16 signext %v) {
+; LA32-LABEL: store_monotonic_i16:
+; LA32:       # %bb.0:
+; LA32-NEXT:    st.h $a0, $a1, 0
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_monotonic_i16:
+; LA64:       # %bb.0:
+; LA64-NEXT:    st.h $a0, $a1, 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i16 %v, ptr %ptr monotonic, align 2
+  ret void
+}
+
+define void @store_monotonic_i32(ptr %ptr, i32 signext %v) {
+; LA32-LABEL: store_monotonic_i32:
+; LA32:       # %bb.0:
+; LA32-NEXT:    st.w $a0, $a1, 0
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_monotonic_i32:
+; LA64:       # %bb.0:
+; LA64-NEXT:    st.w $a0, $a1, 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i32 %v, ptr %ptr monotonic, align 4
+  ret void
+}
+
+define void @store_monotonic_i64(ptr %ptr, i64 %v) {
+; LA32-LABEL: store_monotonic_i64:
+; LA32:       # %bb.0:
+; LA32-NEXT:    addi.w $sp, $sp, -16
+; LA32-NEXT:    .cfi_def_cfa_offset 16
+; LA32-NEXT:    st.w $ra, $sp, 12 # 4-byte Folded Spill
+; LA32-NEXT:    .cfi_offset 1, -4
+; LA32-NEXT:    move $a3, $zero
+; LA32-NEXT:    bl __atomic_store_8
+; LA32-NEXT:    ld.w $ra, $sp, 12 # 4-byte Folded Reload
+; LA32-NEXT:    addi.w $sp, $sp, 16
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_monotonic_i64:
+; LA64:       # %bb.0:
+; LA64-NEXT:    st.d $a0, $a1, 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i64 %v, ptr %ptr monotonic, align 8
+  ret void
+}
+
+define void @store_seq_cst_i8(ptr %ptr, i8 signext %v) {
+; LA32-LABEL: store_seq_cst_i8:
+; LA32:       # %bb.0:
+; LA32-NEXT:    dbar 0
+; LA32-NEXT:    st.b $a0, $a1, 0
+; LA32-NEXT:    dbar 0
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_seq_cst_i8:
+; LA64:       # %bb.0:
+; LA64-NEXT:    dbar 0
+; LA64-NEXT:    st.b $a0, $a1, 0
+; LA64-NEXT:    dbar 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i8 %v, ptr %ptr seq_cst, align 1
+  ret void
+}
+
+define void @store_seq_cst_i16(ptr %ptr, i16 signext %v) {
+; LA32-LABEL: store_seq_cst_i16:
+; LA32:       # %bb.0:
+; LA32-NEXT:    dbar 0
+; LA32-NEXT:    st.h $a0, $a1, 0
+; LA32-NEXT:    dbar 0
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_seq_cst_i16:
+; LA64:       # %bb.0:
+; LA64-NEXT:    dbar 0
+; LA64-NEXT:    st.h $a0, $a1, 0
+; LA64-NEXT:    dbar 0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i16 %v, ptr %ptr seq_cst, align 2
+  ret void
+}
+
+define void @store_seq_cst_i32(ptr %ptr, i32 signext %v) {
+; LA32-LABEL: store_seq_cst_i32:
+; LA32:       # %bb.0:
+; LA32-NEXT:    dbar 0
+; LA32-NEXT:    st.w $a0, $a1, 0
+; LA32-NEXT:    dbar 0
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_seq_cst_i32:
+; LA64:       # %bb.0:
+; LA64-NEXT:    amswap_db.w $zero, $a1, $a0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i32 %v, ptr %ptr seq_cst, align 4
+  ret void
+}
+
+define void @store_seq_cst_i64(ptr %ptr, i64 %v) {
+; LA32-LABEL: store_seq_cst_i64:
+; LA32:       # %bb.0:
+; LA32-NEXT:    addi.w $sp, $sp, -16
+; LA32-NEXT:    .cfi_def_cfa_offset 16
+; LA32-NEXT:    st.w $ra, $sp, 12 # 4-byte Folded Spill
+; LA32-NEXT:    .cfi_offset 1, -4
+; LA32-NEXT:    ori $a3, $zero, 5
+; LA32-NEXT:    bl __atomic_store_8
+; LA32-NEXT:    ld.w $ra, $sp, 12 # 4-byte Folded Reload
+; LA32-NEXT:    addi.w $sp, $sp, 16
+; LA32-NEXT:    jirl $zero, $ra, 0
+;
+; LA64-LABEL: store_seq_cst_i64:
+; LA64:       # %bb.0:
+; LA64-NEXT:    amswap_db.d $zero, $a1, $a0
+; LA64-NEXT:    jirl $zero, $ra, 0
+  store atomic i64 %v, ptr %ptr seq_cst, align 8
+  ret void
+}

llvm/test/Transforms/AtomicExpand/LoongArch/load-store-atomic.ll

 
 define void @store_release_i32(ptr %ptr, i32 signext %v) {
 ; LA32-LABEL: @store_release_i32(
 ; LA32-NEXT:    fence release
 ; LA32-NEXT:    store atomic i32 [[V:%.*]], ptr [[PTR:%.*]] monotonic, align 4
 ; LA32-NEXT:    ret void
 ;
 ; LA64-LABEL: @store_release_i32(
-; LA64-NEXT:    fence release
-; LA64-NEXT:    store atomic i32 [[V:%.*]], ptr [[PTR:%.*]] monotonic, align 4
+; LA64-NEXT:    store atomic i32 [[V:%.*]], ptr [[PTR:%.*]] release, align 4
 ; LA64-NEXT:    ret void
 ;
   store atomic i32 %v, ptr %ptr release, align 4
   ret void
 }
 
 define void @store_release_i64(ptr %ptr, i64 %v) {
 ; LA32-LABEL: @store_release_i64(
 ; LA32-NEXT:    call void @__atomic_store_8(ptr [[PTR:%.*]], i64 [[V:%.*]], i32 3)
 ; LA32-NEXT:    ret void
 ;
 ; LA64-LABEL: @store_release_i64(
-; LA64-NEXT:    fence release
-; LA64-NEXT:    store atomic i64 [[V:%.*]], ptr [[PTR:%.*]] monotonic, align 8
+; LA64-NEXT:    store atomic i64 [[V:%.*]], ptr [[PTR:%.*]] release, align 8
 ; LA64-NEXT:    ret void
 ;
   store atomic i64 %v, ptr %ptr release, align 8
   ret void
 }