amd/common: merge radv_nir_opt_access_speculate() into ac_nir_flag_smem_for_loads()

One shader is negatively affected, but we save 2 entire iterations over every shader.
This effect is also mitigated with the next commits.

Totals from 1 (0.00% of 79839) affected shaders: (Navi48)

Instrs: 947 -> 958 (+1.16%)
CodeSize: 4728 -> 4732 (+0.08%)
Latency: 20678 -> 20723 (+0.22%)
InvThroughput: 2697 -> 2698 (+0.04%)
SClause: 26 -> 27 (+3.85%)
Copies: 139 -> 145 (+4.32%)
Branches: 46 -> 47 (+2.17%)
VALU: 460 -> 463 (+0.65%)
SALU: 201 -> 204 (+1.49%)
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/37843>

src/amd/common/nir/ac_nir_lower_mem_access_bit_sizes.c

@@ -13,12 +13,34 @@
 typedef struct {
    enum amd_gfx_level gfx_level;
    bool use_llvm;
+   bool had_terminate;
 } mem_access_cb_data;
 
 static bool
-use_smem_for_load(nir_builder *b, nir_intrinsic_instr *intrin, void *cb_data_)
+set_smem_access_flags(nir_builder *b, nir_intrinsic_instr *intrin, void *cb_data_)
 {
-   const mem_access_cb_data *cb_data = (mem_access_cb_data *)cb_data_;
+   mem_access_cb_data *cb_data = (mem_access_cb_data *)cb_data_;
+   intrin->instr.pass_flags = 0;
+
+   /* Detect descriptors that are used in top level control flow, and mark all smem users as CAN_SPECULATE. */
+   if (!cb_data->had_terminate) {
+      switch (intrin->intrinsic) {
+      case nir_intrinsic_terminate:
+      case nir_intrinsic_terminate_if:
+         cb_data->had_terminate = true;
+         return false;
+      case nir_intrinsic_load_ubo:
+      case nir_intrinsic_load_ssbo:
+         if (intrin->src[0].ssa->parent_instr->block->cf_node.parent->type != nir_cf_node_function)
+            break;
+         FALLTHROUGH;
+      case nir_intrinsic_load_constant:
+         intrin->src[0].ssa->parent_instr->pass_flags = 1;
+         break;
+      default:
+         break;
+      }
+   }
 
    switch (intrin->intrinsic) {
    case nir_intrinsic_load_ssbo:
@@ -38,7 +60,8 @@ use_smem_for_load(nir_builder *b, nir_intrinsic_instr *intrin, void *cb_data_)
    if (intrin->def.divergent)
       return false;
 
-   enum gl_access_qualifier access = nir_intrinsic_access(intrin);
+   /* Check if this instruction can use SMEM. */
+   const enum gl_access_qualifier access = nir_intrinsic_access(intrin);
    bool glc = access & (ACCESS_VOLATILE | ACCESS_COHERENT);
    bool reorder = nir_intrinsic_can_reorder(intrin) || ((access & ACCESS_NON_WRITEABLE) && !(access & ACCESS_VOLATILE));
    if (!reorder || (glc && cb_data->gfx_level < GFX8))
@@ -48,7 +71,23 @@ use_smem_for_load(nir_builder *b, nir_intrinsic_instr *intrin, void *cb_data_)
       return false;
 
    nir_intrinsic_set_access(intrin, access | ACCESS_SMEM_AMD);
-   return true;
+
+   /* Check if this instruction can be executed speculatively. */
+   if (intrin->src[0].ssa->parent_instr->pass_flags == 1)
+      nir_intrinsic_set_access(intrin, nir_intrinsic_access(intrin) | ACCESS_CAN_SPECULATE);
+
+   return access != nir_intrinsic_access(intrin);
+}
+
+bool
+ac_nir_flag_smem_for_loads(nir_shader *shader, enum amd_gfx_level gfx_level, bool use_llvm)
+{
+   mem_access_cb_data cb_data = {
+      .gfx_level = gfx_level,
+      .use_llvm = use_llvm,
+      .had_terminate = false,
+   };
+   return nir_shader_intrinsics_pass(shader, &set_smem_access_flags, nir_metadata_all, &cb_data);
 }
 
 static nir_mem_access_size_align
@@ -161,16 +200,6 @@ lower_mem_access_cb(nir_intrinsic_op intrin, uint8_t bytes, uint8_t bit_size, ui
    return res;
 }
 
-bool
-ac_nir_flag_smem_for_loads(nir_shader *shader, enum amd_gfx_level gfx_level, bool use_llvm)
-{
-   mem_access_cb_data cb_data = {
-      .gfx_level = gfx_level,
-      .use_llvm = use_llvm,
-   };
-   return nir_shader_intrinsics_pass(shader, &use_smem_for_load, nir_metadata_all, &cb_data);
-}
-
 bool
 ac_nir_lower_mem_access_bit_sizes(nir_shader *shader, enum amd_gfx_level gfx_level, bool use_llvm)
 {

src/amd/vulkan/meson.build

@@ -82,7 +82,6 @@ libradv_files = files(
   'nir/radv_nir_lower_view_index.c',
   'nir/radv_nir_lower_viewport_to_zero.c',
   'nir/radv_nir_lower_vs_inputs.c',
-  'nir/radv_nir_opt_access_speculate.c',
   'nir/radv_nir_opt_fs_builtins.c',
   'nir/radv_nir_opt_tid_function.c',
   'nir/radv_nir_remap_color_attachment.c',

src/amd/vulkan/nir/radv_nir.h

@@ -97,8 +97,6 @@ bool radv_nir_opt_tid_function(nir_shader *shader, const radv_nir_opt_tid_functi
 
 bool radv_nir_opt_fs_builtins(nir_shader *shader, const struct radv_graphics_state_key *gfx_state);
 
-bool radv_nir_opt_access_can_speculate(nir_shader *shader);
-
 bool radv_nir_lower_immediate_samplers(nir_shader *shader, struct radv_device *device,
                                        const struct radv_shader_stage *stage);
 

src/amd/vulkan/nir/radv_nir_opt_access_speculate.c

@@ -1,69 +0,0 @@
-/*
- * Copyright © 2025 Valve Corporation
- *
- * SPDX-License-Identifier: MIT
- */
-
-#include "nir/nir.h"
-#include "nir/nir_builder.h"
-
-#include "radv_nir.h"
-
-static bool
-set_can_speculate(nir_builder *b, nir_intrinsic_instr *intr, void *data)
-{
-   switch (intr->intrinsic) {
-   case nir_intrinsic_load_ubo:
-   case nir_intrinsic_load_ssbo:
-      if (!intr->src[0].ssa->parent_instr->pass_flags)
-         return false;
-      break;
-   case nir_intrinsic_load_constant:
-      break;
-   default:
-      return false;
-   }
-
-   unsigned access = nir_intrinsic_access(intr);
-   if (!(access & ACCESS_SMEM_AMD))
-      return false;
-
-   nir_intrinsic_set_access(intr, access | ACCESS_CAN_SPECULATE);
-   return true;
-}
-
-/* Detect descriptors that are used in top level control flow, and mark all smem users as CAN_SPECULATE. */
-bool
-radv_nir_opt_access_can_speculate(nir_shader *shader)
-{
-   bool had_terminate = false;
-   nir_foreach_function_impl (impl, shader) {
-      nir_foreach_block (block, impl) {
-         bool top_level = block->cf_node.parent->type == nir_cf_node_function;
-         nir_foreach_instr (instr, block) {
-            instr->pass_flags = 0;
-            if (had_terminate)
-               continue;
-
-            if (instr->type != nir_instr_type_intrinsic)
-               continue;
-            nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
-            switch (intr->intrinsic) {
-            case nir_intrinsic_terminate:
-            case nir_intrinsic_terminate_if:
-               had_terminate = true;
-               break;
-            case nir_intrinsic_load_ssbo:
-            case nir_intrinsic_load_ubo:
-               if (top_level)
-                  intr->src[0].ssa->parent_instr->pass_flags = 1;
-               break;
-            default:
-               break;
-            }
-         }
-      }
-   }
-
-   return nir_shader_intrinsics_pass(shader, set_can_speculate, nir_metadata_all, NULL);
-}

src/amd/vulkan/radv_pipeline.c

@@ -355,7 +355,6 @@ radv_postprocess_nir(struct radv_device *device, const struct radv_graphics_stat
 
    nir_divergence_analysis(stage->nir);
    NIR_PASS(_, stage->nir, ac_nir_flag_smem_for_loads, gfx_level, use_llvm);
-   NIR_PASS(_, stage->nir, radv_nir_opt_access_can_speculate);
 
    NIR_PASS(_, stage->nir, nir_lower_memory_model);