feat(shader): link callees of calls whose target is loaded from a table

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Shaders/ShaderRegistry.cpp

@@ -29,9 +29,11 @@
 #ifdef __GLIBC__
 #include <malloc.h>
 #endif
+#include <set>
 #include <sstream>
 #include <stdexcept>
 #include <thread>
+#include "Optimization/ShaderStageInputInfo.hpp"
 #include "RdnaDecoder/RdnaInstructionDecoder.hpp"
 
 namespace AgcDriver::DriverDetail {
@@ -538,14 +540,48 @@ std::vector<std::uint32_t> ReadUserDataCallee(std::uint64_t address, std::uint32
     throw std::runtime_error("its callee at " + hexText(address) + " does not return within " + std::to_string(limitWords * 4u) + " bytes");
 }
 
-std::vector<std::uint32_t> LinkUserDataCallees(const ShaderRecompiler::RecompileRequest& request) {
+std::vector<std::uint8_t> ReadGuestBytes(std::uint64_t address, std::uint64_t bytes, const char* what) {
+    if (address > UINTPTR_MAX - bytes || !GuestMemory::Accessible(reinterpret_cast<const void*>(static_cast<std::uintptr_t>(address)), static_cast<std::size_t>(bytes))) {
+        throw std::runtime_error(std::string("its ") + what + " at " + hexText(address) + " + " + hexText(bytes) + " is not readable");
+    }
+    std::vector<std::uint8_t> data(static_cast<std::size_t>(bytes));
+    GuestMemory::Read(address, std::as_writable_bytes(std::span(data)));
+    return data;
+}
+
+std::vector<ShaderRecompiler::TableCallee> ReadTableCallees(const ShaderRecompiler::RecompileRequest& request, const ShaderRecompiler::TableCall& call) {
+    constexpr std::uint64_t limitBytes = 64u << 20u;
+    constexpr std::size_t limitCallees = 64u;
+    const auto slot = call.userDataRegister - request.context.userDataBaseRegister;
+    const auto pointer = static_cast<std::uint64_t>(request.context.userData[slot]) | static_cast<std::uint64_t>(request.context.userData[slot + 1u]) << 32u;
+    const auto words = ReadGuestBytes(pointer + call.descriptorOffset, 16u, "call table descriptor");
+    ShaderRecompiler::ShaderBufferResource descriptor{};
+    std::memcpy(descriptor.fields.data(), words.data(), words.size());
+    const auto bytes = descriptor.GetSize();
+    if (bytes > limitBytes) throw std::runtime_error("its call table of " + std::to_string(bytes) + " bytes is larger than the " + std::to_string(limitBytes) + " bytes that are searched for callees");
+    const auto table = bytes == 0u ? std::vector<std::uint8_t>{} : ReadGuestBytes(descriptor.Base48(), bytes, "call table");
+    std::set<std::uint64_t> targets;
+    for (std::uint64_t offset = call.entryOffset; offset + 8u <= bytes; offset += call.entryStride) {
+        std::uint64_t target = 0;
+        std::memcpy(&target, table.data() + offset, sizeof(target));
+        if (target != 0u) targets.insert(target);
+        if (call.entryStride == 0u) break;
+    }
+    if (targets.size() > limitCallees) throw std::runtime_error("its call table holds " + std::to_string(targets.size()) + " distinct targets; at most " + std::to_string(limitCallees) + " are linked");
+    std::vector<ShaderRecompiler::TableCallee> callees;
+    for (const auto target : targets) callees.push_back({target, ReadUserDataCallee(target, call.linkRegister)});
+    return callees;
+}
+
+ShaderRecompiler::LinkedCalls LinkUserDataCallees(const ShaderRecompiler::RecompileRequest& request) {
     if (request.shader.stage != ShaderRecompiler::ShaderStage::Compute) return {};
     const auto code = request.shader.code;
     if (std::none_of(code.begin(), code.end(), [](std::uint32_t word) { return (word & 0xff80ff00u) == 0xbe802100u; })) return {};
     const auto program = ShaderRecompiler::RdnaInstructionDecoder{}.Decode(code);
-    const ShaderRecompiler::SwappcInfo swappc{false, request.context.userDataBaseRegister, static_cast<std::uint32_t>(request.context.userData.size())};
+    const ShaderRecompiler::SwappcInfo swappc{false, request.context.userDataBaseRegister, static_cast<std::uint32_t>(request.context.userData.size()), request.context.waveSize};
     const auto calls = ShaderRecompiler::FindUserDataCalls(program, swappc);
-    if (calls.empty()) return {};
+    const auto tableCalls = ShaderRecompiler::FindTableCalls(program, swappc);
+    if (calls.empty() && tableCalls.empty()) return {};
     std::vector<std::vector<std::uint32_t>> callees;
     for (const auto& call : calls) {
         const auto slot = call.userDataRegister - swappc.userDataBaseRegister;
@@ -557,7 +593,15 @@ std::vector<std::uint32_t> LinkUserDataCallees(const ShaderRecompiler::Recompile
                 std::to_string(call.userDataRegister + 1u) + "] at program counter " + hexText(call.programCounter) + ", but " + error.what());
         }
     }
-    return ShaderRecompiler::LinkUserDataCalls(program, calls, callees);
+    std::vector<std::vector<ShaderRecompiler::TableCallee>> tableCallees;
+    for (const auto& call : tableCalls) {
+        try {
+            tableCallees.push_back(ReadTableCallees(request, call));
+        } catch (const std::exception& error) {
+            throw std::runtime_error("AGC driver: compute shader " + hexText(request.shader.codeAddress) + " calls a target loaded from a table at program counter " + hexText(call.programCounter) + ", but " + error.what());
+        }
+    }
+    return ShaderRecompiler::LinkCalls(program, calls, callees, tableCalls, tableCallees);
 }
 
 }
