diff --git a/src/amd/vulkan/nir/radv_nir_lower_abi.c b/src/amd/vulkan/nir/radv_nir_lower_abi.c
index 86b0a5143de..ad3857a5c99 100644
--- a/src/amd/vulkan/nir/radv_nir_lower_abi.c
+++ b/src/amd/vulkan/nir/radv_nir_lower_abi.c
@@ -378,30 +378,28 @@ lower_abi_instr(nir_builder *b, nir_intrinsic_instr *intrin, void *state)
       break;
    }
    case nir_intrinsic_atomic_add_gs_emit_prim_count_amd:
-      nir_gds_atomic_add_amd(b, 32, intrin->src[0].ssa, nir_imm_int(b, RADV_SHADER_QUERY_GS_PRIM_EMIT_OFFSET),
-                             nir_imm_int(b, 0x100));
-      break;
-   case nir_intrinsic_atomic_add_gen_prim_count_amd: {
-      uint32_t offset = stage == MESA_SHADER_MESH ? RADV_SHADER_QUERY_MS_PRIM_GEN_OFFSET
-                                                  : RADV_SHADER_QUERY_PRIM_GEN_OFFSET(nir_intrinsic_stream_id(intrin));
-
-      nir_gds_atomic_add_amd(b, 32, intrin->src[0].ssa, nir_imm_int(b, offset), nir_imm_int(b, 0x100));
-      break;
-   }
+   case nir_intrinsic_atomic_add_gen_prim_count_amd:
    case nir_intrinsic_atomic_add_xfb_prim_count_amd:
-      nir_gds_atomic_add_amd(b, 32, intrin->src[0].ssa,
-                             nir_imm_int(b, RADV_SHADER_QUERY_PRIM_XFB_OFFSET(nir_intrinsic_stream_id(intrin))),
-                             nir_imm_int(b, 0x100));
-      break;
    case nir_intrinsic_atomic_add_shader_invocation_count_amd: {
       uint32_t offset;
 
-      if (stage == MESA_SHADER_MESH) {
-         offset = RADV_SHADER_QUERY_MS_INVOCATION_OFFSET;
-      } else if (stage == MESA_SHADER_TASK) {
-         offset = RADV_SHADER_QUERY_TS_INVOCATION_OFFSET;
+      if (intrin->intrinsic == nir_intrinsic_atomic_add_gs_emit_prim_count_amd) {
+         offset = RADV_SHADER_QUERY_GS_PRIM_EMIT_OFFSET;
+      } else if (intrin->intrinsic == nir_intrinsic_atomic_add_gen_prim_count_amd) {
+         offset = stage == MESA_SHADER_MESH ? RADV_SHADER_QUERY_MS_PRIM_GEN_OFFSET
+                                            : RADV_SHADER_QUERY_PRIM_GEN_OFFSET(nir_intrinsic_stream_id(intrin));
+      } else if (intrin->intrinsic == nir_intrinsic_atomic_add_xfb_prim_count_amd) {
+         offset = RADV_SHADER_QUERY_PRIM_XFB_OFFSET(nir_intrinsic_stream_id(intrin));
       } else {
-         offset = RADV_SHADER_QUERY_GS_INVOCATION_OFFSET;
+         assert(intrin->intrinsic == nir_intrinsic_atomic_add_shader_invocation_count_amd);
+
+         if (stage == MESA_SHADER_MESH) {
+            offset = RADV_SHADER_QUERY_MS_INVOCATION_OFFSET;
+         } else if (stage == MESA_SHADER_TASK) {
+            offset = RADV_SHADER_QUERY_TS_INVOCATION_OFFSET;
+         } else {
+            offset = RADV_SHADER_QUERY_GS_INVOCATION_OFFSET;
+         }
       }
 
       nir_gds_atomic_add_amd(b, 32, intrin->src[0].ssa, nir_imm_int(b, offset), nir_imm_int(b, 0x100));