Use StorePair for AArch64 regular call arg packing

Summary: Extend the AArch64 post-register-allocation call rewrite to use `StorePair` for adjacent 64-bit general-purpose register arguments that overflow onto the regular call stack argument area. This keeps the existing single-store path for FP, immediate, stack, or otherwise unsupported operands.

Reviewed By: DinoV

Differential Revision: D109037584

fbshipit-source-id: cf8f0e044442d0db0ad40dfb3f1b8537ff9f01e6

Author: kddnewton

cinderx/Jit/codegen/autogen.cpp

@@ -1405,7 +1405,8 @@ void translateStorePair(Environ* env, const Instruction* instr) {
       x86::qword_ptr(base, offset + kPointerSize),
       x86::gpq(instr->getInput(3)->getPhyRegister().loc));
 #elif defined(CINDER_AARCH64)
-  auto base = a64::x(instr->getInput(1)->getPhyRegister().loc);
+  auto base_reg = instr->getInput(1)->getPhyRegister();
+  auto base = base_reg == SP ? a64::sp : a64::x(base_reg.loc);
   // stp signed offset range is -512..504. Fall back to two str instructions
   // when the offset is out of range.
   if (Support::isInt7(offset >> 3)) {

cinderx/Jit/lir/postalloc.cpp

@@ -220,6 +220,24 @@ int rewriteRegularFunction(instr_iter_t instr_iter, int base_offset) {
       move->output()->setDataType(operand->dataType());
       move->appendInput(instr->releaseInput(i));
     } else {
+#if defined(CINDER_AARCH64)
+      if (i + 1 < num_inputs) {
+        auto next_arg = instr->getInput(i + 1);
+        if (canStorePairOperand(operand) && canStorePairOperand(next_arg)) {
+          insertStorePairToMemoryLocation(
+              block,
+              instr_iter,
+              arch::reg_stack_pointer_loc,
+              base_offset + stack_arg_size,
+              operand,
+              next_arg);
+          ++i;
+          stack_arg_size += 2 * sizeof(void*);
+          continue;
+        }
+      }
+#endif
+
       insertMoveToMemoryLocation(
           block,
           instr_iter,

cinderx/RuntimeTests/lir_abi_test.cpp

@@ -216,6 +216,42 @@ TEST_F(LIRABITest, TestkCall_PhyReg) {
   translateInstr(Instruction::kCall, makePhyReg());
 }
 
+#if defined(CINDER_AARCH64)
+TEST_F(LIRABITest, TestkStorePair_SPBase) {
+  hir::Function hir_function;
+
+  Environ environ;
+  environ.ctx = getContext();
+  environ.code_rt = environ.ctx->allocateCodeRuntime(
+      hir_function.code.get(),
+      hir_function.builtins.get(),
+      hir_function.globals.get());
+
+  auto code_allocator = std::unique_ptr<ICodeAllocator>(CodeAllocator::make());
+
+  CodeHolder code;
+  code.init(code_allocator->asmJitEnvironment());
+
+  arch::Builder as(&code);
+  environ.as = &as;
+
+  Function function;
+  BasicBlock bb(&function);
+  auto* instr = bb.allocateInstr(
+      Instruction::kStorePair,
+      nullptr,
+      Imm{24},
+      PhyReg{arch::reg_stack_pointer_loc, DataType::k64bit},
+      PhyReg{X25, DataType::k64bit},
+      PhyReg{X20, DataType::k64bit});
+
+  autogen::AutoTranslator::getInstance().translateInstr(&environ, instr);
+
+  EXPECT_EQ(as.finalize(), asmjit::kErrorOk);
+  EXPECT_EQ(code.textSection()->bufferSize(), 4);
+}
+#endif
+
 TEST_F(LIRABITest, TestkCall_FillsCallSiteLiveValueLocations) {
   if constexpr (!kFreeThreadedBuild) {
     GTEST_SKIP() << "Callsite live-value locations are only filled in "

cinderx/RuntimeTests/lir_postalloc_test.cpp

@@ -407,6 +407,44 @@ BB %0
   EXPECT_EQ(store_pairs, 2);
   ASSERT_TRUE(verifyPostRegAllocInvariants(parsed_func.get(), std::cout));
 }
+
+TEST_F(LIRPostAllocRewriteTest, RegularCallArgsUseStorePairForRegisterPairs) {
+  auto lir_input_str = fmt::format(
+      R"(Function:
+BB %0
+                   Call 123456789, {:>9}:Object, {:>9}:Object, {:>9}:Object, {:>9}:Object, {:>9}:Object, {:>9}:Object, {:>9}:Object, {:>9}:Object, {:>9}:Object, {:>9}:Object
+)",
+      X0,
+      X1,
+      X2,
+      X3,
+      X4,
+      X5,
+      X6,
+      X7,
+      X8,
+      X9);
+
+  Parser parser;
+  auto parsed_func = parser.parse(lir_input_str);
+
+  jit::codegen::Environ env;
+  PostRegAllocRewrite rewrite(parsed_func.get(), &env);
+  rewrite.run();
+
+  const Instruction* store_pair = nullptr;
+  for (auto& instr : parsed_func->basicblocks().front()->instructions()) {
+    if (instr->isStorePair()) {
+      ASSERT_EQ(store_pair, nullptr);
+      store_pair = instr.get();
+    }
+  }
+
+  ASSERT_NE(store_pair, nullptr);
+  EXPECT_EQ(
+      store_pair->getInput(1)->getPhyRegister(), arch::reg_stack_pointer_loc);
+  ASSERT_TRUE(verifyPostRegAllocInvariants(parsed_func.get(), std::cout));
+}
 #endif // CINDER_AARCH64
 
 } // namespace jit::lir