nir: add intrinsics for pixel local storage

The pixel local storage load and store instructions keep track of the
format of the pixel local storage variables. This allows drivers to insert
the appropriate conversions on load/store.

Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Reviewed-by: Boris Brezillon <boris.brezillon@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/37110>

src/compiler/nir/nir_gather_info.c

@@ -600,6 +600,7 @@ gather_intrinsic_info(nir_intrinsic_instr *instr, nir_shader *shader)
    case nir_intrinsic_load_per_vertex_output:
    case nir_intrinsic_load_per_view_output:
    case nir_intrinsic_load_per_primitive_output:
+   case nir_intrinsic_load_pixel_local:
       if (shader->info.stage == MESA_SHADER_TESS_CTRL &&
           instr->intrinsic == nir_intrinsic_load_output &&
           !is_patch_special) {
@@ -627,6 +628,10 @@ gather_intrinsic_info(nir_intrinsic_instr *instr, nir_shader *shader)
           !src_is_local_invocation_index(shader, nir_get_io_arrayed_index_src(instr)))
          shader->info.mesh.ms_cross_invocation_output_access |= slot_mask;
 
+      if (shader->info.stage == MESA_SHADER_FRAGMENT &&
+          instr->intrinsic == nir_intrinsic_load_pixel_local)
+         shader->info.fs.accesses_pixel_local_storage = true;
+
       if (shader->info.stage == MESA_SHADER_FRAGMENT &&
           nir_intrinsic_io_semantics(instr).fb_fetch_output)
          shader->info.fs.uses_fbfetch_output = true;
@@ -636,6 +641,7 @@ gather_intrinsic_info(nir_intrinsic_instr *instr, nir_shader *shader)
    case nir_intrinsic_store_per_vertex_output:
    case nir_intrinsic_store_per_view_output:
    case nir_intrinsic_store_per_primitive_output:
