nir: Add a nir_foreach_function_temp_variable helper
[mesa.git] / src / compiler / nir / nir_opt_large_constants.c
index fae4ed27960a57fe9ee95d7f30bf90e59db01360..2575407ca0b5b4e3bb72b745a1a1e2f30e364cd2 100644 (file)
@@ -75,31 +75,28 @@ build_constant_load(nir_builder *b, nir_deref_instr *deref,
    size_align(var->type, &var_size, &var_align);
    assert(var->data.location % var_align == 0);
 
+   UNUSED unsigned deref_size, deref_align;
+   size_align(deref->type, &deref_size, &deref_align);
+
    nir_intrinsic_instr *load =
       nir_intrinsic_instr_create(b->shader, nir_intrinsic_load_constant);
    load->num_components = num_components;
    nir_intrinsic_set_base(load, var->data.location);
    nir_intrinsic_set_range(load, var_size);
+   nir_intrinsic_set_align(load, deref_align, 0);
    load->src[0] = nir_src_for_ssa(nir_build_deref_offset(b, deref, size_align));
    nir_ssa_dest_init(&load->instr, &load->dest,
                      num_components, bit_size, NULL);
    nir_builder_instr_insert(b, &load->instr);
 
    if (load->dest.ssa.bit_size < 8) {
-      /* Booleans are special-cased to be 32-bit
-       *
-       * Ideally, for drivers that can handle 32-bit booleans, we wouldn't
-       * emit the i2b here.  However, at this point, the driver is likely to
-       * still have 1-bit booleans so we need to at least convert bit sizes.
-       * Unfortunately, we don't have a good way to annotate the load as
-       * loading a known boolean value so the optimizer isn't going to be
-       * able to get rid of the conversion.  Some day, we may solve that
-       * problem but not today.
-       */
+      /* Booleans are special-cased to be 32-bit */
       assert(glsl_type_is_boolean(deref->type));
+      assert(deref_size == num_components * 4);
       load->dest.ssa.bit_size = 32;
-      return nir_i2b(b, &load->dest.ssa);
+      return nir_b2b1(b, &load->dest.ssa);
    } else {
+      assert(deref_size == num_components * bit_size / 8);
       return &load->dest.ssa;
    }
 }
@@ -107,6 +104,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 +121,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");
+      }
    }
 }
 
@@ -178,19 +176,17 @@ nir_opt_large_constants(nir_shader *shader,
    /* This pass can only be run once */
    assert(shader->constant_data == NULL && shader->constant_data_size == 0);
 
-   /* The index parameter is unused for local variables so we'll use it for
-    * indexing into our array of variable metadata.
-    */
-   unsigned num_locals = 0;
-   nir_foreach_variable(var, &impl->locals)
-      var->data.index = num_locals++;
+   unsigned num_locals = exec_list_length(&impl->locals);
+   nir_index_vars(shader, impl, nir_var_function_temp);
 
-   if (num_locals == 0)
+   if (num_locals == 0) {
+      nir_shader_preserve_all_metadata(shader);
       return false;
+   }
 
    struct var_info *var_infos = ralloc_array(NULL, struct var_info, num_locals);
-   nir_foreach_variable(var, &impl->locals) {
-      var_infos[var->data.index] = (struct var_info) {
+   nir_foreach_function_temp_variable(var, impl) {
+      var_infos[var->index] = (struct var_info) {
          .var = var,
          .is_constant = true,
          .found_read = false,
@@ -211,10 +207,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:
@@ -222,13 +220,7 @@ nir_opt_large_constants(nir_shader *shader,
             break;
 
          case nir_intrinsic_copy_deref:
-            /* We always assume the src and therefore the dst are not
-             * constants here. Copy and constant propagation passes should
-             * have taken care of this in most cases anyway.
-             */
-            dst_deref = nir_src_as_deref(intrin->src[0]);
-            src_deref = nir_src_as_deref(intrin->src[1]);
-            src_is_const = false;
+            assert(!"Lowering of copy_deref with large constants is prohibited");
             break;
 
          default:
@@ -239,7 +231,7 @@ nir_opt_large_constants(nir_shader *shader,
             nir_variable *var = nir_deref_instr_get_variable(dst_deref);
             assert(var->data.mode == nir_var_function_temp);
 
-            struct var_info *info = &var_infos[var->data.index];
+            struct var_info *info = &var_infos[var->index];
             if (!info->is_constant)
                continue;
 
@@ -255,7 +247,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);
             }
          }
 
@@ -266,7 +259,7 @@ nir_opt_large_constants(nir_shader *shader,
             /* We only consider variables constant if all the reads are
              * dominated by the block that writes to it.
              */
-            struct var_info *info = &var_infos[var->data.index];
+            struct var_info *info = &var_infos[var->index];
             if (!info->is_constant)
                continue;
 
@@ -288,7 +281,7 @@ nir_opt_large_constants(nir_shader *shader,
       struct var_info *info = &var_infos[i];
 
       /* Fix up indices after we sorted. */
-      info->var->data.index = i;
+      info->var->index = i;
 
       if (!info->is_constant)
          continue;
@@ -311,6 +304,7 @@ nir_opt_large_constants(nir_shader *shader,
    }
 
    if (shader->constant_data_size == 0) {
+      nir_shader_preserve_all_metadata(shader);
       ralloc_free(var_infos);
       return false;
    }
@@ -318,7 +312,7 @@ nir_opt_large_constants(nir_shader *shader,
    shader->constant_data = rzalloc_size(shader, shader->constant_data_size);
    for (int i = 0; i < num_locals; i++) {
       struct var_info *info = &var_infos[i];
-      if (!info->duplicate) {
+      if (!info->duplicate && info->is_constant) {
          memcpy((char *)shader->constant_data + info->var->data.location,
                 info->constant_data, info->constant_data_size);
       }
@@ -341,7 +335,7 @@ nir_opt_large_constants(nir_shader *shader,
                continue;
 
             nir_variable *var = nir_deref_instr_get_variable(deref);
-            struct var_info *info = &var_infos[var->data.index];
+            struct var_info *info = &var_infos[var->index];
             if (info->is_constant) {
                b.cursor = nir_after_instr(&intrin->instr);
                nir_ssa_def *val = build_constant_load(&b, deref, size_align);
@@ -359,31 +353,14 @@ nir_opt_large_constants(nir_shader *shader,
                continue;
 
             nir_variable *var = nir_deref_instr_get_variable(deref);
-            struct var_info *info = &var_infos[var->data.index];
+            struct var_info *info = &var_infos[var->index];
             if (info->is_constant) {
                nir_instr_remove(&intrin->instr);
                nir_deref_instr_remove_if_unused(deref);
             }
             break;
          }
-
-         case nir_intrinsic_copy_deref: {
-            nir_deref_instr *deref = nir_src_as_deref(intrin->src[1]);
-            if (deref->mode != nir_var_function_temp)
-               continue;
-
-            nir_variable *var = nir_deref_instr_get_variable(deref);
-            struct var_info *info = &var_infos[var->data.index];
-            if (info->is_constant) {
-               b.cursor = nir_after_instr(&intrin->instr);
-               nir_ssa_def *val = build_constant_load(&b, deref, size_align);
-               nir_store_deref(&b, nir_src_as_deref(intrin->src[0]), val, ~0);
-               nir_instr_remove(&intrin->instr);
-               nir_deref_instr_remove_if_unused(deref);
-            }
-            break;
-         }
-
+         case nir_intrinsic_copy_deref:
          default:
             continue;
          }