llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp

//===- SeedCollection.cpp - Seed collection pass --------------------------===//
//
// 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/Passes/SeedCollection.h"
#include "../../VectorizeOptions.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/SandboxIR/Module.h"
#include "llvm/Transforms/Vectorize/SandboxVectorizer/RegionWithScore.h"
#include "llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.h"
#include "llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h"
#include "llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h"

namespace llvm {

namespace sandboxir {

SeedCollection::SeedCollection(StringRef Pipeline, StringRef AuxArg)
    : FunctionPass("seed-collection"),
      RPM("rpm", Pipeline, SandboxVectorizerPassBuilder::createRegionPass) {
  ArgsRegistry.parse(AuxArg);
  // If the user has not specified seed types default to stores.
  if (!CollectStores && !CollectLoads) {
    CollectStores = true;
  }
}

bool SeedCollection::runOnFunction(Function &F, const Analyses &A) {
  const VectorizeOptions &Opts =
      F.getContext().getLibraryOptions().get<VectorizeOptions>();
  bool Change = false;
  const auto &DL = F.getParent()->getDataLayout();

  // TODO: Start from innermost BBs first
  for (auto &BB : F) {
    SeedCollector SC(&BB, A.getScalarEvolution(), CollectStores, CollectLoads,
                     AllowDiffTypes);
    for (auto &SeedRange : {SC.getStoreSeeds(), SC.getLoadSeeds()}) {
      for (SeedBundle &Seeds : SeedRange) {
        if (Seeds.allUsed())
          continue;
        unsigned FirstUnusedIdx = Seeds.getFirstUnusedElementIdx();
        unsigned ElmBits =
            Utils::getNumBits(VecUtils::getElementType(Utils::getExpectedType(
                                  Seeds[FirstUnusedIdx])),
                              DL);
        unsigned AS = getLoadStoreAddressSpace(Seeds[FirstUnusedIdx]);
        unsigned VecRegBits = Opts.sbvec_vec_reg_bits != 0
                                  ? Opts.sbvec_vec_reg_bits
                                  : A.getTTI().getLoadStoreVecRegBitWidth(AS);

        auto DivideBy2 = [](unsigned Num) {
          auto Floor = VecUtils::getFloorPowerOf2(Num);
          if (Floor == Num)
            return Floor / 2;
          return Floor;
        };
        // Try to create the largest vector supported by the target. If it fails
        // reduce the vector size by half.
        for (unsigned SliceElms = std::min(VecRegBits / ElmBits,
                                           Seeds.getNumUnusedBits() / ElmBits);
             SliceElms >= 2u; SliceElms = DivideBy2(SliceElms)) {
          if (Seeds.allUsed())
            break;
          // Keep trying offsets after FirstUnusedElementIdx, until we vectorize
          // the slice. This could be quite expensive, so we enforce a limit.
          for (unsigned Offset = Seeds.getFirstUnusedElementIdx(),
                        OE = Seeds.size();
               Offset + 1 < OE; Offset += 1) {
            // Seeds are getting used as we vectorize, so skip them.
            if (Seeds.isUsed(Offset))
              continue;
            if (Seeds.allUsed())
              break;

            auto SeedSlice = Seeds.getSlice(Offset, SliceElms * ElmBits,
                                            !Opts.sbvec_allow_non_pow2);
            if (SeedSlice.empty())
              continue;

            assert(SeedSlice.size() >= 2 && "Should have been rejected!");

            // Create a region containing the seed slice.
            auto &Ctx = F.getContext();
            RegionWithScore Rgn(Ctx, A.getTTI());
            Rgn.setAux(SeedSlice);
            // Run the region pass pipeline.
            Change |= RPM.runOnRegion(Rgn, A);
            Rgn.clearAux();
          }
        }
      }
    }
  }
  return Change;
}
} // namespace sandboxir
} // namespace llvm