llvm/lib/Target/LoongArch/LoongArchMemoryBarrierOpt.cpp

//===---- LoongArchMemoryBarrierOpt.cpp - Memory barrier Optimization -----===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
///
/// This pass removes or merges redundant memory barrier instructions.
///
/// - DBAR x + DBAR y -> DBAR (x & y)
/// - DBAR x + AMO_DB -> AMO_DB
/// - DBAR x + AMO    -> AMO_DB
/// - DBAR x + LL     -> LL
/// - AMO_DB + DBAR x -> AMO_DB
/// - AMO    + DBAR x -> AMO_DB
/// - SC     + DBAR x -> SC
///
//===----------------------------------------------------------------------===//

#include "LoongArch.h"
#include "LoongArchInstrInfo.h"
#include "LoongArchSubtarget.h"
#include "llvm/CodeGen/MachineDominators.h"
#include "llvm/CodeGen/MachinePostDominators.h"
#include "llvm/InitializePasses.h"

using namespace llvm;

#define DEBUG_TYPE "loongarch-memory-barrier-opt"
#define LOONGARCH_MEMORY_BARRIER_OPT_NAME                                      \
  "LoongArch Memory Barrier Optimisation pass"

static cl::opt<bool> RequireNoPathBypass(
    "loongarch-require-no-path-bypass",
    cl::desc("Optimize only when no paths bypass either memory barrier"),
    cl::init(true), cl::Hidden);

static cl::opt<bool> MergeAMOWithMB(
    "loongarch-merge-amo-with-dbar",
    cl::desc("Merge AMOs with DBARs into AMO_DB during optimization"),
    cl::init(true), cl::Hidden);

static cl::opt<bool> ReplaceEliminatedMBToNop(
    "loongarch-replace-eliminated-dbar-to-nop",
    cl::desc("Replace eliminated DBARs with NOPs to preserve code layout"),
    cl::init(false), cl::Hidden);

namespace {

#define AMO_CASES                                                              \
  CASE(AMSWAP, B)                                                              \
  CASE(AMSWAP, H)                                                              \
  CASE(AMSWAP, W)                                                              \
  CASE(AMSWAP, D)                                                              \
  CASE(AMADD, B)                                                               \
  CASE(AMADD, H)                                                               \
  CASE(AMADD, W)                                                               \
  CASE(AMADD, D)                                                               \
  CASE(AMAND, W)                                                               \
  CASE(AMAND, D)                                                               \
  CASE(AMOR, W)                                                                \
  CASE(AMOR, D)                                                                \
  CASE(AMXOR, W)                                                               \
  CASE(AMXOR, D)                                                               \
  CASE(AMMAX, W)                                                               \
  CASE(AMMAX, D)                                                               \
  CASE(AMMAX, WU)                                                              \
  CASE(AMMAX, DU)                                                              \
  CASE(AMMIN, W)                                                               \
  CASE(AMMIN, D)                                                               \
  CASE(AMMIN, WU)                                                              \
  CASE(AMMIN, DU)                                                              \
  CASE(AMCAS, B)                                                               \
  CASE(AMCAS, H)                                                               \
  CASE(AMCAS, W)                                                               \
  CASE(AMCAS, D)

static bool isMB(const MachineInstr &MI) {
  return MI.getOpcode() == LoongArch::DBAR;
}

static bool isLL(const MachineInstr &MI) {
  switch (MI.getOpcode()) {
  case LoongArch::LL_W:
  case LoongArch::LL_D:
    return true;
  default:
    return false;
  }
}

static bool isSC(const MachineInstr &MI) {
  switch (MI.getOpcode()) {
  case LoongArch::SC_W:
  case LoongArch::SC_D:
  case LoongArch::SC_Q:
    return true;
  default:
    return false;
  }
}

static std::optional<unsigned> isAM(const MachineInstr &MI) {
#define CASE(Name, Suffix)                                                     \
  case LoongArch::Name##_##Suffix:                                             \
    if (!MergeAMOWithMB)                                                       \
      return std::nullopt;                                                     \
    [[fallthrough]];                                                           \
  case LoongArch::Name##__DB_##Suffix:                                         \
    return LoongArch::Name##__DB_##Suffix;
  switch (MI.getOpcode()) {
    AMO_CASES
  default:
    return std::nullopt;
  }
#undef CASE
}

static bool isSafeToSkip(const MachineInstr &MI) {
  if (MI.mayLoadOrStore())
    return false;
  if (MI.isCall() || MI.isReturn())
    return false;
  if (MI.isInlineAsm())
    return false;
  if (MI.hasUnmodeledSideEffects())
    return isMB(MI);
  return true;
}

struct BarrierHint {
  BarrierHint(unsigned Hint) : Hint(Hint) {}

  bool subsumes(const BarrierHint &O) const { return (Hint & O.Hint) == Hint; }

  BarrierHint merge(const BarrierHint &O) const {
    return BarrierHint(Hint & O.Hint);
  }

