fix(agc): stop optional compute captures at their bounds

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -2398,6 +2398,14 @@ if(BUILD_TESTING)
     endif()
     add_test(NAME agc_compute_pipeline_capture COMMAND agc_compute_pipeline_capture_tests)
     set_tests_properties(agc_compute_pipeline_capture PROPERTIES TIMEOUT 30)
+    foreach(mode IN ITEMS entry-cap byte-cap io invalid)
+        string(REPLACE "-" "_" name "${mode}")
+        add_test(NAME agc_compute_pipeline_capture_optional_${name} COMMAND agc_compute_pipeline_capture_tests --optional-${mode})
+        set_tests_properties(agc_compute_pipeline_capture_optional_${name} PROPERTIES TIMEOUT 30)
+        if(WIN32)
+            set_tests_properties(agc_compute_pipeline_capture_optional_${name} PROPERTIES ENVIRONMENT_MODIFICATION "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
+        endif()
+    endforeach()
     if(WIN32)
         set_tests_properties(agc_compute_pipeline_capture PROPERTIES ENVIRONMENT_MODIFICATION "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
     endif()

core/libs/prx/libSceAgcDriver/Execution/include/ComputePipelineCapture.hpp

@@ -53,6 +53,8 @@ struct ComputePipelineCapture {
     static ComputePipelineCapture Decode(std::span<const std::byte> bytes);
     static ComputePipelineCapture Read(const std::filesystem::path& path);
     std::filesystem::path Write(const std::filesystem::path& directory) const;
+    std::filesystem::path WriteOptional(const std::filesystem::path& directory) const;
+    static bool OptionalEnabled();
     std::string Identity() const;
 };
 

core/libs/prx/libSceAgcDriver/Execution/src/ComputePipelineCapture.cpp

@@ -1,5 +1,7 @@
 #include "prx/libSceAgcDriver/Execution/include/ComputePipelineCapture.hpp"
 #include <algorithm>
+#include <atomic>
+#include <cstdio>
 #include <fstream>
 #include <iomanip>
 #include <limits>
@@ -11,10 +13,52 @@
 namespace AgcDriver {
 namespace {
 
+class CaptureLimit final : public std::runtime_error {
+public:
+    using std::runtime_error::runtime_error;
+};
+
+class CaptureIoFailure final : public std::runtime_error {
+public:
+    using std::runtime_error::runtime_error;
+};
+
+std::atomic<bool> optionalDisabled{false};
+
+void DisableOptionalCapture(const char* message) {
+    if (!optionalDisabled.exchange(true, std::memory_order_relaxed)) std::fprintf(stderr, "[compute-capture] optional capture disabled: %s\n", message);
+}
+
 void RequireCapture(bool condition, const char* message) {
     if (!condition) throw std::runtime_error(std::string("compute pipeline capture: ") + message);
 }
 
+void RequireCaptureIo(bool condition, const char* message) {
+    if (!condition) throw CaptureIoFailure(std::string("compute pipeline capture: ") + message);
+}
+
+void RequireCaptureCapacity(bool condition, const char* message) {
+    if (!condition) throw CaptureLimit(std::string("compute pipeline capture: ") + message);
+}
+
+void CheckCaptureIo(const std::error_code& error, const char* message) {
+    if (error) throw CaptureIoFailure(std::string("compute pipeline capture: ") + message + ": " + error.message());
+}
+
+std::filesystem::file_status CaptureStatus(const std::filesystem::path& path) {
+    std::error_code error;
+    const auto status = std::filesystem::symlink_status(path, error);
+    if (error != std::errc::no_such_file_or_directory) CheckCaptureIo(error, "cannot inspect capture path");
+    return status;
+}
+
+std::uintmax_t CaptureFileSize(const std::filesystem::path& path) {
+    std::error_code error;
+    const auto size = std::filesystem::file_size(path, error);
+    CheckCaptureIo(error, "cannot size capture");
+    return size;
+}
+
 std::string Hash(std::span<const std::byte> bytes) {
     std::uint64_t hash = 14695981039346656037ull;
     for (const auto value : bytes) {
@@ -65,13 +109,13 @@ class Decoder {
 };
 
 std::vector<std::byte> ReadBytes(const std::filesystem::path& path) {
-    const auto size = std::filesystem::file_size(path);
+    const auto size = CaptureFileSize(path);
     RequireCapture(size <= ComputePipelineCapture::MaxFileBytes, "file exceeds 36 MiB");
     std::ifstream input(path, std::ios::binary);
-    RequireCapture(input.good(), "cannot open capture");
+    RequireCaptureIo(input.good(), "cannot open capture");
     std::vector<std::byte> bytes(static_cast<std::size_t>(size));
     input.read(reinterpret_cast<char*>(bytes.data()), static_cast<std::streamsize>(bytes.size()));
-    RequireCapture(input.good() && input.peek() == std::char_traits<char>::eof(), "capture changed while reading");
+    RequireCaptureIo(input.good() && input.peek() == std::char_traits<char>::eof(), "capture changed while reading");
     return bytes;
 }
 
@@ -234,32 +278,65 @@ std::filesystem::path ComputePipelineCapture::Write(const std::filesystem::path&
     const auto path = recordDirectory / "capture.aps5cp";
     static std::mutex mutex;
     std::lock_guard lock(mutex);
-    std::filesystem::create_directories(directory);
-    RequireCapture(!std::filesystem::is_symlink(std::filesystem::symlink_status(directory)), "capture directory is a symbolic link");
-    if (std::filesystem::exists(recordDirectory)) {
-        RequireCapture(std::filesystem::is_directory(recordDirectory) && !std::filesystem::is_symlink(std::filesystem::symlink_status(recordDirectory)) && !std::filesystem::is_symlink(std::filesystem::symlink_status(path)), "existing capture path is not a regular capture");
+    std::error_code error;
+    std::filesystem::create_directories(directory, error);
+    CheckCaptureIo(error, "cannot create capture directory");
+    RequireCapture(!std::filesystem::is_symlink(CaptureStatus(directory)), "capture directory is a symbolic link");
+    const auto recordStatus = CaptureStatus(recordDirectory);
+    if (std::filesystem::exists(recordStatus)) {
+        RequireCapture(std::filesystem::is_directory(recordStatus) && !std::filesystem::is_symlink(recordStatus), "existing capture path is not a regular capture");
+        const auto captureStatus = CaptureStatus(path);
+        RequireCapture(std::filesystem::is_regular_file(captureStatus) && !std::filesystem::is_symlink(captureStatus), "existing capture path is not a regular capture");
         RequireCapture(ReadBytes(path) == bytes, "existing capture differs or is incomplete");
         return path;
     }
     std::uintmax_t totalBytes = 0;
     std::size_t count = 0;
-    for (const auto& entry : std::filesystem::directory_iterator(directory)) {
-        RequireCapture(++count < 128, "capture directory reached 128 entries");
-        RequireCapture(entry.is_directory() && !entry.is_symlink(), "capture directory contains an unexpected entry");
+    auto iterator = std::filesystem::directory_iterator(directory, error);
+    CheckCaptureIo(error, "cannot enumerate capture directory");
+    while (iterator != std::filesystem::directory_iterator{}) {
+        const auto& entry = *iterator;
+        RequireCaptureCapacity(++count < 128, "capture directory reached 128 entries");
+        const auto entryStatus = CaptureStatus(entry.path());
+        RequireCapture(std::filesystem::is_directory(entryStatus) && !std::filesystem::is_symlink(entryStatus), "capture directory contains an unexpected entry");
         const auto capture = entry.path() / "capture.aps5cp";
-        RequireCapture(std::filesystem::is_regular_file(capture) && !std::filesystem::is_symlink(std::filesystem::symlink_status(capture)), "capture directory contains an incomplete record");
-        const auto size = std::filesystem::file_size(capture);
+        const auto captureStatus = CaptureStatus(capture);
+        RequireCapture(std::filesystem::is_regular_file(captureStatus) && !std::filesystem::is_symlink(captureStatus), "capture directory contains an incomplete record");
+        const auto size = CaptureFileSize(capture);
         RequireCapture(size <= MaxFileBytes, "existing capture exceeds size bound");
         totalBytes += size;
+        iterator.increment(error);
+        CheckCaptureIo(error, "cannot advance capture directory");
     }
-    RequireCapture(totalBytes + bytes.size() <= 256u * 1024u * 1024u, "capture directory exceeds 256 MiB");
-    RequireCapture(std::filesystem::create_directory(recordDirectory), "capture directory was concurrently created");
+    RequireCaptureCapacity(totalBytes + bytes.size() <= 256u * 1024u * 1024u, "capture directory exceeds 256 MiB");
+    const bool created = std::filesystem::create_directory(recordDirectory, error);
+    CheckCaptureIo(error, "cannot create capture record directory");
+    RequireCapture(created, "capture directory was concurrently created");
     std::ofstream output(path, std::ios::binary);
-    RequireCapture(output.good(), "cannot create capture");
+    RequireCaptureIo(output.good(), "cannot create capture");
     output.write(reinterpret_cast<const char*>(bytes.data()), static_cast<std::streamsize>(bytes.size()));
     output.close();
-    RequireCapture(output.good(), "cannot finish capture");
+    RequireCaptureIo(output.good(), "cannot finish capture");
     return path;
 }
 
+bool ComputePipelineCapture::OptionalEnabled() {
+    return !optionalDisabled.load(std::memory_order_relaxed);
+}
+
+std::filesystem::path ComputePipelineCapture::WriteOptional(const std::filesystem::path& directory) const {
+    if (!OptionalEnabled()) {
+        Validate();
+        return {};
+    }
+    try {
+        return Write(directory);
+    } catch (const CaptureLimit& error) {
+        DisableOptionalCapture(error.what());
+    } catch (const CaptureIoFailure& error) {
+        DisableOptionalCapture(error.what());
+    }
+    return {};
+}
+
 }

core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp

@@ -4724,7 +4724,7 @@ std::shared_ptr<ComputePipelineObjects> VulkanDevice::createComputePipeline(cons
     const Graphics::PipelineSpecialization specialization(shader);
     const auto subgroupSize = computeSubgroupSize(shader);
     const auto stageFlags = computeStageFlags(shader);
-    if (const char* directory = std::getenv("APS5_DUMP_COMPUTE_PIPELINES"); directory != nullptr && *directory != '\0' && CaptureComputePipeline(shader.spirv.size())) {
+    if (const char* directory = std::getenv("APS5_DUMP_COMPUTE_PIPELINES"); directory != nullptr && *directory != '\0' && ComputePipelineCapture::OptionalEnabled() && CaptureComputePipeline(shader.spirv.size())) {
         ComputePipelineCapture capture;
         capture.spirv = shader.spirv.Copy();
         ComputeDescriptorSetLayout resourcesLayout;
@@ -4747,8 +4747,8 @@ std::shared_ptr<ComputePipelineObjects> VulkanDevice::createComputePipeline(cons
         capture.driverVersion = state->properties.driverVersion;
         capture.apiVersion = state->properties.apiVersion;
         std::copy_n(state->properties.pipelineCacheUUID, VK_UUID_SIZE, capture.pipelineCacheUuid.begin());
-        const auto path = capture.Write(directory);
-        std::fprintf(stderr, "[compute-capture] %s\n", path.string().c_str());
+        const auto path = capture.WriteOptional(directory);
+        if (!path.empty()) std::fprintf(stderr, "[compute-capture] %s\n", path.string().c_str());
     }
     if (profile && shader.spirv.size() > 100000) std::fprintf(stderr, "[dispatch] creating a pipeline for %zu SPIR-V words (program 0x%llx)\n", shader.spirv.size(), static_cast<unsigned long long>(programAddress));
     CreateComputePipeline(state->device, state->DeviceFunction<PFN_vkCreateComputePipelines>("vkCreateComputePipelines"), context.pipelineCache, objects->module, objects->layout, pipelineFlags, stageFlags, specialization.Info(), subgroupSize, &objects->pipeline, context.pipelineCacheSynchronization);

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

@@ -3,6 +3,7 @@
 #include <cstdio>
 #include <fstream>
 #include <stdexcept>
+#include <string_view>
 
 namespace {
 
@@ -21,6 +22,17 @@ void Reject(TAction action, const char* message) {
     Require(rejected, message);
 }
 
+template<typename TAction>
+void RejectWithMessage(TAction action, const char* expected) {
+    bool rejected = false;
+    try {
+        action();
+    } catch (const std::exception& error) {
+        rejected = std::string_view(error.what()).find(expected) != std::string_view::npos;
+    }
+    Require(rejected, "strict capture did not preserve its diagnostic");
+}
+
 class Directory {
 public:
     Directory() {
@@ -163,10 +175,112 @@ void CheckMalformedRecords() {
     invalid([](auto& copy) { copy.sets = {{0, {{}, {}}}}; });
 }
 
+void CheckOptionalEntryCap() {
+    Directory directory;
+    auto capture = AgcDriver::ComputePipelineFixture();
+    std::filesystem::path first;
+    for (std::uint32_t index = 0; index < 128u; ++index) {
+        capture.specialization[0].value = index;
+        const auto path = capture.WriteOptional(directory.path);
+        Require(!path.empty() && AgcDriver::ComputePipelineCapture::OptionalEnabled(), "optional capture stopped before accepting 128 records");
+        if (index == 0u) first = path;
+    }
+    capture.specialization[0].value = 0u;
+    Require(capture.WriteOptional(directory.path) == first, "optional capture did not deduplicate at the entry cap");
+    capture.specialization[0].value = 128u;
+    Require(capture.WriteOptional(directory.path).empty() && !AgcDriver::ComputePipelineCapture::OptionalEnabled(), "optional capture did not stop at the entry cap");
+    RejectWithMessage([&] { capture.Write(directory.path); }, "compute pipeline capture: capture directory reached 128 entries");
+    for (unsigned index = 0; index < 64u; ++index) Require(capture.WriteOptional(directory.path / "after-stop").empty(), "disabled optional capture wrote another file");
+    Require(!std::filesystem::exists(directory.path / "after-stop"), "disabled optional capture touched its output directory");
+    capture.specialization[0].value = 0u;
+    Require(capture.Write(directory.path) == first, "optional stop disabled strict capture deduplication");
+}
+
+void CheckOptionalByteCap() {
+    Directory directory;
+    auto capture = AgcDriver::ComputePipelineFixture();
+    const auto smallBytes = capture.Encode().size();
+    constexpr std::size_t chunkBytes = 32u * 1024u * 1024u;
+    for (std::uint32_t index = 0; index < 8u; ++index) {
+        auto large = capture;
+        large.specialization[0].value = index;
+        const auto bytes = chunkBytes - (index == 7u ? smallBytes : 0u);
+        Require(bytes >= smallBytes && (bytes - smallBytes) % 4u == 0u, "byte-cap fixture size is not representable");
+        large.spirv.insert(large.spirv.end() - 2, (bytes - smallBytes) / 4u, 0x00010000u);
+        Require(std::filesystem::file_size(large.Write(directory.path)) == bytes, "byte-cap fixture size differs");
+    }
+    capture.specialization[0].value = 8u;
+    Require(!capture.WriteOptional(directory.path).empty() && AgcDriver::ComputePipelineCapture::OptionalEnabled(), "optional capture refused exactly 256 MiB");
+    std::uintmax_t total = 0;
+    for (const auto& entry : std::filesystem::directory_iterator(directory.path)) total += std::filesystem::file_size(entry.path() / "capture.aps5cp");
+    Require(total == 256u * 1024u * 1024u, "byte-cap fixture did not reach exactly 256 MiB");
+    capture.specialization[0].value = 9u;
+    Require(capture.WriteOptional(directory.path).empty() && !AgcDriver::ComputePipelineCapture::OptionalEnabled(), "optional capture did not stop beyond the byte cap");
+    RejectWithMessage([&] { capture.Write(directory.path); }, "compute pipeline capture: capture directory exceeds 256 MiB");
+}
+
+void CheckOptionalIo() {
+    Directory directory;
+    const auto blocked = directory.path / "blocked";
+    {
+        std::ofstream output(blocked, std::ios::binary);
+        output.put('x');
+        Require(output.good(), "cannot create output conflict");
+    }
+    const auto capture = AgcDriver::ComputePipelineFixture();
+    Require(capture.WriteOptional(blocked / "captures").empty() && !AgcDriver::ComputePipelineCapture::OptionalEnabled(), "optional capture did not stop after an output I/O error");
+    Require(std::filesystem::file_size(blocked) == 1u, "optional capture changed the conflicting file");
+    RejectWithMessage([&] { capture.Write(blocked / "captures"); }, "compute pipeline capture: cannot create capture directory");
+    Require(capture.WriteOptional(directory.path / "after-stop").empty() && !std::filesystem::exists(directory.path / "after-stop"), "disabled optional capture retried output");
+    auto malformed = capture;
+    malformed.spirv[0] = 0u;
+    Reject([&] { malformed.WriteOptional(directory.path / "after-stop"); }, "disabled optional capture swallowed an invalid pipeline input");
+}
+
+void CheckOptionalInvalid() {
+    Directory directory;
+    auto capture = AgcDriver::ComputePipelineFixture();
+    auto malformed = capture;
+    malformed.spirv[0] = 0u;
+    Reject([&] { malformed.WriteOptional(directory.path / "malformed"); }, "optional capture swallowed an invalid pipeline input");
+    Require(AgcDriver::ComputePipelineCapture::OptionalEnabled() && !std::filesystem::exists(directory.path / "malformed"), "invalid input disabled optional capture or touched its output");
+    const auto path = capture.WriteOptional(directory.path);
+    Require(!path.empty(), "optional capture refused valid input");
+    std::ofstream(path, std::ios::binary | std::ios::trunc).put('x');
+    Reject([&] { capture.WriteOptional(directory.path); }, "optional capture swallowed an existing damaged record");
+    Require(AgcDriver::ComputePipelineCapture::OptionalEnabled() && std::filesystem::file_size(path) == 1u, "damaged record disabled optional capture or was overwritten");
+    Require(std::filesystem::remove(path), "cannot remove owned damaged record");
+    Reject([&] { capture.WriteOptional(directory.path); }, "optional capture swallowed an incomplete existing record");
+    Require(AgcDriver::ComputePipelineCapture::OptionalEnabled() && !std::filesystem::exists(path), "incomplete record disabled optional capture or was replaced");
+    Require(std::filesystem::create_directory(path), "cannot create owned nonregular capture");
+    Reject([&] { capture.WriteOptional(directory.path); }, "optional capture swallowed a nonregular existing capture");
+    Require(AgcDriver::ComputePipelineCapture::OptionalEnabled() && std::filesystem::is_directory(path), "nonregular capture disabled optional capture or was changed");
+    std::filesystem::remove_all(path.parent_path());
+    std::ofstream(path.parent_path(), std::ios::binary).put('x');
+    Reject([&] { capture.WriteOptional(directory.path); }, "optional capture swallowed a non-directory existing record");
+    Require(AgcDriver::ComputePipelineCapture::OptionalEnabled() && std::filesystem::file_size(path.parent_path()) == 1u, "non-directory record disabled optional capture or was changed");
+    Directory unexpected;
+    const auto foreign = unexpected.path / "foreign";
+    std::ofstream(foreign, std::ios::binary).put('x');
+    Reject([&] { capture.WriteOptional(unexpected.path); }, "optional capture swallowed an unexpected directory entry");
+    Require(AgcDriver::ComputePipelineCapture::OptionalEnabled() && std::filesystem::file_size(foreign) == 1u, "unexpected entry disabled optional capture or was changed");
+}
+
 }
 
-int main() {
+int main(int argc, char** argv) {
     try {
+        if (argc == 2) {
+            const std::string_view mode = argv[1];
+            if (mode == "--optional-entry-cap") CheckOptionalEntryCap();
+            else if (mode == "--optional-byte-cap") CheckOptionalByteCap();
+            else if (mode == "--optional-io") CheckOptionalIo();
+            else if (mode == "--optional-invalid") CheckOptionalInvalid();
+            else throw std::runtime_error("unknown optional capture test mode");
+            std::puts("optional compute capture policy passed");
+            return 0;
+        }
+        Require(argc == 1, "unknown capture test arguments");
         CheckRoundTrip();
         CheckMalformedRecords();
         CheckBindingFlags();