jay: smarten predication pass

Merge the empty else optimization, the then-block predication, and the
break-while fusion into a unified "try to predicate each side of an if, peephole
optimizing control flow" optimization. This is simpler and more general.

Totals:
Instrs: 4783809 -> 4775647 (-0.17%)
CodeSize: 70766656 -> 70674064 (-0.13%); split: -0.13%, +0.00%

Totals from 1109 (41.90% of 2647) affected shaders:
Instrs: 4130644 -> 4122482 (-0.20%)
CodeSize: 61180848 -> 61088256 (-0.15%); split: -0.15%, +0.00%

Signed-off-by: Alyssa Rosenzweig <alyssa.rosenzweig@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/41215>

src/intel/compiler/jay/jay_from_nir.c

@@ -2687,7 +2687,7 @@ jay_compile(const struct intel_device_info *devinfo,
             nir->info.bit_sizes_float);
 
    if (!(jay_debug & JAY_DBG_NOOPT)) {
-      JAY_PASS(s, jay_opt_control_flow);
+      JAY_PASS(s, jay_opt_predicate);
    }
 
    JAY_PASS(s, jay_lower_scoreboard);

src/intel/compiler/jay/jay_opt_control_flow.c

@@ -1,137 +0,0 @@
-/*
- * Copyright 2026 Intel Corporation
- * Copyright 2023 Valve Corporation
- * SPDX-License-Identifier: MIT
- */
-
-#include "util/list.h"
-#include "jay_builder.h"
-#include "jay_ir.h"
-#include "jay_opcodes.h"
-#include "jay_private.h"
-
-/*
- * Detect the block "else; endif" and remove the no-op else, effectively
- * removing empty else blocks. Logically, that causes critical edges, so this
- * pass must run late (post-RA).
- */
-static void
-opt_empty_else(jay_block *blk)
-{
-   unsigned i = 0;
-   enum jay_opcode ops[] = { JAY_OPCODE_ELSE, JAY_OPCODE_ENDIF };
-
-   jay_foreach_inst_in_block(blk, I) {
-      if (i >= ARRAY_SIZE(ops) || ops[i++] != I->op)
-         return;
-   }
-
-   if (i == ARRAY_SIZE(ops)) {
-      jay_remove_instruction(jay_first_inst(blk));
-   }
-}
-
-/*
- * Replace short if-statements with predication. Assumes opt_empty_else already
- * ran. TODO: Generalize.
- */
-static void
-opt_predicate(jay_function *f, jay_block *block)
-{
-   jay_inst *if_ = jay_last_inst(block);
-   if (!if_ || if_->op != JAY_OPCODE_IF)
-      return;
-
-   /* If's fallthrough to the then */
-   jay_block *then_block = jay_next_block(block);
-   assert(block->logical_succs[0] == then_block && "successors for if");
-
-   /* We're searching for a single block then, so the next block is else */
-   jay_block *else_block = jay_next_block(then_block);
-   if (block->logical_succs[1] != else_block ||
-       list_length(&then_block->instructions) > 3 ||
-       !list_is_singular(&else_block->instructions))
-      return;
-
-   /* We can only access one flag per instruction, so do not predicate anything
-    * accessing flags. This also ensures the if-condition flag is kept live.
-    *
-    * MIN/MAX turn into SEL which cannot be predicated despite not using flags.
-    *
-    * Predicating NoMask instructions doesn't work if we are electing a nonzero
-    * lane but the NoMask forces lane 0. This should be optimized later.
-    */
-   jay_foreach_inst_in_block(then_block, I) {
-      if (jay_uses_flag(I) ||
-          I->op == JAY_OPCODE_MIN ||
-          I->op == JAY_OPCODE_MAX ||
-          I->op == JAY_OPCODE_CSEL ||
-          jay_is_no_mask(I))
-         return;
-   }
-
-   jay_inst *endif = jay_last_inst(else_block);
-   if (endif->op != JAY_OPCODE_ENDIF)
-      return;
-
-   /* Rewrite with predication */
-   jay_builder b = jay_init_builder(f, jay_after_block(block));
-   assert(if_->predication == JAY_PREDICATED && "if's are always predicated");
-
-   jay_foreach_inst_in_block_safe(then_block, I) {
-      jay_add_predicate(&b, I, *jay_inst_get_predicate(if_));
-   }
-
-   /* Remove the jumps */
-   jay_remove_instruction(if_);
-   jay_remove_instruction(endif);
-}
-
-/*
- * Optimize "(f0) break; while" to "(!f0) while". As break/while appear in
- * different blocks, we optimize the entire function at a time.
- */
-static void
-opt_predicate_while(jay_function *func)
-{
-   jay_inst *prev_break = NULL;
-
-   jay_foreach_block(func, block) {
-      if (list_is_empty(&block->instructions)) {
-         /* Ignore empty blocks */
-      } else if (jay_last_inst(block)->op == JAY_OPCODE_BREAK) {
-         prev_break = jay_last_inst(block);
-      } else if (jay_first_inst(block)->op == JAY_OPCODE_WHILE &&
-                 prev_break &&
-                 prev_break->predication) {
-         assert(!jay_first_inst(block)->predication);
-         jay_inst_get_predicate(prev_break)->negate ^= true;
-
-         jay_remove_instruction(jay_first_inst(block));
-         jay_remove_instruction(prev_break);
-
-         jay_builder b = jay_init_builder(func, jay_before_block(block));
-         jay_builder_insert(&b, prev_break);
-
-         prev_break->op = JAY_OPCODE_WHILE;
-         prev_break = NULL;
-      } else {
-         prev_break = NULL;
-      }
-   }
-}
-
-void
-jay_opt_control_flow(jay_shader *s)
-{
-   jay_foreach_function(s, f) {
-      /* Iterating blocks in reverse lets both opts converge in 1 pass */
-      jay_foreach_block_rev(f, block) {
-         opt_empty_else(block);
-         opt_predicate(f, block);
-      }
-
-      /* Do last: opt_predicate_while depends on both previous optimizations */
-      opt_predicate_while(f);
-   }
-}

