From ca9ff8c8c761e635da55991b4e1d684a8c7944a2 Mon Sep 17 00:00:00 2001
From: Caterina Shablia <caterina.shablia@collabora.com>
Date: Wed, 29 Jan 2025 20:41:09 +0000
Subject: [PATCH] nir: teach nir_lower_bit_size to handle ballot and
 ballot_relaxed

Reviewed-by: Boris Brezillon <boris.brezillon@collabora.com>
Reviewed-by: Eric R. Smith <eric.smith@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/33365>
---
 src/compiler/nir/nir_lower_bit_size.c | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/compiler/nir/nir_lower_bit_size.c b/src/compiler/nir/nir_lower_bit_size.c
index 25f11787888..48245f7048b 100644
--- a/src/compiler/nir/nir_lower_bit_size.c
+++ b/src/compiler/nir/nir_lower_bit_size.c
@@ -131,6 +131,20 @@ lower_intrinsic_instr(nir_builder *b, nir_intrinsic_instr *intrin,
                       unsigned bit_size)
 {
    switch (intrin->intrinsic) {
+   case nir_intrinsic_ballot:
+   case nir_intrinsic_ballot_relaxed: {
+      b->cursor = nir_before_instr(&intrin->instr);
+
+      nir_alu_type type = nir_type_uint;
+      if (intrin->src[0].ssa->bit_size == 1)
+         type = nir_type_bool;
+
+      nir_def *new_src = nir_convert_to_bit_size(b, intrin->src[0].ssa,
+                                                 type, bit_size);
+      nir_src_rewrite(&intrin->src[0], new_src);
+      break;
+   }
+
    case nir_intrinsic_read_invocation:
    case nir_intrinsic_read_first_invocation:
    case nir_intrinsic_shuffle: