[llvm] [SandboxVec][Legality] Scaffolding for Legality (PR #112623)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: vporpo (vporpo)

<details>
<summary>Changes</summary>

This patch adds a LegalityResultWithReason class for describing the reason why legality decided not to vectorize the code.

---
Full diff: https://github.com/llvm/llvm-project/pull/112623.diff


4 Files Affected:

- (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h (+84-3) 
- (modified) llvm/lib/Transforms/Vectorize/CMakeLists.txt (+1) 
- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp (+4) 
- (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp (+19) 


``````````diff
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h
index 50fa56c5b21940..dc82274631a8c1 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h
@@ -13,6 +13,8 @@
 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_LEGALITY_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace llvm::sandboxir {
 
@@ -20,9 +22,38 @@ class LegalityAnalysis;
 class Value;
 
 enum class LegalityResultID {
+  Pack,  ///> Collect scalar values.
   Widen, ///> Vectorize by combining scalars to a vector.
 };
 
+/// The reason for vectorizing or not vectorizing.
+enum class ResultReason {
+  DiffOpcodes,
+  DiffTypes,
+};
+
+#ifndef NDEBUG
+struct ToStr {
+  static const char *getLegalityResultID(LegalityResultID ID) {
+    switch (ID) {
+    case LegalityResultID::Pack:
+      return "Pack";
+    case LegalityResultID::Widen:
+      return "Widen";
+    }
+  }
+
+  static const char *getVecReason(ResultReason Reason) {
+    switch (Reason) {
+    case ResultReason::DiffOpcodes:
+      return "DiffOpcodes";
+    case ResultReason::DiffTypes:
+      return "DiffTypes";
+    }
+  }
+};
+#endif // NDEBUG
+
 /// The legality outcome is represented by a class rather than an enum class
 /// because in some cases the legality checks are expensive and look for a
 /// particular instruction that can be passed along to the vectorizer to avoid
@@ -35,7 +66,34 @@ class LegalityResult {
   friend class LegalityAnalysis;
 
 public:
+  virtual ~LegalityResult() {}
   LegalityResultID getSubclassID() const { return ID; }
+#ifndef NDEBUG
+  virtual void print(raw_ostream &OS) const {
+    OS << ToStr::getLegalityResultID(ID);
+  }
+  LLVM_DUMP_METHOD void dump() const;
+  friend raw_ostream &operator<<(raw_ostream &OS, const LegalityResult &LR) {
+    LR.print(OS);
+    return OS;
+  }
+#endif // NDEBUG
+};
+
+/// Base class for results with reason.
+class LegalityResultWithReason : public LegalityResult {
+  ResultReason Reason;
+  LegalityResultWithReason(LegalityResultID ID, ResultReason Reason)
+      : LegalityResult(ID) {}
+  friend class Pack; // For constructor.
+
+public:
+#ifndef NDEBUG
+  void print(raw_ostream &OS) const override {
+    LegalityResult::print(OS);
+    OS << " Reason: " << ToStr::getVecReason(Reason);
+  }
+#endif
 };
 
 class Widen final : public LegalityResult {
@@ -48,14 +106,37 @@ class Widen final : public LegalityResult {
   }
 };
 
+class Pack final : public LegalityResultWithReason {
+  Pack(ResultReason Reason)
+      : LegalityResultWithReason(LegalityResultID::Pack, Reason) {}
+  friend class LegalityAnalysis; // For constructor.
+
+public:
+  static bool classof(const LegalityResult *From) {
+    return From->getSubclassID() == LegalityResultID::Pack;
+  }
+};
+
 /// Performs the legality analysis and returns a LegalityResult object.
 class LegalityAnalysis {
+  /// Owns the legality result objects created by createLegalityResult().
+  SmallVector<std::unique_ptr<LegalityResult>> ResultPool;
+  /// Checks opcodes, types and other IR-specifics and returns nullopt if it is
+  /// not illegal to vectorize, otherwise it returns a ResultReason object.
+  std::optional<ResultReason>
+  cantVectorizeBasedOnOpcodesAndTypes(ArrayRef<Value *> Bndl);
+
 public:
   LegalityAnalysis() = default;
-  LegalityResult canVectorize(ArrayRef<Value *> Bndl) {
-    // TODO: For now everything is legal.
-    return Widen();
+  /// A LegalityResult factory.
+  template <typename ResultT, typename... ArgsT>
+  ResultT &createLegalityResult(ArgsT... Args) {
+    ResultPool.push_back(std::unique_ptr<ResultT>(new ResultT(Args...)));
+    return cast<ResultT>(*ResultPool.back());
   }
+  /// Checks if it's legal to vectorize the instructions in \p Bndl.
+  /// \Returns a LegalityResult object owned by LegalityAnalysis.
+  LegalityResult &canVectorize(ArrayRef<Value *> Bndl);
 };
 
 } // namespace llvm::sandboxir
diff --git a/llvm/lib/Transforms/Vectorize/CMakeLists.txt b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
index f4e98e576379a4..bffe57c66c4ea4 100644
--- a/llvm/lib/Transforms/Vectorize/CMakeLists.txt
+++ b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
@@ -5,6 +5,7 @@ add_llvm_component_library(LLVMVectorize
   LoopVectorize.cpp
   SandboxVectorizer/DependencyGraph.cpp
   SandboxVectorizer/Interval.cpp
+  SandboxVectorizer/Legality.cpp
   SandboxVectorizer/Passes/BottomUpVec.cpp
   SandboxVectorizer/Passes/RegionsFromMetadata.cpp
   SandboxVectorizer/SandboxVectorizer.cpp
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index 6171d5e52b5869..f11420e47f3e1f 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -50,6 +50,10 @@ void BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl) {
     }
     break;
   }
+  case LegalityResultID::Pack: {
+    // TODO: Unimplemented
+    llvm_unreachable("Unimplemented");
+  }
   }
 }
 
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp
index e16222ddb2d615..26357fec8ab8c5 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp
@@ -55,3 +55,22 @@ define void @foo(ptr %ptr) {
   auto Result = Legality.canVectorize({St0, St1});
   EXPECT_TRUE(isa<sandboxir::Widen>(Result));
 }
+
+TEST_F(LegalityTest, LegalityResultDump) {
+  auto Matches = [](const sandboxir::LegalityResult &Result,
+                    const std::string &ExpectedStr) -> bool {
+    std::string Buff;
+    raw_string_ostream OS(Buff);
+    Result.print(OS);
+    return Buff == ExpectedStr;
+  };
+  sandboxir::LegalityAnalysis Legality;
+  EXPECT_TRUE(
+      Matches(Legality.createLegalityResult<sandboxir::Widen>(), "Widen"));
+  EXPECT_TRUE(Matches(Legality.createLegalityResult<sandboxir::Pack>(
+                          sandboxir::ResultReason::DiffOpcodes),
+                      "Pack Reason: DiffOpcodes"));
+  EXPECT_TRUE(Matches(Legality.createLegalityResult<sandboxir::Pack>(
+                          sandboxir::ResultReason::DiffTypes),
+                      "Pack Reason: DiffTypes"));
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/112623