[VPlan] Add scalar inferencing support for Not and Or insns (#89160)

Fixes #87394.

PR: https://github.com/llvm/llvm-project/pull/89160

2 files changed, 71 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
index ad4ea648cd61..5f93339083f0 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
@@ -35,6 +35,7 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
     CachedTypes[OtherV] = ResTy;
     return ResTy;
   }
+  case Instruction::Or:
   case Instruction::ICmp:
   case VPInstruction::FirstOrderRecurrenceSplice: {
     Type *ResTy = inferScalarType(R->getOperand(0));
@@ -44,6 +45,12 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
     CachedTypes[OtherV] = ResTy;
     return ResTy;
   }
+  case VPInstruction::Not: {
+    Type *ResTy = inferScalarType(R->getOperand(0));
+    assert(IntegerType::get(Ctx, 1) == ResTy &&
+           "unexpected scalar type inferred for operand");
+    return ResTy;
+  }
   case VPInstruction::PtrAdd:
     // Return the type based on the pointer argument (i.e. first operand).
     return inferScalarType(R->getOperand(0));
diff --git a/llvm/test/Transforms/LoopVectorize/vplan-infer-not-or-type.ll b/llvm/test/Transforms/LoopVectorize/vplan-infer-not-or-type.ll
new file mode 100644
index 000000000000..102ef699cb37
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/vplan-infer-not-or-type.ll
@@ -0,0 +1,64 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=loop-vectorize -S | FileCheck %s
+
+; This test used to crash due to missing Or/Not cases in
+; inferScalarTypeForRecipe.
+
+define void @foo(i8 %arg.0, i8 %arg.1) {
+; CHECK-LABEL: define void @foo(
+; CHECK-SAME: i8 [[ARG_0:%.*]], i8 [[ARG_1:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 false, label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
+; CHECK:       vector.ph:
+; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i32 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i32 [[INDEX]], 2
+; CHECK-NEXT:    br i1 true, label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK:       middle.block:
+; CHECK-NEXT:    br i1 true, label [[LOOPEXIT:%.*]], label [[SCALAR_PH]]
+; CHECK:       scalar.ph:
+; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i8 [ 2, [[MIDDLE_BLOCK]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    br label [[LOOP_HEADER:%.*]]
+; CHECK:       loop.header:
+; CHECK-NEXT:    [[INCREMENTOR:%.*]] = phi i8 [ [[ADD:%.*]], [[LATCH:%.*]] ], [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ]
+; CHECK-NEXT:    [[AND:%.*]] = and i8 [[ARG_0]], [[ARG_1]]
+; CHECK-NEXT:    [[EXTRACT_T:%.*]] = trunc i8 [[AND]] to i1
+; CHECK-NEXT:    br i1 [[EXTRACT_T]], label [[LATCH]], label [[INDIRECT_LATCH:%.*]]
+; CHECK:       indirect.latch:
+; CHECK-NEXT:    br label [[LATCH]]
+; CHECK:       latch:
+; CHECK-NEXT:    [[ADD]] = add i8 [[INCREMENTOR]], 1
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[INCREMENTOR]] to i32
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[CONV]], 1
+; CHECK-NEXT:    br i1 [[CMP]], label [[LOOP_HEADER]], label [[LOOPEXIT]], !llvm.loop [[LOOP3:![0-9]+]]
+; CHECK:       loop.exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop.header
+
+loop.header:                                         ; preds = %latch, %entry
+  %incrementor = phi i8 [ %add, %latch ], [ 0, %entry ]
+  %and = and i8 %arg.0, %arg.1
+  %extract.t = trunc i8 %and to i1
+  br i1 %extract.t, label %latch, label %indirect.latch
+
+indirect.latch:                                       ; preds = %loop.header
+  br label %latch
+
+latch:                               ; preds = %loop.header16, %loop.header
+  %add = add i8 %incrementor, 1
+  %conv = zext i8 %incrementor to i32
+  %cmp = icmp ult i32 %conv, 1
+  br i1 %cmp, label %loop.header, label %loop.exit
+
+loop.exit:
+  ret void
+}
+;.
+; CHECK: [[LOOP0]] = distinct !{[[LOOP0]], [[META1:![0-9]+]], [[META2:![0-9]+]]}
+; CHECK: [[META1]] = !{!"llvm.loop.isvectorized", i32 1}
+; CHECK: [[META2]] = !{!"llvm.loop.unroll.runtime.disable"}
+; CHECK: [[LOOP3]] = distinct !{[[LOOP3]], [[META2]], [[META1]]}
+;.