Mesa (main): nir/lower_subgroups: Pad ballot values before bitcasting

[GitLab Mirror]

Module: Mesa
Branch: main
Commit: a195ef123ea31ea89a0994d95b5c1f20462bfb6f
URL:    http://cgit.freedesktop.org/mesa/mesa/commit/?id=a195ef123ea31ea89a0994d95b5c1f20462bfb6f

Author: Jason Ekstrand <jason at jlekstrand.net>
Date:   Thu Jul  8 14:36:15 2021 -0500

nir/lower_subgroups: Pad ballot values before bitcasting

Otherwise, if we cast from a uint32_t to a uint64_t, the bitcast will
fail before we pad.  This happens on Intel.

Fixes: e4e79de2a420 "nir/subgroups: Support > 1 ballot components"
Closes: https://gitlab.freedesktop.org/mesa/mesa/-/issues/5045
Reviewed-by: Connor Abbott <cwabbott0 at gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/11786>

---

 src/compiler/nir/nir_lower_subgroups.c | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/compiler/nir/nir_lower_subgroups.c b/src/compiler/nir/nir_lower_subgroups.c
index 33ba7821752..c809a1783d6 100644
--- a/src/compiler/nir/nir_lower_subgroups.c
+++ b/src/compiler/nir/nir_lower_subgroups.c
@@ -78,8 +78,18 @@ static nir_ssa_def *
 uint_to_ballot_type(nir_builder *b, nir_ssa_def *value,
                     unsigned num_components, unsigned bit_size)
 {
-   value = nir_bitcast_vector(b, value, bit_size);
-   return nir_pad_vector_imm_int(b, value, 0, num_components);
+   assert(util_is_power_of_two_nonzero(num_components));
+   assert(util_is_power_of_two_nonzero(value->num_components));
+
+   /* The ballot type must always have enough bits */
+   unsigned total_bits = bit_size * num_components;
+   assert(total_bits >= value->bit_size * value->num_components);
+
+   /* If the source doesn't have enough bits, zero-pad */
+   if (total_bits > value->bit_size * value->num_components)
+      value = nir_pad_vector_imm_int(b, value, 0, total_bits / value->bit_size);
+
+   return nir_bitcast_vector(b, value, bit_size);
 }
 
 static nir_ssa_def *