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,