All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
StackDeviceMemory.h
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the CC-by-NC license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
#pragma once
#include "DeviceMemory.h"
#include <list>
#include <memory>
#include <unordered_map>
namespace faiss { namespace gpu {
/// Device memory manager that provides temporary memory allocations
/// out of a region of memory
class StackDeviceMemory : public DeviceMemory {
public:
/// Allocate a new region of memory that we manage
explicit StackDeviceMemory(int device, size_t allocPerDevice);
/// Manage a region of memory for a particular device, with or
/// without ownership
StackDeviceMemory(int device, void* p, size_t size, bool isOwner);
~StackDeviceMemory() override;
int getDevice() const override;
DeviceMemoryReservation getMemory(cudaStream_t stream,
size_t size) override;
size_t getSizeAvailable() const override;
std::string toString() const override;
size_t getHighWaterCudaMalloc() const override;
protected:
void returnAllocation(DeviceMemoryReservation& m) override;
protected:
/// Previous allocation ranges and the streams for which
/// synchronization is required
struct Range {
inline Range(char* s, char* e, cudaStream_t str) :
start_(s), end_(e), stream_(str) {
}
// References a memory range [start, end)
char* start_;
char* end_;
cudaStream_t stream_;
};
struct Stack {
/// Constructor that allocates memory via cudaMalloc
Stack(int device, size_t size);
/// Constructor that references a pre-allocated region of memory
Stack(int device, void* p, size_t size, bool isOwner);
~Stack();
/// Returns how much size is available for an allocation without
/// calling cudaMalloc
size_t getSizeAvailable() const;
/// Obtains an allocation; all allocations are guaranteed to be 16
/// byte aligned
char* getAlloc(size_t size, cudaStream_t stream);
/// Returns an allocation
void returnAlloc(char* p, size_t size, cudaStream_t stream);
/// Returns the stack state
std::string toString() const;
/// Returns the high-water mark of cudaMalloc activity
size_t getHighWaterCudaMalloc() const;
/// Device this allocation is on
int device_;
/// Do we own our region of memory?
bool isOwner_;
/// Where our allocation begins and ends
/// [start_, end_) is valid
char* start_;
char* end_;
/// Total size end_ - start_
size_t size_;
/// Stack head within [start, end)
char* head_;
/// List of previous last users of allocations on our stack, for
/// possible synchronization purposes
std::list<Range> lastUsers_;
/// How much cudaMalloc memory is currently outstanding?
size_t mallocCurrent_;
/// What's the high water mark in terms of memory used from the
/// temporary buffer?
size_t highWaterMemoryUsed_;
/// What's the high water mark in terms of memory allocated via
/// cudaMalloc?
size_t highWaterMalloc_;
};
/// Our device
int device_;
/// Memory stack
Stack stack_;
};
} } // namespace
faiss::gpu::StackDeviceMemory::Stack::isOwner_
bool isOwner_
Do we own our region of memory?
Definition: StackDeviceMemory.h:88
faiss::gpu::StackDeviceMemory::getMemory
DeviceMemoryReservation getMemory(cudaStream_t stream, size_t size) override
Definition: StackDeviceMemory.cpp:200
faiss::gpu::StackDeviceMemory::Stack::getHighWaterCudaMalloc
size_t getHighWaterCudaMalloc() const
Returns the high-water mark of cudaMalloc activity.
Definition: StackDeviceMemory.cpp:176
faiss::gpu::StackDeviceMemory::Stack::size_
size_t size_
Total size end_ - start_.
Definition: StackDeviceMemory.h:96
faiss::gpu::StackDeviceMemory::Stack::lastUsers_
std::list< Range > lastUsers_
Definition: StackDeviceMemory.h:103
faiss::gpu::StackDeviceMemory::Stack::Stack
Stack(int device, size_t size)
Constructor that allocates memory via cudaMalloc.
Definition: StackDeviceMemory.cpp:20
faiss::gpu::StackDeviceMemory::Stack::mallocCurrent_
size_t mallocCurrent_
How much cudaMalloc memory is currently outstanding?
Definition: StackDeviceMemory.h:106
faiss::gpu::StackDeviceMemory::stack_
Stack stack_
Memory stack.
Definition: StackDeviceMemory.h:121
faiss::gpu::StackDeviceMemory::Stack::returnAlloc
void returnAlloc(char *p, size_t size, cudaStream_t stream)
Returns an allocation.
Definition: StackDeviceMemory.cpp:130
faiss::gpu::StackDeviceMemory::Stack::head_
char * head_
Stack head within [start, end)
Definition: StackDeviceMemory.h:99
faiss::gpu::StackDeviceMemory::getSizeAvailable
size_t getSizeAvailable() const override
Returns the current size available without calling cudaMalloc.
Definition: StackDeviceMemory.cpp:213
faiss::gpu::StackDeviceMemory::Stack
Definition: StackDeviceMemory.h:59
faiss::gpu::StackDeviceMemory::Stack::getSizeAvailable
size_t getSizeAvailable() const
Definition: StackDeviceMemory.cpp:58
faiss::gpu::StackDeviceMemory::Stack::highWaterMemoryUsed_
size_t highWaterMemoryUsed_
Definition: StackDeviceMemory.h:110
faiss::gpu::StackDeviceMemory::Stack::device_
int device_
Device this allocation is on.
Definition: StackDeviceMemory.h:85
faiss::gpu::StackDeviceMemory
Definition: StackDeviceMemory.h:22
faiss::gpu::StackDeviceMemory::Stack::highWaterMalloc_
size_t highWaterMalloc_
Definition: StackDeviceMemory.h:114
faiss::gpu::StackDeviceMemory::toString
std::string toString() const override
Returns a string containing our current memory manager state.
Definition: StackDeviceMemory.cpp:218
faiss::gpu::StackDeviceMemory::Range
Definition: StackDeviceMemory.h:48
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:46
faiss::gpu::StackDeviceMemory::Stack::toString
std::string toString() const
Returns the stack state.
Definition: StackDeviceMemory.cpp:155
faiss::gpu::StackDeviceMemory::getHighWaterCudaMalloc
size_t getHighWaterCudaMalloc() const override
Definition: StackDeviceMemory.cpp:223
faiss::gpu::StackDeviceMemory::Stack::getAlloc
char * getAlloc(size_t size, cudaStream_t stream)
Definition: StackDeviceMemory.cpp:63
faiss::gpu::StackDeviceMemory::Stack::start_
char * start_
Definition: StackDeviceMemory.h:92
faiss::gpu::StackDeviceMemory::device_
int device_
Our device.
Definition: StackDeviceMemory.h:118
faiss::gpu::DeviceMemoryReservation
Definition: DeviceMemory.h:20
faiss::gpu::StackDeviceMemory::getDevice
int getDevice() const override
Returns the device we are managing memory for.
Definition: StackDeviceMemory.cpp:195
faiss::gpu::StackDeviceMemory::StackDeviceMemory
StackDeviceMemory(int device, size_t allocPerDevice)
Allocate a new region of memory that we manage.
Definition: StackDeviceMemory.cpp:180