Skip to content

Commit

Permalink
[HIP] Add support for multiple active mappings
Browse files Browse the repository at this point in the history 
This patch adds support for mapping the same buffer multiple times with
urEnqueueMemBufferMap.

This is mostly a straight port of oneapi-src#1220.
  • Loading branch information
@frasercrmck
frasercrmck committed May 22, 2024
1 parent d3502dc commit 8422136
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 82 deletions.
46 changes: 20 additions & 26 deletions source/adapters/hip/enqueue.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#include "kernel.hpp"
#include "memory.hpp"
#include "queue.hpp"
#include "ur_api.h"

#include <ur/ur.hpp>

Expand Down Expand Up @@ -1239,27 +1240,20 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferMap(
UR_ASSERT(offset + size <= BufferImpl.getSize(),
UR_RESULT_ERROR_INVALID_SIZE);

ur_result_t Result = UR_RESULT_ERROR_INVALID_OPERATION;
auto MapPtr = BufferImpl.mapToPtr(size, offset, mapFlags);
if (!MapPtr)
return UR_RESULT_ERROR_INVALID_MEM_OBJECT;

const bool IsPinned =
BufferImpl.MemAllocMode == BufferMem::AllocMode::AllocHostPtr;

// Currently no support for overlapping regions
if (BufferImpl.getMapPtr() != nullptr) {
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
}

// Allocate a pointer in the host to store the mapped information
auto HostPtr = BufferImpl.mapToPtr(size, offset, mapFlags);
*ppRetMap = std::get<BufferMem>(hBuffer->Mem).getMapPtr();
if (HostPtr) {
Result = UR_RESULT_SUCCESS;
}
ur_result_t Result = UR_RESULT_SUCCESS;

if (!IsPinned &&
((mapFlags & UR_MAP_FLAG_READ) || (mapFlags & UR_MAP_FLAG_WRITE))) {
// Pinned host memory is already on host so it doesn't need to be read.
Result = urEnqueueMemBufferRead(hQueue, hBuffer, blockingMap, offset, size,
HostPtr, numEventsInWaitList,
MapPtr, numEventsInWaitList,
phEventWaitList, phEvent);
} else {
ScopedContext Active(hQueue->getDevice());
Expand All @@ -1281,6 +1275,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferMap(
}
}

*ppRetMap = MapPtr;

return Result;
}

Expand All @@ -1292,24 +1288,22 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemUnmap(
ur_queue_handle_t hQueue, ur_mem_handle_t hMem, void *pMappedPtr,
uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
ur_event_handle_t *phEvent) {
ur_result_t Result = UR_RESULT_SUCCESS;
UR_ASSERT(hMem->isBuffer(), UR_RESULT_ERROR_INVALID_MEM_OBJECT);
UR_ASSERT(std::get<BufferMem>(hMem->Mem).getMapPtr() != nullptr,
UR_RESULT_ERROR_INVALID_MEM_OBJECT);
UR_ASSERT(std::get<BufferMem>(hMem->Mem).getMapPtr() == pMappedPtr,
UR_RESULT_ERROR_INVALID_MEM_OBJECT);
auto &BufferImpl = std::get<BufferMem>(hMem->Mem);

auto *Map = BufferImpl.getMapDetails(pMappedPtr);
UR_ASSERT(Map != nullptr, UR_RESULT_ERROR_INVALID_MEM_OBJECT);

const bool IsPinned = std::get<BufferMem>(hMem->Mem).MemAllocMode ==
BufferMem::AllocMode::AllocHostPtr;
const bool IsPinned =
BufferImpl.MemAllocMode == BufferMem::AllocMode::AllocHostPtr;

ur_result_t Result = UR_RESULT_SUCCESS;
if (!IsPinned &&
((std::get<BufferMem>(hMem->Mem).getMapFlags() & UR_MAP_FLAG_WRITE) ||
(std::get<BufferMem>(hMem->Mem).getMapFlags() &
UR_MAP_FLAG_WRITE_INVALIDATE_REGION))) {
((Map->getMapFlags() & UR_MAP_FLAG_WRITE) ||
(Map->getMapFlags() & UR_MAP_FLAG_WRITE_INVALIDATE_REGION))) {
// Pinned host memory is only on host so it doesn't need to be written to.
Result = urEnqueueMemBufferWrite(
hQueue, hMem, true, std::get<BufferMem>(hMem->Mem).getMapOffset(),
std::get<BufferMem>(hMem->Mem).getMapSize(), pMappedPtr,
hQueue, hMem, true, Map->getMapOffset(), Map->getMapSize(), pMappedPtr,
numEventsInWaitList, phEventWaitList, phEvent);
} else {
ScopedContext Active(hQueue->getDevice());
Expand All @@ -1331,7 +1325,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemUnmap(
}
}

std::get<BufferMem>(hMem->Mem).unmap(pMappedPtr);
BufferImpl.unmap(pMappedPtr);
return Result;
}

