From aff391bc772242b968d41dee32c3d8914fbbd618 Mon Sep 17 00:00:00 2001
From: Georg Lehmann <dadschoorse@gmail.com>
Date: Wed, 20 Aug 2025 12:43:07 +0200
Subject: [PATCH] nir/uub: handle more reduction ops

Reviewed-by: Rhys Perry <pendingchaos02@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/36874>
---
 src/compiler/nir/nir_range_analysis.c | 48 ++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 9 deletions(-)

diff --git a/src/compiler/nir/nir_range_analysis.c b/src/compiler/nir/nir_range_analysis.c
index 655174a5013..4ba238edfff 100644
--- a/src/compiler/nir/nir_range_analysis.c
+++ b/src/compiler/nir/nir_range_analysis.c
@@ -1751,21 +1751,51 @@ get_intrinsic_uub(struct analysis_state *state, struct uub_query q, uint32_t *re
    case nir_intrinsic_inclusive_scan:
    case nir_intrinsic_exclusive_scan: {
       nir_op op = nir_intrinsic_reduction_op(intrin);
-      unsigned bit_size = q.scalar.def->bit_size;
-      bool exclusive = intrin->intrinsic == nir_intrinsic_exclusive_scan;
 
-      if (op == nir_op_umin || op == nir_op_umax || op == nir_op_imin || op == nir_op_imax) {
+      switch (op) {
+      case nir_op_umin:
+      case nir_op_umax:
+      case nir_op_imax:
+      case nir_op_imin:
+      case nir_op_iand:
+      case nir_op_ior:
+      case nir_op_ixor:
+      case nir_op_iadd:
          if (!q.head.pushed_queries) {
             push_uub_query(state, nir_get_scalar(intrin->src[0].ssa, q.scalar.comp));
             return;
-         } else {
-            *result = src[0];
-            if (exclusive) {
-               uint32_t identity = nir_const_value_as_uint(nir_alu_binop_identity(op, bit_size), bit_size);
-               *result = MAX2(*result, identity);
-            }
          }
+         break;
+      default:
+         return;
       }
+
+      unsigned bit_size = q.scalar.def->bit_size;
+      bool exclusive = intrin->intrinsic == nir_intrinsic_exclusive_scan;
+      switch (op) {
+      case nir_op_umin:
+      case nir_op_umax:
+      case nir_op_imax:
+      case nir_op_imin:
+      case nir_op_iand:
+         *result = src[0];
+         break;
+      case nir_op_ior:
+      case nir_op_ixor:
+         *result = bitmask(util_last_bit64(src[0]));
+         break;
+      case nir_op_iadd:
+         *result = MIN2(*result, (uint64_t)src[0] * (config->max_subgroup_size - exclusive));
+         break;
+      default:
+         UNREACHABLE("unhandled op");
+      }
+
+      if (exclusive) {
+         uint32_t identity = nir_const_value_as_uint(nir_alu_binop_identity(op, bit_size), bit_size);
+         *result = MAX2(*result, identity);
+      }
+
       break;
    }
    case nir_intrinsic_read_first_invocation: