nir/opt_large_constants: Handle store writemasks

This fixes some piglit tests on radeonsi NIR where a varying is
initialized to a constant array in the vertex shader. Varying packing
after nir_lower_io_to_temporaries creates writemasked stores which
persist after pulling the constant initialization down into the fragment
shader.

While we're here, rewrite handle_constant_store() to do the loop over
components outside the switch, so that we don't have to duplicate the
writemask checking for every bitsize.

Fixes: 1235850522c ("nir: Add a large constants optimization pass")
Reviewed-by: Caio Marcelo de Oliveira Filho <[email protected]>
(cherry picked from commit 270fe55256c78ede507d75d4665d73936ea7db31)

1 files changed, 24 insertions, 20 deletions

diff --git a/src/compiler/nir/nir_opt_large_constants.c b/src/compiler/nir/nir_opt_large_constants.c
index 8fd1f0d2e99..d90ca60ff2b 100644
--- a/src/compiler/nir/nir_opt_large_constants.c
+++ b/src/compiler/nir/nir_opt_large_constants.c
@@ -107,6 +107,7 @@ build_constant_load(nir_builder *b, nir_deref_instr *deref,
 static void
 handle_constant_store(void *mem_ctx, struct var_info *info,
                       nir_deref_instr *deref, nir_const_value *val,
+                      unsigned writemask,
                       glsl_type_size_align_func size_align)
 {
    assert(!nir_deref_instr_has_indirect(deref));
@@ -123,35 +124,35 @@ handle_constant_store(void *mem_ctx, struct var_info *info,
    char *dst = (char *)info->constant_data +
                nir_deref_instr_get_const_offset(deref, size_align);
 
-   switch (bit_size) {
-   case 1:
-      /* Booleans are special-cased to be 32-bit */
-      for (unsigned i = 0; i < num_components; i++)
+   for (unsigned i = 0; i < num_components; i++) {
+      if (!(writemask & (1 << i)))
+         continue;
+
+      switch (bit_size) {
+      case 1:
+         /* Booleans are special-cased to be 32-bit */
          ((int32_t *)dst)[i] = -(int)val[i].b;
-      break;
+         break;
 
-   case 8:
-      for (unsigned i = 0; i < num_components; i++)
+      case 8:
          ((uint8_t *)dst)[i] = val[i].u8;
-      break;
+         break;
 
-   case 16:
-      for (unsigned i = 0; i < num_components; i++)
+      case 16:
          ((uint16_t *)dst)[i] = val[i].u16;
-      break;
+         break;
 
-   case 32:
-      for (unsigned i = 0; i < num_components; i++)
+      case 32:
          ((uint32_t *)dst)[i] = val[i].u32;
-      break;
+         break;
 
-   case 64:
-      for (unsigned i = 0; i < num_components; i++)
+      case 64:
          ((uint64_t *)dst)[i] = val[i].u64;
-      break;
+         break;
 
-   default:
-      unreachable("Invalid bit size");
+      default:
+         unreachable("Invalid bit size");
+      }
    }
 }
 
@@ -211,10 +212,12 @@ nir_opt_large_constants(nir_shader *shader,
 
          bool src_is_const = false;
          nir_deref_instr *src_deref = NULL, *dst_deref = NULL;
+         unsigned writemask = 0;
          switch (intrin->intrinsic) {
          case nir_intrinsic_store_deref:
             dst_deref = nir_src_as_deref(intrin->src[0]);
             src_is_const = nir_src_is_const(intrin->src[1]);
+            writemask = nir_intrinsic_write_mask(intrin);
             break;
 
          case nir_intrinsic_load_deref:
@@ -255,7 +258,8 @@ nir_opt_large_constants(nir_shader *shader,
                info->is_constant = false;
             } else {
                nir_const_value *val = nir_src_as_const_value(intrin->src[1]);
-               handle_constant_store(var_infos, info, dst_deref, val, size_align);
+               handle_constant_store(var_infos, info, dst_deref, val, writemask,
+                                     size_align);
             }
          }