[BOLT][AArch64] Enabling Inlining for Memcpy for AArch64 in BOLT

bolt/docs/CommandLineArgumentReference.md

@@ -631,7 +631,7 @@
 
 - `--inline-memcpy`
 
- Inline memcpy using 'rep movsb' instruction (X86-only)
+ Inline memcpy using optimized instruction sequences (X86: 'rep movsb', AArch64: width-optimized register operations)
 
 - `--inline-small-functions`
 

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -14,6 +14,7 @@
 #ifndef BOLT_CORE_MCPLUSBUILDER_H
 #define BOLT_CORE_MCPLUSBUILDER_H
 
+#include "bolt/Core/BinaryBasicBlock.h"
 #include "bolt/Core/MCPlus.h"
 #include "bolt/Core/Relocation.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -1888,13 +1889,38 @@ class MCPlusBuilder {
  return {};
  }
 
+ /// Find memcpy size in bytes by using preceding instructions.
+ /// Returns std::nullopt if size cannot be determined (no-op for most
+ /// targets).
+ virtual std::optional<uint64_t>
+ findMemcpySizeInBytes(const BinaryBasicBlock &BB,
+ BinaryBasicBlock::iterator CallInst) const {
+ return std::nullopt;
+ }
+
  /// Creates inline memcpy instruction. If \p ReturnEnd is true, then return
  /// (dest + n) instead of dest.
  virtual InstructionListType createInlineMemcpy(bool ReturnEnd) const {
  llvm_unreachable("not implemented");
  return {};
  }
 
+ /// Creates size-aware inline memcpy instruction. If \p KnownSize is provided,
+ /// generates optimized code for that specific size. Falls back to regular
+ /// createInlineMemcpy if size is unknown or not needed (e.g. with X86).
+ virtual InstructionListType
+ createInlineMemcpy(bool ReturnEnd, std::optional<uint64_t> KnownSize) const {
+ return createInlineMemcpy(ReturnEnd);
+ }
+
+ /// Extract immediate value from move instruction that sets the given
+ /// register. Returns the immediate value if the instruction is a
+ /// move-immediate to TargetReg.
+ virtual std::optional<uint64_t>
+ extractMoveImmediate(const MCInst &Inst, MCPhysReg TargetReg) const {
+ return std::nullopt;
+ }
+
  /// Create a target-specific relocation out of the \p Fixup.
  /// Note that not every fixup could be converted into a relocation.
  virtual std::optional<Relocation>

bolt/lib/Passes/BinaryPasses.cpp

