shader: support packed D16 XYZW formatted buffer stores

Co-authored-by: fxpw <fxpw@outlook.com>

src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemory.cpp

@@ -815,18 +815,23 @@ void StoreFormattedPrepared(ValueEmitContext& ctx, const IR::Inst& inst,
                              uint32_t components) {
 	if (resource.object_pointer == 0u) return;
 	// RDNA2 formatted stores transfer the entire format, zero-filling missing VGPRs.
+	const auto count = mem.data_bits == 16u ? mem.component_count : components;
 	const auto plan = PrepareFormattedMemory(ctx, inst, mem, resource, info);
 	EmitIfCondition(ctx.state, plan.in_bounds, [&]() {
 		auto packed = ConstantU32(ctx.state, 0);
 		for (uint32_t component = 0; component < info.component_count; component++) {
 			auto value = ConstantU32(ctx.state, 0);
-			if (component < components) {
+			if (component < count) {
+				const auto word = mem.data_bits == 16u ? component / 2u : component;
 				value = data;
 				if (components != 1u) {
 					value = ctx.state.builder.AllocateId();
 					ctx.state.builder.AddFunction(spv::OpCompositeExtract, TypeU32(ctx.state),
-					                              value, data, component);
+					                              value, data, word);
 				}
+				if (mem.data_bits == 16u && (component & 1u) != 0u)
+					value = Binary(ctx.state, spv::OpShiftRightLogical, TypeU32(ctx.state), value,
+					               ConstantU32(ctx.state, 16u));
 			}
 			value = EncodeFormattedStoreComponent(ctx.state, info, component, value, mem.data_bits);
 			const auto bits = info.component_bits[component];

src/graphics/shader/recompiler/frontend/decode/MemoryOps.cpp

@@ -70,6 +70,7 @@ constexpr MemoryOpcodeInfo MUBUF_OPCODE_LIST[] = {
     {0x80u, Opcode::BUFFER_LOAD_FORMAT_D16_X, 1, 16, false, false, true},
     {0x83u, Opcode::BUFFER_LOAD_FORMAT_D16_XYZW, 2, 16, false, false, true},
     {0x84u, Opcode::BUFFER_STORE_FORMAT_D16_X, 1, 16, false, false, true},
+    {0x87u, Opcode::BUFFER_STORE_FORMAT_D16_XYZW, 2, 16, false, false, true},
 };
 
 constexpr MemoryOpcodeInfo MTBUF_OPCODE_LIST[] = {
@@ -307,7 +308,8 @@ void DecodeMubuf(uint32_t pc, std::span<const uint32_t> code, uint32_t word_inde
 	}
 
 	DecodeVectorGpr(vdata, inst.dst);
-	if (inst.opcode == Opcode::BUFFER_LOAD_FORMAT_D16_XYZW) {
+	if (inst.opcode == Opcode::BUFFER_LOAD_FORMAT_D16_XYZW ||
+	    inst.opcode == Opcode::BUFFER_STORE_FORMAT_D16_XYZW) {
 		inst.data_components = 4u;
 	}
 	if (inst.opcode == Opcode::BUFFER_LOAD_FORMAT_D16_X ||

src/graphics/shader/recompiler/frontend/decode/ShaderDecoder.cpp

@@ -637,6 +637,7 @@ std::string InstructionToString(const Instruction& inst) {
 		case Opcode::BUFFER_STORE_FORMAT_XYZ:
 		case Opcode::BUFFER_STORE_FORMAT_XYZW:
 		case Opcode::BUFFER_STORE_FORMAT_D16_X:
+		case Opcode::BUFFER_STORE_FORMAT_D16_XYZW:
 		case Opcode::BUFFER_LOAD_UBYTE:
 		case Opcode::BUFFER_LOAD_UBYTE_D16:
 		case Opcode::BUFFER_LOAD_USHORT:

src/graphics/shader/recompiler/frontend/decode/ShaderDecoder.h

@@ -553,6 +553,7 @@ enum class Opcode {
 	BUFFER_STORE_FORMAT_XYZ,
 	BUFFER_STORE_FORMAT_XYZW,
 	BUFFER_STORE_FORMAT_D16_X,
+	BUFFER_STORE_FORMAT_D16_XYZW,
 	BUFFER_LOAD_UBYTE,
 	BUFFER_LOAD_UBYTE_D16,
 	BUFFER_LOAD_SBYTE,

src/graphics/shader/recompiler/frontend/translate/Memory.cpp

@@ -187,7 +187,7 @@ Decoder::Operand MemorySourceAt(const Decoder::Instruction& decoded, uint32_t in
 	if (decoded.family == Decoder::Family::MUBUF || decoded.family == Decoder::Family::MTBUF) {
 		const bool store_or_atomic =
 		    (decoded.opcode >= Decoder::Opcode::BUFFER_STORE_FORMAT_X &&
-		     decoded.opcode <= Decoder::Opcode::BUFFER_STORE_FORMAT_D16_X) ||
+		     decoded.opcode <= Decoder::Opcode::BUFFER_STORE_FORMAT_D16_XYZW) ||
 		    (decoded.opcode >= Decoder::Opcode::BUFFER_STORE_BYTE &&
 		     decoded.opcode <= Decoder::Opcode::BUFFER_STORE_DWORDX4) ||
 		    (decoded.opcode >= Decoder::Opcode::TBUFFER_STORE_FORMAT_X &&
@@ -1006,6 +1006,7 @@ void Translator::EmitMemory(const Decoder::Instruction& inst) {
 		case Decoder::Opcode::BUFFER_STORE_FORMAT_XYZ:
 		case Decoder::Opcode::BUFFER_STORE_FORMAT_XYZW:
 		case Decoder::Opcode::BUFFER_STORE_FORMAT_D16_X:
+		case Decoder::Opcode::BUFFER_STORE_FORMAT_D16_XYZW:
 		case Decoder::Opcode::TBUFFER_STORE_FORMAT_X:
 		case Decoder::Opcode::TBUFFER_STORE_FORMAT_XY:
 		case Decoder::Opcode::TBUFFER_STORE_FORMAT_XYZ:

src/graphics/shader/recompiler/ir/Program.cpp

@@ -474,11 +474,12 @@ void ValidateProgram(const Program& program, bool require_ssa) {
 					return Fail("indirect buffer requires a scalar, raw 8/16-bit, raw DWORD x1/x2/x3/x4, "
 					            "or formatted X load");
 				}
-				const bool packed_formatted_load =
+				const auto access = BufferAccessOf(inst.GetOpcode());
-				    BufferAccessOf(inst.GetOpcode()) == BufferAccess::Read && memory.formatted &&
+				const bool packed_formatted =
+				    (access == BufferAccess::Read || access == BufferAccess::Write) && memory.formatted &&
 				    memory.data_bits == 16u && memory.component_count == 4u && buffer_components == 2u;
 				if (buffer_components > 1u &&
-				    (!vector_buffer || (memory.data_bits != 32u && !packed_formatted_load) ||
+				    (!vector_buffer || (memory.data_bits != 32u && !packed_formatted) ||
 				     memory.data_dwords != buffer_components || memory.component_index != 0u)) {
 					return Fail(fmt::format("{} has inconsistent native-wide metadata",
 					                        ValueOpcodeName(inst.GetOpcode())));

tests/ShaderRecompilerComputeTests.cpp

@@ -31393,39 +31393,87 @@ std::vector<TestCase> BufferLoadFormatD16XyzwCases() {
   return cases;
 }
 
-std::vector<TestCase> BufferStoreFormatD16XCases() {
+std::vector<TestCase> BufferStoreFormatD16Cases() {
   using F = Prospero::BufferFormat;
   using O = ShaderOpcode;
   struct StoreCase {
     const char *name;
     F format;
-    u32 input;
+    std::array<u32, 2> input;
-    u32 expected;
+    std::array<u32, 4> expected;
+    u32 words = 1;
+    bool wide = false;
   };
   constexpr StoreCase inputs[] = {
-      {"BufferStoreFormatD16XFloat16", F::k16Float, 0x7e63c000u, 0xffffc000u},
+      {"BufferStoreFormatD16XFloat16", F::k16Float, {0x7e63c000u}, {0xffffc000u}},
-      {"BufferStoreFormatD16XFloat32", F::k32Float, 0x7e633800u, 0x3f000000u},
+      {"BufferStoreFormatD16XFloat32", F::k32Float, {0x7e633800u}, {0x3f000000u}},
-      {"BufferStoreFormatD16XUint32", F::k32UInt, 0x7e63fffeu, 0x0000fffeu},
+      {"BufferStoreFormatD16XUint32", F::k32UInt, {0x7e63fffeu}, {0x0000fffeu}},
-      {"BufferStoreFormatD16XSint32", F::k32SInt, 0x7e63fffeu, 0xfffffffeu},
+      {"BufferStoreFormatD16XSint32", F::k32SInt, {0x7e63fffeu}, {0xfffffffeu}},
-      {"BufferStoreFormatD16XUnorm8", F::k8UNorm, 0x7e633800u, 0xffffff80u},
+      {"BufferStoreFormatD16XUnorm8", F::k8UNorm, {0x7e633800u}, {0xffffff80u}},
-      {"BufferStoreFormatD16XSnorm16", F::k16SNorm, 0x7e63b800u, 0xffffc000u},
+      {"BufferStoreFormatD16XSnorm16", F::k16SNorm, {0x7e63b800u}, {0xffffc000u}},
-      {"BufferStoreFormatD16XZeroFill", F::k16_16Float, 0x7e633800u, 0x00003800u},
+      {"BufferStoreFormatD16XZeroFill", F::k16_16Float, {0x7e633800u}, {0x00003800u}},
+      {"BufferStoreFormatD16XyzwFloat16", F::k16_16_16_16Float,
+       {0xc0003c00u, 0x34003800u}, {0xc0003c00u, 0x34003800u}, 2, true},
+      {"BufferStoreFormatD16XyzwFloat32", F::k32_32_32_32Float,
+       {0xc0003c00u, 0x34003800u}, {0x3f800000u, 0xc0000000u, 0x3f000000u, 0x3e800000u}, 4, true},
+      {"BufferStoreFormatD16XyzwUint32", F::k32_32_32_32UInt,
+       {0xffff0001u, 0x80007fffu}, {1u, 0xffffu, 0x7fffu, 0x8000u}, 4, true},
+      {"BufferStoreFormatD16XyzwSint32", F::k32_32_32_32SInt,
+       {0x8000ffffu, 0x7fff0001u}, {0xffffffffu, 0xffff8000u, 1u, 0x7fffu}, 4, true},
+      {"BufferStoreFormatD16XyzwSint8", F::k8_8_8_8SInt,
+       {0xff80ff7fu, 0x0080007fu}, {0x7f7f8080u}, 1, true},
+      {"BufferStoreFormatD16XyzwUnorm8", F::k8_8_8_8UNorm,
+       {0x3800c000u, 0x40003c00u}, {0xffff8000u}, 1, true},
   };
   std::vector<TestCase> cases;
-  for (const auto &input : inputs) {
+  for (const auto &input : inputs)
+  for (const u32 wave : {32u, 64u}) {
+    if (!input.wide && wave == 64u) continue;
     TestCase test;
     test.name = input.name;
-    AppendVMovLiteral(&test.code, 66, input.input);
+    AppendVMovLiteral(&test.code, 66, input.input[0]);
-    AppendVMovU32(&test.code, 65, 1);
+    if (input.wide) {
-    // Match the captured CS824586f35db5c80d instruction at pc 0x484.
+      AppendVMovLiteral(&test.code, 67, input.input[1]);
-    test.code.push_back(EncodeMubuf0(0x84u, 0, true, false));
+      test.code.push_back(EncodeVop2(0x25, 65, InlineU32(1), 0));
-    test.code.push_back(EncodeMubuf1(66, 0, 65));
+      AppendSMovLiteral(&test.code, 126, 0x55555555u);
+      AppendSMovLiteral(&test.code, 127, wave == 64u ? 0x55555555u : 0u);
+    } else AppendVMovU32(&test.code, 65, 1);
+    const auto store = [&] {
+      // XYZW matches captured e21c6000/8002000e with v66:67/v65 and s0:3.
+      test.code.push_back(EncodeMubuf0(input.wide ? 0x87u : 0x84u, 0, true, false, input.wide));
+      test.code.push_back(EncodeMubuf1(66, 0, 65));
+    };
+    store();
+    if (input.wide) {
+      AppendVMovU32(&test.code, 65, wave + 1u);
+      store(); // Active lanes must not write beyond the descriptor range.
+      AppendVMovU32(&test.code, 65, 0u);
+      test.code.push_back(EncodeSop1(0x04, 126, InlineU32(0)));
+      store(); // EXEC-zero must preserve the leading record.
+    }
     AppendEnd(&test.code);
-    test.initial = {0xdeadbeefu, 0xffffffffu, 0xcafef00du};
+    if (input.wide) {
-    test.expected = {0xdeadbeefu, input.expected, 0xcafef00du};
+      test.initial.assign((wave + 2u) * input.words, 0xdeadbeefu);
-    test.user_data = MakeStructuredStorageBufferData(4, 3, false, BufferFormat(input.format));
+      test.expected = test.initial;
+      for (u32 lane = 0; lane + 2u < wave; lane += 2u)
+        std::copy_n(input.expected.begin(), input.words,
+                    test.expected.begin() + (lane + 1u) * input.words);
+      test.storage_buffer_range_bytes = wave * input.words * 4u - 1u;
+      test.has_compute_info = true;
+      test.compute_info.wave_size = wave;
+      test.compute_info.threads_num[0] = wave;
+      test.compute_info.threads_num[1] = test.compute_info.threads_num[2] = 1u;
+      test.compute_info.thread_ids_num = 1u;
+      test.required_spirv = {"Volatile"};
+    } else {
+      test.initial = {0xdeadbeefu, 0xffffffffu, 0xcafef00du};
+      test.expected = {0xdeadbeefu, input.expected[0], 0xcafef00du};
+    }
+    test.user_data = MakeStructuredStorageBufferData(input.words * 4u, input.wide ? wave : 3u,
+                                                    false, BufferFormat(input.format));
     test.has_user_data = true;
-    test.opcodes = {O::V_MOV_B32, O::BUFFER_STORE_FORMAT_D16_X, O::S_ENDPGM};
+    test.opcodes = {O::V_MOV_B32, input.wide ? O::BUFFER_STORE_FORMAT_D16_XYZW :
+                                            O::BUFFER_STORE_FORMAT_D16_X, O::S_ENDPGM};
     cases.push_back(std::move(test));
   }
   return cases;
@@ -40597,7 +40645,7 @@ std::vector<TestCase> MakeCases() {
   for (auto &test : BufferLoadFormatD16XyzwCases()) {
     cases.push_back(std::move(test));
   }
-  for (auto &test : BufferStoreFormatD16XCases()) {
+  for (auto &test : BufferStoreFormatD16Cases()) {
     cases.push_back(std::move(test));
   }
   for (auto &test : FormattedStoreConversionCases()) {
@@ -46271,7 +46319,7 @@ int main(int argc, char **argv) {
     for (const auto &test : BufferLoadFormatD16XyzwCases()) {
       RunCase(&vulkan, test);
     }
-    for (const auto &test : BufferStoreFormatD16XCases()) {
+    for (const auto &test : BufferStoreFormatD16Cases()) {
       RunCase(&vulkan, test);
     }
     for (const auto &test : FormattedStoreConversionCases()) {

tests/shaderCfgTests.cpp

@@ -15375,6 +15375,24 @@ void TestCapturedBufferLoadUbyteD16Decode() {
         "captured BUFFER_LOAD_UBYTE_D16 lost its preserved half, address, or coherence flags");
 }
 
+void TestCapturedBufferStoreD16XyzwDecode() {
+  namespace Decoder = ShaderRecompiler::Decoder;
+  const uint32_t code[] = {0xe21c6000u, 0x8002000eu};
+  Decoder::Instruction inst;
+  Decoder::DecodeInstruction(code, 0, inst);
+  Check(inst.opcode == Decoder::Opcode::BUFFER_STORE_FORMAT_D16_XYZW &&
+            inst.word_count == 2 && inst.data_dwords == 2 && inst.data_bits == 16 &&
+            inst.data_components == 4 && inst.formatted && !inst.typed,
+        "captured D16 XYZW store lost its four components in two packed VGPRs");
+  Check(inst.dst.kind == Decoder::OperandKind::Vgpr && inst.dst.reg == 0 &&
+            inst.dst.sdwa_sel == 6 &&
+            inst.src0.kind == Decoder::OperandKind::Vgpr && inst.src0.reg == 14 &&
+            inst.src1.kind == Decoder::OperandKind::Sgpr && inst.src1.reg == 8 &&
+            inst.src2.kind == Decoder::OperandKind::IntegerInlineConstant && inst.src2.value == 0 &&
+            inst.idxen && !inst.offen && inst.offset == 0 && inst.glc && !inst.dlc && !inst.slc,
+        "captured D16 XYZW store lost its full source VGPRs, index, or coherence flags");
+}
+
 void TestNewShaderRecompilerUnsupportedMemoryDecode() {
   const uint32_t mubuf_unknown[] = {EncodeMubuf0(0x7b), EncodeMubuf1(0, 0, 1),
                                     0xbf810000u};
@@ -16059,6 +16077,7 @@ int main() {
   TestPixelProgramCacheBindingIdentity();
   TestGraphicsPushConstantPlacement();
   TestCapturedBufferLoadUbyteD16Decode();
+  TestCapturedBufferStoreD16XyzwDecode();
   TestNewShaderRecompilerUnsupportedMemoryDecode();
 
   TestNewShaderRecompilerVop3LaneReadDestinationEncoding();