Merge pull request #5791 from Ondra97/codex/fix-draw-mesh-index-cache

fix(agc): validate mesh index descriptors on cached draws

core/libs/prx/libSceAgcDriver/Execution/include/Driver/Dispatch/DispatchCache.hpp

@@ -29,6 +29,7 @@ struct DispatchVariant {
     std::vector<std::uint32_t> words;
     std::uint64_t forgetSerial = 0;
     std::shared_ptr<const ShaderRecompiler::RecompileResult> compiled;
+    std::array<std::uint32_t, 4> meshIndexDescriptor{};
 
     std::shared_ptr<const ShaderSnapshot> shader;
 

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Draw/DrawCache.cpp

@@ -125,7 +125,7 @@ void Driver::insertDrawEntry(std::uint64_t key, std::vector<std::shared_ptr<Disp
     for (std::size_t i = 0; i < fresh.size(); ++i) {
         if (fresh[i] == nullptr) continue;
         auto& variants = replacement->stages[i];
-        const auto present = std::find_if(variants.begin(), variants.end(), [&](const std::shared_ptr<DispatchVariant>& kept) { return kept->pushOffset == fresh[i]->pushOffset && kept->runs == fresh[i]->runs && kept->words == fresh[i]->words; });
+        const auto present = std::find_if(variants.begin(), variants.end(), [&](const std::shared_ptr<DispatchVariant>& kept) { return kept->pushOffset == fresh[i]->pushOffset && kept->meshIndexDescriptor == fresh[i]->meshIndexDescriptor && kept->runs == fresh[i]->runs && kept->words == fresh[i]->words; });
         if (present != variants.end()) {
             ++counters.present;
             fresh[i] = *present;

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Draw/DrawCapture.cpp

@@ -148,6 +148,9 @@ void Driver::cacheDrawStages(bool useDrawEntries, bool drawHit, const Pm4::DrawP
             if (stageCapture.compiled == nullptr || !CacheableResult(*stageCapture.compiled)) continue;
             auto variant = std::make_shared<DispatchVariant>();
             variant->compiled = stageCapture.compiled;
+            if (programs[i].binary.stage == ShaderRecompiler::ShaderStage::Mesh) {
+                std::copy_n(programs[i].userData.begin() + ShaderRecompiler::MeshIndexBufferUserWord, variant->meshIndexDescriptor.size(), variant->meshIndexDescriptor.begin());
+            }
             variant->shader = programs[i].snapshot;
             variant->forgetSerial = stageCapture.forgetSerial;
             variant->pushOffset = stageCapture.pushOffset;

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Draw/DrawLookup.cpp

@@ -1,6 +1,7 @@
 #include "prx/libSceAgcDriver/Execution/include/Driver/Driver.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Diagnostics.hpp"
 #include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
+#include <algorithm>
 
 namespace AgcDriver::DriverDetail {
 
@@ -107,6 +108,11 @@ void Driver::lookupDraw(const Submission& submission, const std::shared_ptr<Vulk
                     for (std::size_t rank = 0; rank < variants.size(); ++rank) {
                         const auto& variant = variants[rank];
                         if (variant->pushOffset != cursor) continue;
+                        if (programs[i].binary.stage == ShaderRecompiler::ShaderStage::Mesh && !std::equal(variant->meshIndexDescriptor.begin(), variant->meshIndexDescriptor.end(), programs[i].userData.begin() + ShaderRecompiler::MeshIndexBufferUserWord)) {
+                            if (!anyLayout) outcome = EntryOutcome::Differing;
+                            anyLayout = true;
+                            continue;
+                        }
                         auto& regions = matchedRegions[i];
                         regions.clear();
                         appendEntryRegions(*variant, regions);

core/libs/prx/libSceAgcDriver/tests/Mesh.cpp

@@ -1,5 +1,6 @@
 #include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp"
 #include "Recompiler.hpp"
+#include "MeshGeometryFixture.hpp"
 #include "execution/VulkanTestDevice.hpp"
 #include "prx/libc/include/GuestAllocations.hpp"
 #include <algorithm>
@@ -24,21 +25,7 @@ constexpr std::uint32_t Height = 128;
 constexpr std::array<std::uint8_t, 4> Background{16, 24, 40, 255};
 alignas(256) std::array<std::byte, Width * Height * 4> Pixels{};
 
-alignas(256) constexpr std::array<std::uint32_t, 104> GeometryCode{
-    0x8f6a9003, 0x94fe6ac1, 0xbf88000b, 0xd7650006, 0x000100c1, 0xd7660006, 0x00020cc1, 0x93ebff03,
-    0x00040018, 0xd7460006, 0x04190c6b, 0x340c0c82, 0xd8340000, 0x00000506, 0xbf8cc07f, 0xbefe04c1,
-    0xbf8a0000, 0x938dff02, 0x00090016, 0x9382ff03, 0x00040018, 0x938cff03, 0x00080008, 0xd7650009,
-    0x000100c1, 0xd7660009, 0x000212c1, 0xd746000a, 0x04250c02, 0x7da8120c, 0xbf88002e, 0x361600ff,
-    0x0000ffff, 0x2c180090, 0x361a02ff, 0x0000ffff, 0xd8d80000, 0x0b00000b, 0xd8d80000, 0x0c00000c,
-    0xd8d80000, 0x0d00000d, 0xbf8cc07f, 0xe0382000, 0x8002100b, 0xe0382010, 0x8002140b, 0xe0382000,
-    0x8002180c, 0xe0382010, 0x80021c0c, 0xe0382000, 0x8002200d, 0xe0382010, 0x8002240d, 0x161c1483,
-    0x161e14ff, 0x00000060, 0xbf8c3f70, 0xdb7c0400, 0x0000100f, 0xdb7c0410, 0x0000140f, 0xdb7c0420,
-    0x0000180f, 0xdb7c0430, 0x00001c0f, 0xdb7c0440, 0x0000200f, 0xdb7c0450, 0x0000240f, 0x4a501c81,
-    0x4a521c82, 0x3450508a, 0x34525294, 0xd772002a, 0x04a6510e, 0xbf8cc07f, 0xbefe04c1, 0xbf8a0000,
-    0x930e830d, 0xbf078002, 0xbf850003, 0x8f0f8c0d, 0x887c0f0e, 0xbf900009, 0x7da8140d, 0xbf880002,
-    0xf8000941, 0x0000002a, 0xbefe04c1, 0x7da8140e, 0xbf88000a, 0x34561485, 0xdbfc0400, 0x2c00002b,
-    0xdbfc0410, 0x3000002b, 0xbf8cc07f, 0xf80008cf, 0x2f2e2d2c, 0xf800020f, 0x33323130, 0xbf810000,
-};
+constexpr auto& GeometryCode = AgcDriver::Test::MeshGeometryCode;
 
 alignas(256) constexpr auto Wave32GeometryCode = [] {
     auto code = GeometryCode;

core/libs/prx/libSceAgcDriver/tests/MeshGeometryFixture.hpp

@@ -0,0 +1,27 @@
+#ifndef CORE_LIBS_PRX_LIBSCEAGCDRIVER_TESTS_MESHGEOMETRYFIXTURE_HPP
+#define CORE_LIBS_PRX_LIBSCEAGCDRIVER_TESTS_MESHGEOMETRYFIXTURE_HPP
+
+#include <array>
+#include <cstdint>
+
+namespace AgcDriver::Test {
+
+alignas(256) inline constexpr std::array<std::uint32_t, 104> MeshGeometryCode{
+    0x8f6a9003, 0x94fe6ac1, 0xbf88000b, 0xd7650006, 0x000100c1, 0xd7660006, 0x00020cc1, 0x93ebff03,
+    0x00040018, 0xd7460006, 0x04190c6b, 0x340c0c82, 0xd8340000, 0x00000506, 0xbf8cc07f, 0xbefe04c1,
+    0xbf8a0000, 0x938dff02, 0x00090016, 0x9382ff03, 0x00040018, 0x938cff03, 0x00080008, 0xd7650009,
+    0x000100c1, 0xd7660009, 0x000212c1, 0xd746000a, 0x04250c02, 0x7da8120c, 0xbf88002e, 0x361600ff,
+    0x0000ffff, 0x2c180090, 0x361a02ff, 0x0000ffff, 0xd8d80000, 0x0b00000b, 0xd8d80000, 0x0c00000c,
+    0xd8d80000, 0x0d00000d, 0xbf8cc07f, 0xe0382000, 0x8002100b, 0xe0382010, 0x8002140b, 0xe0382000,
+    0x8002180c, 0xe0382010, 0x80021c0c, 0xe0382000, 0x8002200d, 0xe0382010, 0x8002240d, 0x161c1483,
+    0x161e14ff, 0x00000060, 0xbf8c3f70, 0xdb7c0400, 0x0000100f, 0xdb7c0410, 0x0000140f, 0xdb7c0420,
+    0x0000180f, 0xdb7c0430, 0x00001c0f, 0xdb7c0440, 0x0000200f, 0xdb7c0450, 0x0000240f, 0x4a501c81,
+    0x4a521c82, 0x3450508a, 0x34525294, 0xd772002a, 0x04a6510e, 0xbf8cc07f, 0xbefe04c1, 0xbf8a0000,
+    0x930e830d, 0xbf078002, 0xbf850003, 0x8f0f8c0d, 0x887c0f0e, 0xbf900009, 0x7da8140d, 0xbf880002,
+    0xf8000941, 0x0000002a, 0xbefe04c1, 0x7da8140e, 0xbf88000a, 0x34561485, 0xdbfc0400, 0x2c00002b,
+    0xdbfc0410, 0x3000002b, 0xbf8cc07f, 0xf80008cf, 0x2f2e2d2c, 0xf800020f, 0x33323130, 0xbf810000,
+};
+
+}
+
+#endif

core/libs/prx/libSceAgcDriver/tests/execution/MeshIndexCache.cpp

@@ -0,0 +1,145 @@
+#include "prx/libSceAgcDriver/Execution/include/Driver/Driver.hpp"
+#include "prx/libSceAgcDriver/Submit/include/Dcb.hpp"
+#include "../MeshGeometryFixture.hpp"
+#include "VulkanTestDevice.hpp"
+#include <array>
+#include <bit>
+#include <cstddef>
+#include <cstdint>
+#include <exception>
+#include <iostream>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace {
+
+using AgcDriver::Graphics::Require;
+using AgcDriver::DriverDetail::DrawMiss;
+
+constexpr std::uint32_t Width = 64;
+constexpr std::uint32_t Height = 64;
+alignas(256) std::array<std::byte, Width * Height * 4> Pixels{};
+alignas(256) constexpr std::array<std::uint16_t, 10> Indices{0, 1, 2, 3, 4, 5, 0, 0, 0, 0};
+alignas(256) constexpr std::array<std::uint16_t, 4> RightIndices{3, 4, 5, 0};
+alignas(256) constexpr std::array<std::array<float, 8>, 6> Vertices{{
+    {-0.9f, -0.8f, 0, 1, 1, 1, 1, 1}, {-0.1f, -0.8f, 0, 1, 1, 1, 1, 1}, {-0.5f, 0.8f, 0, 1, 1, 1, 1, 1},
+    {0.1f, -0.8f, 0, 1, 1, 1, 1, 1}, {0.9f, -0.8f, 0, 1, 1, 1, 1, 1}, {0.5f, 0.8f, 0, 1, 1, 1, 1, 1}
+}};
+alignas(256) constexpr std::array<std::uint32_t, 4> PixelCode{0x7e0e02f2u, 0xf800180fu, 0x07070707u, 0xbf810000u};
+
+struct Header {
+    Shader shader{};
+    ShaderUserData users{};
+};
+
+Header GeometryHeader{};
+Header PixelHeader{};
+
+void Register(Header& header, const std::uint32_t* code, std::size_t size, std::uint8_t type) {
+    header.shader.file_header = 0x34333231u;
+    header.shader.version = 0x18u;
+    header.shader.header_size = sizeof(header);
+    header.shader.shader_size = static_cast<std::uint32_t>(size);
+    header.shader.code = code;
+    header.shader.type = type;
+    header.shader.user_data = &header.users;
+    AgcDriverRegisterShader_nid_postfix(&header.shader);
+}
+
+void Write(std::vector<std::uint32_t>& commands, std::uint32_t opcode, std::uint32_t offset, std::uint32_t value) {
+    commands.insert(commands.end(), {0xc0010000u | (opcode << 8u), offset, value});
+}
+
+std::vector<std::uint32_t> Setup() {
+    const auto color = reinterpret_cast<std::uintptr_t>(Pixels.data());
+    const auto geometry = reinterpret_cast<std::uintptr_t>(AgcDriver::Test::MeshGeometryCode.data());
+    const auto pixel = reinterpret_cast<std::uintptr_t>(PixelCode.data());
+    const auto vertices = reinterpret_cast<std::uintptr_t>(Vertices.data());
+    const std::vector<std::pair<std::uint32_t, std::uint32_t>> context{
+        {0x2d5, 0x20}, {0x1ff, 12}, {0x2ce, 3}, {0x29b, 2}, {0x2ab, 4}, {0x260, 0},
+        {0x1b6, 0}, {0x207, 0}, {0x200, 0}, {0x203, 0x800},
+        {0x2dc, 0xaa00}, {0x2f8, 0}, {0x292, 2}, {0x293, 0},
+        {0x80, 0}, {0x8d, 0}, {0x83, 0xffff}, {0x8c, 0xa},
+        {0x2f9, 0x2d}, {0x313, 0x6000}, {0x30e, 0xffffffff}, {0x30f, 0xffffffff},
+        {0x206, 0x43f}, {0x204, 0x80000}, {0x205, 0x240},
+        {0x8e, 0xf}, {0x8f, 0xf}, {0x202, 0xcc0010},
+        {0x1c4, 0}, {0x1c5, 9}, {0x1c3, 4}, {0x31c, 0x28028},
+        {0x31b, 0}, {0x31d, 0}, {0x3b0, ((Height - 1u) << 14u) | (Width - 1u)},
+        {0x3b8, 0x9000000}, {0x1e0, 0}, {0xc, 0}, {0xd, (Height << 16u) | Width},
+        {0x81, 0x80000000}, {0x82, (Height << 16u) | Width},
+        {0x90, 0x80000000}, {0x91, (Height << 16u) | Width},
+        {0x94, 0x80000000}, {0x95, (Height << 16u) | Width},
+        {0x318, static_cast<std::uint32_t>(color >> 8u)}, {0x390, static_cast<std::uint32_t>(color >> 40u)},
+        {0x10f, std::bit_cast<std::uint32_t>(32.0f)}, {0x110, std::bit_cast<std::uint32_t>(32.0f)},
+        {0x111, std::bit_cast<std::uint32_t>(32.0f)}, {0x112, std::bit_cast<std::uint32_t>(32.0f)},
+        {0x113, std::bit_cast<std::uint32_t>(1.0f)}, {0x114, 0},
+        {0xb4, 0}, {0xb5, std::bit_cast<std::uint32_t>(1.0f)}
+    };
+    const std::vector<std::pair<std::uint32_t, std::uint32_t>> shaders{
+        {0xc8, static_cast<std::uint32_t>(geometry >> 8u)}, {0xc9, static_cast<std::uint32_t>(geometry >> 40u)},
+        {0x8a, 0x600c0000u}, {0x8b, (8u << 19u) | (3u << 16u) | (4u << 1u)},
+        {0x8c, static_cast<std::uint32_t>(vertices)}, {0x8d, static_cast<std::uint32_t>(vertices >> 32u) | (32u << 16u)},
+        {0x8e, static_cast<std::uint32_t>(Vertices.size())}, {0x8f, 0x01016facu},
+        {0x8, static_cast<std::uint32_t>(pixel >> 8u)}, {0x9, static_cast<std::uint32_t>(pixel >> 40u)}, {0xb, 0}
+    };
+    std::vector<std::uint32_t> commands;
+    Write(commands, 0x79, 0x242, 4);
+    Write(commands, 0x79, 0x25b, (12u << 9u) | 4u);
+    for (const auto& [offset, value] : context) Write(commands, 0x69, offset, value);
+    for (const auto& [offset, value] : shaders) Write(commands, 0x76, offset, value);
+    commands.insert(commands.end(), {0xc0002a00u, 0u});
+    return commands;
+}
+
+void Draw(const std::uint16_t* indices, std::uint32_t count, bool left, bool right) {
+    Pixels.fill(std::byte{0});
+    auto commands = Setup();
+    const auto address = reinterpret_cast<std::uintptr_t>(indices);
+    commands.insert(commands.end(), {0xc0042700u, count, static_cast<std::uint32_t>(address), static_cast<std::uint32_t>(address >> 32u), count, 0u});
+    Packet packet{commands.data(), static_cast<std::uint32_t>(commands.size()), 0, {}};
+    Require(sceAgcDriverSubmitDcb(&packet) == 0, "mesh submission failed");
+    AgcDriverWaitIdle_nid_postfix();
+    const auto lit = [](std::uint32_t x) { return std::to_integer<unsigned>(Pixels[(24u * Width + x) * 4u]) > 200u; };
+    Require(lit(16) == left && lit(48) == right, "cached mesh draw used another draw's indices: count=" + std::to_string(count));
+}
+
+}
+
+int main() {
+    try {
+        {
+            const auto device = OpenVulkanTestDevice();
+            if (!device || !TargetHasCapability(device->Target(), spv::CapabilityMeshShadingEXT)) return VulkanTestSkipped;
+        }
+        Register(GeometryHeader, AgcDriver::Test::MeshGeometryCode.data(), sizeof(AgcDriver::Test::MeshGeometryCode), 2);
+        Register(PixelHeader, PixelCode.data(), sizeof(PixelCode), 1);
+        for (unsigned i = 0; i < 3; ++i) Draw(Indices.data(), 3, true, false);
+        const auto warm = AgcDriver::DriverDetail::Driver::Get().DrawCacheCounters();
+        Require(warm.hits != 0, "mesh warm-up did not reach the draw cache");
+        const auto front = static_cast<std::size_t>(DrawMiss::FrontDiffering);
+        const auto layout = static_cast<std::size_t>(DrawMiss::Layout);
+        Draw(Indices.data(), 6, true, true);
+        const auto range = AgcDriver::DriverDetail::Driver::Get().DrawCacheCounters();
+        Require(range.misses[front] == warm.misses[front] + 1 && range.misses[layout] == warm.misses[layout], "an index-range change was not counted as a differing front stage");
+        Draw(Indices.data(), 6, true, true);
+        const auto rangeReuse = AgcDriver::DriverDetail::Driver::Get().DrawCacheCounters();
+        Require(rangeReuse.hits == range.hits + 1, "the changed index range was not retained for cache reuse");
+        Draw(RightIndices.data(), 3, false, true);
+        const auto address = AgcDriver::DriverDetail::Driver::Get().DrawCacheCounters();
+        Require(address.misses[front] == rangeReuse.misses[front] + 1 && address.misses[layout] == rangeReuse.misses[layout], "an index-address change was not counted as a differing front stage");
+        Draw(RightIndices.data(), 3, false, true);
+        const auto addressReuse = AgcDriver::DriverDetail::Driver::Get().DrawCacheCounters();
+        Require(addressReuse.hits == address.hits + 1, "the changed index address was not retained for cache reuse");
+        Draw(Indices.data(), 3, true, false);
+        const auto earlier = AgcDriver::DriverDetail::Driver::Get().DrawCacheCounters();
+        Require(earlier.hits == addressReuse.hits + 1, "returning to an earlier index descriptor did not hit the cache");
+        AgcDriverShutdown_nid_postfix();
+        std::cout << "Cached mesh draws retained their index ranges and addresses\n";
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        AgcDriverShutdown_nid_postfix();
+        return 1;
+    }
+}