amd: lower load_invocation_id in NIR

ACO can't look for it because it's lowered there.

Reviewed-by: Timur Kristóf <timur.kristof@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32782>

src/amd/common/ac_nir.c

@@ -232,6 +232,21 @@ lower_intrinsic_to_arg(nir_builder *b, nir_instr *instr, void *state)
    case nir_intrinsic_load_view_index:
       replacement = ac_nir_load_arg(b, s->args, s->args->view_index);
       break;
+   case nir_intrinsic_load_invocation_id:
+      if (b->shader->info.stage == MESA_SHADER_TESS_CTRL) {
+         replacement = ac_nir_unpack_arg(b, s->args, s->args->tcs_rel_ids, 8, 5);
+      } else if (b->shader->info.stage == MESA_SHADER_GEOMETRY) {
+         if (s->gfx_level >= GFX12) {
+            replacement = ac_nir_unpack_arg(b, s->args, s->args->gs_vtx_offset[0], 27, 5);
+         } else if (s->gfx_level >= GFX10) {
+            replacement = ac_nir_unpack_arg(b, s->args, s->args->gs_invocation_id, 0, 7);
+         } else {
+            replacement = ac_nir_load_arg(b, s->args, s->args->gs_invocation_id);
+         }
+      } else {
+         unreachable("unexpected shader stage");
+      }
+      break;
    default:
       return false;
    }

src/amd/compiler/aco_instruction_selection.cpp

@@ -5375,15 +5375,11 @@ load_input_from_temps(isel_context* ctx, nir_intrinsic_instr* instr, Temp dst)
    if (ctx->shader->info.stage != MESA_SHADER_TESS_CTRL || !ctx->tcs_in_out_eq)
       return false;
 
+   /* This can only be indexing with invocation_id because all other access has been lowered
+    * to load_shared.
+    */
    nir_src* off_src = nir_get_io_offset_src(instr);
-   nir_src* vertex_index_src = nir_get_io_arrayed_index_src(instr);
-   nir_instr* vertex_index_instr = vertex_index_src->ssa->parent_instr;
-   bool can_use_temps =
-      nir_src_is_const(*off_src) && vertex_index_instr->type == nir_instr_type_intrinsic &&
-      nir_instr_as_intrinsic(vertex_index_instr)->intrinsic == nir_intrinsic_load_invocation_id;
-
-   if (!can_use_temps)
-      return false;
+   assert(nir_src_is_const(*off_src));
 
    nir_io_semantics sem = nir_intrinsic_io_semantics(instr);
 
@@ -9009,28 +9005,6 @@ visit_intrinsic(isel_context* ctx, nir_intrinsic_instr* instr)
       bld.copy(Definition(dst), get_arg(ctx, ctx->args->instance_id));
       break;
    }
-   case nir_intrinsic_load_invocation_id: {
-      Temp dst = get_ssa_temp(ctx, &instr->def);
-
-      if (ctx->shader->info.stage == MESA_SHADER_GEOMETRY) {
-         if (ctx->options->gfx_level >= GFX12)
-            bld.vop3(aco_opcode::v_bfe_u32, Definition(dst),
-                     get_arg(ctx, ctx->args->gs_vtx_offset[0]), Operand::c32(27u),
-                     Operand::c32(5u));
-         else if (ctx->options->gfx_level >= GFX10)
-            bld.vop2_e64(aco_opcode::v_and_b32, Definition(dst), Operand::c32(127u),
-                         get_arg(ctx, ctx->args->gs_invocation_id));
-         else
-            bld.copy(Definition(dst), get_arg(ctx, ctx->args->gs_invocation_id));
-      } else if (ctx->shader->info.stage == MESA_SHADER_TESS_CTRL) {
-         bld.vop3(aco_opcode::v_bfe_u32, Definition(dst), get_arg(ctx, ctx->args->tcs_rel_ids),
-                  Operand::c32(8u), Operand::c32(5u));
-      } else {
-         unreachable("Unsupported stage for load_invocation_id");
-      }
-
-      break;
-   }
    case nir_intrinsic_load_primitive_id: {
       Temp dst = get_ssa_temp(ctx, &instr->def);
 

src/amd/compiler/aco_instruction_selection_setup.cpp

@@ -578,7 +578,6 @@ init_context(isel_context* ctx, nir_shader* shader)
                case nir_intrinsic_shared_atomic:
                case nir_intrinsic_shared_atomic_swap:
                case nir_intrinsic_load_scratch:
-               case nir_intrinsic_load_invocation_id:
                case nir_intrinsic_load_primitive_id:
                case nir_intrinsic_load_typed_buffer_amd:
                case nir_intrinsic_load_buffer_amd:

src/amd/llvm/ac_nir_to_llvm.c

@@ -2953,18 +2953,6 @@ static bool visit_intrinsic(struct ac_nir_context *ctx, nir_intrinsic_instr *ins
    case nir_intrinsic_load_vertex_id_zero_base:
       result = ctx->abi->vertex_id_replaced ? ctx->abi->vertex_id_replaced : ctx->abi->vertex_id;
       break;
-   case nir_intrinsic_load_invocation_id:
-      assert(ctx->stage == MESA_SHADER_TESS_CTRL || ctx->stage == MESA_SHADER_GEOMETRY);
-      if (ctx->stage == MESA_SHADER_TESS_CTRL) {
-         result = ac_unpack_param(&ctx->ac, ac_get_arg(&ctx->ac, ctx->args->tcs_rel_ids), 8, 5);
-      } else if (ctx->ac.gfx_level >= GFX12) {
-         result = ac_unpack_param(&ctx->ac, ac_get_arg(&ctx->ac, ctx->args->gs_vtx_offset[0]), 27, 5);
-      } else if (ctx->ac.gfx_level >= GFX10) {
-         result = ac_unpack_param(&ctx->ac, ac_get_arg(&ctx->ac, ctx->args->gs_invocation_id), 0, 7);
-      } else {
-         result = ac_get_arg(&ctx->ac, ctx->args->gs_invocation_id);
-      }
-      break;
    case nir_intrinsic_load_primitive_id:
       if (ctx->stage == MESA_SHADER_GEOMETRY) {
          result = ac_get_arg(&ctx->ac, ctx->args->gs_prim_id);