diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h index 286d971ff4851..f55dbf753dbcc 100644 --- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h +++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h @@ -25,6 +25,11 @@ class SeedCollection final : public FunctionPass { /// The PM containing the pipeline of region passes. RegionPassManager RPM; + /// The auxiliary argument passed to the pass that tells us that we should + /// collect seeds of different types. + static constexpr StringRef DiffTypesArgStr = "enable-diff-types"; + /// Collect seeds of different types. + bool AllowDiffTypes = false; public: SeedCollection(StringRef Pipeline); diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h index d32bfbaf7a4c8..76752872f4d5c 100644 --- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h +++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h @@ -262,6 +262,35 @@ class VecUtils { return Pack; } + static constexpr const char AuxArgBeginToken = '('; + static constexpr const char AuxArgEndToken = ')'; + /// \Returns the auxiliary pass argument by parsing \p Args. The auxiliary + /// pass argument is enclosed in parentheses and should be before any other + /// argument. + static StringRef getAuxPassArg(StringRef Args) { + if (Args.empty()) + return Args; + if (Args[0] != AuxArgBeginToken) + return StringRef(); + // We found the Begin token, so look for the End token. + size_t EndIdx = Args.find(AuxArgEndToken); + if (EndIdx == StringRef::npos) { + errs() << "Missing '" << AuxArgEndToken << "' in '" << Args << "' !\n"; + exit(1); + } + assert(EndIdx >= 1 && "Expected at index 1 or later!"); + return Args.substr(1, EndIdx - 1); + } + + /// \Returns \p Args with the auxiliary argument removed. + /// For example: "(foo)bar1,bar2" returns "bar1,bar2". + static StringRef stripAuxPassArg(StringRef Args) { + StringRef Aux = getAuxPassArg(Args); + if (Aux.empty()) + return Args; + return Args.slice(Aux.size() + 2, StringRef::npos); + } + #ifndef NDEBUG /// Helper dump function for debugging. LLVM_DUMP_METHOD static void dump(ArrayRef Bndl); diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp index ddb4a1e154a18..e87e6426d3be6 100644 --- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp +++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp @@ -32,9 +32,20 @@ cl::opt CollectSeeds( "list of '" StoreSeedsDef "' and '" LoadSeedsDef "'.")); namespace sandboxir { + SeedCollection::SeedCollection(StringRef Pipeline) : FunctionPass("seed-collection"), - RPM("rpm", Pipeline, SandboxVectorizerPassBuilder::createRegionPass) {} + RPM("rpm", VecUtils::stripAuxPassArg(Pipeline), + SandboxVectorizerPassBuilder::createRegionPass) { + StringRef AuxArg = VecUtils::getAuxPassArg(Pipeline); + if (!AuxArg.empty()) { + if (DiffTypesArgStr != DiffTypesArgStr) { + errs() << "SeedCollection only supports '" << DiffTypesArgStr << "' !\n"; + exit(1); + } + AllowDiffTypes = true; + } +} bool SeedCollection::runOnFunction(Function &F, const Analyses &A) { bool Change = false; @@ -50,7 +61,8 @@ bool SeedCollection::runOnFunction(Function &F, const Analyses &A) { // TODO: Start from innermost BBs first for (auto &BB : F) { - SeedCollector SC(&BB, A.getScalarEvolution(), CollectStores, CollectLoads); + SeedCollector SC(&BB, A.getScalarEvolution(), CollectStores, CollectLoads, + AllowDiffTypes); for (SeedBundle &Seeds : SC.getStoreSeeds()) { unsigned ElmBits = Utils::getNumBits(VecUtils::getElementType(Utils::getExpectedType( diff --git a/llvm/test/Transforms/SandboxVectorizer/seed_collection.ll b/llvm/test/Transforms/SandboxVectorizer/seed_collection.ll new file mode 100644 index 0000000000000..0fc1fa04a081f --- /dev/null +++ b/llvm/test/Transforms/SandboxVectorizer/seed_collection.ll @@ -0,0 +1,19 @@ +; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5 +; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -disable-output -sbvec-passes="seed-collection<(enable-diff-types)print-region>" %s | FileCheck %s + +define void @collect_seeds_of_diff_types(ptr %ptr) { +; CHECK: Aux: +; CHECK-NEXT: store i8 %ld0, ptr %ptr0, {{.*}} +; CHECK-NEXT: store i16 %ld1, ptr %ptr1, {{.*}} +; CHECK-NEXT: store i8 %ld2, ptr %ptr2, {{.*}} + %ptr0 = getelementptr i8, ptr %ptr, i32 0 + %ptr1 = getelementptr i16, ptr %ptr, i32 1 + %ptr2 = getelementptr i8, ptr %ptr, i32 3 + %ld0 = load i8, ptr %ptr0 + %ld1 = load i16, ptr %ptr1 + %ld2 = load i8, ptr %ptr2 + store i8 %ld0, ptr %ptr0 + store i16 %ld1, ptr %ptr1 + store i8 %ld2, ptr %ptr2 + ret void +} diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp index 2bfea6908305c..566c31446ec20 100644 --- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp +++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp @@ -617,3 +617,30 @@ define void @foo(i8 %v0, i8 %v1) { EXPECT_FALSE(sandboxir::VecUtils::matchPack(NotPack)); } } + +TEST_F(VecUtilsTest, GetAuxPassArg) { + // Check no aux argument. + EXPECT_EQ(sandboxir::VecUtils::getAuxPassArg("no aux arg"), ""); + // Check an illegal aux argument. + EXPECT_EQ(sandboxir::VecUtils::getAuxPassArg("illegal (arg) other"), ""); + // Check a valid argument. + EXPECT_EQ(sandboxir::VecUtils::getAuxPassArg("(some arg)other stuff"), + "some arg"); + // Missing token. + EXPECT_DEBUG_DEATH(sandboxir::VecUtils::getAuxPassArg("(arg other"), + "Missing.*"); +} + +TEST_F(VecUtilsTest, StripAuxPassArg) { + // Check no aux argument. + EXPECT_EQ(sandboxir::VecUtils::stripAuxPassArg("no aux arg"), "no aux arg"); + // Check a legal aux argument. + EXPECT_EQ(sandboxir::VecUtils::stripAuxPassArg("(legal aux arg)foo bar"), + "foo bar"); + // Check an illegal aux argument. + EXPECT_EQ(sandboxir::VecUtils::stripAuxPassArg("illegal (arg) other"), + "illegal (arg) other"); + // Missing token. + EXPECT_DEBUG_DEATH(sandboxir::VecUtils::stripAuxPassArg("(arg other"), + "Missing.*"); +}