Expand Down
125 changes: 69 additions & 56 deletions source/adapters/hip/memory.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,53 +9,72 @@
//===----------------------------------------------------------------------===//
#pragma once

#include "common.hpp"
#include "context.hpp"
#include "event.hpp"
#include <cassert>
#include <memory>
#include <unordered_map>
#include <variant>

#include "common.hpp"

ur_result_t allocateMemObjOnDeviceIfNeeded(ur_mem_handle_t,
const ur_device_handle_t);
ur_result_t migrateMemoryToDeviceIfNeeded(ur_mem_handle_t,
const ur_device_handle_t);

// Handler for plain, pointer-based HIP allocations
struct BufferMem {
struct BufferMap {
/// Size of the active mapped region.
size_t MapSize;
/// Offset of the active mapped region.
size_t MapOffset;
/// Original flags for the mapped region
ur_map_flags_t MapFlags;
/// Allocated host memory used exclusively for this map.
std::shared_ptr<unsigned char[]> MapMem;

BufferMap(size_t MapSize, size_t MapOffset, ur_map_flags_t MapFlags)
: MapSize(MapSize), MapOffset(MapOffset), MapFlags(MapFlags),
MapMem(nullptr) {}

BufferMap(size_t MapSize, size_t MapOffset, ur_map_flags_t MapFlags,
std::unique_ptr<unsigned char[]> &&MapMem)
: MapSize(MapSize), MapOffset(MapOffset), MapFlags(MapFlags),
MapMem(std::move(MapMem)) {}

size_t getMapSize() const noexcept { return MapSize; }

size_t getMapOffset() const noexcept { return MapOffset; }

ur_map_flags_t getMapFlags() const noexcept { return MapFlags; }
};

/** AllocMode
* Classic: Just a normal buffer allocated on the device via hip malloc
* UseHostPtr: Use an address on the host for the device
* CopyIn: The data for the device comes from the host but the host
pointer is not available later for re-use
* AllocHostPtr: Uses pinned-memory allocation
*/
enum class AllocMode { Classic, UseHostPtr, CopyIn, AllocHostPtr };

using native_type = hipDeviceptr_t;

// If this allocation is a sub-buffer (i.e., a view on an existing
// allocation), this is the pointer to the parent handler structure
ur_mem_handle_t Parent = nullptr;
// Outer mem holding this struct in variant
ur_mem_handle_t OuterMemStruct;

/// Pointer associated with this device on the host
void *HostPtr;
/// Size of the allocation in bytes
size_t Size;
/// Size of the active mapped region.
size_t MapSize;
/// Offset of the active mapped region.
size_t MapOffset;
/// Pointer to the active mapped region, if any
void *MapPtr;
/// Original flags for the mapped region
ur_map_flags_t MapFlags;
/// A map that contains all the active mappings for this buffer.
std::unordered_map<void *, BufferMap> PtrToBufferMap;

/** AllocMode
* Classic: Just a normal buffer allocated on the device via hip malloc
* UseHostPtr: Use an address on the host for the device
* CopyIn: The data for the device comes from the host but the host
pointer is not available later for re-use
* AllocHostPtr: Uses pinned-memory allocation
*/
enum class AllocMode {
Classic,
UseHostPtr,
CopyIn,
AllocHostPtr
} MemAllocMode;
AllocMode MemAllocMode;

private:
// Vector of HIP pointers
Expand All @@ -65,10 +84,8 @@ struct BufferMem {
BufferMem(ur_context_handle_t Context, ur_mem_handle_t OuterMemStruct,
AllocMode Mode, void *HostPtr, size_t Size)
: OuterMemStruct{OuterMemStruct}, HostPtr{HostPtr}, Size{Size},
MapSize{0}, MapOffset{0}, MapPtr{nullptr}, MapFlags{UR_MAP_FLAG_WRITE},
MemAllocMode{Mode}, Ptrs(Context->Devices.size(), native_type{0}){};

BufferMem(const BufferMem &Buffer) = default;
PtrToBufferMap{}, MemAllocMode{Mode},
Ptrs(Context->Devices.size(), native_type{0}) {};

// This will allocate memory on device if there isn't already an active
// allocation on the device
Expand Down Expand Up @@ -98,45 +115,41 @@ struct BufferMem {

size_t getSize() const noexcept { return Size; }

void *getMapPtr() const noexcept { return MapPtr; }

size_t getMapSize() const noexcept { return MapSize; }

size_t getMapOffset() const noexcept { return MapOffset; }
BufferMap *getMapDetails(void *Map) {
auto details = PtrToBufferMap.find(Map);
if (details != PtrToBufferMap.end()) {
return &details->second;
}
return nullptr;
}

/// Returns a pointer to data visible on the host that contains
/// the data on the device associated with this allocation.
/// The offset is used to index into the HIP allocation.
///
void *mapToPtr(size_t Size, size_t Offset, ur_map_flags_t Flags) noexcept {
assert(MapPtr == nullptr);
MapSize = Size;
MapOffset = Offset;
MapFlags = Flags;
if (HostPtr) {
MapPtr = static_cast<char *>(HostPtr) + Offset;
void *mapToPtr(size_t MapSize, size_t MapOffset,
ur_map_flags_t MapFlags) noexcept {
void *MapPtr = nullptr;
if (HostPtr == nullptr) {
/// If HostPtr is invalid, we need to create a Mapping that owns its own
/// memory on the host.
auto MapMem = std::make_unique<unsigned char[]>(MapSize);
MapPtr = MapMem.get();
PtrToBufferMap.insert(
{MapPtr, BufferMap(MapSize, MapOffset, MapFlags, std::move(MapMem))});
} else {
// TODO: Allocate only what is needed based on the offset
MapPtr = static_cast<void *>(malloc(this->getSize()));
/// However, if HostPtr already has valid memory (e.g. pinned allocation),
/// we can just use that memory for the mapping.
MapPtr = static_cast<char *>(HostPtr) + MapOffset;
PtrToBufferMap.insert({MapPtr, BufferMap(MapSize, MapOffset, MapFlags)});
}
return MapPtr;
}

/// Detach the allocation from the host memory.
void unmap(void *) noexcept {
void unmap(void *MapPtr) noexcept {
assert(MapPtr != nullptr);

if (MapPtr != HostPtr) {
free(MapPtr);
}
MapPtr = nullptr;
MapSize = 0;
MapOffset = 0;
}

ur_map_flags_t getMapFlags() const noexcept {
assert(MapPtr != nullptr);
return MapFlags;
PtrToBufferMap.erase(MapPtr);
}

ur_result_t clear() {
Expand Down Expand Up @@ -414,7 +427,7 @@ struct ur_mem_handle_t_ {
HaveMigratedToDeviceSinceLastWrite(Context->Devices.size(), false),
Mem{std::in_place_type<BufferMem>, Ctxt, this, Mode, HostPtr, Size} {
urContextRetain(Context);
};
}

// Subbuffer constructor
ur_mem_handle_t_(ur_mem Parent, size_t SubBufferOffset)
Expand All @@ -435,7 +448,7 @@ struct ur_mem_handle_t_ {
}
}
urMemRetain(Parent);
};
}

/// Constructs the UR mem handler for an Image object
ur_mem_handle_t_(ur_context Ctxt, ur_mem_flags_t MemFlags,
Expand Down

0 comments on commit 8422136

Please sign in to comment.