Mesa (main): nir/lower_subgroups: Handle down-casts in uint_to_ballot_type

[GitLab Mirror]

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

Author: Jason Ekstrand <jason at jlekstrand.net>
Date:   Tue Jul 20 17:14:56 2021 -0500

nir/lower_subgroups: Handle down-casts in uint_to_ballot_type

This is required for Zink where the API ballot type is a uint64_t and
the "hardware" ballot type is uvec4.

Reviewed-By: Mike Blumenkrantz <michael.blumenkrantz at gmail.com>
Reviewed-by: Connor Abbott <cwabbott0 at gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/11989>

---

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

diff --git a/src/compiler/nir/nir_lower_subgroups.c b/src/compiler/nir/nir_lower_subgroups.c
index c809a1783d6..ecac878325d 100644
--- a/src/compiler/nir/nir_lower_subgroups.c
+++ b/src/compiler/nir/nir_lower_subgroups.c
@@ -81,15 +81,26 @@ uint_to_ballot_type(nir_builder *b, nir_ssa_def *value,
    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);
+   value = nir_bitcast_vector(b, value, bit_size);
+
+   /* If the source has too many components, truncate.  This can happen if,
+    * for instance, we're implementing GL_ARB_shader_ballot or
+    * VK_EXT_shader_subgroup_ballot which have 64-bit ballot values on an
+    * architecture with a native 128-bit uvec4 ballot.  This comes up in Zink
+    * for OpenGL on Vulkan.  It's the job of the driver calling this lowering
+    * pass to ensure that it's restricted subgroup sizes sufficiently that we
+    * have enough ballot bits.
+    */
+   if (value->num_components > num_components)
+      value = nir_channels(b, value, BITFIELD_MASK(num_components));
+
+   return value;
 }
 
 static nir_ssa_def *