From 8368a972943d5265b22eab86aa17c9a3cfe3406b Mon Sep 17 00:00:00 2001
From: Gert Wollny <gert.wollny@collabora.com>
Date: Tue, 30 Jan 2024 21:34:43 +0100
Subject: [PATCH] r600: handle indirect access to kcache 14 and 15

r600 can't handle indirect access to kcache 14 and 15, so
if the shader has more than 14 UBOs and there is indirect UBO
access load the values from kcache 14 (and 15) directly and do
a bcsel based on the buffer id to return the right value.

v2: - replace superfluous check with an assert (Triang3l)
    - change the lowering pass to work on load_ubo

Closes: https://gitlab.freedesktop.org/mesa/mesa/-/issues/10112

Signed-off-by: Gert Wollny <gert.wollny@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27372>
---
 src/gallium/drivers/r600/sfn/sfn_nir.cpp      |  1 +
 .../drivers/r600/sfn/sfn_nir_lower_alu.cpp    | 48 +++++++++++++++++++
 .../drivers/r600/sfn/sfn_nir_lower_alu.h      |  3 ++
 3 files changed, 52 insertions(+)

diff --git a/src/gallium/drivers/r600/sfn/sfn_nir.cpp b/src/gallium/drivers/r600/sfn/sfn_nir.cpp
index 70f10c83deb..724ef461e11 100644
--- a/src/gallium/drivers/r600/sfn/sfn_nir.cpp
+++ b/src/gallium/drivers/r600/sfn/sfn_nir.cpp
@@ -762,6 +762,7 @@ r600_lower_and_optimize_nir(nir_shader *sh,
       ((sh->info.bit_sizes_float | sh->info.bit_sizes_int) & 64);
 
    r600::sort_uniforms(sh);
+   NIR_PASS_V(sh, r600_nir_fix_kcache_indirect_access);
 
    while (optimize_once(sh))
       ;
diff --git a/src/gallium/drivers/r600/sfn/sfn_nir_lower_alu.cpp b/src/gallium/drivers/r600/sfn/sfn_nir_lower_alu.cpp
index 002aaf17cc0..bb5abc8ed59 100644
--- a/src/gallium/drivers/r600/sfn/sfn_nir_lower_alu.cpp
+++ b/src/gallium/drivers/r600/sfn/sfn_nir_lower_alu.cpp
@@ -127,6 +127,47 @@ LowerSinCos::lower(nir_instr *instr)
       return nir_fcos_amd(b, normalized);
 }
 
+class FixKcacheIndirectRead : public NirLowerInstruction {
+private:
+   bool filter(const nir_instr *instr) const override;
+   nir_def *lower(nir_instr *instr) override;
+};
+
+bool FixKcacheIndirectRead::filter(const nir_instr *instr) const
+{
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   auto intr = nir_instr_as_intrinsic(instr);
+   if (intr->intrinsic != nir_intrinsic_load_ubo)
+      return false;
+
+   return nir_src_as_const_value(intr->src[0]) == nullptr;
+}
+
+nir_def *FixKcacheIndirectRead::lower(nir_instr *instr)
+{
+   auto intr = nir_instr_as_intrinsic(instr);
+   assert(nir_src_as_const_value(intr->src[0]) == nullptr);
+
+   nir_def *result = &intr->def;
+   for (unsigned i = 14; i < b->shader->info.num_ubos; ++i) {
+      auto test_bufid = nir_imm_int(b, i);
+      auto direct_value =
+	    nir_load_ubo(b, intr->num_components,
+			 intr->def.bit_size,
+			 test_bufid,
+			 intr->src[1].ssa);
+      auto direct_load = nir_instr_as_intrinsic(direct_value->parent_instr);
+      nir_intrinsic_copy_const_indices(direct_load, intr);
+      result = nir_bcsel(b,
+			 nir_ieq(b, test_bufid, intr->src[0].ssa),
+	                 direct_value,
+	                 result);
+   }
+   return result;
+}
+
 } // namespace r600
 
 bool
@@ -140,3 +181,10 @@ r600_nir_lower_trigen(nir_shader *shader, amd_gfx_level gfx_level)
 {
    return r600::LowerSinCos(gfx_level).run(shader);
 }
+
+bool
+r600_nir_fix_kcache_indirect_access(nir_shader *shader)
+{
+   return shader->info.num_ubos > 14 ?
+	    r600::FixKcacheIndirectRead().run(shader) : false;
+}
diff --git a/src/gallium/drivers/r600/sfn/sfn_nir_lower_alu.h b/src/gallium/drivers/r600/sfn/sfn_nir_lower_alu.h
index 5dbe7c75770..bb419cbfbe8 100644
--- a/src/gallium/drivers/r600/sfn/sfn_nir_lower_alu.h
+++ b/src/gallium/drivers/r600/sfn/sfn_nir_lower_alu.h
@@ -36,4 +36,7 @@ r600_nir_lower_pack_unpack_2x16(nir_shader *shader);
 bool
 r600_nir_lower_trigen(nir_shader *shader, enum amd_gfx_level gfx_level);
 
+bool
+r600_nir_fix_kcache_indirect_access(nir_shader *shader);
+
 #endif // SFN_NIR_LOWER_ALU_H