From 0badd0547d083a1679d85406655b653afdf6a300 Mon Sep 17 00:00:00 2001
From: Enrico Galli <enrico.galli@intel.com>
Date: Wed, 30 Mar 2022 08:25:22 -0700
Subject: [PATCH] microsoft/spirv_to_dxil: Add pass to lower dynamic accesses
 on ubo[1]

Reviewed-by: Boris Brezillon <boris.brezillon@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/14988>
---
 src/microsoft/compiler/dxil_nir.c           | 58 +++++++++++++++++++++
 src/microsoft/compiler/dxil_nir.h           |  1 +
 src/microsoft/spirv_to_dxil/spirv_to_dxil.c |  1 +
 3 files changed, 60 insertions(+)

diff --git a/src/microsoft/compiler/dxil_nir.c b/src/microsoft/compiler/dxil_nir.c
index 5102a222c45..407cc2490e0 100644
--- a/src/microsoft/compiler/dxil_nir.c
+++ b/src/microsoft/compiler/dxil_nir.c
@@ -27,6 +27,7 @@
 #include "nir_deref.h"
 #include "nir_to_dxil.h"
 #include "util/u_math.h"
+#include "vulkan/vulkan_core.h"
 
 static void
 cl_type_size_align(const struct glsl_type *type, unsigned *size,
@@ -1851,3 +1852,60 @@ dxil_reassign_driver_locations(nir_shader* s, nir_variable_mode modes,
    }
    return result;
 }
+
+static bool
+lower_ubo_array_one_to_static(struct nir_builder *b, nir_instr *inst,
+                              void *cb_data)
+{
+   if (inst->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(inst);
+
+   if (intrin->intrinsic != nir_intrinsic_load_vulkan_descriptor)
+      return false;
+
+   nir_variable *var =
+      nir_get_binding_variable(b->shader, nir_chase_binding(intrin->src[0]));
+
+   if (!var)
+      return false;
+
+   if (!glsl_type_is_array(var->type) || glsl_array_size(var->type) != 1)
+      return false;
+
+   nir_intrinsic_instr *index = nir_src_as_intrinsic(intrin->src[0]);
+   /* We currently do not support reindex */
+   assert(index && index->intrinsic == nir_intrinsic_vulkan_resource_index);
+
+   if (nir_src_is_const(index->src[0]) && nir_src_as_uint(index->src[0]) == 0)
+      return false;
+
+   if (nir_intrinsic_desc_type(index) != VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER)
+      return false;
+
+   b->cursor = nir_instr_remove(&index->instr);
+
+   // Indexing out of bounds on array of UBOs is considered undefined
+   // behavior. Therefore, we just hardcode all the index to 0.
+   uint8_t bit_size = index->dest.ssa.bit_size;
+   nir_ssa_def *zero = nir_imm_intN_t(b, 0, bit_size);
+   nir_ssa_def *dest =
+      nir_vulkan_resource_index(b, index->num_components, bit_size, zero,
+                                .desc_set = nir_intrinsic_desc_set(index),
+                                .binding = nir_intrinsic_binding(index),
+                                .desc_type = nir_intrinsic_desc_type(index));
+
+   nir_ssa_def_rewrite_uses(&index->dest.ssa, dest);
+
+   return true;
+}
+
+bool
+dxil_nir_lower_ubo_array_one_to_static(nir_shader *s)
+{
+   bool progress = nir_shader_instructions_pass(
+      s, lower_ubo_array_one_to_static, nir_metadata_none, NULL);
+
+   return progress;
+}
diff --git a/src/microsoft/compiler/dxil_nir.h b/src/microsoft/compiler/dxil_nir.h
index 2ba474736e5..fe4ecf56aed 100644
--- a/src/microsoft/compiler/dxil_nir.h
+++ b/src/microsoft/compiler/dxil_nir.h
@@ -70,6 +70,7 @@ dxil_reassign_driver_locations(nir_shader* s, nir_variable_mode modes,
 
 void dxil_nir_split_tess_ctrl(nir_shader *nir, nir_function **patch_const_func);
 bool dxil_nir_fixup_tess_level_for_domain(nir_shader *nir);
+bool dxil_nir_lower_ubo_array_one_to_static(nir_shader *s);
 
 #ifdef __cplusplus
 }
diff --git a/src/microsoft/spirv_to_dxil/spirv_to_dxil.c b/src/microsoft/spirv_to_dxil/spirv_to_dxil.c
index 36dcaac4ffd..d02b6e1cbc1 100644
--- a/src/microsoft/spirv_to_dxil/spirv_to_dxil.c
+++ b/src/microsoft/spirv_to_dxil/spirv_to_dxil.c
@@ -714,6 +714,7 @@ spirv_to_dxil(const uint32_t *words, size_t word_count,
    NIR_PASS_V(nir, dxil_nir_lower_loads_stores_to_dxil);
    NIR_PASS_V(nir, dxil_nir_split_typed_samplers);
    NIR_PASS_V(nir, dxil_nir_lower_bool_input);
+   NIR_PASS_V(nir, dxil_nir_lower_ubo_array_one_to_static);
    NIR_PASS_V(nir, nir_opt_dce);
    NIR_PASS_V(nir, nir_remove_dead_variables, nir_var_uniform, NULL);