[SandboxVec][Legality] Per opcode checks (llvm#114145)

This patch adds more opcode-specific legality checks.

Author: vporpo

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/Analysis/ScalarEvolution.h"
+#include "llvm/IR/DataLayout.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -33,6 +35,9 @@ enum class ResultReason {
   DiffTypes,
   DiffMathFlags,
   DiffWrapFlags,
+  NotConsecutive,
+  Unimplemented,
+  Infeasible,
 };
 
 #ifndef NDEBUG
@@ -59,6 +64,12 @@ struct ToStr {
       return "DiffMathFlags";
     case ResultReason::DiffWrapFlags:
       return "DiffWrapFlags";
+    case ResultReason::NotConsecutive:
+      return "NotConsecutive";
+    case ResultReason::Unimplemented:
+      return "Unimplemented";
+    case ResultReason::Infeasible:
+      return "Infeasible";
     }
     llvm_unreachable("Unknown ResultReason enum");
   }
@@ -142,8 +153,12 @@ class LegalityAnalysis {
   std::optional<ResultReason>
   notVectorizableBasedOnOpcodesAndTypes(ArrayRef<Value *> Bndl);
 
+  ScalarEvolution &SE;
+  const DataLayout &DL;
+
 public:
-  LegalityAnalysis() = default;
+  LegalityAnalysis(ScalarEvolution &SE, const DataLayout &DL)
+      : SE(SE), DL(DL) {}
   /// A LegalityResult factory.
   template <typename ResultT, typename... ArgsT>
   ResultT &createLegalityResult(ArgsT... Args) {

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h

@@ -24,7 +24,7 @@ namespace llvm::sandboxir {
 
 class BottomUpVec final : public FunctionPass {
   bool Change = false;
-  LegalityAnalysis Legality;
+  std::unique_ptr<LegalityAnalysis> Legality;
   void vectorizeRec(ArrayRef<Value *> Bndl);
   void tryVectorize(ArrayRef<Value *> Seeds);
 

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

@@ -12,7 +12,10 @@
 #ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_VECUTILS_H
 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_VECUTILS_H
 
+#include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/IR/DataLayout.h"
 #include "llvm/SandboxIR/Type.h"
+#include "llvm/SandboxIR/Utils.h"
 
 namespace llvm::sandboxir {
 
@@ -29,6 +32,40 @@ class VecUtils {
   static Type *getElementType(Type *Ty) {
     return Ty->isVectorTy() ? cast<FixedVectorType>(Ty)->getElementType() : Ty;
   }
+
+  /// \Returns true if \p I1 and \p I2 are load/stores accessing consecutive
+  /// memory addresses.
+  template <typename LoadOrStoreT>
+  static bool areConsecutive(LoadOrStoreT *I1, LoadOrStoreT *I2,
+                             ScalarEvolution &SE, const DataLayout &DL) {
+    static_assert(std::is_same<LoadOrStoreT, LoadInst>::value ||
+                      std::is_same<LoadOrStoreT, StoreInst>::value,
+                  "Expected Load or Store!");
+    auto Diff = Utils::getPointerDiffInBytes(I1, I2, SE);
+    if (!Diff)
+      return false;
+    int ElmBytes = Utils::getNumBits(I1) / 8;
+    return *Diff == ElmBytes;
+  }
+
+  template <typename LoadOrStoreT>
+  static bool areConsecutive(ArrayRef<Value *> &Bndl, ScalarEvolution &SE,
+                             const DataLayout &DL) {
+    static_assert(std::is_same<LoadOrStoreT, LoadInst>::value ||
+                      std::is_same<LoadOrStoreT, StoreInst>::value,
+                  "Expected Load or Store!");
+    assert(isa<LoadOrStoreT>(Bndl[0]) && "Expected Load or Store!");
+    auto *LastLS = cast<LoadOrStoreT>(Bndl[0]);
+    for (Value *V : drop_begin(Bndl)) {
+      assert(isa<LoadOrStoreT>(V) &&
+             "Unimplemented: we only support StoreInst!");
+      auto *LS = cast<LoadOrStoreT>(V);
+      if (!VecUtils::areConsecutive(LastLS, LS, SE, DL))
+        return false;
+      LastLS = LS;
+    }
+    return true;
+  }
 };
 
 } // namespace llvm::sandboxir

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

@@ -70,7 +70,109 @@ LegalityAnalysis::notVectorizableBasedOnOpcodesAndTypes(
     }
   }
 
-  // TODO: Missing checks
+  // Now we need to do further checks for specific opcodes.
+  switch (Opcode) {
+  case Instruction::Opcode::ZExt:
+  case Instruction::Opcode::SExt:
+  case Instruction::Opcode::FPToUI:
+  case Instruction::Opcode::FPToSI:
+  case Instruction::Opcode::FPExt:
+  case Instruction::Opcode::PtrToInt:
+  case Instruction::Opcode::IntToPtr:
+  case Instruction::Opcode::SIToFP:
+  case Instruction::Opcode::UIToFP:
+  case Instruction::Opcode::Trunc:
+  case Instruction::Opcode::FPTrunc:
+  case Instruction::Opcode::BitCast: {
+    // We have already checked that they are of the same opcode.
+    assert(all_of(Bndl,
+                  [Opcode](Value *V) {
+                    return cast<Instruction>(V)->getOpcode() == Opcode;
+                  }) &&
+           "Different opcodes, should have early returned!");
+    // But for these opcodes we should also check the operand type.
+    Type *FromTy0 = Utils::getExpectedType(I0->getOperand(0));
+    if (any_of(drop_begin(Bndl), [FromTy0](Value *V) {
+          return Utils::getExpectedType(cast<User>(V)->getOperand(0)) !=
+                 FromTy0;
+        }))
+      return ResultReason::DiffTypes;
+    return std::nullopt;
+  }
+  case Instruction::Opcode::FCmp:
+  case Instruction::Opcode::ICmp: {
+    // We need the same predicate..
+    auto Pred0 = cast<CmpInst>(I0)->getPredicate();
+    bool Same = all_of(Bndl, [Pred0](Value *V) {
+      return cast<CmpInst>(V)->getPredicate() == Pred0;
+    });
+    if (Same)
+      return std::nullopt;
+    return ResultReason::DiffOpcodes;
+  }
+  case Instruction::Opcode::Select:
+  case Instruction::Opcode::FNeg:
+  case Instruction::Opcode::Add:
+  case Instruction::Opcode::FAdd:
+  case Instruction::Opcode::Sub:
+  case Instruction::Opcode::FSub:
+  case Instruction::Opcode::Mul:
+  case Instruction::Opcode::FMul:
+  case Instruction::Opcode::FRem:
+  case Instruction::Opcode::UDiv:
+  case Instruction::Opcode::SDiv:
+  case Instruction::Opcode::FDiv:
+  case Instruction::Opcode::URem:
+  case Instruction::Opcode::SRem:
+  case Instruction::Opcode::Shl:
+  case Instruction::Opcode::LShr:
+  case Instruction::Opcode::AShr:
+  case Instruction::Opcode::And:
+  case Instruction::Opcode::Or:
+  case Instruction::Opcode::Xor:
+    return std::nullopt;
+  case Instruction::Opcode::Load:
+    if (VecUtils::areConsecutive<LoadInst>(Bndl, SE, DL))
+      return std::nullopt;
+    return ResultReason::NotConsecutive;
+  case Instruction::Opcode::Store:
+    if (VecUtils::areConsecutive<StoreInst>(Bndl, SE, DL))
+      return std::nullopt;
+    return ResultReason::NotConsecutive;
+  case Instruction::Opcode::PHI:
+    return ResultReason::Unimplemented;
+  case Instruction::Opcode::Opaque:
+    return ResultReason::Unimplemented;
+  case Instruction::Opcode::Br:
+  case Instruction::Opcode::Ret:
+  case Instruction::Opcode::AddrSpaceCast:
+  case Instruction::Opcode::InsertElement:
+  case Instruction::Opcode::InsertValue:
+  case Instruction::Opcode::ExtractElement:
+  case Instruction::Opcode::ExtractValue:
+  case Instruction::Opcode::ShuffleVector:
+  case Instruction::Opcode::Call:
+  case Instruction::Opcode::GetElementPtr:
+  case Instruction::Opcode::Switch:
+    return ResultReason::Unimplemented;
+  case Instruction::Opcode::VAArg:
+  case Instruction::Opcode::Freeze:
+  case Instruction::Opcode::Fence:
+  case Instruction::Opcode::Invoke:
+  case Instruction::Opcode::CallBr:
+  case Instruction::Opcode::LandingPad:
+  case Instruction::Opcode::CatchPad:
+  case Instruction::Opcode::CleanupPad:
+  case Instruction::Opcode::CatchRet:
+  case Instruction::Opcode::CleanupRet:
+  case Instruction::Opcode::Resume:
+  case Instruction::Opcode::CatchSwitch:
+  case Instruction::Opcode::AtomicRMW:
+  case Instruction::Opcode::AtomicCmpXchg:
+  case Instruction::Opcode::Alloca:
+  case Instruction::Opcode::Unreachable:
+    return ResultReason::Infeasible;
+  }
 
   return std::nullopt;
 }

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

@@ -11,6 +11,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/SandboxIR/Function.h"
 #include "llvm/SandboxIR/Instruction.h"
+#include "llvm/SandboxIR/Module.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.h"
 
 namespace llvm::sandboxir {
@@ -40,7 +41,7 @@ static SmallVector<Value *, 4> getOperand(ArrayRef<Value *> Bndl,
 }
 
 void BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl) {
-  const auto &LegalityRes = Legality.canVectorize(Bndl);
+  const auto &LegalityRes = Legality->canVectorize(Bndl);
   switch (LegalityRes.getSubclassID()) {
   case LegalityResultID::Widen: {
     auto *I = cast<Instruction>(Bndl[0]);
@@ -60,6 +61,8 @@ void BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl) {
 void BottomUpVec::tryVectorize(ArrayRef<Value *> Bndl) { vectorizeRec(Bndl); }
 
 bool BottomUpVec::runOnFunction(Function &F, const Analyses &A) {
+  Legality = std::make_unique<LegalityAnalysis>(A.getScalarEvolution(),
+                                                F.getParent()->getDataLayout());
   Change = false;
   // TODO: Start from innermost BBs first
   for (auto &BB : F) {