/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_COMPILER_TF2XLA_ALLOCATOR_H_
#define TENSORFLOW_COMPILER_TF2XLA_ALLOCATOR_H_

#include <cstddef>

#include "absl/types/span.h"
#include "xla/backends/cpu/buffer_allocation_info.h"

namespace tensorflow {

// AlignedBufferBytes returns the sum of the size of each buffer in
// `buffer_infos`, skipping constants, on-stack buffers and, if
// allocate_entry_params is false, entry parameters.  There are `n` entries in
// `buffer_infos`.  Each buffer is aligned to Align() byte boundaries.
size_t AlignedBufferBytes(
    absl::Span<const xla::cpu::BufferAllocationInfo> buffers,
    bool allocate_entry_params);

// MallocContiguousBuffers allocates buffers for use by the entry point
// generated by tfcompile.  There are `n` entries in `buffer_infos`.  If
// `annotate_initialized` is set, the allocated memory will be annotated as
// having been initialized - this is useful when allocating temporary buffers.
// If allocate_entry_params is true then allocates temp buffers and entry
// parameters, otherwise allocated only temp buffers.  Slots in `bufs`
// corresponding to unallocated buffers are set to nullptr.
//
// A single contiguous block of memory is allocated, and portions of it are
// parceled out into `bufs`, which must have space for `n` entries.  Returns
// the head of the allocated contiguous block, which should be passed to
// FreeContiguous when the buffers are no longer in use.
void* MallocContiguousBuffers(
    absl::Span<const xla::cpu::BufferAllocationInfo> buffers,
    bool allocate_entry_params, void** bufs, bool annotate_initialized);

// FreeContiguous frees the contiguous block of memory allocated by
// MallocContiguousBuffers.
void FreeContiguous(void* contiguous);

}  // namespace tensorflow

#endif  // TENSORFLOW_COMPILER_TF2XLA_ALLOCATOR_H_