+   case nir_intrinsic_store_pixel_local:
       if (shader->info.stage == MESA_SHADER_TESS_CTRL &&
           instr->intrinsic == nir_intrinsic_store_output &&
           !is_patch_special) {
@@ -675,6 +681,10 @@ gather_intrinsic_info(nir_intrinsic_instr *instr, nir_shader *shader)
       if (shader->info.stage == MESA_SHADER_FRAGMENT &&
           nir_intrinsic_io_semantics(instr).dual_source_blend_index)
          shader->info.fs.color_is_dual_source = true;
+
+      if (shader->info.stage == MESA_SHADER_FRAGMENT &&
+          instr->intrinsic == nir_intrinsic_store_pixel_local)
+         shader->info.fs.accesses_pixel_local_storage = true;
       break;
 
    case nir_intrinsic_load_color0:

src/compiler/nir/nir_intrinsics.py

@@ -1254,6 +1254,8 @@ load("ssbo", [-1, 1], [ACCESS, ALIGN_MUL, ALIGN_OFFSET, OFFSET_SHIFT], [CAN_ELIM
 load("ssbo_address", [1, 1], [], [CAN_ELIMINATE, CAN_REORDER])
 # src[] = { offset }.
 load("output", [1], [BASE, RANGE, COMPONENT, DEST_TYPE, IO_SEMANTICS], flags=[CAN_ELIMINATE])
+# src[] = { offset }.
+load("pixel_local", [1], [BASE, RANGE, COMPONENT, DEST_TYPE, FORMAT, IO_SEMANTICS], flags=[CAN_ELIMINATE])
 # src[] = { vertex, offset }.
 load("per_vertex_output", [1, 1], [BASE, RANGE, COMPONENT, DEST_TYPE, IO_SEMANTICS], [CAN_ELIMINATE])
 # src[] = { view_index, offset }.
@@ -1303,6 +1305,8 @@ def store(name, srcs, indices=[], flags=[]):
 
 # src[] = { value, offset }.
 store("output", [1], [BASE, RANGE, WRITE_MASK, COMPONENT, SRC_TYPE, IO_SEMANTICS, IO_XFB, IO_XFB2])
+# src[] = { value, offset }.
+store("pixel_local", [1], [BASE, RANGE, WRITE_MASK, COMPONENT, FORMAT, SRC_TYPE, IO_SEMANTICS])
 # src[] = { value, vertex, offset }.
 store("per_vertex_output", [1, 1], [BASE, RANGE, WRITE_MASK, COMPONENT, SRC_TYPE, IO_SEMANTICS])
 # src[] = { value, view_index, offset }.

src/compiler/nir/nir_lower_amul.c

@@ -153,11 +153,13 @@ lower_intrinsic(lower_state *state, nir_intrinsic_instr *intr)
    case nir_intrinsic_global_atomic_swap:
    case nir_intrinsic_load_global_constant:
    case nir_intrinsic_load_global:
+   case nir_intrinsic_load_pixel_local:
       /* just assume that 24b is not sufficient: */
       lower_large_src(&intr->src[0], state);
       return;
 
    case nir_intrinsic_store_global:
+   case nir_intrinsic_store_pixel_local:
       /* just assume that 24b is not sufficient: */
       lower_large_src(&intr->src[1], state);
       return;

src/compiler/nir/nir_lower_io.c

@@ -386,6 +386,11 @@ emit_load(struct lower_io_state *state,
    case nir_var_uniform:
       op = nir_intrinsic_load_uniform;
       break;
+   case nir_var_mem_pixel_local_in:
+   case nir_var_mem_pixel_local_inout:
+      assert(!array_index);
+      op = nir_intrinsic_load_pixel_local;
+      break;
    default:
       UNREACHABLE("Unknown variable mode");
    }
@@ -416,6 +421,11 @@ emit_load(struct lower_io_state *state,
 
    nir_intrinsic_set_dest_type(load, dest_type);
 
+   if (op == nir_intrinsic_load_pixel_local) {
+      assert(var && var->data.image.format != PIPE_FORMAT_NONE);
+      nir_intrinsic_set_format(load, var->data.image.format);
+   }
+
    if (load->intrinsic != nir_intrinsic_load_uniform) {
       int location = var->data.location;
       unsigned num_slots = get_number_of_slots(state, var);
@@ -535,9 +545,11 @@ emit_store(struct lower_io_state *state, nir_def *data,
 {
    nir_builder *b = &state->builder;
 
-   assert(var->data.mode == nir_var_shader_out);
    nir_intrinsic_op op;
-   if (!array_index)
+   if (var->data.mode == nir_var_mem_pixel_local_out ||
+       var->data.mode == nir_var_mem_pixel_local_inout)
+      op = nir_intrinsic_store_pixel_local;
+   else if (!array_index)
       op = nir_intrinsic_store_output;
    else if (var->data.per_view)
       op = nir_intrinsic_store_per_view_output;
@@ -602,6 +614,10 @@ emit_store(struct lower_io_state *state, nir_def *data,
    semantics.per_view = var->data.per_view;
 
    nir_intrinsic_set_io_semantics(store, semantics);
+   if (op == nir_intrinsic_store_pixel_local) {
+      assert(var && var->data.image.format != PIPE_FORMAT_NONE);
+      nir_intrinsic_set_format(store, var->data.image.format);
+   }
 
    nir_builder_instr_insert(b, &store->instr);
 }
@@ -904,7 +920,8 @@ nir_lower_io_impl(nir_function_impl *impl,
                   _mesa_hash_string, _mesa_key_string_equal);
 
    ASSERTED nir_variable_mode supported_modes =
-      nir_var_shader_in | nir_var_shader_out | nir_var_uniform;
+      nir_var_shader_in | nir_var_shader_out | nir_var_uniform |
+      nir_var_any_pixel_local;
    assert(!(modes & ~supported_modes));
 
    nir_foreach_block(block, impl) {
@@ -951,6 +968,7 @@ nir_get_io_offset_src_number(const nir_intrinsic_instr *instr)
    case nir_intrinsic_load_input:
    case nir_intrinsic_load_per_primitive_input:
    case nir_intrinsic_load_output:
+   case nir_intrinsic_load_pixel_local:
    case nir_intrinsic_load_shared:
    case nir_intrinsic_load_task_payload:
    case nir_intrinsic_load_uniform:
@@ -991,6 +1009,7 @@ nir_get_io_offset_src_number(const nir_intrinsic_instr *instr)
    case nir_intrinsic_load_interpolated_input:
    case nir_intrinsic_load_global_amd:
    case nir_intrinsic_store_output:
+   case nir_intrinsic_store_pixel_local:
    case nir_intrinsic_store_shared:
    case nir_intrinsic_store_task_payload:
    case nir_intrinsic_store_global:

src/compiler/nir/nir_opt_dead_cf.c

@@ -241,6 +241,7 @@ node_is_dead(nir_cf_node *node)
             case nir_intrinsic_load_shared:
             case nir_intrinsic_load_shared2_amd:
             case nir_intrinsic_load_output:
+            case nir_intrinsic_load_pixel_local:
             case nir_intrinsic_load_per_vertex_output:
             case nir_intrinsic_load_per_view_output:
                /* Same as above loads. */

src/compiler/nir/nir_print.c

@@ -1484,6 +1484,11 @@ print_intrinsic_instr(nir_intrinsic_instr *instr, print_state *state)
             mode = nir_var_shader_out;
             break;
 
+         case nir_intrinsic_load_pixel_local:
+         case nir_intrinsic_store_pixel_local:
+            mode = nir_var_mem_pixel_local_inout;
+            break;
+
          default:
             break;
          }