intel: Use hardware generated compute shader local invocation IDs

Reviewed-by: Ivan Briano <ivan.briano@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27167>

src/intel/compiler/brw_fs.h

@@ -147,6 +147,8 @@ struct cs_thread_payload : public thread_payload {
 
    void load_subgroup_id(const brw::fs_builder &bld, fs_reg &dest) const;
 
+   fs_reg local_invocation_id[3];
+
 protected:
    fs_reg subgroup_id_;
 };

src/intel/compiler/brw_fs_nir.cpp

@@ -4447,14 +4447,15 @@ fs_nir_emit_cs_intrinsic(nir_to_brw_state &ntb,
       s.cs_payload().load_subgroup_id(bld, dest);
       break;
 
-   case nir_intrinsic_load_local_invocation_id: {
-      fs_reg val = ntb.system_values[SYSTEM_VALUE_LOCAL_INVOCATION_ID];
-      assert(val.file != BAD_FILE);
-      dest.type = val.type;
+   case nir_intrinsic_load_local_invocation_id:
+      /* This is only used for hardware generated local IDs. */
+      assert(cs_prog_data->generate_local_id);
+
+      dest.type = BRW_REGISTER_TYPE_UD;
+
       for (unsigned i = 0; i < 3; i++)
-         bld.MOV(offset(dest, bld, i), offset(val, bld, i));
+         bld.MOV(offset(dest, bld, i), s.cs_payload().local_invocation_id[i]);
       break;
-   }
 
    case nir_intrinsic_load_workgroup_id:
    case nir_intrinsic_load_workgroup_id_zero_base: {

src/intel/compiler/brw_fs_thread_payload.cpp

@@ -471,13 +471,31 @@ fs_thread_payload::fs_thread_payload(const fs_visitor &v,
 
 cs_thread_payload::cs_thread_payload(const fs_visitor &v)
 {
+   struct brw_cs_prog_data *prog_data = brw_cs_prog_data(v.prog_data);
+
+   unsigned r = reg_unit(v.devinfo);
+
    /* See nir_setup_uniforms for subgroup_id in earlier versions. */
-   if (v.devinfo->verx10 >= 125)
+   if (v.devinfo->verx10 >= 125) {
       subgroup_id_ = brw_ud1_grf(0, 2);
 
-   /* TODO: Fill out uses_btd_stack_ids automatically */
-   num_regs = (1 + brw_cs_prog_data(v.prog_data)->uses_btd_stack_ids) *
-              reg_unit(v.devinfo);
+      for (int i = 0; i < 3; i++) {
+         if (prog_data->generate_local_id & (1 << i)) {
+            local_invocation_id[i] = brw_uw8_grf(r, 0);
+            r += reg_unit(v.devinfo);
+            if (v.devinfo->ver < 20 && v.dispatch_width == 32)
+               r += reg_unit(v.devinfo);
+         } else {
+            local_invocation_id[i] = brw_imm_uw(0);
+         }
+      }
+
+      /* TODO: Fill out uses_btd_stack_ids automatically */
+      if (prog_data->uses_btd_stack_ids)
+         r += reg_unit(v.devinfo);
+   }
+
+   num_regs = r;
 }
 
 void

src/intel/compiler/brw_nir_lower_cs_intrinsics.c

@@ -28,6 +28,7 @@ struct lower_intrinsics_state {
    nir_shader *nir;
    nir_function_impl *impl;
    bool progress;
+   bool hw_generated_local_id;
    nir_builder builder;
 };
 
@@ -189,8 +190,12 @@ lower_cs_intrinsics_convert_block(struct lower_intrinsics_state *state,
 
       nir_def *sysval;
       switch (intrinsic->intrinsic) {
-      case nir_intrinsic_load_local_invocation_index:
-      case nir_intrinsic_load_local_invocation_id: {
+      case nir_intrinsic_load_local_invocation_id:
+         if (state->hw_generated_local_id)
+            continue;
+
+         FALLTHROUGH;
+      case nir_intrinsic_load_local_invocation_index: {
          if (!local_index && !nir->info.workgroup_size_variable) {
             const uint16_t *ws = nir->info.workgroup_size;
             if (ws[0] * ws[1] * ws[2] == 1) {
@@ -209,6 +214,21 @@ lower_cs_intrinsics_convert_block(struct lower_intrinsics_state *state,
                continue;
             }
 
+            if (state->hw_generated_local_id) {
+               nir_def *local_id_vec = nir_load_local_invocation_id(b);
+               nir_def *local_id[3] = { nir_channel(b, local_id_vec, 0),
+                                        nir_channel(b, local_id_vec, 1),
+                                        nir_channel(b, local_id_vec, 2) };
+               nir_def *size_x = nir_imm_int(b, nir->info.workgroup_size[0]);
+               nir_def *size_y = nir_imm_int(b, nir->info.workgroup_size[1]);
+
+               sysval = nir_imul(b, local_id[2], nir_imul(b, size_x, size_y));
+               sysval = nir_iadd(b, sysval, nir_imul(b, local_id[1], size_x));
+               sysval = nir_iadd(b, sysval, local_id[0]);
+               local_index = sysval;
+               break;
+            }
+
             /* First time we are using those, so let's calculate them. */
             assert(!local_id);
             compute_local_index_id(b, nir, &local_index, &local_id);
@@ -283,6 +303,7 @@ brw_nir_lower_cs_intrinsics(nir_shader *nir,
 
    struct lower_intrinsics_state state = {
       .nir = nir,
+      .hw_generated_local_id = false,
    };
 
    /* Constraints from NV_compute_shader_derivatives. */
@@ -300,6 +321,38 @@ brw_nir_lower_cs_intrinsics(nir_shader *nir,
       }
    }
 
+   if (devinfo->verx10 >= 125 && prog_data &&
+       nir->info.stage == MESA_SHADER_COMPUTE &&
+       nir->info.cs.derivative_group != DERIVATIVE_GROUP_QUADS &&
+       !nir->info.workgroup_size_variable &&
+       util_is_power_of_two_nonzero(nir->info.workgroup_size[0]) &&
+       util_is_power_of_two_nonzero(nir->info.workgroup_size[1])) {
+
+      state.hw_generated_local_id = true;
+
+      /* TODO: more heuristics about 1D/SLM access vs. 2D access */
+      bool linear =
+         BITSET_TEST(nir->info.system_values_read,
+                     SYSTEM_VALUE_LOCAL_INVOCATION_INDEX) ||
+         (nir->info.workgroup_size[1] == 1 &&
+          nir->info.workgroup_size[2] == 1) ||
+         (nir->info.num_images == 0 && nir->info.num_textures == 0);
+
+      prog_data->walk_order =
+         linear ? BRW_WALK_ORDER_XYZ : BRW_WALK_ORDER_YXZ;
+
+      /* nir_lower_compute_system_values will replace any references to
+       * SYSTEM_VALUE_LOCAL_INVOCATION_ID vector components with zero for
+       * any dimension where the workgroup size is 1, so we can skip
+       * generating those.  However, the hardware can only generate
+       * X, XY, or XYZ - it can't skip earlier components.
+       */
+      prog_data->generate_local_id =
+         (nir->info.workgroup_size[0] > 1 ? WRITEMASK_X   : 0) |
+         (nir->info.workgroup_size[1] > 1 ? WRITEMASK_XY  : 0) |
+         (nir->info.workgroup_size[2] > 1 ? WRITEMASK_XYZ : 0);
+   }
+
    nir_foreach_function_impl(impl, nir) {
       state.impl = impl;
       lower_cs_intrinsics_convert_impl(&state);