nir: Add src/dest num_components helpers
[mesa.git] / src / compiler / nir / nir_lower_atomics_to_ssbo.c
index 27fe50bbc22d573b61fd951c9e933bcd8d7b2732..934ae81d750595af5b9427b44e7b2f9801d8c6d7 100644 (file)
@@ -59,6 +59,7 @@ lower_instr(nir_intrinsic_instr *instr, unsigned ssbo_offset, nir_builder *b)
    case nir_intrinsic_ssbo_atomic_comp_swap:
    case nir_intrinsic_store_ssbo:
    case nir_intrinsic_load_ssbo:
+   case nir_intrinsic_get_buffer_size:
       /* easy case, keep same opcode and just remap SSBO buffer index: */
       op = instr->intrinsic;
       idx_src = (op == nir_intrinsic_store_ssbo) ? 1 : 0;
@@ -115,7 +116,7 @@ lower_instr(nir_intrinsic_instr *instr, unsigned ssbo_offset, nir_builder *b)
       /* remapped to ssbo_atomic_add: { buffer_idx, offset, +1 } */
       temp = nir_imm_int(b, +1);
       new_instr->src[0] = nir_src_for_ssa(buffer);
-      new_instr->src[1] = instr->src[0];
+      nir_src_copy(&new_instr->src[1], &instr->src[0], new_instr);
       new_instr->src[2] = nir_src_for_ssa(temp);
       break;
    case nir_intrinsic_atomic_counter_dec:
@@ -123,21 +124,21 @@ lower_instr(nir_intrinsic_instr *instr, unsigned ssbo_offset, nir_builder *b)
       /* NOTE semantic difference so we adjust the return value below */
       temp = nir_imm_int(b, -1);
       new_instr->src[0] = nir_src_for_ssa(buffer);
-      new_instr->src[1] = instr->src[0];
+      nir_src_copy(&new_instr->src[1], &instr->src[0], new_instr);
       new_instr->src[2] = nir_src_for_ssa(temp);
       break;
    case nir_intrinsic_atomic_counter_read:
       /* remapped to load_ssbo: { buffer_idx, offset } */
       new_instr->src[0] = nir_src_for_ssa(buffer);
-      new_instr->src[1] = instr->src[0];
+      nir_src_copy(&new_instr->src[1], &instr->src[0], new_instr);
       break;
    default:
       /* remapped to ssbo_atomic_x: { buffer_idx, offset, data, (compare)? } */
       new_instr->src[0] = nir_src_for_ssa(buffer);
-      new_instr->src[1] = instr->src[0];
-      new_instr->src[2] = instr->src[1];
+      nir_src_copy(&new_instr->src[1], &instr->src[0], new_instr);
+      nir_src_copy(&new_instr->src[2], &instr->src[1], new_instr);
       if (op == nir_intrinsic_ssbo_atomic_comp_swap)
-         new_instr->src[3] = instr->src[2];
+         nir_src_copy(&new_instr->src[3], &instr->src[2], new_instr);
       break;
    }
 
@@ -155,9 +156,22 @@ lower_instr(nir_intrinsic_instr *instr, unsigned ssbo_offset, nir_builder *b)
       nir_ssa_def_rewrite_uses(&instr->dest.ssa, nir_src_for_ssa(&new_instr->dest.ssa));
    }
 
+   /* we could be replacing an intrinsic with fixed # of dest num_components
+    * with one that has variable number.  So best to take this from the dest:
+    */
+   new_instr->num_components = instr->dest.ssa.num_components;
+
    return true;
 }
 
+static bool
+is_atomic_uint(const struct glsl_type *type)
+{
+   if (glsl_get_base_type(type) == GLSL_TYPE_ARRAY)
+      return is_atomic_uint(glsl_get_array_element(type));
+   return glsl_get_base_type(type) == GLSL_TYPE_ATOMIC_UINT;
+}
+
 bool
 nir_lower_atomics_to_ssbo(nir_shader *shader, unsigned ssbo_offset)
 {
@@ -184,7 +198,7 @@ nir_lower_atomics_to_ssbo(nir_shader *shader, unsigned ssbo_offset)
       /* replace atomic_uint uniforms with ssbo's: */
       unsigned replaced = 0;
       nir_foreach_variable_safe(var, &shader->uniforms) {
-         if (glsl_get_base_type(var->type) == GLSL_TYPE_ATOMIC_UINT) {
+         if (is_atomic_uint(var->type)) {
             exec_node_remove(&var->node);
 
             if (replaced & (1 << var->data.binding))