[SandboxVec][Legality] Check wrap flags (#113975)

Author: vporpo

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

@@ -32,6 +32,7 @@ enum class ResultReason {
   DiffOpcodes,
   DiffTypes,
   DiffMathFlags,
+  DiffWrapFlags,
 };
 
 #ifndef NDEBUG
@@ -56,6 +57,8 @@ struct ToStr {
       return "DiffTypes";
     case ResultReason::DiffMathFlags:
       return "DiffMathFlags";
+    case ResultReason::DiffWrapFlags:
+      return "DiffWrapFlags";
     }
     llvm_unreachable("Unknown ResultReason enum");
   }

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

@@ -55,6 +55,21 @@ LegalityAnalysis::notVectorizableBasedOnOpcodesAndTypes(
       return ResultReason::DiffMathFlags;
   }
 
+  // TODO: Allow vectorization by using common flags.
+  // For now Pack if they don't have the same wrap flags.
+  bool CanHaveWrapFlags =
+      isa<OverflowingBinaryOperator>(I0) || isa<TruncInst>(I0);
+  if (CanHaveWrapFlags) {
+    bool NUW0 = I0->hasNoUnsignedWrap();
+    bool NSW0 = I0->hasNoSignedWrap();
+    if (any_of(drop_begin(Bndl), [NUW0, NSW0](auto *V) {
+          return cast<Instruction>(V)->hasNoUnsignedWrap() != NUW0 ||
+                 cast<Instruction>(V)->hasNoSignedWrap() != NSW0;
+        })) {
+      return ResultReason::DiffWrapFlags;
+    }
+  }
+
   // TODO: Missing checks
 
   return std::nullopt;

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

@@ -29,7 +29,7 @@ struct LegalityTest : public testing::Test {
 
 TEST_F(LegalityTest, Legality) {
   parseIR(C, R"IR(
-define void @foo(ptr %ptr, <2 x float> %vec2, <3 x float> %vec3, i8 %arg, float %farg0, float %farg1) {
+define void @foo(ptr %ptr, <2 x float> %vec2, <3 x float> %vec3, i8 %arg, float %farg0, float %farg1, i64 %v0, i64 %v1) {
   %gep0 = getelementptr float, ptr %ptr, i32 0
   %gep1 = getelementptr float, ptr %ptr, i32 1
   %gep3 = getelementptr float, ptr %ptr, i32 3
@@ -42,6 +42,8 @@ define void @foo(ptr %ptr, <2 x float> %vec2, <3 x float> %vec3, i8 %arg, float
   store i8 %arg, ptr %gep1
   %fadd0 = fadd float %farg0, %farg0
   %fadd1 = fadd fast float %farg1, %farg1
+  %trunc0 = trunc nuw nsw i64 %v0 to i8
+  %trunc1 = trunc nsw i64 %v1 to i8
   ret void
 }
 )IR");
@@ -62,6 +64,8 @@ define void @foo(ptr %ptr, <2 x float> %vec2, <3 x float> %vec3, i8 %arg, float
   auto *StI8 = cast<sandboxir::StoreInst>(&*It++);
   auto *FAdd0 = cast<sandboxir::BinaryOperator>(&*It++);
   auto *FAdd1 = cast<sandboxir::BinaryOperator>(&*It++);
+  auto *Trunc0 = cast<sandboxir::TruncInst>(&*It++);
+  auto *Trunc1 = cast<sandboxir::TruncInst>(&*It++);
 
   sandboxir::LegalityAnalysis Legality;
   const auto &Result = Legality.canVectorize({St0, St1});
@@ -98,6 +102,13 @@ define void @foo(ptr %ptr, <2 x float> %vec2, <3 x float> %vec3, i8 %arg, float
     EXPECT_EQ(cast<sandboxir::Pack>(Result).getReason(),
               sandboxir::ResultReason::DiffMathFlags);
   }
+  {
+    // Check DiffWrapFlags
+    const auto &Result = Legality.canVectorize({Trunc0, Trunc1});
+    EXPECT_TRUE(isa<sandboxir::Pack>(Result));
+    EXPECT_EQ(cast<sandboxir::Pack>(Result).getReason(),
+              sandboxir::ResultReason::DiffWrapFlags);
+  }
 }
 
 #ifndef NDEBUG
@@ -124,5 +135,8 @@ TEST_F(LegalityTest, LegalityResultDump) {
   EXPECT_TRUE(Matches(Legality.createLegalityResult<sandboxir::Pack>(
                           sandboxir::ResultReason::DiffMathFlags),
                       "Pack Reason: DiffMathFlags"));
+  EXPECT_TRUE(Matches(Legality.createLegalityResult<sandboxir::Pack>(
+                          sandboxir::ResultReason::DiffWrapFlags),
+                      "Pack Reason: DiffWrapFlags"));
 }
 #endif // NDEBUG