From 5a8808086a9815b67d4822cd20cf1e8a7a05f002 Mon Sep 17 00:00:00 2001 From: nmzik Date: Thu, 8 Oct 2026 03:14:25 +0400 Subject: [PATCH] ampr: honor shared counter waits and completion writes --- src/libs/libAmpr.cpp | 69 ++++++++++++++++++++++++--------- tests/KernelFileSystemTests.cpp | 59 ++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 18 deletions(-) diff --git a/src/libs/libAmpr.cpp b/src/libs/libAmpr.cpp index ebb32c402..3c2fc5852 100644 --- a/src/libs/libAmpr.cpp +++ b/src/libs/libAmpr.cpp @@ -1,4 +1,5 @@ #include "common/abi.h" +#include "common/assert.h" #include "common/dateTime.h" #include "common/file.h" #include "common/logging/log.h" @@ -649,10 +650,15 @@ struct CommandBufferState { uint64_t address = 0; uint64_t value = 0; }; - struct WaitAddressCommand { - uint64_t address; + struct WaitCommand { + uint64_t source; uint64_t value; uint8_t compare; + bool is_counter; + }; + struct WriteCounterCommand { + uint8_t index; + uint32_t value; }; struct AmmMapCommand { AmmCommandKind kind = AmmCommandKind::MapAuto; @@ -664,7 +670,7 @@ struct CommandBufferState { uint8_t gpu_mask_id = 0; }; using CommandData = std::variant; + AmmMapCommand, WaitCommand, WriteCounterCommand>; struct Command { uint64_t record_offset; CommandData data; @@ -689,6 +695,8 @@ static_assert(sizeof(AmmUsageStatsData) == 0x18); static std::mutex g_command_buffer_mutex; static std::unordered_map g_command_buffers; static std::unordered_map g_command_buffer_aliases; +// AMM and APR share 256 32-bit counters; only the first 128 are application-writable. +static std::array, 256> g_counters {}; using CommandBufferIterator = std::unordered_map::iterator; static void RegisterCommandBufferAliasLocked(uint64_t command_buffer, uint64_t buffer) { @@ -1191,6 +1199,9 @@ static int ExecuteCommand(const CommandBufferState::Command& entry) { const auto& command = *payload; std::atomic_ref(*reinterpret_cast(command.address)) .store(command.value, std::memory_order_release); + } else if (const auto* command = + std::get_if(&entry.data)) { + g_counters[command->index].store(command->value, std::memory_order_release); } else if (const auto* payload = std::get_if(&entry.data)) { const auto& command = *payload; int result = OK; @@ -1219,17 +1230,20 @@ struct Submission { uint32_t id = 0; }; -static bool IsWaitSatisfied(const CommandBufferState::WaitAddressCommand& wait) { - const uint64_t value = - std::atomic_ref(*reinterpret_cast(wait.address)).load(std::memory_order_acquire); +static bool IsWaitSatisfied(const CommandBufferState::WaitCommand& wait) { + const uint64_t value = wait.is_counter + ? g_counters[wait.source].load(std::memory_order_acquire) + : std::atomic_ref(*reinterpret_cast(wait.source)) + .load(std::memory_order_acquire); + const uint64_t sign = uint64_t {1} << (wait.is_counter ? 31 : 63); switch (wait.compare) { case 0: return value == wait.value; case 1: return value > wait.value; case 2: return value < wait.value; case 3: return value != wait.value; - case 4: return static_cast(value - wait.value) >= 0; - case 5: return static_cast(value) > static_cast(wait.value); - default: return static_cast(value) < static_cast(wait.value); + case 4: return ((value - wait.value) & sign) == 0; + case 5: return (value ^ sign) > (wait.value ^ sign); + default: return (value ^ sign) < (wait.value ^ sign); } } @@ -1263,7 +1277,7 @@ private: pending = true; auto& submission = queue.front(); if (submission.cursor < submission.commands.size()) { - const auto* wait = std::get_if( + const auto* wait = std::get_if( &submission.commands[submission.cursor].data); if (wait != nullptr && !IsWaitSatisfied(*wait)) { continue; @@ -1911,18 +1925,27 @@ static int KYTY_SYSV_ABI CommandBufferWaitOnAddress(void* command_b return LibKernel::KERNEL_ERROR_EINVAL; } return AppendCommandRecord(reinterpret_cast(command_buffer), - CommandBufferState::WaitAddressCommand { - reinterpret_cast(address), value, compare}, + CommandBufferState::WaitCommand { + reinterpret_cast(address), value, compare, false}, 0x20) ? OK : LibKernel::KERNEL_ERROR_EBUSY; } -static int KYTY_SYSV_ABI CommandBufferWaitOnCounter(void* command_buffer, uint8_t, uint8_t, - uint64_t, uint8_t, uint8_t, uint64_t, uint8_t) { +static int KYTY_SYSV_ABI CommandBufferWaitOnCounter(void* command_buffer, uint8_t index, + uint8_t access, uint64_t value, uint8_t compare, + uint8_t mask_op, uint64_t, uint8_t flush) { PRINT_NAME(); - return AppendNoOpCommand(command_buffer, 0x20); + if (command_buffer == nullptr || access > 7 || compare > 6 || mask_op > 1 || flush > 1) { + return LibKernel::KERNEL_ERROR_EINVAL; + } + EXIT_NOT_IMPLEMENTED(access != 1 || mask_op != 0 || flush != 0 || value > UINT32_MAX); + return AppendCommandRecord(reinterpret_cast(command_buffer), + CommandBufferState::WaitCommand {index, value, compare, true}, + 0x20) + ? OK + : LibKernel::KERNEL_ERROR_EBUSY; } static int AppendWriteAddressCommand(void* command_buffer, volatile uint64_t* address, @@ -1954,11 +1977,21 @@ static int KYTY_SYSV_ABI CommandBufferWriteAddressOnCompletion(void* return AppendWriteAddressCommand(command_buffer, address, value); } -static int KYTY_SYSV_ABI CommandBufferWriteCounter(void* command_buffer, uint8_t, uint8_t, uint64_t, - uint8_t, uint32_t) { +static int KYTY_SYSV_ABI CommandBufferWriteCounter(void* command_buffer, uint8_t index, + uint8_t access, uint64_t value, + uint8_t operation, uint32_t mode) { PRINT_NAME(); - return AppendNoOpCommand(command_buffer, 0x20); + if (command_buffer == nullptr || index >= 128 || access > 7 || operation > 4 || mode > 1) { + return LibKernel::KERNEL_ERROR_EINVAL; + } + EXIT_NOT_IMPLEMENTED(access != 1 || operation != 0 || mode != 0); + return AppendCommandRecord(reinterpret_cast(command_buffer), + CommandBufferState::WriteCounterCommand { + index, static_cast(value)}, + 0x20) + ? OK + : LibKernel::KERNEL_ERROR_EBUSY; } static int KYTY_SYSV_ABI CommandBufferWriteCounterOnCompletion(void* command_buffer, uint8_t, diff --git a/tests/KernelFileSystemTests.cpp b/tests/KernelFileSystemTests.cpp index 9c138726c..b0aeb86f9 100644 --- a/tests/KernelFileSystemTests.cpp +++ b/tests/KernelFileSystemTests.cpp @@ -712,6 +712,10 @@ void CheckAmprOrdering(Loader::SymbolDatabase &symbols, uint32_t file_id) { using Offset = uint64_t (KYTY_SYSV_ABI *)(void *); using WaitAddress = int (KYTY_SYSV_ABI *)(void *, volatile uint64_t *, uint64_t, uint8_t, uint8_t); + using WaitCounter = int (KYTY_SYSV_ABI *)(void *, uint8_t, uint8_t, uint64_t, + uint8_t, uint8_t, uint64_t, uint8_t); + using WriteCounter = int (KYTY_SYSV_ABI *)(void *, uint8_t, uint8_t, uint64_t, + uint8_t, uint32_t); using WriteAddress = int (KYTY_SYSV_ABI *)(void *, volatile uint64_t *, uint64_t); using ReadFile = int (KYTY_SYSV_ABI *)(void *, uint64_t, uint64_t, uint32_t, void *, uint64_t, uint64_t); @@ -728,6 +732,8 @@ void CheckAmprOrdering(Loader::SymbolDatabase &symbols, uint32_t file_id) { const auto set_buffer = reinterpret_cast(find("N-FSPA4S3nI")); const auto offset = reinterpret_cast(find("GnxKOHEawhk")); const auto wait_address = reinterpret_cast(find("DLfoNxTFNVk")); + const auto wait_counter = reinterpret_cast(find("cQb8Zr8Q0Y0")); + const auto write_counter = reinterpret_cast(find("jK+yuYCI7MA")); const auto write_address = reinterpret_cast(find("sJXyWHjP-F8")); const auto read_file = reinterpret_cast(find("mQ16-QdKv7k")); const auto write_event = reinterpret_cast(find("H896Pt-yB4I")); @@ -801,6 +807,59 @@ void CheckAmprOrdering(Loader::SymbolDatabase &symbols, uint32_t file_id) { "submission wait observes the completed result"); } } + constexpr std::array counter_comparisons { + Comparison{1, 0, 0, 1}, Comparison{1, 0, 0, 1}, + Comparison{4, 0x7ffffffe, 0x7fffffff, 0x80000000}, + Comparison{5, 0xffffffff, 0, 1}}; + for (const auto &comparison : counter_comparisons) { + reset(1); + auto *producer = buffers[0].header.data(); + auto *consumer = buffers[1].header.data(); + auto *lower = buffers[2].header.data(); + constexpr uint8_t Counter = 127; + constexpr auto Invalid = Libs::LibKernel::KERNEL_ERROR_EINVAL; + Check(write_counter(producer, 128, 1, 1, 0, 0) == Invalid && + write_counter(producer, Counter, 8, 1, 0, 0) == Invalid && + write_counter(producer, Counter, 1, 1, 5, 0) == Invalid && + write_counter(producer, Counter, 1, 1, 0, 2) == Invalid && + wait_counter(consumer, Counter, 8, 0, 1, 0, 0, 0) == Invalid && + wait_counter(consumer, Counter, 1, 0, 7, 0, 0, 0) == Invalid && + wait_counter(consumer, Counter, 1, 0, 1, 2, 0, 0) == Invalid && + wait_counter(consumer, Counter, 1, 0, 1, 0, 0, 2) == Invalid && + offset(producer) == 0 && offset(consumer) == 0, + "invalid counter operations fail without appending a command"); + uint64_t retired = 0, lower_done = 0; + std::array output {}; + Result result {1234, 5678}; + std::array ids {}; + if (comparison.blocked != 0) { + auto *initializer = buffers[3].header.data(); + Check(write_counter(initializer, Counter, 1, comparison.blocked, 0, 0) == OK && + submit_amm(buffers[3].data.data(), static_cast(offset(initializer)), + 0, &ids[3]) == OK && wait_amm(ids[3]) == OK, + "initialize shared counter before dependent submissions"); + } + Check(read_file(producer, reinterpret_cast(&buffers[0].header[3]), + reinterpret_cast(&buffers[0].header[4]), file_id, + output.data(), output.size(), 0) == OK && + write_counter(producer, Counter, 1, comparison.released, 0, 0) == OK && + wait_counter(consumer, Counter, 1, comparison.reference, comparison.compare, + 0, 0, 0) == OK && + write_counter(consumer, Counter, 1, 0, 0, 0) == OK && + write_address(consumer, &retired, 1) == OK && + write_address(lower, &lower_done, 1) == OK, + "build APR completion and dependent AMM counter reset"); + Check(submit_amm(buffers[1].data.data(), static_cast(offset(consumer)), + 0, &ids[1]) == OK && + submit_amm(buffers[2].data.data(), static_cast(offset(lower)), + 1, &ids[2]) == OK && wait_amm(ids[2]) == OK && + lower_done == 1 && retired == 0, + "counter wait blocks AMM retirement while its lower priority progresses"); + Check(submit_apr(producer, 3, &result, &ids[0]) == OK && wait_amm(ids[1]) == OK && + wait_apr(ids[0]) == OK && result.result == OK && retired == 1 && + std::memcmp(output.data(), "APR", 3) == 0, + "shared counter completion releases AMM only after APR read and resets for reuse"); + } reset(1); namespace EventQueue = Libs::LibKernel::EventQueue; EventQueue::KernelEqueue queue = EventQueue::KERNEL_EQUEUE_INVALID;