/*

* Copyright (c) Meta Platforms, Inc. and affiliates.

* All rights reserved.

*

* This source code is licensed under the BSD-style license found in the

* LICENSE file in the root directory of this source tree.

*/

#pragma once

#include <stdio.h>

#include <cinttypes>

#include <cstdint>

#include <c10/util/safe_numerics.h>

#include <executorch/runtime/core/memory_allocator.h>

#include <executorch/runtime/platform/log.h>

#include <executorch/runtime/platform/platform.h>

namespace executorch {

namespace ET_RUNTIME_NAMESPACE {

namespace internal {

/**

* PlatformMemoryAllocator is a memory allocator that uses a linked list to

* manage allocated nodes. It overrides the allocate method of MemoryAllocator

* using the PAL fallback allocator method `et_pal_allocate`.

*/

class PlatformMemoryAllocator final : public MemoryAllocator {

private:

// We allocate a little more than requested and use that memory as a node in

// a linked list, pushing the allocated buffers onto a list that's iterated

// and freed when the KernelRuntimeContext is destroyed.

struct AllocationNode {

void* data;

AllocationNode* next;

};

AllocationNode* head_ = nullptr;

public:

PlatformMemoryAllocator() : MemoryAllocator(0, nullptr) {}

void* allocate(size_t size, size_t alignment = kDefaultAlignment) override {

if (!isPowerOf2(alignment)) {

ET_LOG(Error, "Alignment %zu is not a power of 2", alignment);

return nullptr;

}

// Check for overflow before computing total allocation size.

// Allocate enough for the node, data, and alignment bump (at most

// alignment - 1 extra bytes to align the data pointer).

size_t alloc_size = 0;

if (c10::add_overflows(sizeof(AllocationNode), size, &alloc_size) ||

c10::add_overflows(alloc_size, alignment - 1, &alloc_size)) {

ET_LOG(

Error,

"Allocation size overflow: size %zu, alignment %zu",

size,

alignment);

return nullptr;

}

void* node_memory = runtime::pal_allocate(alloc_size);

// If allocation failed, log message and return nullptr.

if (node_memory == nullptr) {

ET_LOG(Error, "Failed to allocate %zu bytes", alloc_size);

return nullptr;

}

// Compute data pointer.

uint8_t* data_ptr =

reinterpret_cast<uint8_t*>(node_memory) + sizeof(AllocationNode);

// Align the data pointer.

void* aligned_data_ptr = alignPointer(data_ptr, alignment);

// Assert that the alignment didn't overflow the allocated memory.

ET_DCHECK_MSG(

reinterpret_cast<uintptr_t>(aligned_data_ptr) + size <=

reinterpret_cast<uintptr_t>(node_memory) + alloc_size,

"aligned_data_ptr %p + size %zu > node_memory %p + alloc_size %zu",

aligned_data_ptr,

size,

node_memory,

alloc_size);

// Construct the node.

AllocationNode* new_node = reinterpret_cast<AllocationNode*>(node_memory);

new_node->data = aligned_data_ptr;

new_node->next = head_;

head_ = new_node;

// Return the aligned data pointer.

return head_->data;

}

void reset() override {

AllocationNode* current = head_;

while (current != nullptr) {

AllocationNode* next = current->next;

runtime::pal_free(current);

current = next;

}

head_ = nullptr;

}

~PlatformMemoryAllocator() override {

reset();

}

private:

// Disable copy and move.

PlatformMemoryAllocator(const PlatformMemoryAllocator&) = delete;

PlatformMemoryAllocator& operator=(const PlatformMemoryAllocator&) = delete;

PlatformMemoryAllocator(PlatformMemoryAllocator&&) noexcept = delete;

PlatformMemoryAllocator& operator=(PlatformMemoryAllocator&&) noexcept =

delete;

};

} // namespace internal

} // namespace ET_RUNTIME_NAMESPACE

} // namespace executorch