From 3fb732ad7dade313a22d2628a50977a00460757b Mon Sep 17 00:00:00 2001
From: Rebecca Mckeever <rebecca.mckeever@collabora.com>
Date: Fri, 9 Dec 2022 16:32:57 -0600
Subject: [PATCH] nvk: Lower base_workgroup_id

Lower load_base_workgroup_id intrinsics to a load_ubo using NIR.

Signed-off-by: Rebecca Mckeever <rebecca.mckeever@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24326>
---
 .../vulkan/nvk_nir_lower_descriptors.c        | 24 +++++++++++++++++++
 src/nouveau/vulkan/nvk_shader.c               |  1 +
 2 files changed, 25 insertions(+)

diff --git a/src/nouveau/vulkan/nvk_nir_lower_descriptors.c b/src/nouveau/vulkan/nvk_nir_lower_descriptors.c
index dc66647ee23..ad607476d36 100644
--- a/src/nouveau/vulkan/nvk_nir_lower_descriptors.c
+++ b/src/nouveau/vulkan/nvk_nir_lower_descriptors.c
@@ -151,6 +151,27 @@ lower_num_workgroups(nir_builder *b, nir_intrinsic_instr *load,
    return true;
 }
 
+static bool
+lower_load_base_workgroup_id(nir_builder *b, nir_intrinsic_instr *load,
+                             const struct lower_descriptors_ctx *ctx)
+{
+   const uint32_t root_table_offset =
+      nvk_root_descriptor_offset(cs.base_group);
+
+   b->cursor = nir_instr_remove(&load->instr);
+
+   nir_ssa_def *val = nir_load_ubo(b, 3, 32,
+                                   nir_imm_int(b, 0),
+                                   nir_imm_int(b, root_table_offset),
+                                   .align_mul = 4,
+                                   .align_offset = 0,
+                                   .range = root_table_offset + 3 * 4);
+
+   nir_ssa_def_rewrite_uses(&load->dest.ssa, val);
+
+   return true;
+}
+
 static bool
 lower_load_push_constant(nir_builder *b, nir_intrinsic_instr *load,
                          const struct lower_descriptors_ctx *ctx)
@@ -246,6 +267,9 @@ lower_intrin(nir_builder *b, nir_intrinsic_instr *intrin,
    case nir_intrinsic_load_num_workgroups:
       return lower_num_workgroups(b, intrin, ctx);
 
+   case nir_intrinsic_load_base_workgroup_id:
+      return lower_load_base_workgroup_id(b, intrin, ctx);
+
    case nir_intrinsic_load_push_constant:
       return lower_load_push_constant(b, intrin, ctx);
 
diff --git a/src/nouveau/vulkan/nvk_shader.c b/src/nouveau/vulkan/nvk_shader.c
index 41039e177bf..936aaf743ec 100644
--- a/src/nouveau/vulkan/nvk_shader.c
+++ b/src/nouveau/vulkan/nvk_shader.c
@@ -282,6 +282,7 @@ nvk_lower_nir(struct nvk_device *device, nir_shader *nir,
    }
 
    nir_lower_compute_system_values_options csv_options = {
+      .has_base_workgroup_id = true,
    };
    NIR_PASS(_, nir, nir_lower_compute_system_values, &csv_options);