From 91fd5c94965b4077490c6bfcc9aa4b0e4146b38a Mon Sep 17 00:00:00 2001
From: Tamar Christina <tamar.christina@arm.com>
Date: Wed, 10 Jan 2024 14:31:02 +0000
Subject: [PATCH] middle-end: correctly identify the edge taken when condition
 is true. [PR113287]

The vectorizer needs to know during early break vectorization whether the edge
that will be taken if the condition is true stays or leaves the loop.

This is because the code assumes that if you take the true branch you exit the
loop.  If you don't exit the loop it has to generate a different condition.

Basically it uses this information to decide whether it's generating a
"any element" or an "all element" check.

Bootstrapped Regtested on aarch64-none-linux-gnu, x86_64-pc-linux-gnu
and no issues with --enable-lto --with-build-config=bootstrap-O3
--enable-checking=release,yes,rtl,extra.

gcc/ChangeLog:

	PR tree-optimization/113287
	* tree-vect-stmts.cc (vectorizable_early_exit): Check the flags on edge
	instead of using BRANCH_EDGE to determine true edge.

gcc/testsuite/ChangeLog:

	PR tree-optimization/113287
	* gcc.dg/vect/vect-early-break_100-pr113287.c: New test.
	* gcc.dg/vect/vect-early-break_99-pr113287.c: New test.
---
 .../vect/vect-early-break_100-pr113287.c      | 35 +++++++++++++++++++
 .../vect/vect-early-break_99-pr113287.c       | 32 +++++++++++++++++
 gcc/tree-vect-stmts.cc                        |  9 +++--
 3 files changed, 74 insertions(+), 2 deletions(-)
 create mode 100644 gcc/testsuite/gcc.dg/vect/vect-early-break_100-pr113287.c
 create mode 100644 gcc/testsuite/gcc.dg/vect/vect-early-break_99-pr113287.c

diff --git a/gcc/testsuite/gcc.dg/vect/vect-early-break_100-pr113287.c b/gcc/testsuite/gcc.dg/vect/vect-early-break_100-pr113287.c
new file mode 100644
index 000000000000..f908e5bc6077
--- /dev/null
+++ b/gcc/testsuite/gcc.dg/vect/vect-early-break_100-pr113287.c
@@ -0,0 +1,35 @@
+/* { dg-add-options vect_early_break } */
+/* { dg-require-effective-target vect_early_break } */
+/* { dg-require-effective-target vect_int } */
+/* { dg-require-effective-target bitint } */
+
+__attribute__((noipa)) void
+bar (unsigned long *p)
+{
+  __builtin_memset (p, 0, 142 * sizeof (unsigned long));
+  p[17] = 0x50000000000UL;
+}
+
+__attribute__((noipa)) int
+foo (void)
+{
+  unsigned long r[142];
+  bar (r);
+  unsigned long v = ((long) r[0] >> 31);
+  if (v + 1 > 1)
+    return 1;
+  for (unsigned long i = 1; i <= 140; ++i)
+    if (r[i] != v)
+      return 1;
+  unsigned long w = r[141];
+  if ((unsigned long) (((long) (w << 60)) >> 60) != v)
+    return 1;
+  return 0;
+}
+
+int
+main ()
+{
+  if (foo () != 1)
+    __builtin_abort ();
+}
diff --git a/gcc/testsuite/gcc.dg/vect/vect-early-break_99-pr113287.c b/gcc/testsuite/gcc.dg/vect/vect-early-break_99-pr113287.c
new file mode 100644
index 000000000000..b92a8a268d80
--- /dev/null
+++ b/gcc/testsuite/gcc.dg/vect/vect-early-break_99-pr113287.c
@@ -0,0 +1,32 @@
+/* { dg-add-options vect_early_break } */
+/* { dg-require-effective-target vect_early_break } */
+/* { dg-require-effective-target vect_int } */
+/* { dg-require-effective-target bitint } */
+
+_BitInt(998) b;
+char c;
+char d;
+char e;
+char f;
+char g;
+char h;
+char i;
+char j;
+
+void
+foo(char y, _BitInt(9020) a, char *r)
+{
+  char x = __builtin_mul_overflow_p(a << sizeof(a), y, 0);
+  x += c + d + e + f + g + h + i + j + b;
+  *r = x;
+}
+
+int
+main(void)
+{
+  char x;
+  foo(5, 5, &x);
+  if (x != 1)
+    __builtin_abort();
+  return 0;
+}
diff --git a/gcc/tree-vect-stmts.cc b/gcc/tree-vect-stmts.cc
index bdbf08c9d298..cabd4e3ae864 100644
--- a/gcc/tree-vect-stmts.cc
+++ b/gcc/tree-vect-stmts.cc
@@ -12870,13 +12870,18 @@ vectorizable_early_exit (vec_info *vinfo, stmt_vec_info stmt_info,
      rewrite conditions to always be a comparison against 0.  To do this it
      sometimes flips the edges.  This is fine for scalar,  but for vector we
      then have to flip the test, as we're still assuming that if you take the
-     branch edge that we found the exit condition.  */
+     branch edge that we found the exit condition.  i.e. we need to know whether
+     we are generating a `forall` or an `exist` condition.  */
   auto new_code = NE_EXPR;
   auto reduc_optab = ior_optab;
   auto reduc_op = BIT_IOR_EXPR;
   tree cst = build_zero_cst (vectype);
+  edge exit_true_edge = EDGE_SUCC (gimple_bb (cond_stmt), 0);
+  if (exit_true_edge->flags & EDGE_FALSE_VALUE)
+    exit_true_edge = EDGE_SUCC (gimple_bb (cond_stmt), 1);
+  gcc_assert (exit_true_edge->flags & EDGE_TRUE_VALUE);
   if (flow_bb_inside_loop_p (LOOP_VINFO_LOOP (loop_vinfo),
-			     BRANCH_EDGE (gimple_bb (cond_stmt))->dest))
+			     exit_true_edge->dest))
     {
       new_code = EQ_EXPR;
       reduc_optab = and_optab;
-- 
2.47.2