[llvm] [SandboxVec][Legality] Reject non-instructions (PR #113190)

[via llvm-commits]

https://github.com/vporpo created https://github.com/llvm/llvm-project/pull/113190

None

>From 5dd53704ff1a82714116b1fc67fb4d39fea6f76f Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Thu, 17 Oct 2024 09:13:00 -0700
Subject: [PATCH] [SandboxVec][Legality] Reject non-instructions

---
 .../Vectorize/SandboxVectorizer/Legality.h     | 10 +++++++++-
 .../Vectorize/SandboxVectorizer/Legality.cpp   | 18 +++++++++++++++++-
 .../SandboxVectorizer/Passes/BottomUpVec.cpp   |  2 +-
 .../SandboxVectorizer/LegalityTest.cpp         | 13 ++++++++++++-
 4 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h
index 233abf3efd64e1..11d72ff6b771e6 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h
@@ -28,6 +28,7 @@ enum class LegalityResultID {
 
 /// The reason for vectorizing or not vectorizing.
 enum class ResultReason {
+  NotInstructions,
   DiffOpcodes,
   DiffTypes,
 };
@@ -45,6 +46,8 @@ struct ToStr {
 
   static const char *getVecReason(ResultReason Reason) {
     switch (Reason) {
+    case ResultReason::NotInstructions:
+      return "NotInstructions";
     case ResultReason::DiffOpcodes:
       return "DiffOpcodes";
     case ResultReason::DiffTypes:
@@ -65,6 +68,10 @@ class LegalityResult {
   LegalityResult(LegalityResultID ID) : ID(ID) {}
   friend class LegalityAnalysis;
 
+  /// We shouldn't need copies.
+  LegalityResult(const LegalityResult &) = delete;
+  LegalityResult &operator=(const LegalityResult &) = delete;
+
 public:
   virtual ~LegalityResult() {}
   LegalityResultID getSubclassID() const { return ID; }
@@ -88,6 +95,7 @@ class LegalityResultWithReason : public LegalityResult {
   friend class Pack; // For constructor.
 
 public:
+  ResultReason getReason() const { return Reason; }
 #ifndef NDEBUG
   void print(raw_ostream &OS) const override {
     LegalityResult::print(OS);
@@ -136,7 +144,7 @@ class LegalityAnalysis {
   }
   /// Checks if it's legal to vectorize the instructions in \p Bndl.
   /// \Returns a LegalityResult object owned by LegalityAnalysis.
-  LegalityResult &canVectorize(ArrayRef<Value *> Bndl);
+  const LegalityResult &canVectorize(ArrayRef<Value *> Bndl);
 };
 
 } // namespace llvm::sandboxir
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Legality.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Legality.cpp
index 0e2cd83c37b0cd..f1c4577cece78a 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Legality.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Legality.cpp
@@ -7,11 +7,15 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h"
+#include "llvm/SandboxIR/Instruction.h"
+#include "llvm/SandboxIR/Utils.h"
 #include "llvm/SandboxIR/Value.h"
 #include "llvm/Support/Debug.h"
 
 namespace llvm::sandboxir {
 
+#define DEBUG_TYPE "SBVec:Legality"
+
 #ifndef NDEBUG
 void LegalityResult::dump() const {
   print(dbgs());
@@ -26,7 +30,19 @@ LegalityAnalysis::notVectorizableBasedOnOpcodesAndTypes(
   return std::nullopt;
 }
 
-LegalityResult &LegalityAnalysis::canVectorize(ArrayRef<Value *> Bndl) {
+static void dumpBndl(ArrayRef<Value *> Bndl) {
+  for (auto *V : Bndl)
+    dbgs() << *V << "\n";
+}
+
+const LegalityResult &LegalityAnalysis::canVectorize(ArrayRef<Value *> Bndl) {
+  // If Bndl contains values other than instructions, we need to Pack.
+  if (any_of(Bndl, [](auto *V) { return !isa<Instruction>(V); })) {
+    LLVM_DEBUG(dbgs() << "Not vectorizing: Not Instructions:\n";
+               dumpBndl(Bndl););
+    return createLegalityResult<Pack>(ResultReason::NotInstructions);
+  }
+
   if (auto ReasonOpt = notVectorizableBasedOnOpcodesAndTypes(Bndl))
     return createLegalityResult<Pack>(*ReasonOpt);
 
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index f11420e47f3e1f..ede41cd661b559 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -40,7 +40,7 @@ static SmallVector<Value *, 4> getOperand(ArrayRef<Value *> Bndl,
 }
 
 void BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl) {
-  auto LegalityRes = Legality.canVectorize(Bndl);
+  const auto &LegalityRes = Legality.canVectorize(Bndl);
   switch (LegalityRes.getSubclassID()) {
   case LegalityResultID::Widen: {
     auto *I = cast<Instruction>(Bndl[0]);
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp
index 76e5a5ce5aed92..56c6bf5f1ef1f5 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp
@@ -52,8 +52,16 @@ define void @foo(ptr %ptr) {
   auto *St1 = cast<sandboxir::StoreInst>(&*It++);
 
   sandboxir::LegalityAnalysis Legality;
-  auto Result = Legality.canVectorize({St0, St1});
+  const auto &Result = Legality.canVectorize({St0, St1});
   EXPECT_TRUE(isa<sandboxir::Widen>(Result));
+
+  {
+    // Check NotInstructions
+    auto &Result = Legality.canVectorize({F, St0});
+    EXPECT_TRUE(isa<sandboxir::Pack>(Result));
+    EXPECT_EQ(cast<sandboxir::Pack>(Result).getReason(),
+              sandboxir::ResultReason::NotInstructions);
+  }
 }
 
 #ifndef NDEBUG
@@ -68,6 +76,9 @@ TEST_F(LegalityTest, LegalityResultDump) {
   sandboxir::LegalityAnalysis Legality;
   EXPECT_TRUE(
       Matches(Legality.createLegalityResult<sandboxir::Widen>(), "Widen"));
+  EXPECT_TRUE(Matches(Legality.createLegalityResult<sandboxir::Pack>(
+                          sandboxir::ResultReason::NotInstructions),
+                      "Pack Reason: NotInstructions"));
   EXPECT_TRUE(Matches(Legality.createLegalityResult<sandboxir::Pack>(
                           sandboxir::ResultReason::DiffOpcodes),
                       "Pack Reason: DiffOpcodes"));