From 28ad600b4349b57a8f37e8a22a605cc5201775c4 Mon Sep 17 00:00:00 2001
From: Francisco Casas <fcasas@codeweavers.com>
Date: Fri, 8 Nov 2024 15:57:13 -0300
Subject: [PATCH] vkd3d-shader/hlsl: Store SM4 jumps in the vsir program.

---
 libs/vkd3d-shader/hlsl_codegen.c | 36 +++++++++++++++++++++++
 libs/vkd3d-shader/tpf.c          | 49 ++++++--------------------------
 2 files changed, 45 insertions(+), 40 deletions(-)

diff --git a/libs/vkd3d-shader/hlsl_codegen.c b/libs/vkd3d-shader/hlsl_codegen.c
index e8512806..e649ce4f 100644
--- a/libs/vkd3d-shader/hlsl_codegen.c
+++ b/libs/vkd3d-shader/hlsl_codegen.c
@@ -9280,6 +9280,37 @@ static bool sm4_generate_vsir_instr_resource_load(struct hlsl_ctx *ctx,
     }
 }
 
+static bool sm4_generate_vsir_instr_jump(struct hlsl_ctx *ctx,
+        struct vsir_program *program, const struct hlsl_ir_jump *jump)
+{
+    const struct hlsl_ir_node *instr = &jump->node;
+    struct vkd3d_shader_instruction *ins;
+
+    switch (jump->type)
+    {
+        case HLSL_IR_JUMP_BREAK:
+            return generate_vsir_add_program_instruction(ctx, program, &instr->loc, VKD3DSIH_BREAK, 0, 0);
+
+        case HLSL_IR_JUMP_CONTINUE:
+            return generate_vsir_add_program_instruction(ctx, program, &instr->loc, VKD3DSIH_CONTINUE, 0, 0);
+
+        case HLSL_IR_JUMP_DISCARD_NZ:
+            if (!(ins = generate_vsir_add_program_instruction(ctx, program, &instr->loc, VKD3DSIH_DISCARD, 0, 1)))
+                return false;
+            ins->flags = VKD3D_SHADER_CONDITIONAL_OP_NZ;
+
+            vsir_src_from_hlsl_node(&ins->src[0], ctx, jump->condition.node, VKD3DSP_WRITEMASK_ALL);
+            return true;
+
+        case HLSL_IR_JUMP_RETURN:
+            vkd3d_unreachable();
+
+        default:
+            hlsl_fixme(ctx, &jump->node.loc, "Jump type %s.", hlsl_jump_type_to_string(jump->type));
+            return false;
+    }
+}
+
 static void sm4_generate_vsir_block(struct hlsl_ctx *ctx, struct hlsl_block *block, struct vsir_program *program)
 {
     struct vkd3d_string_buffer *dst_type_string;
@@ -9340,6 +9371,11 @@ static void sm4_generate_vsir_block(struct hlsl_ctx *ctx, struct hlsl_block *blo
                     replace_instr_with_last_vsir_instr(ctx, program, instr);
                 break;
 
+            case HLSL_IR_JUMP:
+                if (sm4_generate_vsir_instr_jump(ctx, program, hlsl_ir_jump(instr)))
+                    replace_instr_with_last_vsir_instr(ctx, program, instr);
+                break;
+
             case HLSL_IR_STORE:
                 if (sm4_generate_vsir_instr_store(ctx, program, hlsl_ir_store(instr)))
                     replace_instr_with_last_vsir_instr(ctx, program, instr);
diff --git a/libs/vkd3d-shader/tpf.c b/libs/vkd3d-shader/tpf.c
index c580bb3b..acedcd83 100644
--- a/libs/vkd3d-shader/tpf.c
+++ b/libs/vkd3d-shader/tpf.c
@@ -4804,42 +4804,6 @@ static void write_sm4_if(struct tpf_compiler *tpf, const struct hlsl_ir_if *iff)
     write_sm4_instruction(tpf, &instr);
 }
 
-static void write_sm4_jump(const struct tpf_compiler *tpf, const struct hlsl_ir_jump *jump)
-{
-    struct sm4_instruction instr = {0};
-
-    switch (jump->type)
-    {
-        case HLSL_IR_JUMP_BREAK:
-            instr.opcode = VKD3D_SM4_OP_BREAK;
-            break;
-
-        case HLSL_IR_JUMP_CONTINUE:
-            instr.opcode = VKD3D_SM4_OP_CONTINUE;
-            break;
-
-        case HLSL_IR_JUMP_DISCARD_NZ:
-        {
-            instr.opcode = VKD3D_SM4_OP_DISCARD;
-            instr.extra_bits = VKD3D_SM4_CONDITIONAL_NZ;
-
-            memset(&instr.srcs[0], 0, sizeof(*instr.srcs));
-            instr.src_count = 1;
-            sm4_src_from_node(tpf, &instr.srcs[0], jump->condition.node, VKD3DSP_WRITEMASK_ALL);
-            break;
-        }
-
-        case HLSL_IR_JUMP_RETURN:
-            vkd3d_unreachable();
-
-        default:
-            hlsl_fixme(tpf->ctx, &jump->node.loc, "Jump type %s.", hlsl_jump_type_to_string(jump->type));
-            return;
-    }
-
-    write_sm4_instruction(tpf, &instr);
-}
-
 static void write_sm4_loop(struct tpf_compiler *tpf, const struct hlsl_ir_loop *loop)
 {
     struct sm4_instruction instr =
@@ -4953,6 +4917,12 @@ static void tpf_simple_instruction(struct tpf_compiler *tpf, const struct vkd3d_
         modifier->u.aoffimmi.w = ins->texel_offset.w;
     }
 
+    if (ins->opcode == VKD3DSIH_DISCARD)
+    {
+        if (ins->flags == VKD3D_SHADER_CONDITIONAL_OP_NZ)
+            instr.extra_bits = VKD3D_SM4_CONDITIONAL_NZ;
+    }
+
     write_sm4_instruction(tpf, &instr);
 }
 
@@ -5002,6 +4972,9 @@ static void tpf_handle_instruction(struct tpf_compiler *tpf, const struct vkd3d_
 
         case VKD3DSIH_ADD:
         case VKD3DSIH_AND:
+        case VKD3DSIH_BREAK:
+        case VKD3DSIH_CONTINUE:
+        case VKD3DSIH_DISCARD:
         case VKD3DSIH_DIV:
         case VKD3DSIH_DP2:
         case VKD3DSIH_DP3:
@@ -5119,10 +5092,6 @@ static void write_sm4_block(struct tpf_compiler *tpf, const struct hlsl_block *b
                 write_sm4_if(tpf, hlsl_ir_if(instr));
                 break;
 
-            case HLSL_IR_JUMP:
-                write_sm4_jump(tpf, hlsl_ir_jump(instr));
-                break;
-
             case HLSL_IR_LOOP:
                 write_sm4_loop(tpf, hlsl_ir_loop(instr));
                 break;