Skip to content

Commit

Permalink
umf providers
Browse files Browse the repository at this point in the history
  • Loading branch information
bratpiorka committed Aug 22, 2023
1 parent 59d5fe4 commit 83f170c
Show file tree
Hide file tree
Showing 5 changed files with 312 additions and 32 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,29 @@ extern "C" {

typedef struct umf_memory_provider_t *umf_memory_provider_handle_t;

typedef enum umf_device_type_t {
UMF_DEVICE_TYPE_INVALID = -1,
UMF_DEVICE_TYPE_NUMA,
UMF_DEVICE_TYPE_GPU,
} umf_device_type_t;

typedef struct umf_memory_provider_config_t {
umf_device_type_t type;

union {
struct {
// TODO
size_t id;
} numa;

struct {
char address[12]; // in format "0000:00:00.0"
} gpu;

// other partition types ...
};
} umf_memory_provider_config_t;

///
/// \brief Creates new memory provider.
/// \param ops instance of umf_memory_provider_ops_t
Expand All @@ -39,14 +62,12 @@ void umfMemoryProviderDestroy(umf_memory_provider_handle_t hProvider);

// TODO comment
enum umf_result_t
umfMemoryProviderRegister(struct umf_memory_provider_ops_t *ops);

enum umf_result_t
umfMemoryProvidersRegistryGet(struct umf_memory_provider_ops_t *providers,
size_t *numProviders);
umfMemoryProviderRegister(umf_device_type_t type,
struct umf_memory_provider_ops_t *ops);

const struct umf_memory_provider_ops_t *
umfMemoryProvidersRegistryGetOps(char *name);
enum umf_result_t umfMemoryProvidersCreateFromType(
umf_device_type_t type, const umf_memory_provider_config_t *config,
umf_memory_provider_handle_t *hProvider /* out */);

///
/// \brief Allocates size bytes of uninitialized storage from memory provider
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
extern "C" {
#endif

struct umf_memory_provider_config_t;

/// This structure comprises function pointers used by corresponding
/// umfMemoryProvider* calls. Each memory provider implementation should
/// initialize all function pointers.
Expand All @@ -26,6 +28,8 @@ struct umf_memory_provider_ops_t {
/// Should be initialized using UMF_VERSION_CURRENT
uint32_t version;

void *priv;

///
/// \brief Initializes memory provider.
/// \param params provider-specific params
Expand All @@ -51,7 +55,7 @@ struct umf_memory_provider_ops_t {
enum umf_result_t (*purge_lazy)(void *provider, void *ptr, size_t size);
enum umf_result_t (*purge_force)(void *provider, void *ptr, size_t size);
const char *(*get_name)(void *provider);
bool (*supports_device)(const char *name);
bool (*supports_device)(const struct umf_memory_provider_config_t *config);
};

#ifdef __cplusplus
Expand Down
57 changes: 33 additions & 24 deletions source/common/unified_malloc_framework/src/memory_provider.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

#include <algorithm>
#include <cstring>
#include <memory>
#include <string>
#include <vector>

Expand All @@ -25,7 +26,12 @@ struct umf_memory_provider_t {
void *provider_priv;
};

std::vector<struct umf_memory_provider_ops_t> globalProviders;
// TODO here I use the ptr to vector because the system calls
// globalProviders destructor twice - why?
typedef std::pair<umf_device_type_t, struct umf_memory_provider_ops_t>
umf_provider_desc_pair_t;

std::vector<umf_provider_desc_pair_t> *globalProviders;

enum umf_result_t
umfMemoryProviderCreate(const struct umf_memory_provider_ops_t *ops,
Expand Down Expand Up @@ -54,40 +60,43 @@ umfMemoryProviderCreate(const struct umf_memory_provider_ops_t *ops,
return UMF_RESULT_SUCCESS;
}

enum umf_result_t umfMemoryProviderRegister(umf_memory_provider_ops_t *ops) {
enum umf_result_t umfMemoryProviderRegister(umf_device_type_t type,
umf_memory_provider_ops_t *ops) {

if (globalProviders == NULL) {
// TODO this is never freed
globalProviders = new std::vector<umf_provider_desc_pair_t>;
}

// TODO check if this provider isn't already registered
globalProviders.push_back(*ops);
globalProviders->push_back(std::make_pair(type, *ops));

return UMF_RESULT_SUCCESS;
}

enum umf_result_t
umfMemoryProvidersRegistryGet(umf_memory_provider_ops_t *providers,
size_t *numProviders) {

if (providers == NULL) {
*numProviders = globalProviders.size();
} else {
memcpy(providers, globalProviders.data(),
sizeof(umf_memory_provider_ops_t) * *numProviders);
enum umf_result_t umfMemoryProvidersCreateFromType(
umf_device_type_t type, const umf_memory_provider_config_t *config,
umf_memory_provider_handle_t *hProvider /* out */) {
if (globalProviders == NULL) {
return UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC;
}

return UMF_RESULT_SUCCESS;
}

// TODO rename ;)
const umf_memory_provider_ops_t *umfMemoryProvidersRegistryGetOps(char *name) {
auto it = std::find_if(
std::begin(globalProviders), std::end(globalProviders),
[&](auto &ops) { return std::strcmp(ops.get_name(NULL), name) == 0; });
if (hProvider != NULL) {
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
}

if (it != globalProviders.end()) {
return &(*it);
for (size_t i = 0; i < globalProviders->size(); i++) {
if (globalProviders->at(i).first == type &&
globalProviders->at(i).second.supports_device(config)) {
enum umf_result_t ret = umfMemoryProviderCreate(
&globalProviders->at(i).second, (void *)config, hProvider);
if (ret == UMF_RESULT_SUCCESS) {
return ret;
}
}
}

// else
return NULL;
return UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC;
}

void umfMemoryProviderDestroy(umf_memory_provider_handle_t hProvider) {
Expand Down
230 changes: 230 additions & 0 deletions source/common/ur_memory_provider.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,230 @@
/*
*
* Copyright (C) 2023 Intel Corporation
*
* Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
* See LICENSE.TXT
* SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
*
*/

#include <assert.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <umf.h>
#include <umf/memory_provider.h>

#include "ur_api.h"

struct ur_provider_config_t {
// NOTE: when the context is not NULL, create/get a provider instance
// based on it
ur_context_handle_t context;

// when the context is NULL, create/get a provider instance based
// on model, dev and PCI
char *model_name;
char *pci;

// type of USM allocations
ur_usm_type_t usm_type;
};

typedef struct ur_provider_priv_t {
// TODO there will be always single dev per provider instance?
ur_context_handle_t context;
ur_device_handle_t device;
ur_usm_type_t usm_type;
} ur_provider_priv_t;

enum umf_result_t ur_initialize(void *params, void **priv_ptr) {
urInit(0, 0);

if (priv_ptr == NULL) {
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
}

if (params == NULL) {
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
}

*(struct ur_provider_priv_t **)priv_ptr =
(struct ur_provider_priv_t *)malloc(sizeof(struct ur_provider_priv_t));

ur_provider_priv_t *priv = (struct ur_provider_priv_t *)*priv_ptr;
if (priv == NULL) {
return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
}

priv->usm_type = ((struct ur_provider_config_t *)params)->usm_type;

if (((struct ur_provider_config_t *)params)->context) {
priv->context = ((struct ur_provider_config_t *)params)->context;
// get the devices list from the context
size_t num_devices = 0;
urContextGetInfo(priv->context, UR_CONTEXT_INFO_NUM_DEVICES,
sizeof(size_t), &num_devices, NULL);
// assume there will be always single device per provider instance
assert(num_devices == 1); // TODO report error
urContextGetInfo(priv->context, UR_CONTEXT_INFO_DEVICES,
sizeof(ur_device_handle_t), &priv->device, NULL);
assert(priv->device != NULL); // TODO report error

return UMF_RESULT_SUCCESS;
} else if (((struct ur_provider_config_t *)params)->pci) {
// find a device that matches this name

// NOTE: we browse all platforms here because right now
// there is a single UR provider for all platforms

uint32_t adapterCount = 0;
urAdapterGet(0, NULL, &adapterCount);
ur_adapter_handle_t *adapters = (ur_adapter_handle_t *)malloc(
sizeof(ur_adapter_handle_t) * adapterCount);
urAdapterGet(adapterCount, adapters, NULL);

uint32_t platformCount = 0;
urPlatformGet(adapters, adapterCount, 1, NULL, &platformCount);
ur_platform_handle_t *platforms = (ur_platform_handle_t *)malloc(
sizeof(ur_platform_handle_t) * platformCount);
urPlatformGet(adapters, adapterCount, platformCount, platforms, NULL);

for (size_t pid = 0; pid < platformCount; pid++) {
uint32_t deviceCount = 0;
urDeviceGet(platforms[pid], UR_DEVICE_TYPE_GPU, 0, NULL,
&deviceCount);
ur_device_handle_t *devices = (ur_device_handle_t *)malloc(
sizeof(ur_device_handle_t) * deviceCount);
urDeviceGet(platforms[pid], UR_DEVICE_TYPE_GPU, deviceCount,
devices, NULL);

for (size_t did = 0; did < deviceCount; did++) {
static const size_t DEVICE_INFO_MAX_LEN = 1024;

char *device_pci = (char *)malloc(DEVICE_INFO_MAX_LEN);
memset(device_pci, 0, DEVICE_INFO_MAX_LEN);
urDeviceGetInfo(devices[did], UR_DEVICE_INFO_PCI_ADDRESS,
DEVICE_INFO_MAX_LEN - 1, device_pci, NULL);

if (strcmp(((struct ur_provider_config_t *)params)->pci,
device_pci) == 0) {
ur_context_handle_t ctx = NULL;
urContextCreate(1, &devices[did], NULL, &ctx);
priv->context = ctx;
priv->device = devices[did];

free(device_pci);
return UMF_RESULT_SUCCESS;
}
}
}
}

// else
return UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC;
}

void ur_finalize(void *provider) { free(provider); }

static enum umf_result_t ur_alloc(void *provider, size_t size, size_t alignment,
void **resultPtr) {
struct ur_provider_priv_t *config = (struct ur_provider_priv_t *)provider;
enum umf_result_t result = UMF_RESULT_SUCCESS;

// TODO check errors
assert(config->context);

switch (config->usm_type) {
case UR_USM_TYPE_HOST:
// TODO ur_usm_desc_t
urUSMHostAlloc(config->context, NULL, NULL, size, resultPtr);
break;
case UR_USM_TYPE_DEVICE:
urUSMDeviceAlloc(config->context, config->device, NULL, NULL, size,
resultPtr);
break;
case UR_USM_TYPE_SHARED:
urUSMSharedAlloc(config->context, config->device, NULL, NULL, size,
resultPtr);
break;
default:
assert(0);
}

return result;
}

static enum umf_result_t ur_free(void *provider, void *ptr, size_t size) {
struct ur_provider_priv_t *config = (struct ur_provider_priv_t *)provider;

// TODO check errors
urUSMFree(config->context, ptr);

// TODO - size?

return UMF_RESULT_SUCCESS;
}

void ur_get_last_native_error(void *provider, const char **ppMessage,
int32_t *pError) {
// TODO
}

enum umf_result_t ur_get_min_page_size(void *provider, void *ptr,
size_t *pageSize) {
*pageSize = 1024; // TODO call urVirtualMemGranularityGetInfo here
return UMF_RESULT_SUCCESS;
}

const char *ur_get_name(void *provider) { return "USM"; }

bool ur_supports_device(const struct umf_memory_provider_config_t *config) {
if (config->type != UMF_DEVICE_TYPE_GPU) {
return false;
}

urInit(0, 0);

uint32_t adapterCount = 0;
urAdapterGet(0, NULL, &adapterCount);
ur_adapter_handle_t *adapters = (ur_adapter_handle_t *)malloc(
sizeof(ur_adapter_handle_t) * adapterCount);
urAdapterGet(adapterCount, adapters, NULL);

uint32_t platformCount = 0;
urPlatformGet(adapters, adapterCount, 1, NULL, &platformCount);
ur_platform_handle_t *platforms = (ur_platform_handle_t *)malloc(
sizeof(ur_platform_handle_t) * platformCount);
urPlatformGet(adapters, adapterCount, platformCount, platforms, NULL);

for (uint32_t pid = 0; pid < platformCount; pid++) {
uint32_t deviceCount = 0;
urDeviceGet(platforms[pid], UR_DEVICE_TYPE_GPU, 0, NULL, &deviceCount);
ur_device_handle_t *devices = (ur_device_handle_t *)malloc(
sizeof(ur_device_handle_t) * deviceCount);
urDeviceGet(platforms[pid], UR_DEVICE_TYPE_GPU, deviceCount, devices,
NULL);

for (uint32_t did = 0; did < deviceCount; did++) {
static const size_t DEVICE_INFO_MAX_LEN = 1024;

char *device_name = (char *)malloc(DEVICE_INFO_MAX_LEN);
memset(device_name, 0, DEVICE_INFO_MAX_LEN);
urDeviceGetInfo(devices[did], UR_DEVICE_INFO_NAME,
DEVICE_INFO_MAX_LEN - 1, device_name, NULL);

if (memcpy((void *)config->gpu.address, device_name,
sizeof(config->gpu.address)) == 0) {
return true;
}
}
}

// no match
return false;
}
Loading

0 comments on commit 83f170c

Please sign in to comment.