Refactor to match memset (NFC)

AZero13 · AZero13 · commit ddcb8134a0b2 · 2024-10-03T12:18:03.000-04:00
Refactored the memset and memcpy codegen to share the alignment-determining code.
diff --git a/llvm/lib/Target/X86/X86SelectionDAGInfo.cpp b/llvm/lib/Target/X86/X86SelectionDAGInfo.cpp
@@ -28,6 +28,23 @@ static cl::opt<bool>
     UseFSRMForMemcpy("x86-use-fsrm-for-memcpy", cl::Hidden, cl::init(false),
                      cl::desc("Use fast short rep mov in memcpy lowering"));
 
+/// Returns the best type to use with repmovs/repstos depending on alignment.
+static MVT getOptimalRepType(const X86Subtarget &Subtarget, Align Alignment) {
+  uint64_t Align = Alignment.value();
+  assert((Align != 0) && "Align is normalized");
+  assert(isPowerOf2_64(Align) && "Align is a power of 2");
+  switch (Align) {
+  case 1:
+    return MVT::i8;
+  case 2:
+    return MVT::i16;
+  case 4:
+    return MVT::i32;
+  default:
+    return Subtarget.is64Bit() ? MVT::i64 : MVT::i32;
+  }
+}
+
 bool X86SelectionDAGInfo::isBaseRegConflictPossible(
     SelectionDAG &DAG, ArrayRef<MCPhysReg> ClobberSet) const {
   // We cannot use TRI->hasBasePointer() until *after* we select all basic
@@ -44,102 +61,99 @@ bool X86SelectionDAGInfo::isBaseRegConflictPossible(
   return llvm::is_contained(ClobberSet, TRI->getBaseRegister());
 }
 
-SDValue X86SelectionDAGInfo::EmitTargetCodeForMemset(
-    SelectionDAG &DAG, const SDLoc &dl, SDValue Chain, SDValue Dst, SDValue Val,
-    SDValue Size, Align Alignment, bool isVolatile, bool AlwaysInline,
-    MachinePointerInfo DstPtrInfo) const {
-  // If to a segment-relative address space, use the default lowering.
-  if (DstPtrInfo.getAddrSpace() >= 256)
-    return SDValue();
+/// Emit a single REP STOSB instruction for a particular constant size.
+static SDValue emitRepstos(const X86Subtarget &Subtarget, SelectionDAG &DAG,
+                           const SDLoc &dl, SDValue Chain, SDValue Dst,
+                           SDValue Val, SDValue Size, MVT AVT) {
+  const bool Use64BitRegs = Subtarget.isTarget64BitLP64();
+  unsigned AX = X86::AL;
+  switch (AVT.getSizeInBits()) {
+  case 8:
+    AX = X86::AL;
+    break;
+  case 16:
+    AX = X86::AX;
+    break;
+  case 32:
+    AX = X86::EAX;
+    break;
+  default:
+    AX = X86::RAX;
+    break;
+  }
 
-  // If the base register might conflict with our physical registers, bail out.
-  const MCPhysReg ClobberSet[] = {X86::RCX, X86::RAX, X86::RDI,
-                                  X86::ECX, X86::EAX, X86::EDI};
-  if (isBaseRegConflictPossible(DAG, ClobberSet))
-    return SDValue();
+  const unsigned CX = Use64BitRegs ? X86::RCX : X86::ECX;
+  const unsigned DI = Use64BitRegs ? X86::RDI : X86::EDI;
 
-  ConstantSDNode *ConstantSize = dyn_cast<ConstantSDNode>(Size);
-  const X86Subtarget &Subtarget =
-      DAG.getMachineFunction().getSubtarget<X86Subtarget>();
+  SDValue InGlue;
+  Chain = DAG.getCopyToReg(Chain, dl, AX, Val, InGlue);
+  InGlue = Chain.getValue(1);
+  Chain = DAG.getCopyToReg(Chain, dl, CX, Size, InGlue);
+  InGlue = Chain.getValue(1);
+  Chain = DAG.getCopyToReg(Chain, dl, DI, Dst, InGlue);
+  InGlue = Chain.getValue(1);
+
+  SDVTList Tys = DAG.getVTList(MVT::Other, MVT::Glue);
+  SDValue Ops[] = {Chain, DAG.getValueType(AVT), InGlue};
+  return DAG.getNode(X86ISD::REP_STOS, dl, Tys, Ops);
+}
+
+/// Returns a REP STOS instruction, possibly with a few load/stores to implement
+/// a constant size memory set. In some cases where we know REP MOVS is
+/// inefficient we return an empty SDValue so the calling code can either
+/// generate a store sequence or call the runtime memset function.
+static SDValue emitConstantSizeRepstos(SelectionDAG &DAG,
+                                       const X86Subtarget &Subtarget,
+                                       const SDLoc &dl, SDValue Chain,
+                                       SDValue Dst, SDValue Val, uint64_t Size,
+                                       EVT SizeVT, Align Alignment,
+                                       bool isVolatile, bool AlwaysInline,
+                                       MachinePointerInfo DstPtrInfo) {
+
+  if (Size > Subtarget.getMaxInlineSizeThreshold())
+    return SDValue();
 
   // If not DWORD aligned or size is more than the threshold, call the library.
   // The libc version is likely to be faster for these cases. It can use the
   // address value and run time information about the CPU.
-  if (Alignment < Align(4) || !ConstantSize ||
-      ConstantSize->getZExtValue() > Subtarget.getMaxInlineSizeThreshold())
+  if (Alignment < Align(4))
     return SDValue();
 
-  uint64_t SizeVal = ConstantSize->getZExtValue();
-  SDValue InGlue;
-  EVT AVT;
-  SDValue Count;
-  unsigned BytesLeft = 0;
+  MVT BlockType = MVT::i8;
+  uint64_t BlockCount = Size;
+  uint64_t BytesLeft = 0;
   if (auto *ValC = dyn_cast<ConstantSDNode>(Val)) {
-    unsigned ValReg;
-    uint64_t Val = ValC->getZExtValue() & 255;
-
-    // If the value is a constant, then we can potentially use larger sets.
-    if (Alignment >= Align(4)) {
-      // DWORD aligned
-      AVT = MVT::i32;
-      ValReg = X86::EAX;
-      Val = (Val << 8)  | Val;
-      Val = (Val << 16) | Val;
-      if (Subtarget.is64Bit() && Alignment >= Align(8)) { // QWORD aligned
-        AVT = MVT::i64;
-        ValReg = X86::RAX;
-        Val = (Val << 32) | Val;
-      }
-    } else if (Alignment == Align(2)) {
-      // WORD aligned
-      AVT = MVT::i16;
-      ValReg = X86::AX;
-      Val = (Val << 8) | Val;
-    } else {
-      // Byte aligned
-      AVT = MVT::i8;
-      ValReg = X86::AL;
-      Count = DAG.getIntPtrConstant(SizeVal, dl);
-    }
-
-    if (AVT.bitsGT(MVT::i8)) {
-      unsigned UBytes = AVT.getSizeInBits() / 8;
-      Count = DAG.getIntPtrConstant(SizeVal / UBytes, dl);
-      BytesLeft = SizeVal % UBytes;
-    }
-
-    Chain = DAG.getCopyToReg(Chain, dl, ValReg, DAG.getConstant(Val, dl, AVT),
-                             InGlue);
-    InGlue = Chain.getValue(1);
-  } else {
-    AVT = MVT::i8;
-    Count  = DAG.getIntPtrConstant(SizeVal, dl);
-    Chain  = DAG.getCopyToReg(Chain, dl, X86::AL, Val, InGlue);
-    InGlue = Chain.getValue(1);
-  }
+    BlockType = getOptimalRepType(Subtarget, Alignment);
+    uint64_t Value = ValC->getZExtValue() & 255;
+    const uint64_t BlockBits = BlockType.getSizeInBits();
 
-  bool Use64BitRegs = Subtarget.isTarget64BitLP64();
-  Chain = DAG.getCopyToReg(Chain, dl, Use64BitRegs ? X86::RCX : X86::ECX,
-                           Count, InGlue);
-  InGlue = Chain.getValue(1);
-  Chain = DAG.getCopyToReg(Chain, dl, Use64BitRegs ? X86::RDI : X86::EDI,
-                           Dst, InGlue);
-  InGlue = Chain.getValue(1);
+    if (BlockBits >= 16)
+      Value = (Value << 8) | Value;
 
-  SDVTList Tys = DAG.getVTList(MVT::Other, MVT::Glue);
-  SDValue Ops[] = {Chain, DAG.getValueType(AVT), InGlue};
-  SDValue RepStos = DAG.getNode(X86ISD::REP_STOS, dl, Tys, Ops);
+    if (BlockBits >= 32)
+      Value = (Value << 16) | Value;
+
+    if (BlockBits >= 64)
+      Value = (Value << 32) | Value;
+
+    const uint64_t BlockBytes = BlockBits / 8;
+    BlockCount = Size / BlockBytes;
+    BytesLeft = Size % BlockBytes;
+    Val = DAG.getConstant(Value, dl, BlockType);
+  }
 
+  SDValue RepStos =
+      emitRepstos(Subtarget, DAG, dl, Chain, Dst, Val,
+                  DAG.getIntPtrConstant(BlockCount, dl), BlockType);
   /// RepStos can process the whole length.
   if (BytesLeft == 0)
     return RepStos;
 
   // Handle the last 1 - 7 bytes.
   SmallVector<SDValue, 4> Results;
   Results.push_back(RepStos);
-  unsigned Offset = SizeVal - BytesLeft;
+  unsigned Offset = Size - BytesLeft;
   EVT AddrVT = Dst.getValueType();
-  EVT SizeVT = Size.getValueType();
 
   Results.push_back(
       DAG.getMemset(Chain, dl,
@@ -152,6 +166,31 @@ SDValue X86SelectionDAGInfo::EmitTargetCodeForMemset(
   return DAG.getNode(ISD::TokenFactor, dl, MVT::Other, Results);
 }
 
+SDValue X86SelectionDAGInfo::EmitTargetCodeForMemset(
+    SelectionDAG &DAG, const SDLoc &dl, SDValue Chain, SDValue Dst, SDValue Val,
+    SDValue Size, Align Alignment, bool isVolatile, bool AlwaysInline,
+    MachinePointerInfo DstPtrInfo) const {
+  // If to a segment-relative address space, use the default lowering.
+  if (DstPtrInfo.getAddrSpace() >= 256)
+    return SDValue();
+
+  // If the base register might conflict with our physical registers, bail out.
+  const MCPhysReg ClobberSet[] = {X86::RCX, X86::RAX, X86::RDI,
+                                  X86::ECX, X86::EAX, X86::EDI};
+  if (isBaseRegConflictPossible(DAG, ClobberSet))
+    return SDValue();
+
+  ConstantSDNode *ConstantSize = dyn_cast<ConstantSDNode>(Size);
+  if (!ConstantSize)
+    return SDValue();
+
+  const X86Subtarget &Subtarget =
+      DAG.getMachineFunction().getSubtarget<X86Subtarget>();
+  return emitConstantSizeRepstos(
+      DAG, Subtarget, dl, Chain, Dst, Val, ConstantSize->getZExtValue(),
+      Size.getValueType(), Alignment, isVolatile, AlwaysInline, DstPtrInfo);
+}
+
 /// Emit a single REP MOVS{B,W,D,Q} instruction.
 static SDValue emitRepmovs(const X86Subtarget &Subtarget, SelectionDAG &DAG,
                            const SDLoc &dl, SDValue Chain, SDValue Dst,
@@ -182,24 +221,6 @@ static SDValue emitRepmovsB(const X86Subtarget &Subtarget, SelectionDAG &DAG,
                      DAG.getIntPtrConstant(Size, dl), MVT::i8);
 }
 
-/// Returns the best type to use with repmovs depending on alignment.
-static MVT getOptimalRepmovsType(const X86Subtarget &Subtarget,
-                                 Align Alignment) {
-  uint64_t Align = Alignment.value();
-  assert((Align != 0) && "Align is normalized");
-  assert(isPowerOf2_64(Align) && "Align is a power of 2");
-  switch (Align) {
-  case 1:
-    return MVT::i8;
-  case 2:
-    return MVT::i16;
-  case 4:
-    return MVT::i32;
-  default:
-    return Subtarget.is64Bit() ? MVT::i64 : MVT::i32;
-  }
-}
-
 /// Returns a REP MOVS instruction, possibly with a few load/stores to implement
 /// a constant size memory copy. In some cases where we know REP MOVS is
 /// inefficient we return an empty SDValue so the calling code can either
@@ -222,10 +243,10 @@ static SDValue emitConstantSizeRepmov(
   assert(!Subtarget.hasERMSB() && "No efficient RepMovs");
   /// We assume runtime memcpy will do a better job for unaligned copies when
   /// ERMS is not present.
-  if (!AlwaysInline && (Alignment.value() & 3) != 0)
+  if (!AlwaysInline && (Alignment < Align(4)))
     return SDValue();
 
-  const MVT BlockType = getOptimalRepmovsType(Subtarget, Alignment);
+  const MVT BlockType = getOptimalRepType(Subtarget, Alignment);
   const uint64_t BlockBytes = BlockType.getSizeInBits() / 8;
   const uint64_t BlockCount = Size / BlockBytes;
   const uint64_t BytesLeft = Size % BlockBytes;
@@ -239,7 +260,7 @@ static SDValue emitConstantSizeRepmov(
 
   assert(BytesLeft && "We have leftover at this point");
 
-  /// In case we optimize for size we use repmovsb even if it's less efficient
+  /// In case we optimize for size, we use repmovsb even if it's less efficient
   /// so we can save the loads/stores of the leftover.
   if (DAG.getMachineFunction().getFunction().hasMinSize())
     return emitRepmovsB(Subtarget, DAG, dl, Chain, Dst, Src, Size);
@@ -282,7 +303,7 @@ SDValue X86SelectionDAGInfo::EmitTargetCodeForMemcpy(
   if (UseFSRMForMemcpy && Subtarget.hasFSRM())
     return emitRepmovs(Subtarget, DAG, dl, Chain, Dst, Src, Size, MVT::i8);
 
-  /// Handle constant sizes,
+  /// Handle constant sizes
   if (ConstantSDNode *ConstantSize = dyn_cast<ConstantSDNode>(Size))
     return emitConstantSizeRepmov(DAG, Subtarget, dl, Chain, Dst, Src,
                                   ConstantSize->getZExtValue(),
