llvm/lib/Transforms/Vectorize/SandboxVectorizer/VecUtils.cpp

//===- VecUtils.cpp -------------------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h"
#include "../VectorizeOptions.h"

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/SandboxIR/Instruction.h"
#include "llvm/Transforms/Vectorize/SandboxVectorizer/Debug.h"
#include "llvm/Transforms/Vectorize/SandboxVectorizer/InstrMaps.h"

namespace llvm::sandboxir {

static SmallVector<unsigned, 2> getOperandIndicesInUser(User *U, Value *Op) {
  SmallVector<unsigned, 2> OpIdxVec;
  for (unsigned Idx : seq<unsigned>(U->getNumOperands()))
    if (U->getOperand(Idx) == Op)
      OpIdxVec.push_back(Idx);
  return OpIdxVec;
}

static std::optional<BundleTy>
getMatchingBundle(ArrayRef<Value *> Bndl, const InstrMaps &IMaps, Value *Seed,
                  Instruction *SeedUserInst,
                  SmallPtrSet<Instruction *, 4> &Claimed) {
  SmallVector<unsigned, 2> OpIdxVec0 =
      getOperandIndicesInUser(SeedUserInst, Seed);
  assert(!OpIdxVec0.empty() && "U0 does not use Seed!");
  BundleTy NextUserBndl;
  NextUserBndl.push_back(SeedUserInst);
  Claimed.insert(SeedUserInst);
  for (Value *V : drop_begin(Bndl)) {
    Instruction *Match = nullptr;
    for (User *U : V->users()) {
      auto *UI = dyn_cast<Instruction>(U);
      if (!UI || IMaps.isVectorized(UI) || Claimed.contains(UI) ||
          UI->getOpcode() != SeedUserInst->getOpcode() ||
          UI->getType() != SeedUserInst->getType() ||
          UI->getParent() != SeedUserInst->getParent() ||
          getOperandIndicesInUser(UI, V) != OpIdxVec0)
        continue;

      Match = UI;
      break;
    }
    if (!Match)
      return std::nullopt;
    NextUserBndl.push_back(Match);
  }

  for (auto *I : NextUserBndl)
    Claimed.insert(cast<Instruction>(I));
  return NextUserBndl;
}

SmallVector<BundleTy>
VecUtils::getNextUserBundles(ArrayRef<Value *> Bndl, const InstrMaps &IMaps,
                             SmallPtrSet<Instruction *, 4> &Claimed) {
  SmallVector<BundleTy> Bundles;
  if (Bndl.empty())
    return Bundles;

  Value *V0 = Bndl[0];
  DenseSet<User *> SeenUsers;
  const VectorizeOptions &Opts =
      V0->getContext().getLibraryOptions().get<VectorizeOptions>();
  // For each user U0 of lane 0, try to form a bundle of matching users across
  // all lanes. Cap the number of users considered to bound compilation time,
  // since each one may trigger an O(Bndl.size()) search across the other
  // lanes' users.
  for (User *U0 : V0->users()) {
    if (SeenUsers.size() >= Opts.sbvec_max_users_to_consider)
      break;
    if (!SeenUsers.insert(U0).second)
      continue;
    auto *UI0 = dyn_cast<Instruction>(U0);
    if (!UI0 || IMaps.isVectorized(UI0) || Claimed.contains(UI0))
      continue;
    std::optional<BundleTy> NextUserBndl =
        getMatchingBundle(Bndl, IMaps, V0, UI0, Claimed);
    if (NextUserBndl)
      Bundles.emplace_back(std::move(*NextUserBndl));
  }
  return Bundles;
}

unsigned VecUtils::getFloorPowerOf2(unsigned Num) {
  if (Num == 0)
    return Num;
  unsigned Mask = Num;
  Mask >>= 1;
  for (unsigned ShiftBy = 1; ShiftBy < sizeof(Num) * 8; ShiftBy <<= 1)
    Mask |= Mask >> ShiftBy;
  return Num & ~Mask;
}

template <typename T>
void VecUtils::DeadInstructionMorgue::collectPotentiallyDeadInstrs(
    BndlRef<T *> Bndl) {
  for (T *V : Bndl) {
    assert(isa<Instruction>(V) && "Only works with instructions");
    DeadInstrCandidates.insert(cast<Instruction>(V));
  }
  // Also collect the GEPs of vectorized loads and stores.
  auto Opcode = cast<Instruction>(Bndl[0])->getOpcode();
  switch (Opcode) {
  case Instruction::Opcode::Load: {
    for (T *V : drop_begin(Bndl))
      if (auto *Ptr =
              dyn_cast<Instruction>(cast<LoadInst>(V)->getPointerOperand()))
        DeadInstrCandidates.insert(Ptr);
    break;
  }
  case Instruction::Opcode::Store: {
    for (T *V : drop_begin(Bndl))
      if (auto *Ptr =
              dyn_cast<Instruction>(cast<StoreInst>(V)->getPointerOperand()))
        DeadInstrCandidates.insert(Ptr);
    break;
  }
  default:
    break;
  }
}

template void
    VecUtils::DeadInstructionMorgue::collectPotentiallyDeadInstrs<Value>(
        BndlRef<Value *>);
template LLVM_EXPORT_TEMPLATE void
    VecUtils::DeadInstructionMorgue::collectPotentiallyDeadInstrs<Instruction>(
        BndlRef<Instruction *>);

void VecUtils::DeadInstructionMorgue::tryEraseDeadInstrs() {
  DenseMap<BasicBlock *, SmallVector<Instruction *>> SortedDeadInstrCandidates;
  // The dead instrs could span BBs, so we need to collect and sort them per BB.
  for (auto *V : DeadInstrCandidates) {
    auto *DeadI = cast<Instruction>(V);
    SortedDeadInstrCandidates[DeadI->getParent()].push_back(DeadI);
  }
  for (auto &Pair : SortedDeadInstrCandidates)
    sort(Pair.second,
         [](Instruction *I1, Instruction *I2) { return I1->comesBefore(I2); });
  for (const auto &Pair : SortedDeadInstrCandidates) {
    for (Instruction *I : reverse(Pair.second)) {
      if (I->hasNUses(0)) {
        // Erase the dead instructions bottom-to-top.
        LLVM_DEBUG(dbgs() << DEBUG_PREFIX << "Erase dead: " << *I << "\n");
        I->eraseFromParent();
      }
    }
  }
  DeadInstrCandidates.clear();
}

#ifndef NDEBUG
template <typename T> static void dumpImpl(ArrayRef<T *> Bndl) {
  for (auto [Idx, V] : enumerate(Bndl))
    dbgs() << Idx << "." << *V << "\n";
}
void VecUtils::dump(ArrayRef<Value *> Bndl) { dumpImpl(Bndl); }
void VecUtils::dump(ArrayRef<Instruction *> Bndl) { dumpImpl(Bndl); }

template <typename T> void BndlRef<T>::dump() const {
  print(dbgs());
  dbgs() << "\n";
}
// Explicit instantiation for commonly used types.
template class BndlRef<Instruction *>;
template class BndlRef<Value *>;

#endif // NDEBUG

} // namespace llvm::sandboxir