  static inline bool isValid(unsigned Hint) { return (Hint & ~0x1f) == 0; }

  unsigned Hint;
};

struct InstBarrier {
  InstBarrier(MachineInstr &MI)
      : MI(&MI), Pre(0), Post(0), OpcAMDB(0), IsMB(false), IsAM(false) {
    if (isMB(MI)) {
      unsigned Hint = MI.getOperand(0).getImm();
      if (!BarrierHint::isValid(Hint))
        return;
      IsMB = true;
      Pre = Post = BarrierHint(Hint);
    } else if (isLL(MI)) {
      IsAM = true;
      Pre = BarrierHint(0b10000);
      Post = BarrierHint(0b11111);
    } else if (isSC(MI)) {
      IsAM = true;
      Pre = BarrierHint(0b11111);
      Post = BarrierHint(0b10000);
    } else if (auto R = isAM(MI)) {
      IsAM = true;
      OpcAMDB = *R;
      Pre = Post = BarrierHint(0b10000);
    }
  }

  MachineInstr *MI;
  BarrierHint Pre;
  BarrierHint Post;
  StringRef OpName;
  StringRef Operands;
  unsigned OpcAMDB;
  bool IsMB;
  bool IsAM;
};

class LoongArchMemoryBarrierOpt : public MachineFunctionPass {
public:
  static char ID;

  LoongArchMemoryBarrierOpt() : MachineFunctionPass(ID) {}

  StringRef getPassName() const override {
    return LOONGARCH_MEMORY_BARRIER_OPT_NAME;
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<MachineDominatorTreeWrapperPass>();
    AU.addPreserved<MachineDominatorTreeWrapperPass>();
    AU.addRequired<MachinePostDominatorTreeWrapperPass>();
    AU.addPreserved<MachinePostDominatorTreeWrapperPass>();
    MachineFunctionPass::getAnalysisUsage(AU);
  }

  bool runOnMachineFunction(MachineFunction &Fn) override;

private:
  enum : unsigned {
    CandidateA = 1u << 0,
    CandidateB = 1u << 1,
  };

  unsigned resolveBarrierRedundancy(const MachineInstr *A,
                                    const MachineInstr *B) const;
  bool eliminateRedundantBarrier(InstBarrier &IA, InstBarrier &IB) const;