@@ -567,9 +611,12 @@ ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& s
     const auto code = std::span(snapshot.code).subspan(codeOffset);
     require(request.shader.code.data() == code.data() && request.shader.code.size() == code.size(), "prepared invocation does not refer to registered code");
     auto linked = LinkUserDataCallees(request);
-    if (linked.empty()) linked = LinkGraphicsUserDataTail(request, snapshot);
+    if (linked.code.empty()) linked.code = LinkGraphicsUserDataTail(request, snapshot);
     auto source = request;
-    if (!linked.empty()) source.shader.code = linked;
+    if (!linked.code.empty()) {
+        source.shader.code = linked.code;
+        source.context.linkedCallFallbacks = linked.fallbacks;
+    }
     auto invocationRequest = source;
     struct PreparedKeyStorage {};
     auto& key = HostThreadLocal<std::vector<std::uint64_t>, PreparedKeyStorage>();
@@ -579,7 +626,7 @@ ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& s
     for (const auto& entry : snapshot.prepared->entries) {
         if (entry.codeOffset != codeOffset) continue;
         invocationRequest.shader.code = ShaderRecompiler::GetPreparedCode(*entry.handle);
-        if (!linked.empty() && !std::ranges::equal(invocationRequest.shader.code, linked)) continue;
+        if (!linked.code.empty() && !std::ranges::equal(invocationRequest.shader.code, linked.code)) continue;
         if (auto invocation = ShaderRecompiler::PreparedShaderInvocation::TryCreate(invocationRequest, entry.handle, key)) return std::move(*invocation);
     }
     ReportPreparedAtUse(snapshot, source);

core/libs/prx/libSceAgcDriver/Graphics/src/BdaResources.cpp

@@ -275,6 +275,10 @@ void BdaResources::CheckFault() const {
         throw std::runtime_error(message.str());
     }
     std::ostringstream message;
