From 065ef0c5e4adcc34dbaffa3e71445b823290b7e7 Mon Sep 17 00:00:00 2001
From: Conor McCarthy <cmccarthy@codeweavers.com>
Date: Tue, 23 Apr 2024 21:56:26 +1000
Subject: [PATCH] vkd3d-shader/spirv: Implement the WAVE_IS_FIRST_LANE
 instruction.

---
 libs/vkd3d-shader/spirv.c            | 21 +++++++++++++++++++++
 tests/hlsl/wave-ops-uint.shader_test |  4 ++--
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/libs/vkd3d-shader/spirv.c b/libs/vkd3d-shader/spirv.c
index fd2babc3..015dbe90 100644
--- a/libs/vkd3d-shader/spirv.c
+++ b/libs/vkd3d-shader/spirv.c
@@ -1768,6 +1768,13 @@ static uint32_t vkd3d_spirv_build_op_group_nonuniform_ballot_bit_count(struct vk
             result_type, vkd3d_spirv_get_op_scope_subgroup(builder), group_op, val_id);
 }
 
+static uint32_t vkd3d_spirv_build_op_group_nonuniform_elect(struct vkd3d_spirv_builder *builder)
+{
+    vkd3d_spirv_enable_capability(builder, SpvCapabilityGroupNonUniform);
+    return vkd3d_spirv_build_op_tr1(builder, &builder->function_stream, SpvOpGroupNonUniformElect,
+            vkd3d_spirv_get_op_type_bool(builder), vkd3d_spirv_get_op_scope_subgroup(builder));
+}
+
 static uint32_t vkd3d_spirv_build_op_glsl_std450_tr1(struct vkd3d_spirv_builder *builder,
         enum GLSLstd450 op, uint32_t result_type, uint32_t operand)
 {
@@ -9904,6 +9911,17 @@ static void spirv_compiler_emit_wave_bit_count(struct spirv_compiler *compiler,
     spirv_compiler_emit_store_dst(compiler, dst, val_id);
 }
 
+static void spirv_compiler_emit_wave_is_first_lane(struct spirv_compiler *compiler,
+        const struct vkd3d_shader_instruction *instruction)
+{
+    struct vkd3d_spirv_builder *builder = &compiler->spirv_builder;
+    const struct vkd3d_shader_dst_param *dst = instruction->dst;
+    uint32_t val_id;
+
+    val_id = vkd3d_spirv_build_op_group_nonuniform_elect(builder);
+    spirv_compiler_emit_store_dst(compiler, dst, val_id);
+}
+
 /* This function is called after declarations are processed. */
 static void spirv_compiler_emit_main_prolog(struct spirv_compiler *compiler)
 {
@@ -10272,6 +10290,9 @@ static int spirv_compiler_handle_instruction(struct spirv_compiler *compiler,
         case VKD3DSIH_WAVE_ALL_BIT_COUNT:
             spirv_compiler_emit_wave_bit_count(compiler, instruction);
             break;
+        case VKD3DSIH_WAVE_IS_FIRST_LANE:
+            spirv_compiler_emit_wave_is_first_lane(compiler, instruction);
+            break;
         case VKD3DSIH_DCL:
         case VKD3DSIH_DCL_HS_MAX_TESSFACTOR:
         case VKD3DSIH_DCL_INPUT_CONTROL_POINT_COUNT:
diff --git a/tests/hlsl/wave-ops-uint.shader_test b/tests/hlsl/wave-ops-uint.shader_test
index 2c817366..9bae40ca 100644
--- a/tests/hlsl/wave-ops-uint.shader_test
+++ b/tests/hlsl/wave-ops-uint.shader_test
@@ -24,7 +24,7 @@ void main(uint id : SV_GroupIndex)
 }
 
 [test]
-todo dispatch 4 1 1
+dispatch 4 1 1
 probe uav 1 (0) rui (1)
 probe uav 1 (1) rui (0)
 probe uav 1 (2) rui (0)
@@ -46,7 +46,7 @@ void main(uint id : SV_GroupIndex)
 }
 
 [test]
-todo dispatch 4 1 1
+dispatch 4 1 1
 probe uav 1 (0) rui (0)
 probe uav 1 (1) rui (1)
 probe uav 1 (2) rui (0)