Optimize load pairs

Summary: Adds a post-alloc pass that re-writes two adjacent loads into a single load pair instruction. This is not part of the main post alloc pass because we only need to run it once everything else is done and we're not going to benefit from running it multiple times to hit a fix point.

Reviewed By: alexmalyshev

Differential Revision: D114958913

fbshipit-source-id: 72f896af39153480078a6605ba4a29641e555a00

Author: DinoV

cinderx/Jit/codegen/gen_asm.cpp

@@ -882,6 +882,15 @@ void* NativeGenerator::getVectorcallEntry() {
       "Post Reg Alloc Rewrite",
       post_rewrite.run())
 
+#if defined(CINDER_AARCH64)
+  // Peepholes go last, once nothing else will add to or remove from the
+  // instruction stream.
+  COMPILE_TIMER(
+      getFunction()->compilation_phase_timer,
+      "Post Reg Alloc Peephole",
+      runPostRegAllocPeephole(lir_func.get()))
+#endif
+
   JIT_LOGIF(
       getConfig().log.dump_lir,
       "LIR for {} after postalloc rewrites:\n{}",

cinderx/Jit/lir/postalloc.cpp

@@ -124,6 +124,151 @@ void insertStorePairToMemoryLocation(
 }
 #endif
 
+#if defined(CINDER_AARCH64)
+// A 64-bit general-purpose move between a register and memory.
+struct PairCandidate {
+  bool is_load;
+  // Frame slots report the frame pointer; their offset is the slot location.
+  PhyLocation base;
+  int32_t offset;
+  PhyLocation reg;
+
+  bool isAdjacentLoad(const PairCandidate& other) const {
+    return is_load == other.is_load && base == other.base &&
+        std::abs(
+            static_cast<int64_t>(offset) -
+            static_cast<int64_t>(other.offset)) == kPointerSize;
+  }
+};
+
+std::optional<std::pair<PhyLocation, int32_t>> pairMemoryLocation(
+    const Operand* operand) {
+  if (operand->isStack()) {
+    return std::make_pair(
+        codegen::arch::reg_frame_pointer_loc, operand->getStackSlot().loc);
+  }
+  if (operand->isInd()) {
+    MemoryIndirect* ind = operand->getMemoryIndirect();
+    // An index register would need scaling that ldp/stp can't express.
+    if (ind->getIndexRegOperand() != nullptr) {
+      return std::nullopt;
+    }
+    const Operand* base = ind->getBaseRegOperand();
+    if (base == nullptr || !base->isReg()) {
+      return std::nullopt;
+    }
+    return std::make_pair(base->getPhyRegister(), ind->getOffset());
+  }
+  return std::nullopt;
+}
+
+std::optional<PairCandidate> describePairCandidate(const Instruction* instr) {
+  if (!instr->isMove() || instr->getNumInputs() != 1 ||
+      instr->getNumOutputs() != 1) {
+    return std::nullopt;
+  }
+
+  const Operand* out = instr->output();
+  const Operand* in = instr->getInput(0);
+
+  // ldp/stp only come in general-purpose and floating-point flavours that
+  // can't be mixed, and only the 64-bit width is worth matching here. The
+  // access width is decided by the memory operand for a store and by the
+  // destination register for a load; requiring both to be 64-bit covers it.
+  if (out->isFp() || in->isFp() || out->sizeInBits() != 64 ||
+      in->sizeInBits() != 64) {
+    return std::nullopt;
+  }
+
+  if (out->isReg()) {
+    auto loc = pairMemoryLocation(in);
+    if (!loc.has_value()) {
+      return std::nullopt;
+    }
+    return PairCandidate{true, loc->first, loc->second, out->getPhyRegister()};
+  }
+
+  if (in->isReg()) {
+    auto loc = pairMemoryLocation(out);
+    if (!loc.has_value()) {
+      return std::nullopt;
+    }
+    return PairCandidate{false, loc->first, loc->second, in->getPhyRegister()};
+  }
+
+  return std::nullopt;
+}
+
+// Merge adjacent 64-bit loads or stores of neighbouring memory into ldp/stp.
+//
+// Spill traffic is the main source of these: the register allocator emits one
+// move per slot, and consecutive slots are a pointer apart. Argument loading
+// off the vectorcall array has the same shape.
+//
+// Only directly adjacent instructions are considered, which keeps this honest
+// about ordering without needing an aliasing check — nothing runs in between,
+// so the only reordering is between the two accesses themselves, and they
+// cover disjoint memory.
+void pairAdjacentMemoryOps(BasicBlock* block) {
+  auto& instrs = block->instructions();
+
+  for (auto it = instrs.begin(); it != instrs.end();) {
+    auto second = std::next(it);
+    if (second == instrs.end()) {
+      break;
+    }
+
+    auto first_desc = describePairCandidate(it->get());
+    auto second_desc = describePairCandidate(second->get());
+    if (!first_desc.has_value() || !second_desc.has_value() ||
+        !first_desc->isAdjacentLoad(*second_desc)) {
+      ++it;
+      continue;
+    }
+
+    // ldp/stp always take the lower address first.
+    bool in_order = first_desc->offset < second_desc->offset;
+    const PairCandidate& low = in_order ? *first_desc : *second_desc;
+    const PairCandidate& high = in_order ? *second_desc : *first_desc;
+
+    if (low.is_load) {
+      // ldp with a repeated destination, or with a destination that is also
+      // the base, is architecturally unpredictable. The unmerged pair would
+      // also have fed the first load's result into the second's address.
+      if (low.reg == high.reg || low.reg == low.base || high.reg == low.base) {
+        ++it;
+        continue;
+      }
+    }
+
+    auto offset = static_cast<uint64_t>(static_cast<int64_t>(low.offset));
+    // describePairCandidate filtered out floating point and non-64 bit regs
+    if (low.is_load) {
+      block->allocateInstrBefore(
+          it,
+          Instruction::kLoadPair,
+          OutPhyReg{low.reg, DataType::k64bit},
+          Imm{offset},
+          PhyReg{low.base, DataType::k64bit},
+          PhyReg{high.reg, DataType::k64bit});
+    } else {
+      block->allocateInstrBefore(
+          it,
+          Instruction::kStorePair,
+          Imm{offset},
+          PhyReg{low.base, DataType::k64bit},
+          PhyReg{low.reg, DataType::k64bit},
+          PhyReg{high.reg, DataType::k64bit});
+    }
+
+    auto next = std::next(second);
+    block->removeInstr(it);
+    block->removeInstr(second);
+    it = next;
+  }
+}
+#endif
+
 int rewriteRegularFunction(instr_iter_t instr_iter, int base_offset) {
   auto instr = instr_iter->get();
   auto block = instr->basicBlock();
@@ -1628,4 +1773,12 @@ void PostRegAllocRewrite::registerRewrites() {
 #endif
 }
 
+#if defined(CINDER_AARCH64)
+void runPostRegAllocPeephole(Function* func) {
+  for (auto& block : func->basicBlocks()) {
+    pairAdjacentMemoryOps(block);
+  }
+}
+#endif
+
 } // namespace cinderx::jit::lir

cinderx/Jit/lir/postalloc.h

@@ -18,4 +18,11 @@ class PostRegAllocRewrite : public Rewrite {
   void registerRewrites();
 };
 
+#if defined(CINDER_AARCH64)
+
+// Peephole rewrites run once the instruction stream is otherwise final.
+void runPostRegAllocPeephole(Function* func);
+
+#endif
+
 } // namespace cinderx::jit::lir

