diff --git a/libs/vkd3d-shader/hlsl_codegen.c b/libs/vkd3d-shader/hlsl_codegen.c
index f380195e3..100af3f7b 100644
--- a/libs/vkd3d-shader/hlsl_codegen.c
+++ b/libs/vkd3d-shader/hlsl_codegen.c
@@ -13574,6 +13574,19 @@ static void process_entry_function(struct hlsl_ctx *ctx,
     lower_ir(ctx, lower_matrix_swizzles, body);
     lower_ir(ctx, lower_index_loads, body);
 
+    if (entry_func->return_var)
+    {
+        if (profile->type == VKD3D_SHADER_TYPE_GEOMETRY)
+            hlsl_error(ctx, &entry_func->loc, VKD3D_SHADER_ERROR_HLSL_INCOMPATIBLE_PROFILE,
+                    "Geometry shaders cannot return values.");
+        else if (entry_func->return_var->data_type->class != HLSL_CLASS_STRUCT
+                && !entry_func->return_var->semantic.name)
+            hlsl_error(ctx, &entry_func->loc, VKD3D_SHADER_ERROR_HLSL_MISSING_SEMANTIC,
+                    "Entry point \"%s\" is missing a return value semantic.", entry_func->func->name);
+
+        append_output_var_copy(ctx, entry_func, entry_func->return_var);
+    }
+
     for (i = 0; i < entry_func->parameters.count; ++i)
     {
         var = entry_func->parameters.vars[i];
@@ -13678,18 +13691,9 @@ static void process_entry_function(struct hlsl_ctx *ctx,
             }
         }
     }
+
     if (entry_func->return_var)
     {
-        if (profile->type == VKD3D_SHADER_TYPE_GEOMETRY)
-            hlsl_error(ctx, &entry_func->loc, VKD3D_SHADER_ERROR_HLSL_INCOMPATIBLE_PROFILE,
-                    "Geometry shaders cannot return values.");
-        else if (entry_func->return_var->data_type->class != HLSL_CLASS_STRUCT
-                && !entry_func->return_var->semantic.name)
-            hlsl_error(ctx, &entry_func->loc, VKD3D_SHADER_ERROR_HLSL_MISSING_SEMANTIC,
-                    "Entry point \"%s\" is missing a return value semantic.", entry_func->func->name);
-
-        append_output_var_copy(ctx, entry_func, entry_func->return_var);
-
         if (profile->type == VKD3D_SHADER_TYPE_HULL && !ctx->is_patch_constant_func)
             ctx->output_control_point_type = entry_func->return_var->data_type;
     }
diff --git a/tests/hlsl/point-sprite.shader_test b/tests/hlsl/point-sprite.shader_test
index 8a413259c..a5c644c15 100644
--- a/tests/hlsl/point-sprite.shader_test
+++ b/tests/hlsl/point-sprite.shader_test
@@ -33,16 +33,16 @@ clear rtv 0  0 0 0 0
 point-size 30.0 1.0 64.0
 point-sprite off
 draw point list 1
-probe (310, 230) rgba (1, 0, 1, 1)
-probe (330, 230) rgba (1, 0, 1, 1)
-probe (310, 250) rgba (1, 0, 1, 1)
-probe (330, 250) rgba (1, 0, 1, 1)
+todo(sm>=4 & sm<6) probe (310, 230) f32(1, 0, 1, 1)
+todo(sm>=4 & sm<6) probe (330, 230) f32(1, 0, 1, 1)
+todo(sm>=4 & sm<6) probe (310, 250) f32(1, 0, 1, 1)
+todo(sm>=4 & sm<6) probe (330, 250) f32(1, 0, 1, 1)
 
 clear rtv 0  0 0 0 0
 point-size 30.0 1.0 64.0
 point-sprite on
 draw point list 1
-probe (310, 230) rgba (1, 1, 0, 1)
-probe (330, 230) rgba (0, 1, 1, 1)
-probe (310, 250) rgba (1, 0, 0, 1)
-probe (330, 250) rgba (1, 0, 1, 1)
+bug(mvk & sm>=4 & sm<6) probe (310, 230) f32(1, 1, 0, 1)
+bug(mvk & sm>=4 & sm<6) probe (330, 230) f32(0, 1, 1, 1)
+bug(mvk & sm>=4 & sm<6) probe (310, 250) f32(1, 0, 0, 1)
+bug(mvk & sm>=4 & sm<6) probe (330, 250) f32(1, 0, 1, 1)
diff --git a/tests/hlsl/return-semantics.shader_test b/tests/hlsl/return-semantics.shader_test
index b3cb303e0..8ef4cd410 100644
--- a/tests/hlsl/return-semantics.shader_test
+++ b/tests/hlsl/return-semantics.shader_test
@@ -39,4 +39,4 @@ float4 main(float cc : SEMC, float aa : SEMA, float bb[2] : SEMB) : sv_target
 
 [test]
 bug(mvk & sm>=6) draw quad
-todo(sm<6) probe (0, 0) f32(300, 100, 200, 201)
+probe (0, 0) f32(300, 100, 200, 201)