Merge branch 'main' into fixup/alias-analysis-atomic-op

Author: jinhuang1102

.github/CODEOWNERS

@@ -93,9 +93,9 @@
 /mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp @MaheshRavishankar @nicolasvasilache
 /mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp @dcaballe @MaheshRavishankar @nicolasvasilache
 /mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp @MaheshRavishankar @nicolasvasilache
-/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp @hanhanW @nicolasvasilache
-/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @dcaballe @hanhanW @nicolasvasilache
-/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @banach-space @dcaballe @hanhanW @nicolasvasilache @Groverkss
+/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp @nicolasvasilache
+/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @dcaballe @nicolasvasilache
+/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @banach-space @dcaballe @nicolasvasilache @Groverkss
 
 # MemRef Dialect in MLIR.
 /mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp @MaheshRavishankar @nicolasvasilache
@@ -112,16 +112,16 @@
 /mlir/include/mlir/Dialect/Vector @banach-space @dcaballe @nicolasvasilache @Groverkss
 /mlir/include/mlir/Dialect/Vector/IR @kuhar
 /mlir/lib/Dialect/Vector @banach-space @dcaballe @nicolasvasilache @Groverkss
-/mlir/lib/Dialect/Vector/Transforms/* @banach-space @dcaballe @hanhanW @nicolasvasilache
+/mlir/lib/Dialect/Vector/Transforms/* @banach-space @dcaballe @nicolasvasilache
 /mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp @banach-space @dcaballe @MaheshRavishankar @nicolasvasilache
-/mlir/**/*EmulateNarrowType* @dcaballe @hanhanW
+/mlir/**/*EmulateNarrowType* @dcaballe
 
 # Presburger library in MLIR
 /mlir/**/*Presburger* @Groverkss @Superty
 
 # Tensor Dialect in MLIR.
-/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp @hanhanW @nicolasvasilache
-/mlir/lib/Dialect/Tensor/Transforms/* @hanhanW @nicolasvasilache
+/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp @nicolasvasilache
+/mlir/lib/Dialect/Tensor/Transforms/* @nicolasvasilache
 
 # Transform Dialect in MLIR.
 /mlir/include/mlir/Dialect/Transform/* @ftynse @nicolasvasilache @rolfmorel

bolt/include/bolt/Core/BinaryContext.h

@@ -1113,7 +1113,7 @@ class BinaryContext {
     return FragmentClasses.isEquivalent(LHS, RHS);
   }
 
-  /// Add interprocedural reference for \p Function to \p Address
+  /// Add interprocedural branch reference from \p Function to \p Address.
   void addInterproceduralReference(BinaryFunction *Function, uint64_t Address) {
     InterproceduralReferences.push_back({Function, Address});
   }
@@ -1128,7 +1128,8 @@ class BinaryContext {
   /// argument is false.
   bool handleAArch64Veneer(uint64_t Address, bool MatchOnly = false);
 
-  /// Resolve inter-procedural dependencies from
+  /// Resolve inter-procedural branch dependencies discovered during
+  /// disassembly.
   void processInterproceduralReferences();
 
   /// Skip functions with all parent and child fragments transitively.

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -538,11 +538,6 @@ class MCPlusBuilder {
     llvm_unreachable("not implemented");
   }
 
-  virtual void createDirectBranch(MCInst &Inst, const MCSymbol *Target,
-                                  MCContext *Ctx) {
-    llvm_unreachable("not implemented");
-  }
-
   virtual MCPhysReg getX86R11() const { llvm_unreachable("not implemented"); }
 
   virtual unsigned getShortBranchOpcode(unsigned Opcode) const {

bolt/lib/Core/BinaryContext.cpp

@@ -1438,8 +1438,6 @@ void BinaryContext::processInterproceduralReferences() {
       continue;
     }
 
-    // Check if address falls in function padding space - this could be
-    // unmarked data in code. In this case adjust the padding space size.
     ErrorOr<BinarySection &> Section = getSectionForAddress(Address);
     assert(Section && "cannot get section for referenced address");
 
@@ -1451,7 +1449,7 @@ void BinaryContext::processInterproceduralReferences() {
     if (SectionName == ".plt" || SectionName == ".plt.got")
       continue;
 
-    // Check if it is aarch64 veneer written at Address
+    // Check if it is aarch64 veneer written at Address.
     if (isAArch64() && handleAArch64Veneer(Address))
       continue;
 
@@ -1463,6 +1461,8 @@ void BinaryContext::processInterproceduralReferences() {
       exit(1);
     }
 
+    // Check if the address falls into the function padding space - this could
+    // be an unmarked data in code. In this case, adjust the padding space size.
     TargetFunction = getBinaryFunctionContainingAddress(Address,
                                                         /*CheckPastEnd=*/false,
                                                         /*UseMaxSize=*/true);

