ampr: honor shared counter waits and completion writes

This commit is contained in:
nmzik
2026-10-08 03:14:25 +04:00
parent 861dd76d7f
commit 5a8808086a
2 changed files with 110 additions and 18 deletions
+51 -18
View File
@@ -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<ReadFileCommand, KernelEventCommand, WriteAddressCommand,
AmmMapCommand, WaitAddressCommand>;
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<uint64_t, CommandBufferState> g_command_buffers;
static std::unordered_map<uint64_t, uint64_t> g_command_buffer_aliases;
// AMM and APR share 256 32-bit counters; only the first 128 are application-writable.
static std::array<std::atomic<uint32_t>, 256> g_counters {};
using CommandBufferIterator = std::unordered_map<uint64_t, CommandBufferState>::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<uint64_t*>(command.address))
.store(command.value, std::memory_order_release);
} else if (const auto* command =
std::get_if<CommandBufferState::WriteCounterCommand>(&entry.data)) {
g_counters[command->index].store(command->value, std::memory_order_release);
} else if (const auto* payload = std::get_if<CommandBufferState::AmmMapCommand>(&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<uint64_t*>(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<uint64_t*>(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<int64_t>(value - wait.value) >= 0;
case 5: return static_cast<int64_t>(value) > static_cast<int64_t>(wait.value);
default: return static_cast<int64_t>(value) < static_cast<int64_t>(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<CommandBufferState::WaitAddressCommand>(
const auto* wait = std::get_if<CommandBufferState::WaitCommand>(
&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<uint64_t>(command_buffer),
CommandBufferState::WaitAddressCommand {
reinterpret_cast<uint64_t>(address), value, compare},
CommandBufferState::WaitCommand {
reinterpret_cast<uint64_t>(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<uint64_t>(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<uint64_t>(command_buffer),
CommandBufferState::WriteCounterCommand {
index, static_cast<uint32_t>(value)},
0x20)
? OK
: LibKernel::KERNEL_ERROR_EBUSY;
}
static int KYTY_SYSV_ABI CommandBufferWriteCounterOnCompletion(void* command_buffer, uint8_t,
+59
View File
@@ -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<SetBuffer>(find("N-FSPA4S3nI"));
const auto offset = reinterpret_cast<Offset>(find("GnxKOHEawhk"));
const auto wait_address = reinterpret_cast<WaitAddress>(find("DLfoNxTFNVk"));
const auto wait_counter = reinterpret_cast<WaitCounter>(find("cQb8Zr8Q0Y0"));
const auto write_counter = reinterpret_cast<WriteCounter>(find("jK+yuYCI7MA"));
const auto write_address = reinterpret_cast<WriteAddress>(find("sJXyWHjP-F8"));
const auto read_file = reinterpret_cast<ReadFile>(find("mQ16-QdKv7k"));
const auto write_event = reinterpret_cast<WriteKernelEvent>(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<char, 3> output {};
Result result {1234, 5678};
std::array<uint32_t, 4> 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<uint32_t>(offset(initializer)),
0, &ids[3]) == OK && wait_amm(ids[3]) == OK,
"initialize shared counter before dependent submissions");
}
Check(read_file(producer, reinterpret_cast<uint64_t>(&buffers[0].header[3]),
reinterpret_cast<uint64_t>(&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<uint32_t>(offset(consumer)),
0, &ids[1]) == OK &&
submit_amm(buffers[2].data.data(), static_cast<uint32_t>(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;