src/intel/compiler/jay/jay_opt_predicate.c

@@ -0,0 +1,121 @@
+/*
+ * Copyright 2026 Intel Corporation
+ * Copyright 2023 Valve Corporation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "jay_builder.h"
+#include "jay_ir.h"
+#include "jay_opcodes.h"
+#include "jay_private.h"
+
+static bool
+predicate_block(jay_builder *b,
+                jay_block *block,
+                jay_def condition,
+                signed limit)
+{
+   /* We can only access one flag per instruction, so do not predicate anything
+    * accessing flags. This also ensures the if-condition flag is kept live.
+    *
+    * A few opcodes can't be predicated due to ISA restrictions.
+    *
+    * Predicating NoMask instructions doesn't work if we are electing a nonzero
+    * lane but the NoMask forces lane 0. This should be optimized later.
+    */
+   jay_foreach_inst_in_block(block, I) {
+      if (jay_uses_flag(I) ||
+          (I->op == JAY_OPCODE_MIN || I->op == JAY_OPCODE_MAX) ||
+          I->op == JAY_OPCODE_CSEL ||
+          jay_is_no_mask(I) ||
+          (--limit) < 0)
+         return false;
+   }
+
+   /* Hoist everything but branches. Branches cannot be hoisted because we
+    * cannot have a branch in the middle of a block. Since they appear last,
+    * this rule does not cause any nontrivial reordering. Branches make the
+    * block considered unpredicatable so we don't remove the control flow ops.
+    */
+   jay_foreach_inst_in_block_safe(block, I) {
+      if (I->op != JAY_OPCODE_ENDIF && I->op != JAY_OPCODE_ELSE) {
+         I = jay_add_predicate(b, I, condition);
+
+         if (I->op == JAY_OPCODE_BREAK) {
+            return false;
+         }
+
+         jay_remove_instruction(I);
+         jay_builder_insert(b, I);
+      }
+   }
+
+   return true;
+}
+
+/*
+ * Replace short if-statements with predication.
+ */
+static void
+predicate_if(jay_function *f, jay_block *if_block, jay_inst *if_)
+{
+   /* If's fallthrough to the then and branch to the else */
+   jay_block *then_block = if_block->logical_succs[0],
+             *else_block = if_block->logical_succs[1];
+   assert(then_block == jay_next_block(if_block) && "successors for if");
+
+   jay_builder b = jay_init_builder(f, jay_before_inst(if_));
+   jay_inst *endif = jay_last_inst(else_block);
+   jay_def pred = *jay_inst_get_predicate(if_);
+
+   /* Else has a higher limit to account for else/endif ops */
+   bool no_then = else_block == jay_next_block(then_block) &&
+                  predicate_block(&b, then_block, pred, 3);
+   bool no_else = endif->op == JAY_OPCODE_ENDIF &&
+                  predicate_block(&b, else_block, jay_negate(pred), 5);
+
+   if (no_then && no_else) {
+      jay_remove_instruction(if_);
+      jay_remove_instruction(endif);
+   } else if (no_then) {
+      /* if (x) {} else { ... } -----> if (!x) { .... } */
+      jay_inst_get_predicate(if_)->negate ^= true;
+   }
+
+   if (no_then || no_else) {
+      assert(jay_first_inst(else_block)->op == JAY_OPCODE_ELSE);
+      jay_remove_instruction(jay_first_inst(else_block));
+   }
+
+   /* Optimize "if (f0) { break }" and "if (f0) { break } while", leaving the
+    * control flow graph intact for global data flow analysis.
+    */
+   if (!no_then && no_else) {
+      jay_inst *brk = jay_first_inst(then_block);
+      jay_inst *whl = jay_first_inst(jay_next_block(else_block));
+
+      if (brk && brk->op == JAY_OPCODE_BREAK) {
+         jay_remove_instruction(if_);
+         jay_remove_instruction(endif);
+
+         if (whl && whl->op == JAY_OPCODE_WHILE && brk->predication) {
+            jay_add_predicate(&b, whl,
+                              jay_negate(*jay_inst_get_predicate(brk)));
+            jay_remove_instruction(brk);
+         }
+      }
+   }
+}
+
+static void
+pass(jay_function *f)
+{
+   jay_foreach_block_rev(f, block) {
+      jay_inst *if_ = jay_last_inst(block);
+      if (if_ && if_->op == JAY_OPCODE_IF) {
+         predicate_if(f, block, if_);
+      }
+   }
+}
+
+JAY_DEFINE_FUNCTION_PASS(jay_opt_predicate, pass)

src/intel/compiler/jay/jay_private.h

@@ -66,7 +66,7 @@ jay_validate_ra(jay_function *func)
 void jay_opt_propagate_forwards(jay_shader *s);
 void jay_opt_propagate_backwards(jay_shader *s);
 void jay_opt_dead_code(jay_shader *s);
-void jay_opt_control_flow(jay_shader *s);
+void jay_opt_predicate(jay_shader *s);
 
 void jay_lower_pre_ra(jay_shader *s);
 void jay_lower_post_ra(jay_shader *s);

src/intel/compiler/jay/meson.build

@@ -58,7 +58,7 @@ libintel_compiler_jay_files = files(
   'jay_lower_spill.c',
   'jay_nir.c',
   'jay_opt_dead_code.c',
-  'jay_opt_control_flow.c',
+  'jay_opt_predicate.c',
   'jay_opt_propagate.c',
   'jay_print.c',
   'jay_private.h',