[SOL] add support for (pseudo) atomics to SBF (#23)

alessandrod · LucasSte · commit 2d40886aa43e · 2024-08-19T11:53:36.000-03:00
Lower atomic operations to their regular non-atomic equivalents. Lowering for
all operations except atomic fence is done at DAG legalization time. Fences are
removed at instruction emission time.
diff --git a/llvm/lib/Target/BPF/BPFISelLowering.cpp b/llvm/lib/Target/BPF/BPFISelLowering.cpp
@@ -81,10 +81,30 @@ BPFTargetLowering::BPFTargetLowering(const TargetMachine &TM,
   setOperationAction(ISD::STACKSAVE, MVT::Other, Expand);
   setOperationAction(ISD::STACKRESTORE, MVT::Other, Expand);
 
-  // Set unsupported atomic operations as Custom so
-  // we can emit better error messages than fatal error
-  // from selectiondag.
-  for (auto VT : {MVT::i8, MVT::i16, MVT::i32}) {
+  for (auto VT : {MVT::i8, MVT::i16, MVT::i32, MVT::i32, MVT::i64}) {
+    if (Subtarget->isSolana()) {
+      // Implement custom lowering for all atomic operations
+      setOperationAction(ISD::ATOMIC_SWAP, VT, Custom);
+      setOperationAction(ISD::ATOMIC_CMP_SWAP_WITH_SUCCESS, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_ADD, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_AND, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_MAX, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_MIN, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_NAND, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_OR, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_SUB, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_UMAX, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_UMIN, VT, Custom);
+      setOperationAction(ISD::ATOMIC_LOAD_XOR, VT, Custom);
+      continue;
+    }
+
+    if (VT == MVT::i64) {
+      continue;
+    }
+
+    // Set unsupported atomic operations as Custom so we can emit better error
+    // messages than fatal error from selectiondag.
     if (VT == MVT::i32) {
       if (STI.getHasAlu32())
         continue;
@@ -222,7 +242,17 @@ bool BPFTargetLowering::allowsMisalignedMemoryAccesses(
   return isSolana;
 }
 
-bool BPFTargetLowering::isOffsetFoldingLegal(const GlobalAddressSDNode *GA) const {
+bool BPFTargetLowering::lowerAtomicStoreAsStoreSDNode(
+    const StoreInst &SI) const {
+  return Subtarget->isSolana();
+}
+
+bool BPFTargetLowering::lowerAtomicLoadAsLoadSDNode(const LoadInst &LI) const {
+  return Subtarget->isSolana();
+}
+
+bool BPFTargetLowering::isOffsetFoldingLegal(
+    const GlobalAddressSDNode *GA) const {
   return false;
 }
 
@@ -318,6 +348,21 @@ void BPFTargetLowering::ReplaceNodeResults(
   case ISD::ATOMIC_LOAD_XOR:
   case ISD::ATOMIC_SWAP:
   case ISD::ATOMIC_CMP_SWAP_WITH_SUCCESS:
+  case ISD::ATOMIC_LOAD_ADD:
+  case ISD::ATOMIC_LOAD_AND:
+  case ISD::ATOMIC_LOAD_MAX:
+  case ISD::ATOMIC_LOAD_MIN:
+  case ISD::ATOMIC_LOAD_NAND:
+  case ISD::ATOMIC_LOAD_OR:
+  case ISD::ATOMIC_LOAD_SUB:
+  case ISD::ATOMIC_LOAD_UMAX:
+  case ISD::ATOMIC_LOAD_UMIN:
+  case ISD::ATOMIC_LOAD_XOR:
+    if (Subtarget->isSolana()) {
+      // We do lowering during legalization, see LowerOperation()
+      return;
+    }
+
     if (HasAlu32 || Opcode == ISD::ATOMIC_LOAD_ADD)
       Msg = "unsupported atomic operation, please use 32/64 bit version";
     else
@@ -346,6 +391,23 @@ SDValue BPFTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
     return LowerSDIVSREM(Op, DAG);
   case ISD::DYNAMIC_STACKALLOC:
     return LowerDYNAMIC_STACKALLOC(Op, DAG);
+  case ISD::ATOMIC_SWAP:
+  case ISD::ATOMIC_CMP_SWAP_WITH_SUCCESS:
+  case ISD::ATOMIC_LOAD_ADD:
+  case ISD::ATOMIC_LOAD_AND:
+  case ISD::ATOMIC_LOAD_MAX:
+  case ISD::ATOMIC_LOAD_MIN:
+  case ISD::ATOMIC_LOAD_NAND:
+  case ISD::ATOMIC_LOAD_OR:
+  case ISD::ATOMIC_LOAD_SUB:
+  case ISD::ATOMIC_LOAD_UMAX:
+  case ISD::ATOMIC_LOAD_UMIN:
+  case ISD::ATOMIC_LOAD_XOR:
+    return LowerATOMICRMW(Op, DAG);
+  case ISD::DYNAMIC_STACKALLOC:
+    report_fatal_error("Unsupported dynamic stack allocation");
+  default:
+    llvm_unreachable("unimplemented atomic operand");
   }
 }
 
@@ -444,7 +506,6 @@ SDValue BPFTargetLowering::LowerFormalArguments(
     fail(DL, DAG, "functions with VarArgs or StructRet are not supported");
   }
 
-
   return Chain;
 }
 
@@ -788,6 +849,114 @@ SDValue BPFTargetLowering::LowerSELECT_CC(SDValue Op, SelectionDAG &DAG) const {
   return DAG.getNode(BPFISD::SELECT_CC, DL, VTs, Ops);
 }
 
+SDValue BPFTargetLowering::LowerATOMICRMW(SDValue Op, SelectionDAG &DAG) const {
+  SDLoc DL(Op);
+  AtomicSDNode *AN = cast<AtomicSDNode>(Op);
+  assert(AN && "Expected custom lowering of an atomic load node");
+
+  SDValue Chain = AN->getChain();
+  SDValue Ptr = AN->getBasePtr();
+  EVT PtrVT = AN->getMemoryVT();
+  EVT RetVT = Op.getValueType();
+
+  // Load the current value
+  SDValue Load =
+      DAG.getExtLoad(ISD::EXTLOAD, DL, RetVT, Chain, Ptr, MachinePointerInfo(),
+                     PtrVT, AN->getAlignment());
+  Chain = Load.getValue(1);
+
+  // Most ops return the current value, except CMP_SWAP_WITH_SUCCESS see below
+  SDValue Ret = Load;
+  SDValue RetFlag;
+
+  // Val contains the new value we want to set. For CMP_SWAP, Cmp contains the
+  // expected current value.
+  SDValue Cmp, Val;
+  if (AN->isCompareAndSwap()) {
+    Cmp = Op.getOperand(2);
+    Val = Op.getOperand(3);
+
+    // The Cmp value must match the pointer type
+    EVT CmpVT = Cmp->getValueType(0);
+    if (CmpVT != RetVT) {
+      Cmp = RetVT.bitsGT(CmpVT) ? DAG.getNode(ISD::SIGN_EXTEND, DL, RetVT, Cmp)
+                                : DAG.getNode(ISD::TRUNCATE, DL, RetVT, Cmp);
+    }
+  } else {
+    Val = AN->getVal();
+  }
+
+  // The new value type must match the pointer type
+  EVT ValVT = Val->getValueType(0);
+  if (ValVT != RetVT) {
+    Val = RetVT.bitsGT(ValVT) ? DAG.getNode(ISD::SIGN_EXTEND, DL, RetVT, Val)
+                              : DAG.getNode(ISD::TRUNCATE, DL, RetVT, Val);
+    ValVT = Val->getValueType(0);
+  }
+
+  SDValue NewVal;
+  switch (Op.getOpcode()) {
+  case ISD::ATOMIC_SWAP:
+    NewVal = Val;
+    break;
+  case ISD::ATOMIC_CMP_SWAP_WITH_SUCCESS: {
+    EVT RetFlagVT = AN->getValueType(1);
+    NewVal = DAG.getSelectCC(DL, Load, Cmp, Val, Load, ISD::SETEQ);
+    RetFlag = DAG.getSelectCC(
+        DL, Load, Cmp, DAG.getBoolConstant(true, DL, RetFlagVT, RetFlagVT),
+        DAG.getBoolConstant(false, DL, RetFlagVT, RetFlagVT), ISD::SETEQ);
+    break;
+  }
+  case ISD::ATOMIC_LOAD_ADD:
+    NewVal = DAG.getNode(ISD::ADD, DL, ValVT, Load, Val);
+    break;
+  case ISD::ATOMIC_LOAD_SUB:
+    NewVal = DAG.getNode(ISD::SUB, DL, ValVT, Load, Val);
+    break;
+  case ISD::ATOMIC_LOAD_AND:
+    NewVal = DAG.getNode(ISD::AND, DL, ValVT, Load, Val);
+    break;
+  case ISD::ATOMIC_LOAD_NAND: {
+    NewVal =
+        DAG.getNOT(DL, DAG.getNode(ISD::AND, DL, ValVT, Load, Val), ValVT);
+    break;
+  }
+  case ISD::ATOMIC_LOAD_OR:
+    NewVal = DAG.getNode(ISD::OR, DL, ValVT, Load, Val);
+    break;
+  case ISD::ATOMIC_LOAD_XOR:
+    NewVal = DAG.getNode(ISD::XOR, DL, ValVT, Load, Val);
+    break;
+  case ISD::ATOMIC_LOAD_MIN:
+    NewVal = DAG.getNode(ISD::SMIN, DL, ValVT, Load, Val);
+    break;
+  case ISD::ATOMIC_LOAD_UMIN:
+    NewVal = DAG.getNode(ISD::UMIN, DL, ValVT, Load, Val);
+    break;
+  case ISD::ATOMIC_LOAD_MAX:
+    NewVal = DAG.getNode(ISD::SMAX, DL, ValVT, Load, Val);
+    break;
+  case ISD::ATOMIC_LOAD_UMAX:
+    NewVal = DAG.getNode(ISD::UMAX, DL, ValVT, Load, Val);
+    break;
+  default:
+    llvm_unreachable("unknown atomicrmw op");
+  }
+
+  Chain =
+      DAG.getTruncStore(Chain, DL, NewVal, Ptr, MachinePointerInfo(), PtrVT);
+
+  if (RetFlag) {
+    // CMP_SWAP_WITH_SUCCESS returns {value, success, chain}
+    Ret = DAG.getMergeValues({Ret, RetFlag, Chain}, DL);
+  } else {
+    // All the other ops return {value, chain}
+    Ret = DAG.getMergeValues({Ret, Chain}, DL);
+  }
+
+  return Ret;
+}
+
 const char *BPFTargetLowering::getTargetNodeName(unsigned Opcode) const {
   switch ((BPFISD::NodeType)Opcode) {
   case BPFISD::FIRST_NUMBER:
@@ -897,6 +1066,7 @@ BPFTargetLowering::EmitInstrWithCustomInserter(MachineInstr &MI,
                        Opc == BPF::Select_32_64);
 
   bool isMemcpyOp = Opc == BPF::MEMCPY;
+  bool isAtomicFence = Opc == BPF::ATOMIC_FENCE;
 
 #ifndef NDEBUG
   bool isSelectRIOp = (Opc == BPF::Select_Ri ||
@@ -911,6 +1081,12 @@ BPFTargetLowering::EmitInstrWithCustomInserter(MachineInstr &MI,
   if (isMemcpyOp)
     return EmitInstrWithCustomInserterMemcpy(MI, BB);
 
+  if (isAtomicFence) {
+    // this is currently a nop
+    MI.eraseFromParent();
+    return BB;
+  }
+
   bool is32BitCmp = (Opc == BPF::Select_32 ||
                      Opc == BPF::Select_32_64 ||
                      Opc == BPF::Select_Ri_32 ||
diff --git a/llvm/lib/Target/BPF/BPFISelLowering.h b/llvm/lib/Target/BPF/BPFISelLowering.h
@@ -71,6 +71,9 @@ class BPFTargetLowering : public TargetLowering {
 
   MVT getScalarShiftAmountTy(const DataLayout &, EVT) const override;
 
+  bool lowerAtomicStoreAsStoreSDNode(const StoreInst &SI) const override;
+  bool lowerAtomicLoadAsLoadSDNode(const LoadInst &LI) const override;
+
 private:
   // Control Instruction Selection Features
   bool HasAlu32;
@@ -83,6 +86,7 @@ class BPFTargetLowering : public TargetLowering {
   SDValue LowerBR_CC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSELECT_CC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerGlobalAddress(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerATOMICRMW(SDValue Op, SelectionDAG &DAG) const;
 
   // Lower the result values of a call, copying them out of physregs into vregs
   SDValue LowerCallResult(SDValue Chain, SDValue InGlue,
diff --git a/llvm/lib/Target/BPF/BPFInstrInfo.td b/llvm/lib/Target/BPF/BPFInstrInfo.td
@@ -68,6 +68,7 @@ class ImmediateAsmOperand<string name> : AsmOperandClass {
 }
 
 def SImm16AsmOperand : ImmediateAsmOperand<"SImm16">;
+def BPFSubtargetSolana : Predicate<"Subtarget->isSolana()">;
 
 def brtarget : Operand<OtherVT> {
   let PrintMethod = "printBrTargetOperand";
@@ -873,6 +874,14 @@ def : Pat<(atomic_load_sub_32 ADDRri:$addr, GPR32:$val),
 def : Pat<(atomic_load_sub_64 ADDRri:$addr, GPR:$val),
           (XFADDD ADDRri:$addr, (NEG_64 GPR:$val))>;
 
+let Predicates = [BPFSubtargetSolana], usesCustomInserter = 1, isCodeGenOnly = 1 in {
+  def ATOMIC_FENCE : Pseudo<
+    (outs),
+    (ins),
+    "#atomic_fence",
+    [(atomic_fence timm, timm)]>;
+}
+
 // Atomic Exchange
 class XCHG<BPFWidthModifer SizeOp, string OpcodeStr, PatFrag OpNode>
     : TYPE_LD_ST<BPF_ATOMIC.Value, SizeOp.Value,
diff --git a/llvm/test/CodeGen/BPF/atomics_sbf.ll b/llvm/test/CodeGen/BPF/atomics_sbf.ll
