From 71a3d55e8cc5f7cf9096c0990d4dd5961c8cef7d Mon Sep 17 00:00:00 2001
From: Elizabeth Figura <zfigura@codeweavers.com>
Date: Mon, 10 Jun 2024 17:06:54 -0500
Subject: [PATCH] vkd3d-shader/hlsl: Implement the GetRenderTargetSampleCount()
 intrinsic.

---
 libs/vkd3d-shader/hlsl.c                  |  1 +
 libs/vkd3d-shader/hlsl.h                  |  1 +
 libs/vkd3d-shader/hlsl.y                  | 15 +++++++++++++
 libs/vkd3d-shader/tpf.c                   | 27 +++++++++++++++++++++++
 tests/hlsl/rt-get-sample-info.shader_test |  4 ++--
 5 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/libs/vkd3d-shader/hlsl.c b/libs/vkd3d-shader/hlsl.c
index de00aa1b..acf50869 100644
--- a/libs/vkd3d-shader/hlsl.c
+++ b/libs/vkd3d-shader/hlsl.c
@@ -2793,6 +2793,7 @@ const char *debug_hlsl_expr_op(enum hlsl_ir_expr_op op)
     static const char *const op_names[] =
     {
         [HLSL_OP0_VOID]         = "void",
+        [HLSL_OP0_RASTERIZER_SAMPLE_COUNT] = "GetRenderTargetSampleCount",
 
         [HLSL_OP1_ABS]          = "abs",
         [HLSL_OP1_BIT_NOT]      = "~",
diff --git a/libs/vkd3d-shader/hlsl.h b/libs/vkd3d-shader/hlsl.h
index 4262d36c..0f78f0b5 100644
--- a/libs/vkd3d-shader/hlsl.h
+++ b/libs/vkd3d-shader/hlsl.h
@@ -638,6 +638,7 @@ struct hlsl_ir_switch
 enum hlsl_ir_expr_op
 {
     HLSL_OP0_VOID,
+    HLSL_OP0_RASTERIZER_SAMPLE_COUNT,
 
     HLSL_OP1_ABS,
     HLSL_OP1_BIT_NOT,
diff --git a/libs/vkd3d-shader/hlsl.y b/libs/vkd3d-shader/hlsl.y
index 61a98e7a..7b058a65 100644
--- a/libs/vkd3d-shader/hlsl.y
+++ b/libs/vkd3d-shader/hlsl.y
@@ -4688,6 +4688,20 @@ static bool intrinsic_d3dcolor_to_ubyte4(struct hlsl_ctx *ctx,
     return true;
 }
 
+static bool intrinsic_GetRenderTargetSampleCount(struct hlsl_ctx *ctx,
+        const struct parse_initializer *params, const struct vkd3d_shader_location *loc)
+{
+    struct hlsl_ir_node *operands[HLSL_MAX_OPERANDS] = {0};
+    struct hlsl_ir_node *expr;
+
+    if (!(expr = hlsl_new_expr(ctx, HLSL_OP0_RASTERIZER_SAMPLE_COUNT,
+            operands, hlsl_get_scalar_type(ctx, HLSL_TYPE_UINT), loc)))
+        return false;
+    hlsl_block_add_instr(params->instrs, expr);
+
+    return true;
+}
+
 static const struct intrinsic_function
 {
     const char *name;
@@ -4700,6 +4714,7 @@ intrinsic_functions[] =
 {
     /* Note: these entries should be kept in alphabetical order. */
     {"D3DCOLORtoUBYTE4",                    1, true,  intrinsic_d3dcolor_to_ubyte4},
+    {"GetRenderTargetSampleCount",          0, true,  intrinsic_GetRenderTargetSampleCount},
     {"abs",                                 1, true,  intrinsic_abs},
     {"acos",                                1, true,  intrinsic_acos},
     {"all",                                 1, true,  intrinsic_all},
diff --git a/libs/vkd3d-shader/tpf.c b/libs/vkd3d-shader/tpf.c
index ca7cdfd5..0c4d93dc 100644
--- a/libs/vkd3d-shader/tpf.c
+++ b/libs/vkd3d-shader/tpf.c
@@ -5042,6 +5042,25 @@ static void write_sm4_store_uav_typed(const struct tpf_writer *tpf, const struct
     write_sm4_instruction(tpf, &instr);
 }
 
+static void write_sm4_rasterizer_sample_count(const struct tpf_writer *tpf, const struct hlsl_ir_node *dst)
+{
+    struct sm4_instruction instr;
+
+    memset(&instr, 0, sizeof(instr));
+    instr.opcode = VKD3D_SM4_OP_SAMPLE_INFO;
+    instr.extra_bits |= VKD3DSI_SAMPLE_INFO_UINT << VKD3D_SM4_INSTRUCTION_FLAGS_SHIFT;
+
+    sm4_dst_from_node(&instr.dsts[0], dst);
+    instr.dst_count = 1;
+
+    instr.srcs[0].reg.type = VKD3DSPR_RASTERIZER;
+    instr.srcs[0].reg.dimension = VSIR_DIMENSION_VEC4;
+    instr.srcs[0].swizzle = VKD3D_SHADER_SWIZZLE(X, X, X, X);
+    instr.src_count = 1;
+
+    write_sm4_instruction(tpf, &instr);
+}
+
 static void write_sm4_expr(const struct tpf_writer *tpf, const struct hlsl_ir_expr *expr)
 {
     const struct hlsl_ir_node *arg1 = expr->operands[0].node;
@@ -5057,6 +5076,14 @@ static void write_sm4_expr(const struct tpf_writer *tpf, const struct hlsl_ir_ex
 
     switch (expr->op)
     {
+        case HLSL_OP0_RASTERIZER_SAMPLE_COUNT:
+            if (tpf->ctx->profile->type == VKD3D_SHADER_TYPE_PIXEL && hlsl_version_ge(tpf->ctx, 4, 1))
+                write_sm4_rasterizer_sample_count(tpf, &expr->node);
+            else
+                hlsl_error(tpf->ctx, &expr->node.loc, VKD3D_SHADER_ERROR_HLSL_INCOMPATIBLE_PROFILE,
+                        "GetRenderTargetSampleCount() can only be used from a pixel shader using version 4.1 or higher.");
+            break;
+
         case HLSL_OP1_ABS:
             switch (dst_type->e.numeric.type)
             {
diff --git a/tests/hlsl/rt-get-sample-info.shader_test b/tests/hlsl/rt-get-sample-info.shader_test
index c552f17b..4168e769 100644
--- a/tests/hlsl/rt-get-sample-info.shader_test
+++ b/tests/hlsl/rt-get-sample-info.shader_test
@@ -6,14 +6,14 @@ shader model >= 4.1
 format r32g32b32a32 float
 size (2dms, 4, 640, 480)
 
-[pixel shader todo]
+[pixel shader]
 float4 main() : sv_target
 {
     return float4(GetRenderTargetSampleCount(), 0, 0, 0);
 }
 
 [test]
-todo(sm<6) draw quad
+todo(glsl) draw quad
 probe (0, 0) rgba (4.0, 0.0, 0.0, 0.0)