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

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

Author: vporpo

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h

@@ -13,16 +13,47 @@
 #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 {
 
 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), Reason(Reason) {}
+  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 a ResultReason
+  /// object if not vectorizable, or nullptr otherwise.
+  std::optional<ResultReason>
+  notVectorizableBasedOnOpcodesAndTypes(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

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

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

@@ -0,0 +1,39 @@
+//===- Legality.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/Legality.h"
+#include "llvm/SandboxIR/Value.h"
+#include "llvm/Support/Debug.h"
+
+namespace llvm::sandboxir {
+
+#ifndef NDEBUG
+void LegalityResult::dump() const {
+  print(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
+
+std::optional<ResultReason>
+LegalityAnalysis::notVectorizableBasedOnOpcodesAndTypes(
+    ArrayRef<Value *> Bndl) {
+  // TODO: Unimplemented.
+  return std::nullopt;
+}
+
+LegalityResult &LegalityAnalysis::canVectorize(ArrayRef<Value *> Bndl) {
+  if (auto ReasonOpt = notVectorizableBasedOnOpcodesAndTypes(Bndl))
+    return createLegalityResult<Pack>(*ReasonOpt);
+
+  // TODO: Check for existing vectors containing values in Bndl.
+
+  // TODO: Check with scheduler.
+
+  return createLegalityResult<Widen>();
+}
+} // namespace llvm::sandboxir

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");
+  }
   }
 }
 

llvm/unittests/Transforms/Vectorize/SandboxVectorizer/LegalityTest.cpp

@@ -55,3 +55,24 @@ define void @foo(ptr %ptr) {
   auto Result = Legality.canVectorize({St0, St1});
   EXPECT_TRUE(isa<sandboxir::Widen>(Result));
 }
+
+#ifndef NDEBUG
+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"));
+}
+#endif // NDEBUG