shader: implement BUFFER_ATOMIC_CMPSWAP (#357)

shader: implement RDNA2 buffer atomic compare swap

Co-authored-by: nmzik <Nmzik@mail.ru>

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

@@ -614,8 +614,9 @@ void FormattedStore(ValueEmitContext& ctx, const IR::Inst& inst, const IR::Memor
 	});
 }
 
-uint32_t AtomicOpcode(IR::ValueOpcode opcode) {
+uint32_t SpirvAtomicOpcode(IR::ValueOpcode opcode) {
 	switch (opcode) {
+		case IR::ValueOpcode::BufferAtomicCmpSwap32: return OpAtomicCompareExchange;
 		case IR::ValueOpcode::BufferAtomicSwap32:
 		case IR::ValueOpcode::BufferAtomicSwap64:
 		case IR::ValueOpcode::SharedAtomicSwap32: return OpAtomicExchange;
@@ -642,23 +643,37 @@ uint32_t AtomicOpcode(IR::ValueOpcode opcode) {
 	}
 }
 
-uint32_t EmitAtomic(ValueEmitContext& ctx, const IR::Inst& inst, const IR::MemoryInfo& mem,
+uint32_t EmitAtomicOperation(ValueEmitContext& ctx, const IR::Inst& inst, uint32_t pointer,
-                    bool return_value) {
+                             uint32_t scope) {
+	const auto old = ctx.state.builder.AllocateId();
+	if (inst.GetOpcode() == IR::ValueOpcode::BufferAtomicCmpSwap32) {
+		const auto desired    = ctx.Arg(inst, inst.NumArgs() - 3);
+		const auto comparator = ctx.Arg(inst, inst.NumArgs() - 2);
+		ctx.state.builder.AddFunction(
+		    {OpAtomicCompareExchange, TypeU32(ctx.state), old, pointer, ConstantU32(ctx.state, scope),
+		     ConstantU32(ctx.state, MemorySemanticsNone),
+		     ConstantU32(ctx.state, MemorySemanticsNone), desired, comparator});
+	} else {
+		const auto value = ctx.Arg(inst, inst.NumArgs() - 2);
+		ctx.state.builder.AddFunction(
+		    {SpirvAtomicOpcode(inst.GetOpcode()), TypeU32(ctx.state), old, pointer,
+		     ConstantU32(ctx.state, scope), ConstantU32(ctx.state, MemorySemanticsNone), value});
+	}
+	return old;
+}
+
+uint32_t EmitAtomic(ValueEmitContext& ctx, const IR::Inst& inst, const IR::MemoryInfo& mem) {
 	const auto result =
 	    EmitValueOrZeroIfCondition(ctx.state, ctx.Arg(inst, inst.NumArgs() - 1), [&]() {
 		    const auto access = PrepareMemoryElement(ctx, mem, DwordIndex(ctx, inst, mem));
 		    return EmitValueOrZeroIfCondition(
 		        ctx.state, EmitMemoryElementInBounds(ctx.state, access.resource, access.index),
 		        [&]() {
-			        const auto value = ctx.Arg(inst, inst.NumArgs() - 2);
-			        const auto old   = ctx.state.builder.AllocateId();
 			        const auto scope =
 			            mem.kind == IR::ResourceKind::Lds ? ScopeWorkgroup : ScopeDevice;
-			        ctx.state.builder.AddFunction(
+			        const auto pointer =
-			            {AtomicOpcode(inst.GetOpcode()), TypeU32(ctx.state), old,
+			            EmitMemoryElementPointer(ctx.state, access.resource, access.index);
-			             EmitMemoryElementPointer(ctx.state, access.resource, access.index),
+			        const auto old = EmitAtomicOperation(ctx, inst, pointer, scope);
-			             ConstantU32(ctx.state, scope), ConstantU32(ctx.state, MemorySemanticsNone),
-			             value});
 			        if (mem.kind == IR::ResourceKind::Lds) {
 				        const auto semantics =
 				            MemorySemanticsAcquireRelease | MemorySemanticsWorkgroupMemory;
@@ -671,7 +686,7 @@ uint32_t EmitAtomic(ValueEmitContext& ctx, const IR::Inst& inst, const IR::Memor
 			        return old;
 		        });
 	    });
-	return return_value ? result : ConstantU32(ctx.state, 0);
+	return result;
 }
 
 uint32_t EmitBufferAtomic64(ValueEmitContext& ctx, const IR::Inst& inst,
@@ -692,7 +707,7 @@ uint32_t EmitBufferAtomic64(ValueEmitContext& ctx, const IR::Inst& inst,
 			                                 ctx.Arg(inst, inst.NumArgs() - 2));
 			        const auto old   = state.builder.AllocateId();
 			        state.builder.AddFunction(
-			            {AtomicOpcode(inst.GetOpcode()), TypeScalarU64(state), old,
+			            {SpirvAtomicOpcode(inst.GetOpcode()), TypeScalarU64(state), old,
 			             EmitStorageBufferElementPointer(
 			                 state, resource, index, TypeStorageBufferU64ElementPointer(state)),
 			             ConstantU32(state, ScopeDevice),
@@ -1210,13 +1225,13 @@ bool EmitValueMemory(ValueEmitContext& ctx, const IR::Inst& inst) {
 			StoreWord(ctx, inst, mem);
 		return true;
 	}
-	const auto atomic_opcode = AtomicOpcode(op);
+	const auto atomic_opcode = SpirvAtomicOpcode(op);
 	if (atomic_opcode != 0 && inst.GetType() == IR::Type::U64) {
 		ctx.Define(inst, EmitBufferAtomic64(ctx, inst, ctx.Memory(inst)));
 		return true;
 	}
 	if (atomic_opcode != 0) {
-		ctx.Define(inst, EmitAtomic(ctx, inst, ctx.Memory(inst), true));
+		ctx.Define(inst, EmitAtomic(ctx, inst, ctx.Memory(inst)));
 		return true;
 	}
 	if (op == IR::ValueOpcode::BufferAtomicFMin32 || op == IR::ValueOpcode::BufferAtomicFMax32) {

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

@@ -47,6 +47,7 @@ constexpr MemoryOpcodeInfo MUBUF_OPCODE_LIST[] = {
     {0x1eu, Opcode::BUFFER_STORE_DWORDX4, 4, 32},
     {0x1fu, Opcode::BUFFER_STORE_DWORDX3, 3, 32},
     {0x30u, Opcode::BUFFER_ATOMIC_SWAP, 1, 32},
+    {0x31u, Opcode::BUFFER_ATOMIC_CMPSWAP, 1, 32},
     {0x32u, Opcode::BUFFER_ATOMIC_ADD, 1, 32},
     {0x33u, Opcode::BUFFER_ATOMIC_SUB, 1, 32},
     {0x35u, Opcode::BUFFER_ATOMIC_SMIN, 1, 32},

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

@@ -592,6 +592,7 @@ std::string InstructionToString(const Instruction& inst) {
 		case Opcode::TBUFFER_STORE_FORMAT_XYZ:
 		case Opcode::TBUFFER_STORE_FORMAT_XYZW:
 		case Opcode::BUFFER_ATOMIC_SWAP:
+		case Opcode::BUFFER_ATOMIC_CMPSWAP:
 		case Opcode::BUFFER_ATOMIC_SWAP_X2:
 		case Opcode::BUFFER_ATOMIC_ADD:
 		case Opcode::BUFFER_ATOMIC_SUB:

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

@@ -455,6 +455,7 @@ enum class Opcode {
 	TBUFFER_STORE_FORMAT_XYZ,
 	TBUFFER_STORE_FORMAT_XYZW,
 	BUFFER_ATOMIC_SWAP,
+	BUFFER_ATOMIC_CMPSWAP,
 	BUFFER_ATOMIC_SWAP_X2,
 	BUFFER_ATOMIC_ADD,
 	BUFFER_ATOMIC_SUB,

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

@@ -523,12 +523,23 @@ bool Translator::BUFFER_ATOMIC(const Decoder::Instruction& inst, IR::ValueOpcode
 	const auto resource = GetBufferResource(memory);
 	const auto address  = ReadBufferAddress(inst, 1);
 	const auto data_src = MemorySourceAt(inst, 0);
-	const IR::Value data =
+	const auto flags    = AddMemoryInfo(memory, inst.pc);
-	    memory.data_dwords == 2u ? IR::Value(ReadU64(data_src)) : IR::Value(ReadU32(data_src));
+	IR::Value  result;
-	const auto result = ir.Emit(opcode,
+	if (opcode == IR::ValueOpcode::BufferAtomicCmpSwap32) {
-	                            {resource, address.index, address.offset, address.soffset, data,
+		const auto desired    = ReadU32(data_src);
-	                             ir.GetExec()},
+		const auto comparator = ReadU32(OffsetOperand(data_src, 1u));
-	                            AddMemoryInfo(memory, inst.pc));
+		result = ir.Emit(opcode,
+		                 {resource, address.index, address.offset, address.soffset, desired,
+		                  comparator, ir.GetExec()},
+		                 flags);
+	} else {
+		const IR::Value value =
+		    memory.data_dwords == 2u ? IR::Value(ReadU64(data_src)) : IR::Value(ReadU32(data_src));
+		result = ir.Emit(opcode,
+		                 {resource, address.index, address.offset, address.soffset, value,
+		                  ir.GetExec()},
+		                 flags);
+	}
 	if (inst.glc) {
 		WriteOperand(inst.dst, result);
 	}
@@ -922,6 +933,8 @@ bool Translator::EmitMemory(const Decoder::Instruction& inst) {
 
 		case Decoder::Opcode::BUFFER_ATOMIC_SWAP:
 			return BUFFER_ATOMIC(inst, IR::ValueOpcode::BufferAtomicSwap32);
+		case Decoder::Opcode::BUFFER_ATOMIC_CMPSWAP:
+			return BUFFER_ATOMIC(inst, IR::ValueOpcode::BufferAtomicCmpSwap32);
 		case Decoder::Opcode::BUFFER_ATOMIC_SWAP_X2:
 			return BUFFER_ATOMIC(inst, IR::ValueOpcode::BufferAtomicSwap64);
 		case Decoder::Opcode::BUFFER_ATOMIC_ADD:

src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.cpp

@@ -111,6 +111,7 @@ BufferAccess BufferAccessOf(ValueOpcode opcode) {
 		case ValueOpcode::StoreBufferU32x3:
 		case ValueOpcode::StoreBufferU32x4: return BufferAccess::Write;
 		case ValueOpcode::BufferAtomicSwap32:
+		case ValueOpcode::BufferAtomicCmpSwap32:
 		case ValueOpcode::BufferAtomicSwap64:
 		case ValueOpcode::BufferAtomicIAdd32:
 		case ValueOpcode::BufferAtomicISub32:

src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.inc

@@ -222,6 +222,7 @@ VALUE_OPCODE(StoreBufferU32x2, Void, BufferResource, U32, U32, U32, U32x2, U1)
 VALUE_OPCODE(StoreBufferU32x3, Void, BufferResource, U32, U32, U32, U32x3, U1)
 VALUE_OPCODE(StoreBufferU32x4, Void, BufferResource, U32, U32, U32, U32x4, U1)
 VALUE_OPCODE(BufferAtomicSwap32, U32, BufferResource, U32, U32, U32, U32, U1)
+VALUE_OPCODE(BufferAtomicCmpSwap32, U32, BufferResource, U32, U32, U32, U32, U32, U1)
 VALUE_OPCODE(BufferAtomicSwap64, U64, BufferResource, U32, U32, U32, U64, U1)
 VALUE_OPCODE(BufferAtomicIAdd32, U32, BufferResource, U32, U32, U32, U32, U1)
 VALUE_OPCODE(BufferAtomicISub32, U32, BufferResource, U32, U32, U32, U32, U1)

tests/ShaderRecompilerComputeTests.cpp

@@ -12226,6 +12226,7 @@ CoverageClass ClassifyOpcode(ShaderOpcode opcode,
   case Opcode::TBUFFER_STORE_FORMAT_XYZ:
   case Opcode::TBUFFER_STORE_FORMAT_XYZW:
   case Opcode::BUFFER_ATOMIC_SWAP:
+  case Opcode::BUFFER_ATOMIC_CMPSWAP:
   case Opcode::BUFFER_ATOMIC_ADD:
   case Opcode::BUFFER_ATOMIC_SUB:
   case Opcode::BUFFER_ATOMIC_SMIN:
@@ -19450,6 +19451,52 @@ TestCase BufferAtomicVariants() {
        O::S_ENDPGM}};
 }
 
+TestCase BufferAtomicCmpSwapExactRaw() {
+  using O = ShaderOpcode;
+
+  std::vector<u32> code;
+  const auto Store = [&](u32 value, u32 index) {
+    AppendVMovU32(&code, 31, index * sizeof(u32));
+    code.push_back(EncodeMubuf0(0x1cu));
+    code.push_back(EncodeMubuf1(value, 8, 31));
+  };
+  AppendVMovU32(&code, 4, 1);
+  AppendVMovLiteral(&code, 1, 0x22222222u);
+  AppendVMovLiteral(&code, 2, 0x11111111u);
+  code.push_back(0xe0c46000u);
+  code.push_back(0x80080104u);
+  Store(1, 2);
+
+  AppendVMovLiteral(&code, 1, 0x33333333u);
+  AppendVMovLiteral(&code, 2, 0x44444444u);
+  code.push_back(0xe0c46000u);
+  code.push_back(0x80080104u);
+  Store(1, 3);
+  Store(2, 4);
+
+  AppendVMovLiteral(&code, 1, 0x55555555u);
+  AppendVMovLiteral(&code, 2, 0x22222222u);
+  code.push_back(EncodeMubuf0(0x31u, 0, true, false));
+  code.push_back(0x80080104u);
+  Store(1, 5);
+  AppendEnd(&code);
+
+  TestCase test;
+  test.name = "BufferAtomicCmpSwapExactRaw";
+  test.code = std::move(code);
+  test.initial = {0xaaaaaaaau, 0x11111111u, 0, 0, 0, 0};
+  test.expected = {0xaaaaaaaau, 0x55555555u, 0x11111111u,
+                   0x22222222u, 0x44444444u, 0x55555555u};
+  test.opcodes = {O::V_MOV_B32, O::BUFFER_ATOMIC_CMPSWAP,
+                  O::BUFFER_STORE_DWORD, O::S_ENDPGM};
+  test.required_spirv = {"OpAtomicCompareExchange"};
+  const auto descriptor = MakeStructuredStorageBufferData(
+      sizeof(u32), static_cast<u32>(test.initial.size()));
+  std::copy_n(descriptor.begin(), 4, test.user_data.begin() + 32);
+  test.has_user_data = true;
+  return test;
+}
+
 TestCase BufferAtomicGlc0DoesNotReturnOldValue() {
   using O = ShaderOpcode;
 
@@ -21654,6 +21701,7 @@ std::vector<TestCase> MakeCases() {
   AddCase(DsBpermuteCapturedExecOffsetAndWrap);
   AddCase(DsBpermuteWave64UsesIndependentHalves);
   AddCase(BufferAtomicVariants);
+  AddCase(BufferAtomicCmpSwapExactRaw);
   AddCase(BufferAtomicGlc0DoesNotReturnOldValue);
   AddCase(BufferAtomicFMinExactRawGlcModes);
   AddCase(BufferAtomicFMinSpecialValues);
@@ -25841,6 +25889,11 @@ int main(int argc, char **argv) {
     RunCase(&vulkan, BufferOffsetsUsePackedLaneAndStorageFallback());
     return 0;
   }
+  if (argc == 2 && std::strcmp(argv[1], "--buffer-cmpswap-only") == 0) {
+    VulkanHarness vulkan;
+    RunCase(&vulkan, BufferAtomicCmpSwapExactRaw());
+    return 0;
+  }
   if (argc == 2 && std::strcmp(argv[1], "--mapped-range-only") == 0) {
     VulkanHarness vulkan;
     vulkan.CheckGpuMappedRangeLifecycle();

tests/shaderCfgTests.cpp

@@ -8557,40 +8557,47 @@ void TestNewShaderRecompilerBufferLoadsGuardedByExec() {
 }
 
 void TestNewShaderRecompilerBufferAtomicsGuardedByBounds() {
-  const uint32_t shader[] = {
+  struct Case {
-      EncodeMubuf0(0x32, 0, true, true),
+    std::array<uint32_t, 3> shader;
-      EncodeMubuf1(0, 0, 1), // buffer_atomic_add
+    const char *decoded;
-      EncodeSopp(0x01),
+    const char *ir;
+    const char *spirv;
+  };
+  const Case cases[] = {
+      {{EncodeMubuf0(0x32, 0, true, true), EncodeMubuf1(0, 0, 1),
+        EncodeSopp(0x01)},
+       "BUFFER_ATOMIC_ADD", "BufferAtomicIAdd32", "OpAtomicIAdd"},
+      {{0xe0c46000u, 0x80080104u, EncodeSopp(0x01)},
+       "BUFFER_ATOMIC_CMPSWAP", "BufferAtomicCmpSwap32",
+       "OpAtomicCompareExchange"},
   };
-
   auto options = MakeCompileOptions(ShaderType::Compute);
   options.dump_ir = true;
 
-  auto result = RecompileForTest(shader, options);
+  for (const auto &test : cases) {
-  Check(Common::ContainsStr(result.decoded_dump, "BUFFER_ATOMIC_ADD"),
+    const auto result = RecompileForTest(test.shader, options);
-        "buffer atomic bounds regression did not decode MUBUF atomic");
+    Check(Common::ContainsStr(result.decoded_dump, test.decoded),
-  CheckSpirvBinaryValidates(result.spirv);
+          "buffer atomic bounds regression did not decode MUBUF atomic");
+    Check(Common::ContainsStr(result.ir_dump, test.ir),
+          "buffer atomic did not lower to its native IR opcode");
+    CheckSpirvBinaryValidates(result.spirv);
 
-  const auto source = DisassembleSpirvBinary(result.spirv);
+    const auto source = DisassembleSpirvBinary(result.spirv);
-  const auto array_length =
+    const auto array_length =
-      Common::FindIndex(source, std::string("OpArrayLength"), 0);
+        Common::FindIndex(source, std::string("OpArrayLength"), 0);
-  const auto bounds_branch = Common::FindIndex(
+    const auto bounds_branch = Common::FindIndex(
-      source, std::string("OpBranchConditional"), array_length);
+        source, std::string("OpBranchConditional"), array_length);
-  const auto atomic = Common::FindIndex(source, std::string("OpAtomicIAdd"), 0);
+    const auto atomic = Common::FindIndex(source, std::string(test.spirv), 0);
-  const auto memory_barrier =
+    const auto memory_barrier =
-      Common::FindIndex(source, std::string("OpMemoryBarrier"), atomic);
+        Common::FindIndex(source, std::string("OpMemoryBarrier"), atomic);
-  Check(array_length != Common::FIND_INVALID_INDEX,
+    Check(array_length != Common::FIND_INVALID_INDEX &&
-        "buffer atomic SPIR-V lacks storage buffer array-length bounds check");
+              bounds_branch != Common::FIND_INVALID_INDEX &&
-  Check(bounds_branch != Common::FIND_INVALID_INDEX,
+              atomic != Common::FIND_INVALID_INDEX &&
-        "buffer atomic SPIR-V lacks storage buffer bounds branch");
+              memory_barrier != Common::FIND_INVALID_INDEX,
-  Check(atomic != Common::FIND_INVALID_INDEX,
+          "buffer atomic SPIR-V lacks its bounds guard, operation, or barrier");
-        "buffer atomic SPIR-V lacks atomic operation");
+    Check(bounds_branch < atomic && atomic < memory_barrier,
-  Check(memory_barrier != Common::FIND_INVALID_INDEX,
+          "buffer atomic bounds guard, operation, and barrier are misordered");
-        "buffer atomic SPIR-V lacks memory barrier after atomic operation");
+  }
-  Check(bounds_branch < atomic,
-        "buffer atomic was emitted before bounds guard");
-  Check(atomic < memory_barrier,
-        "buffer atomic memory barrier was emitted before atomic");
 }
 
 void TestCapturedBufferAtomicsX2() {