+    if (report.reason == ShaderRecompiler::BdaAbi::FaultReason::UnlinkedCall) {
+        message << "shader stage " << report.stage << " called address 0x" << std::hex << report.address << " at program counter 0x" << report.instruction << ", which no linked callee of that call matches";
+        throw std::runtime_error(message.str());
+    }
     message << "BDA access failed: address=0x" << std::hex << report.address << " instruction=0x" << report.instruction << std::dec << " bytes=" << report.bytes << " stage=" << report.stage << " reason=" << static_cast<std::uint32_t>(report.reason);
     if (report.reason == ShaderRecompiler::BdaAbi::FaultReason::Permission && table != nullptr) {
         const auto bytes = table->Bytes();

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

@@ -2836,7 +2836,13 @@ void verifyPixelRequestSerialization() {
     minimal.context.pixel->layerExported = true;
     minimal.context.pixel->ditheredAlphaToCoverage = true;
     const auto encoded = serializer.Serialize(minimal);
-    require(requestPrefix(encoded, 8u) == "NVNQQRkAAAA=", "new requests did not use serialization version 25");
+    require(requestPrefix(encoded, 8u) == "NVNQQRoAAAA=", "new requests did not use serialization version 26");
+    const std::array<std::uint32_t, 2> fallbacks{0x40u, 0x1ac4u};
+    auto withFallbacks = minimal;
+    withFallbacks.context.linkedCallFallbacks = fallbacks;
+    const auto fallbackRequest = serializer.Deserialize(serializer.Serialize(withFallbacks));
+    require(std::ranges::equal(fallbackRequest.request.context.linkedCallFallbacks, fallbacks), "the linked call fallbacks did not round-trip");
+    require(serializer.Deserialize(encoded).request.context.linkedCallFallbacks.empty(), "a request without linked call fallbacks gained some");
     require(serializer.Deserialize(encoded).request.context.pixel->rasterPrimitive == RasterPrimitive::Rectangle && serializer.Deserialize(encoded).request.context.pixel->ditheredAlphaToCoverage, "the pixel primitive class or dither mode did not round-trip");
     require(serializer.Deserialize(encoded).request.context.pixel->layerExported, "the pixel layer export did not round-trip");
     const auto legacy16 = serializer.Deserialize(
@@ -2948,7 +2954,7 @@ void verifyPixelRequestSerialization() {
     for (std::size_t bytes = 0; bytes < 23u; ++bytes) {
         expectFailure([&] { static_cast<void>(serializer.Deserialize(requestPrefix(encoded, mappingOffset + bytes))); }, "truncated data", "a truncated pixel mapping, packing, dual-source, alpha-to-coverage, stencil reference export flag, primitive class, layer export or dither flag was accepted");
     }
-    for (const auto unsupported : {"NVNQQQAAAAA=", "NVNQQRoAAAA="}) {
+    for (const auto unsupported : {"NVNQQQAAAAA=", "NVNQQRsAAAA="}) {
         expectFailure([&] { static_cast<void>(serializer.Deserialize(unsupported)); }, "serialization version", "an unsupported request version was accepted");
     }
     static constexpr std::string_view version18Request{
@@ -3067,6 +3073,7 @@ void verifyVertexInfoVersion19() {
     while (flagOffset < legacyPixel.size() && legacyPixel[flagOffset] == exportedBytes[flagOffset]) ++flagOffset;
     require(flagOffset < legacyPixel.size(), "the pixel layer export flag is not serialized");
     legacyPixel.erase(legacyPixel.begin() + static_cast<std::ptrdiff_t>(flagOffset), legacyPixel.begin() + static_cast<std::ptrdiff_t>(flagOffset + 2u));
+    legacyPixel.resize(legacyPixel.size() - sizeof(std::uint64_t));
     for (const auto version : {19u, 20u, 21u, 22u}) {
         auto versionPixel = legacyPixel;
         versionPixel[4] = static_cast<std::uint8_t>(version);

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

@@ -0,0 +1,164 @@
+#include "prx/libSceAgcDriver/Execution/include/Driver.hpp"
+#include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/Draw.hpp"
+#include "SceShaders.hpp"
+#include "VulkanTestDevice.hpp"
+#include <array>
+#include <cstdint>
+#include <cstdio>
+#include <iostream>
+#include <span>
+#include <stdexcept>
+#include <string>
+#include <vector>
+#ifdef _WIN32
+#include <io.h>
+#else
+#include <unistd.h>
+#endif
+
+namespace {
+
+using AgcDriver::Graphics::Require;
+
+constexpr std::uint32_t Threads = 64;
+constexpr std::uint32_t Untouched = 0xcafebabeu;
+
+alignas(256) constexpr std::array<std::uint32_t, 64> Caller{
+    0x7e020280u, 0xf4080400u, 0xfa000000u, 0x8f148402u, 0xbf8cc07fu, 0xf4280208u, 0x28000000u, 0xbf8cc07fu,
+    0xbe8e0308u, 0xbe8f0309u, 0xbf068103u, 0xbe8e210eu, 0x34040082u, 0xe0701000u, 0x80010102u, 0xbf810000u,
+    0xbf9f0000u, 0x30306c73u, 0x00000040u, 0x00000008u};
+alignas(256) constexpr std::array<std::uint32_t, 3> SelectSixteen{0x85159091u, 0x7e020215u, 0xbe80200eu};
+alignas(256) constexpr std::array<std::uint32_t, 3> SelectThirtyTwo{0x8515a0a1u, 0x7e020215u, 0xbe80200eu};
+alignas(16) std::array<std::uint32_t, 8> Table{};
+alignas(16) std::array<std::uint32_t, 4> Descriptors{};
+alignas(256) std::array<std::uint32_t, Threads> Output{};
+
+void Fill() {
+    const auto first = reinterpret_cast<std::uintptr_t>(SelectSixteen.data());
+    const auto second = reinterpret_cast<std::uintptr_t>(SelectThirtyTwo.data());
+    Table = {static_cast<std::uint32_t>(first), static_cast<std::uint32_t>(first >> 32u), 0u, 0u, static_cast<std::uint32_t>(second), static_cast<std::uint32_t>(second >> 32u), 0u, 0u};
+    const auto table = reinterpret_cast<std::uintptr_t>(Table.data());
+    Descriptors = {static_cast<std::uint32_t>(table), static_cast<std::uint32_t>((table >> 32u) & 0xffffu) | (16u << 16u), 2u, 0x31016facu};
+}
+
+std::array<std::uint32_t, 8> UserData(std::uint32_t index, bool scc) {
+    const auto descriptors = reinterpret_cast<std::uintptr_t>(Descriptors.data());
+    const auto output = reinterpret_cast<std::uintptr_t>(Output.data());
+    return {static_cast<std::uint32_t>(descriptors), static_cast<std::uint32_t>(descriptors >> 32u), index, scc ? 1u : 0u,
+        static_cast<std::uint32_t>(output), static_cast<std::uint32_t>((output >> 32u) & 0xffffu), static_cast<std::uint32_t>(sizeof(Output)), 0x31016facu};
+}
+
+std::vector<std::uint32_t> Commands(const std::array<std::uint32_t, 8>& data, bool wave32) {
+    std::vector<std::uint32_t> words;
+    const auto registers = [&](std::uint32_t first, std::span<const std::uint32_t> values) {
+        words.push_back(0xc0007600u | (static_cast<std::uint32_t>(values.size()) << 16u));
+        words.push_back(first);
+        words.insert(words.end(), values.begin(), values.end());
+    };
+    const std::array<std::uint32_t, 3> shape{Threads, 1, 1};
+    registers(0x207, shape);
+    const auto address = reinterpret_cast<std::uintptr_t>(Caller.data());
+    const std::array<std::uint32_t, 2> program{static_cast<std::uint32_t>(address >> 8u), static_cast<std::uint32_t>(address >> 40u)};
+    registers(0x20c, program);
+    const std::array<std::uint32_t, 1> resource{static_cast<std::uint32_t>(data.size()) << 1u};
+    registers(0x213, resource);
+    registers(0x240, data);
+    words.insert(words.end(), {0xc0031500u, 1, 1, 1, 0x41u | (wave32 ? 0x8000u : 0u)});
+    return words;
+}
+
+class StderrCapture {
+public:
+    StderrCapture() {
+        std::fflush(stderr);
+        file = std::tmpfile();
+        Require(file != nullptr, "cannot open a temporary file for stderr");
+#ifdef _WIN32
+        saved = _dup(_fileno(stderr));
+        _dup2(_fileno(file), _fileno(stderr));
+#else
+        saved = dup(fileno(stderr));
+        dup2(fileno(file), fileno(stderr));
+#endif
+    }
+    std::string Finish() {
+        std::fflush(stderr);
+#ifdef _WIN32
+        _dup2(saved, _fileno(stderr));
+        _close(saved);
+#else
+        dup2(saved, fileno(stderr));
+        close(saved);
+#endif
+        std::rewind(file);
+        std::string text;
+        char buffer[4096];
+        for (std::size_t read = 0; (read = std::fread(buffer, 1, sizeof(buffer), file)) != 0;) text.append(buffer, read);
+        std::fclose(file);
+        std::fwrite(text.data(), 1, text.size(), stderr);
+        return text;
+    }
+
+private:
+    std::FILE* file = nullptr;
+    int saved = -1;
+};
+
+std::string Dispatch(std::uint32_t index, bool scc, bool wave32) {
+    Output.fill(Untouched);
+    const auto words = Commands(UserData(index, scc), wave32);
+    Packet packet{const_cast<std::uint32_t*>(words.data()), static_cast<std::uint32_t>(words.size()), 0, {}};
+    StderrCapture log;
+    std::string failure;
+    try {
+        AgcDriver::Submit(&packet, 0);
+        AgcDriverWaitIdle_nid_postfix();
+    } catch (const std::exception& error) {
+        failure = error.what();
+    }
+    AgcDriver::GuestMemory::FlushGpuWrites(reinterpret_cast<std::uintptr_t>(Output.data()), sizeof(Output));
+    return log.Finish() + failure;
+}
+
+void Check(std::uint32_t index, bool scc, bool wave32, std::uint32_t expected, const std::string& what) {
+    const auto log = Dispatch(index, scc, wave32);
+    Require(log.find("no linked callee") == std::string::npos, what + ": a linked callee was reported missing: " + log);
+    for (std::uint32_t lane = 0; lane < Threads; ++lane) {
+        Require(Output[lane] == expected, what + (wave32 ? " (wave32)" : " (wave64)") + ": lane " + std::to_string(lane) + " is " + std::to_string(Output[lane]) + ", expected " + std::to_string(expected));
+    }
+}
+
+}
+
+int main() {
+    try {
+        {
+            const auto device = OpenVulkanTestDevice();
+            if (!device) return VulkanTestSkipped;
+        }
+        Fill();
+        for (const bool wave32 : {false, true}) {
+            Check(0u, false, wave32, 16u, "the first record's callee with SCC 0");
+            Check(0u, true, wave32, 17u, "the first record's callee with SCC 1");
+            Check(1u, false, wave32, 32u, "the second record's callee with SCC 0");
+            Check(1u, true, wave32, 33u, "the second record's callee with SCC 1");
+            const auto log = Dispatch(2u, false, wave32);
+            Require(log.find("called address 0x0 at program counter") != std::string::npos && log.find("which no linked callee of that call matches") != std::string::npos,
+                std::string("a record past the table did not report an unlinked call") + (wave32 ? " (wave32)" : " (wave64)") + ": " + log);
+            for (std::uint32_t lane = 0; lane < Threads; ++lane) {
+                Require(Output[lane] == Untouched, "a lane stored after the unlinked call: lane " + std::to_string(lane) + " is " + std::to_string(Output[lane]));
+            }
+        }
+        AgcDriverShutdown_nid_postfix();
+        std::puts("compute table call tests passed");
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        try {
+            AgcDriverShutdown_nid_postfix();
+        } catch (const std::exception&) {
+        }
+        return 1;
+    }
+}