cinderx/RuntimeTests/lir_postalloc_test.cpp

@@ -445,6 +445,95 @@ BB %0
       store_pair->getInput(1)->getPhyRegister(), arch::reg_stack_pointer_loc);
   ASSERT_TRUE(verifyPostRegAllocInvariants(parsed_func.get(), std::cout));
 }
+TEST_F(LIRPostAllocRewriteTest, AdjacentFrameSlotStoresBecomeStorePair) {
+  Function func;
+  auto* bb = func.allocateBasicBlock();
+
+  bb->allocateInstr(
+      Instruction::kMove,
+      nullptr,
+      OutStk{PhyLocation(-24, 64), DataType::k64bit},
+      PhyReg{X0, DataType::k64bit});
+  bb->allocateInstr(
+      Instruction::kMove,
+      nullptr,
+      OutStk{PhyLocation(-16, 64), DataType::k64bit},
+      PhyReg{X1, DataType::k64bit});
+
+  jit::codegen::Environ env;
+  PostRegAllocRewrite rewrite(&func, &env);
+  rewrite.run();
+  runPostRegAllocPeephole(&func);
+
+  auto instrs = collectInstrs(*bb);
+  ASSERT_EQ(instrs.size(), 1);
+  ASSERT_TRUE(instrs[0]->isStorePair());
+  // The lower address comes first, so the slot at -24 supplies the offset.
+  EXPECT_EQ(static_cast<int32_t>(instrs[0]->getInput(0)->getConstant()), -24);
+  EXPECT_EQ(
+      instrs[0]->getInput(1)->getPhyRegister(), arch::reg_frame_pointer_loc);
+  EXPECT_EQ(instrs[0]->getInput(2)->getPhyRegister(), X0);
+  EXPECT_EQ(instrs[0]->getInput(3)->getPhyRegister(), X1);
+}
+
+TEST_F(LIRPostAllocRewriteTest, DescendingFrameSlotLoadsBecomeLoadPair) {
+  Function func;
+  auto* bb = func.allocateBasicBlock();
+
+  // Written high address first, so the pair has to swap the register order.
+  bb->allocateInstr(
+      Instruction::kMove,
+      nullptr,
+      OutPhyReg{X0, DataType::k64bit},
+      Stk{PhyLocation(-16, 64), DataType::k64bit});
+  bb->allocateInstr(
+      Instruction::kMove,
+      nullptr,
+      OutPhyReg{X1, DataType::k64bit},
+      Stk{PhyLocation(-24, 64), DataType::k64bit});
+
+  jit::codegen::Environ env;
+  PostRegAllocRewrite rewrite(&func, &env);
+  rewrite.run();
+  runPostRegAllocPeephole(&func);
+
+  auto instrs = collectInstrs(*bb);
+  ASSERT_EQ(instrs.size(), 1);
+  ASSERT_TRUE(instrs[0]->isLoadPair());
+  EXPECT_EQ(static_cast<int32_t>(instrs[0]->getInput(0)->getConstant()), -24);
+  EXPECT_EQ(instrs[0]->output()->getPhyRegister(), X1);
+  EXPECT_EQ(instrs[0]->getInput(2)->getPhyRegister(), X0);
+}
+
+// ldp with a destination that is also its base is unpredictable, and the
+// unmerged form would have fed the first load's result into the second's
+// address, so this pair has to be left alone.
+TEST_F(LIRPostAllocRewriteTest, LoadPairSkippedWhenDestinationIsBase) {
+  Function func;
+  auto* bb = func.allocateBasicBlock();
+
+  bb->allocateInstr(
+      Instruction::kMove,
+      nullptr,
+      OutPhyReg{X2, DataType::k64bit},
+      Ind(X2, static_cast<int32_t>(0)));
+  bb->allocateInstr(
+      Instruction::kMove,
+      nullptr,
+      OutPhyReg{X3, DataType::k64bit},
+      Ind(X2, static_cast<int32_t>(8)));
+
+  jit::codegen::Environ env;
+  PostRegAllocRewrite rewrite(&func, &env);
+  rewrite.run();
+  runPostRegAllocPeephole(&func);
+
+  auto instrs = collectInstrs(*bb);
+  ASSERT_EQ(instrs.size(), 2);
+  EXPECT_FALSE(instrs[0]->isLoadPair());
+  EXPECT_FALSE(instrs[1]->isLoadPair());
+}
+
 #endif // CINDER_AARCH64
 
 } // namespace cinderx::jit::lir