[VPlan] Add scalar inferencing support for Not and Or insns

[patrick-rivos]

Fixes #87394.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Patrick O'Neill (patrick-rivos)

Changes

Fixes #87394.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp (+3-1)
• (added) llvm/test/Transforms/LoopVectorize/vplan-infer-not-or-type.ll (+61)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
index 130fb04f586e75..d38e5d23907174 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,8 +45,9 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPInstruction *R) {
     CachedTypes[OtherV] = ResTy;
     return ResTy;
   }
+  case VPInstruction::Not:
   case VPInstruction::PtrAdd:
-    // Return the type based on the pointer argument (i.e. first operand).
+    // Return the type based on the pointer/not argument (i.e. first operand).
     return inferScalarType(R->getOperand(0));
   default:
     break;
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 00000000000000..488269c2eb6155
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/vplan-infer-not-or-type.ll
@@ -0,0 +1,61 @@
+; 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 i32 @foo() {
+; CHECK-LABEL: define i32 @foo() {
+; 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 i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[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 [[FOR_COND_CLEANUP_LOOPEXIT:%.*]], label [[SCALAR_PH]]
+; CHECK:       scalar.ph:
+; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ 2, [[MIDDLE_BLOCK]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.cond.cleanup.loopexit:
+; CHECK-NEXT:    ret i32 0
+; CHECK:       for.body:
+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[INDVARS_IV_NEXT:%.*]], [[COND_END:%.*]] ]
+; CHECK-NEXT:    [[ZEXT_0:%.*]] = zext i1 false to i64
+; CHECK-NEXT:    [[DEAD_INSN:%.*]] = trunc i64 [[ZEXT_0]] to i16
+; CHECK-NEXT:    br i1 false, label [[COND_FALSE:%.*]], label [[COND_END]]
+; CHECK:       cond.false:
+; CHECK-NEXT:    br label [[COND_END]]
+; CHECK:       cond.end:
+; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add i64 [[INDVARS_IV]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[INDVARS_IV]], 1
+; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_COND_CLEANUP_LOOPEXIT]], !llvm.loop [[LOOP3:![0-9]+]]
+;
+entry:
+  br label %for.body
+
+for.cond.cleanup.loopexit:                        ; preds = %cond.end
+  ret i32 0
+
+for.body:                                         ; preds = %cond.end, %entry
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %cond.end ]
+  %zext.0 = zext i1 false to i64
+  %dead.insn = trunc i64 %zext.0 to i16
+  br i1 false, label %cond.false, label %cond.end
+
+cond.false:                                       ; preds = %for.body
+  br label %cond.end
+
+cond.end:                                         ; preds = %cond.false, %for.body
+  %indvars.iv.next = add i64 %indvars.iv, 1
+  %cmp = icmp ult i64 %indvars.iv, 1
+  br i1 %cmp, label %for.body, label %for.cond.cleanup.loopexit
+}
+;.
+; 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]]}
+;.

[patrick-rivos]

VPlan so @fhahn
Please let me know if this isn't the right approach and I can dig deeper.
There are likely other missing cases but those can be added as the fuzzer finds them.

[fhahn]

Thanks for the patch!

[fhahn]

To make the test more robust w.r.t. future changes, would is it possible to reproduce the issue with a non-loop invariant/constant zext here?

[patrick-rivos]

Thanks for the review!

I split out the not and added an assert. Thankfully the fuzzer has quite a few finds of this ice so I was able to find another testcase that doesn't have any dead variables. The loop-constant seems to be a bit more challenging to get rid of.
I'll check the remaining testcases again just to make sure I didn't miss one.

EDIT: I don't know how I missed it - this can be resolved by passing in 2 values into the function and adding those together. I'll send another fixup shortly. IIUC still an invariant but no longer a constant.

[fhahn]

LGTM, thanks!

[patrick-rivos]

Thank you! I don't have write permission to llvm - could you land this for me?