bolt/lib/Passes/Instrumentation.cpp

@@ -305,12 +305,9 @@ void Instrumentation::instrumentIndirectTarget(BinaryBasicBlock &BB,
                  : IndCallHandlerExitBBFunction->getSymbol(),
       IndCallSiteID, &*BC.Ctx);
 
-  if (!BC.isAArch64()) {
-    Iter = BB.eraseInstruction(Iter);
-    Iter = insertInstructions(CounterInstrs, BB, Iter);
-    --Iter;
-  } else
-    Iter = insertInstructions(CounterInstrs, BB, Iter);
+  Iter = BB.eraseInstruction(Iter);
+  Iter = insertInstructions(CounterInstrs, BB, Iter);
+  --Iter;
 }
 
 bool Instrumentation::instrumentOneTarget(

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -48,14 +48,14 @@ static cl::opt<bool> NoLSEAtomics(
 
 namespace {
 
-[[maybe_unused]] static void getSystemFlag(MCInst &Inst, MCPhysReg RegName) {
+static void getSystemFlag(MCInst &Inst, MCPhysReg RegName) {
   Inst.setOpcode(AArch64::MRS);
   Inst.clear();
   Inst.addOperand(MCOperand::createReg(RegName));
   Inst.addOperand(MCOperand::createImm(AArch64SysReg::NZCV));
 }
 
-[[maybe_unused]] static void setSystemFlag(MCInst &Inst, MCPhysReg RegName) {
+static void setSystemFlag(MCInst &Inst, MCPhysReg RegName) {
   Inst.setOpcode(AArch64::MSR);
   Inst.clear();
   Inst.addOperand(MCOperand::createImm(AArch64SysReg::NZCV));
@@ -2114,14 +2114,6 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
       convertJmpToTailCall(Inst);
   }
 
-  void createDirectBranch(MCInst &Inst, const MCSymbol *Target,
-                          MCContext *Ctx) override {
-    Inst.setOpcode(AArch64::B);
-    Inst.clear();
-    Inst.addOperand(MCOperand::createExpr(getTargetExprFor(
-        Inst, MCSymbolRefExpr::create(Target, *Ctx), *Ctx, 0)));
-  }
-
   bool analyzeBranch(InstructionIterator Begin, InstructionIterator End,
                      const MCSymbol *&TBB, const MCSymbol *&FBB,
                      MCInst *&CondBranch,
@@ -2479,14 +2471,21 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
   }
 
   InstructionListType createInstrumentedIndCallHandlerExitBB() const override {
+    InstructionListType Insts(5);
     // Code sequence for instrumented indirect call handler:
-    //   ret
-
-    InstructionListType Insts;
-
-    Insts.emplace_back();
-    createReturn(Insts.back());
-
+    //   msr  nzcv, x1
+    //   ldp  x0, x1, [sp], #16
+    //   ldr  x16, [sp], #16
+    //   ldp  x0, x1, [sp], #16
+    //   br   x16
+    setSystemFlag(Insts[0], AArch64::X1);
+    createPopRegisters(Insts[1], AArch64::X0, AArch64::X1);
+    // Here we load address of the next function which should be called in the
+    // original binary to X16 register. Writing to X16 is permitted without
+    // needing to restore.
+    loadReg(Insts[2], AArch64::X16, AArch64::SP);
+    createPopRegisters(Insts[3], AArch64::X0, AArch64::X1);
+    createIndirectBranch(Insts[4], AArch64::X16, 0);
     return Insts;
   }
 
@@ -2562,59 +2561,39 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
                                                      MCSymbol *HandlerFuncAddr,
                                                      int CallSiteID,
                                                      MCContext *Ctx) override {
+    InstructionListType Insts;
     // Code sequence used to enter indirect call instrumentation helper:
-    //   stp x0, x1, [sp, #-16]! createPushRegisters  (1)
-    //   mov target, x0  convertIndirectCallToLoad -> orr x0 target xzr
+    //   stp x0, x1, [sp, #-16]! createPushRegisters
+    //   mov target x0  convertIndirectCallToLoad -> orr x0 target xzr
     //   mov x1 CallSiteID createLoadImmediate ->
     //   movk    x1, #0x0, lsl #48
     //   movk    x1, #0x0, lsl #32
     //   movk    x1, #0x0, lsl #16
     //   movk    x1, #0x0
-    //   stp x0, x30, [sp, #-16]!    (2)
+    //   stp x0, x1, [sp, #-16]!
+    //   bl *HandlerFuncAddr createIndirectCall ->
     //   adr x0 *HandlerFuncAddr -> adrp + add
-    //   blr x0   (__bolt_instr_ind_call_handler_func)
-    //   ldp x0, x30, [sp], #16   (2)
-    //   mov x0, target ; move target address to used register
-    //   ldp x0, x1, [sp], #16   (1)
-
-    InstructionListType Insts;
+    //   blr x0
     Insts.emplace_back();
-    createPushRegisters(Insts.back(), getIntArgRegister(0),
-                        getIntArgRegister(1));
+    createPushRegisters(Insts.back(), AArch64::X0, AArch64::X1);
     Insts.emplace_back(CallInst);
-    convertIndirectCallToLoad(Insts.back(), getIntArgRegister(0));
+    convertIndirectCallToLoad(Insts.back(), AArch64::X0);
     InstructionListType LoadImm =
         createLoadImmediate(getIntArgRegister(1), CallSiteID);
     Insts.insert(Insts.end(), LoadImm.begin(), LoadImm.end());
     Insts.emplace_back();
-    createPushRegisters(Insts.back(), getIntArgRegister(0), AArch64::LR);
+    createPushRegisters(Insts.back(), AArch64::X0, AArch64::X1);
     Insts.resize(Insts.size() + 2);
-    InstructionListType Addr = materializeAddress(
-        HandlerFuncAddr, Ctx, CallInst.getOperand(0).getReg());
+    InstructionListType Addr =
+        materializeAddress(HandlerFuncAddr, Ctx, AArch64::X0);
     assert(Addr.size() == 2 && "Invalid Addr size");
     std::copy(Addr.begin(), Addr.end(), Insts.end() - Addr.size());
-
     Insts.emplace_back();
-    createIndirectCallInst(Insts.back(), false,
-                           CallInst.getOperand(0).getReg());
+    createIndirectCallInst(Insts.back(), isTailCall(CallInst), AArch64::X0);
 
-    Insts.emplace_back();
-    createPopRegisters(Insts.back(), getIntArgRegister(0), AArch64::LR);
-
-    // move x0 to indirect call register
-    Insts.emplace_back();
-    Insts.back().setOpcode(AArch64::ORRXrs);
-    Insts.back().insert(Insts.back().begin(),
-                        MCOperand::createReg(CallInst.getOperand(0).getReg()));
-    Insts.back().insert(Insts.back().begin() + 1,
-                        MCOperand::createReg(AArch64::XZR));
-    Insts.back().insert(Insts.back().begin() + 2,
-                        MCOperand::createReg(getIntArgRegister(0)));
-    Insts.back().insert(Insts.back().begin() + 3, MCOperand::createImm(0));
-
-    Insts.emplace_back();
-    createPopRegisters(Insts.back(), getIntArgRegister(0),
-                       getIntArgRegister(1));
+    // Carry over metadata including tail call marker if present.
+    stripAnnotations(Insts.back());
+    moveAnnotations(std::move(CallInst), Insts.back());
 
     return Insts;
   }
@@ -2623,53 +2602,43 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
   createInstrumentedIndCallHandlerEntryBB(const MCSymbol *InstrTrampoline,
                                           const MCSymbol *IndCallHandler,
                                           MCContext *Ctx) override {
-    // Code sequence used to check whether InstrTrampoline was initialized
+    // Code sequence used to check whether InstrTampoline was initialized
     // and call it if so, returns via IndCallHandler
-    //   adrp    x0, InstrTrampoline
-    //   ldr     x0, [x0, #lo12:InstrTrampoline]
+    //   stp     x0, x1, [sp, #-16]!
+    //   mrs     x1, nzcv
+    //   adr     x0, InstrTrampoline -> adrp + add
+    //   ldr     x0, [x0]
     //   subs    x0, x0, #0x0
     //   b.eq    IndCallHandler
     //   str     x30, [sp, #-16]!
     //   blr     x0
     //   ldr     x30, [sp], #16
     //   b       IndCallHandler
     InstructionListType Insts;
-
-    // load handler address
-    MCInst InstAdrp;
-    InstAdrp.setOpcode(AArch64::ADRP);
-    InstAdrp.addOperand(MCOperand::createReg(getIntArgRegister(0)));
-    InstAdrp.addOperand(MCOperand::createImm(0));
-    setOperandToSymbolRef(InstAdrp, /* OpNum */ 1, InstrTrampoline,
-                          /* Addend */ 0, Ctx, ELF::R_AARCH64_ADR_GOT_PAGE);
-    Insts.emplace_back(InstAdrp);
-
-    MCInst InstLoad;
-    InstLoad.setOpcode(AArch64::LDRXui);
-    InstLoad.addOperand(MCOperand::createReg(getIntArgRegister(0)));
-    InstLoad.addOperand(MCOperand::createReg(getIntArgRegister(0)));
-    InstLoad.addOperand(MCOperand::createImm(0));
-    setOperandToSymbolRef(InstLoad, /* OpNum */ 2, InstrTrampoline,
-                          /* Addend */ 0, Ctx, ELF::R_AARCH64_LD64_GOT_LO12_NC);
-    Insts.emplace_back(InstLoad);
-
-    InstructionListType CmpJmp =
-        createCmpJE(getIntArgRegister(0), 0, IndCallHandler, Ctx);
-    Insts.insert(Insts.end(), CmpJmp.begin(), CmpJmp.end());
-
     Insts.emplace_back();
-    storeReg(Insts.back(), AArch64::LR, getSpRegister(/*Size*/ 8));
-
+    createPushRegisters(Insts.back(), AArch64::X0, AArch64::X1);
+    Insts.emplace_back();
+    getSystemFlag(Insts.back(), getIntArgRegister(1));
+    Insts.emplace_back();
+    Insts.emplace_back();
+    InstructionListType Addr =
+        materializeAddress(InstrTrampoline, Ctx, AArch64::X0);
+    std::copy(Addr.begin(), Addr.end(), Insts.end() - Addr.size());
+    assert(Addr.size() == 2 && "Invalid Addr size");
+    Insts.emplace_back();
+    loadReg(Insts.back(), AArch64::X0, AArch64::X0);
+    InstructionListType cmpJmp =
+        createCmpJE(AArch64::X0, 0, IndCallHandler, Ctx);
+    Insts.insert(Insts.end(), cmpJmp.begin(), cmpJmp.end());
+    Insts.emplace_back();
+    storeReg(Insts.back(), AArch64::LR, AArch64::SP);
     Insts.emplace_back();
     Insts.back().setOpcode(AArch64::BLR);
-    Insts.back().addOperand(MCOperand::createReg(getIntArgRegister(0)));
-
+    Insts.back().addOperand(MCOperand::createReg(AArch64::X0));
     Insts.emplace_back();
-    loadReg(Insts.back(), AArch64::LR, getSpRegister(/*Size*/ 8));
-
+    loadReg(Insts.back(), AArch64::LR, AArch64::SP);
     Insts.emplace_back();
-    createDirectBranch(Insts.back(), IndCallHandler, Ctx);
-
+    createDirectCall(Insts.back(), IndCallHandler, Ctx, /*IsTailCall*/ true);
     return Insts;
   }
 

bolt/runtime/instr.cpp

@@ -1691,12 +1691,9 @@ instrumentIndirectCall(uint64_t Target, uint64_t IndCallID) {
 extern "C" __attribute((naked)) void __bolt_instr_indirect_call()
 {
 #if defined(__aarch64__)
-  // the target address is placed on stack
-  // the identifier of the indirect call site is placed in X1 register
-
   // clang-format off
   __asm__ __volatile__(SAVE_ALL
-                       "ldr x0, [sp, #272]\n"
+                       "ldp x0, x1, [sp, #288]\n"
                        "bl instrumentIndirectCall\n"
                        RESTORE_ALL
                        "ret\n"
@@ -1731,12 +1728,9 @@ extern "C" __attribute((naked)) void __bolt_instr_indirect_call()
 extern "C" __attribute((naked)) void __bolt_instr_indirect_tailcall()
 {
 #if defined(__aarch64__)
-  // the target address is placed on stack
-  // the identifier of the indirect call site is placed in X1 register
-
   // clang-format off
   __asm__ __volatile__(SAVE_ALL
-                       "ldr x0, [sp, #272]\n"
+                       "ldp x0, x1, [sp, #288]\n"
                        "bl instrumentIndirectCall\n"
                        RESTORE_ALL
                        "ret\n"

bolt/runtime/sys_aarch64.h

@@ -18,12 +18,10 @@
   "stp x24, x25, [sp, #-16]!\n"                                                \
   "stp x26, x27, [sp, #-16]!\n"                                                \
   "stp x28, x29, [sp, #-16]!\n"                                                \
-  "mrs x29, nzcv\n"                                                            \
-  "stp x29, x30, [sp, #-16]!\n"
+  "str x30, [sp,#-16]!\n"
 // Mirrors SAVE_ALL
 #define RESTORE_ALL                                                            \
-  "ldp x29, x30, [sp], #16\n"                                                  \
-  "msr nzcv, x29\n"                                                            \
+  "ldr x30, [sp], #16\n"                                                       \
   "ldp x28, x29, [sp], #16\n"                                                  \
   "ldp x26, x27, [sp], #16\n"                                                  \
   "ldp x24, x25, [sp], #16\n"                                                  \