@@ -1843,7 +1843,7 @@ Error StripRepRet::runOnFunctions(BinaryContext &BC) {
 }
 
 Error InlineMemcpy::runOnFunctions(BinaryContext &BC) {
- if (!BC.isX86())
+ if (!BC.isX86() && !BC.isAArch64())
  return Error::success();
 
  uint64_t NumInlined = 0;
@@ -1866,8 +1866,16 @@ Error InlineMemcpy::runOnFunctions(BinaryContext &BC) {
  const bool IsMemcpy8 = (CalleeSymbol->getName() == "_memcpy8");
  const bool IsTailCall = BC.MIB->isTailCall(Inst);
 
+ // Extract size from preceding instructions (AArch64 only).
+ // Pattern: MOV X2, #nb-bytes; BL memcpy src, dest, X2.
+ std::optional<uint64_t> KnownSize =
+ BC.MIB->findMemcpySizeInBytes(BB, II);
+
+ if (BC.isAArch64() && (!KnownSize.has_value() || *KnownSize > 64))
+ continue;
+
  const InstructionListType NewCode =
- BC.MIB->createInlineMemcpy(IsMemcpy8);
+ BC.MIB->createInlineMemcpy(IsMemcpy8, KnownSize);
  II = BB.replaceInstruction(II, NewCode);
  std::advance(II, NewCode.size() - 1);
  if (IsTailCall) {

bolt/lib/Rewrite/BinaryPassManager.cpp

@@ -247,7 +247,9 @@ static cl::opt<bool> Stoke("stoke", cl::desc("turn on the stoke analysis"),
 
 static cl::opt<bool> StringOps(
  "inline-memcpy",
- cl::desc("inline memcpy using 'rep movsb' instruction (X86-only)"),
+ cl::desc(
+ "inline memcpy using size-specific optimized instructions "
+ "(X86: 'rep movsb', AArch64: width-optimized register operations)"),
  cl::cat(BoltOptCategory));
 
 static cl::opt<bool> StripRepRet(

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -2597,6 +2597,120 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
  getInstructionSize(const MCInst &Inst) const override {
  return 4;
  }
+
+ std::optional<uint64_t>
+ extractMoveImmediate(const MCInst &Inst, MCPhysReg TargetReg) const override {
+ // Match MOVZ instructions (both X and W register variants) with no shift.
+ if ((Inst.getOpcode() == AArch64::MOVZXi ||
+ Inst.getOpcode() == AArch64::MOVZWi) &&
+ Inst.getOperand(2).getImm() == 0 &&
+ getAliases(TargetReg)[Inst.getOperand(0).getReg()])
+ return Inst.getOperand(1).getImm();
+ return std::nullopt;
+ }
+
+ std::optional<uint64_t>
+ findMemcpySizeInBytes(const BinaryBasicBlock &BB,
+ BinaryBasicBlock::iterator CallInst) const override {
+ MCPhysReg SizeReg = getIntArgRegister(2);
+ if (SizeReg == getNoRegister())
+ return std::nullopt;
+
+ BitVector WrittenRegs(RegInfo->getNumRegs());
+ const BitVector &SizeRegAliases = getAliases(SizeReg);
+
+ for (auto InstIt = BB.begin(); InstIt != CallInst; ++InstIt) {
+ const MCInst &Inst = *InstIt;
+ WrittenRegs.reset();
+ getWrittenRegs(Inst, WrittenRegs);
+
+ if (WrittenRegs.anyCommon(SizeRegAliases))
+ return extractMoveImmediate(Inst, SizeReg);
+ }
+ return std::nullopt;
+ }
+
+ InstructionListType
+ createInlineMemcpy(bool ReturnEnd,
+ std::optional<uint64_t> KnownSize) const override {
+ assert(KnownSize.has_value() &&
+ "AArch64 memcpy inlining requires known size");
+ InstructionListType Code;
+ uint64_t Size = *KnownSize;
+
+ generateSizeSpecificMemcpy(Code, Size);
+
+ // If _memcpy8, adjust X0 to return dest+size instead of dest.
+ if (ReturnEnd)
+ Code.emplace_back(MCInstBuilder(AArch64::ADDXri)
+ .addReg(AArch64::X0)
+ .addReg(AArch64::X0)
+ .addImm(Size)
+ .addImm(0));
+ return Code;
+ }
+
+ InstructionListType generateSizeSpecificMemcpy(InstructionListType &Code,
+ uint64_t Size) const {
+ auto AddLoadStorePair = [&](unsigned LoadOpc, unsigned StoreOpc,
+ unsigned Reg, unsigned Offset = 0) {
+ Code.emplace_back(MCInstBuilder(LoadOpc)
+ .addReg(Reg)
+ .addReg(AArch64::X1)
+ .addImm(Offset));
+ Code.emplace_back(MCInstBuilder(StoreOpc)
+ .addReg(Reg)
+ .addReg(AArch64::X0)
+ .addImm(Offset));
+ };
+
+ // Generate optimal instruction sequences based on exact size.
+ switch (Size) {
+ case 1:
+ AddLoadStorePair(AArch64::LDRBBui, AArch64::STRBBui, AArch64::W9);
+ break;
+ case 2:
+ AddLoadStorePair(AArch64::LDRHHui, AArch64::STRHHui, AArch64::W9);
+ break;
+ case 4:
+ AddLoadStorePair(AArch64::LDRWui, AArch64::STRWui, AArch64::W9);
+ break;
+ case 8:
+ AddLoadStorePair(AArch64::LDRXui, AArch64::STRXui, AArch64::X9);
+ break;
+ case 16:
+ AddLoadStorePair(AArch64::LDRQui, AArch64::STRQui, AArch64::Q16);
+ break;
+ case 32:
+ AddLoadStorePair(AArch64::LDRQui, AArch64::STRQui, AArch64::Q16, 0);
+ AddLoadStorePair(AArch64::LDRQui, AArch64::STRQui, AArch64::Q17, 1);
+ break;
+
+ default:
+ if (Size <= 64) {
+ // For sizes up to 64 bytes, greedily use the largest possible loads.
+ uint64_t Remaining = Size;
+ uint64_t Offset = 0;
+
+ const std::array<std::tuple<uint64_t, unsigned, unsigned, unsigned>, 5>
+ LoadStoreOps = {
+ {{16, AArch64::LDRQui, AArch64::STRQui, AArch64::Q16},
+ {8, AArch64::LDRXui, AArch64::STRXui, AArch64::X9},
+ {4, AArch64::LDRWui, AArch64::STRWui, AArch64::W9},
+ {2, AArch64::LDRHHui, AArch64::STRHHui, AArch64::W9},
+ {1, AArch64::LDRBBui, AArch64::STRBBui, AArch64::W9}}};
+
+ for (const auto &[OpSize, LoadOp, StoreOp, TempReg] : LoadStoreOps)
+ while (Remaining >= OpSize) {
+ AddLoadStorePair(LoadOp, StoreOp, TempReg, Offset / OpSize);
+ Remaining -= OpSize;
+ Offset += OpSize;
+ }
+ }
+ break;
+ }
+ return Code;
+ }
 };
 
 } // end anonymous namespace