wehub-resource-sync
57 lines
2.4 KiB
C++
57 lines
2.4 KiB
C++
/* 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_
|