chromium/third_party/xnnpack/src/src/xnnpack/memory-planner.h

// Copyright 2020 Google LLC
//
// 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 <stddef.h>
#include <stdint.h>

#include "xnnpack.h"
#include "xnnpack/allocator.h"
#include "xnnpack/common.h"

#ifdef __cplusplus
extern "C" {
#endif

struct xnn_usage_record { … };

// Track the memory allocation in a memory arena for a runtime.
struct xnn_value_allocation_tracker { … };

// Initialize the memory allocation tracker for xnn_values.
XNN_INTERNAL void xnn_init_value_allocation_tracker(struct xnn_value_allocation_tracker* tracker,
                                                    const struct xnn_runtime* runtime);

inline static void xnn_release_value_allocation_tracker(struct xnn_value_allocation_tracker* tracker) { … }

// Add a to-be-allocated xnn_value (referred by 'value_id') of size 'tensor_size' to the allocation tracker.
// Note: this function assumes 'value_id's added in increasing order for simplicity as it's called inside a loop
// iterating over 'runtime->values'.
XNN_INTERNAL void xnn_add_value_allocation_tracker(struct xnn_value_allocation_tracker* tracker,
                                                   uint32_t value_id, size_t tensor_size);

// Add a value to represent operator workspace. This is a temporary buffer that is only used during the invocation of
// operator.
XNN_INTERNAL void xnn_add_operator_workspace_allocation_tracker(
  struct xnn_value_allocation_tracker* tracker,
  uint32_t operator_workspace_value_id,
  size_t tensor_size,
  uint32_t opdata_id);

// Mark value_id as reusing the memory that is allocated to another reuse_value_id. No memory is then
// allocated to value_id. The usage record of reuse_value_id needs to be expanded to include al consumers of value_id,
// indicated by new_last_node.
XNN_INTERNAL void xnn_mark_tensor_as_reuse(
  struct xnn_value_allocation_tracker* tracker,
  uint32_t value_id,
  uint32_t reuse_value_id,
  uint32_t new_last_node);

// Plan the exact the memory allocation for intermediate tensors according to the xnn_value allocation tracker.
XNN_INTERNAL void xnn_plan_value_allocation_tracker(struct xnn_value_allocation_tracker* tracker);

#ifdef __cplusplus
}  // extern "C"
#endif