  MachineFunction *MF;
  const MachineDominatorTree *MDT;
  const MachinePostDominatorTree *MPDT;
};

static bool checkAllPathSafe(const MachineBasicBlock *MBBA,
                             const MachineBasicBlock *MBBB, bool IsAToB) {
  const MachineBasicBlock *Start = IsAToB ? MBBA : MBBB;
  const MachineBasicBlock *End = IsAToB ? MBBB : MBBA;

  SmallVector<const MachineBasicBlock *, 16> Worklist;
  DenseSet<const MachineBasicBlock *> Visited;

  Worklist.push_back(Start);
  Visited.insert(Start);

  while (!Worklist.empty()) {
    const MachineBasicBlock *BB = Worklist.pop_back_val();

    if (BB == End)
      continue;

    if (BB != Start) {
      for (const MachineInstr &MI : *BB) {
        if (!isSafeToSkip(MI))
          return false;
      }
    }

    if (IsAToB) {
      for (const MachineBasicBlock *Succ : BB->successors()) {
        if (Visited.insert(Succ).second)
          Worklist.push_back(Succ);
      }
    } else {
      for (const MachineBasicBlock *Pred : BB->predecessors()) {
        if (Visited.insert(Pred).second)
          Worklist.push_back(Pred);
      }
    }
  }

  return true;
}

// Returns a bitmask indicating removal candidates: A (bit 1) and B (bit 2).
unsigned LoongArchMemoryBarrierOpt::resolveBarrierRedundancy(
    const MachineInstr *A, const MachineInstr *B) const {
  const MachineBasicBlock *MBBA = A->getParent();
  const MachineBasicBlock *MBBB = B->getParent();

  if (MBBA == MBBB) {
    /* A -> B */
    for (auto It = std::next(A->getIterator()); It != MBBA->end(); ++It) {
      if (It == B->getIterator())
        return CandidateA | CandidateB;
      if (!isSafeToSkip(*It))
        return 0;
    }
    return 0;
  }

  // Cross-block walk
  bool ADomB = MDT->dominates(MBBA, MBBB);
  bool BPostDomA = MPDT->dominates(MBBB, MBBA);
  unsigned Mask = 0;
  if (!ADomB && !BPostDomA)
    return 0;

  /* A -> MBBA->end() */
  for (auto It = std::next(A->getIterator()); It != MBBA->end(); ++It)
    if (!isSafeToSkip(*It))
      return 0;
  /* B -> MBBB->begin() */
  for (auto It = MBBB->begin(); It != B->getIterator(); ++It)
    if (!isSafeToSkip(*It))
      return 0;

  /* MBBA -> MBBB */
  if (BPostDomA)
    if (checkAllPathSafe(MBBA, MBBB, true /*IsAToB*/))
      Mask |= CandidateA;

  /* MBBB -> MBBA */
  if (ADomB)
    if (checkAllPathSafe(MBBA, MBBB, false /*IsAToB*/))
      Mask |= CandidateB;

  return Mask;
}

// Update DBAR hint
static void updateMB(InstBarrier &I, BarrierHint Hint, MachineFunction *MF) {
  assert(I.IsMB && "Unexpected!");
  I.Pre = I.Post = Hint;
  I.MI->getOperand(0).setImm(Hint.Hint);
}

// Replace AMO to AMO_DB
static void replaceAM(InstBarrier &I, MachineFunction *MF) {
  if (!I.IsAM)
    return;
  if (I.OpcAMDB) {
    auto &ST = MF->getSubtarget<LoongArchSubtarget>();
    I.MI->setDesc(ST.getInstrInfo()->get(I.OpcAMDB));
  }
}

bool LoongArchMemoryBarrierOpt::eliminateRedundantBarrier(
    InstBarrier &IA, InstBarrier &IB) const {
  MachineInstr *A = IA.MI;
  MachineInstr *B = IB.MI;

  if (!A || !B)
    return false; // Already erased
  if (A == B)
    return false;

  unsigned Mask = resolveBarrierRedundancy(A, B);
  if (!Mask)
    return false;

  auto eraseOrReplaceWithNop = [&](MachineInstr *MI) {
    if (ReplaceEliminatedMBToNop) {
      auto &ST = MF->getSubtarget<LoongArchSubtarget>();
      BuildMI(*MI->getParent(), MI->getIterator(), MI->getDebugLoc(),
              ST.getInstrInfo()->get(LoongArch::ANDI), LoongArch::R0)
          .addReg(LoongArch::R0)
          .addImm(0);
    }
    MI->eraseFromParent();
  };

  // A        B
  // DBAR x + DBAR y -> DBAR (x & y)
  // DBAR x + AMO_DB -> AMO_DB
  // DBAR x + AMO    -> AMO_DB
  // DBAR x + LL     -> LL
  if ((Mask & CandidateA) && IA.IsMB) {
    if (!IB.Pre.subsumes(IA.Post)) {
      if (!IB.IsMB || (RequireNoPathBypass && !(Mask & CandidateB)))
        return false;
      updateMB(IB, IB.Pre.merge(IA.Post), MF);
    }
    replaceAM(IB, MF);
    eraseOrReplaceWithNop(A);
    IA.MI = nullptr;
    return true;
  }

  // A        B
  // DBAR x + DBAR y -> DBAR (x & y)
  // AMO_DB + DBAR x -> AMO_DB
  // AMO    + DBAR x -> AMO_DB
  // SC     + DBAR x -> SC
  if ((Mask & CandidateB) && IB.IsMB) {
    if (!IA.Post.subsumes(IB.Pre)) {
      if (!IA.IsMB || (RequireNoPathBypass && !(Mask & CandidateA)))
        return false;
      updateMB(IA, IA.Post.merge(IB.Pre), MF);
    }
    replaceAM(IA, MF);
    eraseOrReplaceWithNop(B);
    IB.MI = nullptr;
    return true;
  }

  return false;
}

bool LoongArchMemoryBarrierOpt::runOnMachineFunction(MachineFunction &Fn) {
  if (skipFunction(Fn.getFunction()))
    return false;

  MF = &Fn;
  MDT = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
  MPDT = &getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree();

  SmallVector<InstBarrier, 32> Sites;
  bool Changed = false;

  for (MachineBasicBlock &MBB : Fn)
    for (MachineInstr &MI : MBB) {
      InstBarrier IB(MI);
      if (IB.IsMB || IB.IsAM)
        Sites.push_back(IB);
    }

  for (size_t a = 0; a < Sites.size(); ++a) {
    for (size_t b = a + 1; b < Sites.size(); ++b) {
      InstBarrier &IA = Sites[a];
      InstBarrier &IB = Sites[b];
      Changed |= eliminateRedundantBarrier(IA, IB);
      Changed |= eliminateRedundantBarrier(IB, IA);
    }
  }

  return Changed;
}
} // namespace

char LoongArchMemoryBarrierOpt::ID = 0;
INITIALIZE_PASS_BEGIN(LoongArchMemoryBarrierOpt, DEBUG_TYPE,
                      LOONGARCH_MEMORY_BARRIER_OPT_NAME, false, false)
INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
INITIALIZE_PASS_DEPENDENCY(MachinePostDominatorTreeWrapperPass)
INITIALIZE_PASS_END(LoongArchMemoryBarrierOpt, DEBUG_TYPE,
                    LOONGARCH_MEMORY_BARRIER_OPT_NAME, false, false)

FunctionPass *llvm::createLoongArchMemoryBarrierOptPass() {
  return new LoongArchMemoryBarrierOpt();
}