Skip to content

UCT/ZE/COPY: Initialize memtype cache when opening MD #12069

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tvegas1 merged 6 commits into from
Oct 30, 2026
Merged

UCT/ZE/COPY: Initialize memtype cache when opening MD #12069

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0157e44
UCT/ZE/COPY: Initialize memtype cache when opening MD
yafshar Oct 9, 2026
f0a7cf5
UCT/ZE/COPY: Probe memtype cache creation and share alloc helper in test
yafshar Oct 9, 2026
5f6c3db
UCT/ZE/COPY: Hold the ze_copy MD in an RAII handle in tests
yafshar Oct 10, 2026
2d0f0e9
Merge branch 'master' into fix/ze-copy-memtype-cache-eager
yafshar Oct 19, 2026
bffc39c
UCT/ZE/COPY: Clarify memtype cache handler comment
yafshar Oct 19, 2026
6abaccd
UCT/ZE/COPY: Reuse test_md fixture in ze_copy MD tests
yafshar Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/uct/ze/copy/ze_copy_md.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -341,6 +341,9 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name,
md->super.component = &uct_ze_copy_component;

*md_p = (uct_md_h)md;

/* Install the memtype cache handler before Level Zero allocations. */
ucs_memtype_cache_global_create();
return UCS_OK;
}

Expand Down
219 changes: 80 additions & 139 deletions test/gtest/uct/ze/test_ze_copy_md.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,13 @@
* See file LICENSE for terms.
*/

#include <common/test.h>
#include <common/mem_buffer.h>
#include <uct/test_md.h>

extern "C" {
#include <uct/api/uct.h>
#include <uct/ze/base/ze_base.h>
#include <uct/ze/copy/ze_copy_md.h>
#include <ucs/memory/memtype_cache.h>
}


Expand All @@ -17,187 +18,127 @@ extern "C" {
* and memory-type detection. The ze_copy MD is the path used by UCP for
* VRAM<->host staging on Intel XPU.
*/
class test_ze_copy_md : public ucs::test {
class test_ze_copy_md : public test_md {
protected:
void SetUp() override
/* Allocate ZE_DEVICE memory through the MD, which advertises it */
static void alloc_device(uct_md_h md, size_t *length_p, void **addr_p,
uct_mem_h *memh_p)
{
ucs::test::SetUp();
if (uct_ze_base_init() != ZE_RESULT_SUCCESS) {
UCS_TEST_SKIP_R("Level Zero runtime not available");
}
if (uct_ze_base_get_num_devices() == 0) {
UCS_TEST_SKIP_R("No Level Zero devices available");
}
}

static uct_md_h open_first_md()
{
uct_md_resource_desc_t *resources = NULL;
unsigned num = 0;
ucs_status_t status;

status = uct_ze_copy_component.query_md_resources(
&uct_ze_copy_component, &resources, &num);
EXPECT_UCS_OK(status);
if (num == 0) {
ucs_free(resources);
return NULL;
}

uct_md_config_t *md_config = NULL;
status = uct_md_config_read(&uct_ze_copy_component, NULL, NULL,
&md_config);
EXPECT_UCS_OK(status);

uct_md_h md = NULL;
status = uct_md_open(&uct_ze_copy_component, resources[0].md_name,
md_config, &md);
uct_config_release(md_config);
ucs_free(resources);

return (status == UCS_OK) ? md : NULL;
ASSERT_UCS_OK(md->ops->mem_alloc(md, length_p, addr_p,
UCS_MEMORY_TYPE_ZE_DEVICE,
UCS_SYS_DEVICE_ID_UNKNOWN, 0,
"test_ze_copy", memh_p));
}
};


UCS_TEST_F(test_ze_copy_md, component_is_registered) {
EXPECT_STREQ("ze_cpy", uct_ze_copy_component.name);
}


UCS_TEST_F(test_ze_copy_md, query_md_resources_returns_at_least_one) {
uct_md_resource_desc_t *resources = NULL;
unsigned num = 0;

EXPECT_UCS_OK(
uct_ze_copy_component.query_md_resources(&uct_ze_copy_component,
&resources, &num));
EXPECT_GE(num, 1u);
ucs_free(resources);
}


UCS_TEST_F(test_ze_copy_md, open_close_md) {
uct_md_h md = open_first_md();
if (md == NULL) {
UCS_TEST_SKIP_R("Could not open ze_cpy MD on this system");
}

auto *ze_md = ucs_derived_of(md, uct_ze_copy_md_t);
EXPECT_TRUE(ze_md->ze_device != NULL);
EXPECT_TRUE(ze_md->ze_context != NULL);

uct_md_close(md);
UCS_TEST_P(test_ze_copy_md, open_close_md) {
auto *ze_md = ucs_derived_of(md(), uct_ze_copy_md_t);
EXPECT_NE(nullptr, ze_md->ze_device);
EXPECT_NE(nullptr, ze_md->ze_context);
}


UCS_TEST_F(test_ze_copy_md, md_attr_advertises_alloc_caps) {
uct_md_h md = open_first_md();
if (md == NULL) {
UCS_TEST_SKIP_R("Could not open ze_cpy MD on this system");
}

uct_md_attr_v2_t attr;
attr.field_mask = UINT64_MAX;
EXPECT_UCS_OK(uct_md_query_v2(md, &attr));

UCS_TEST_P(test_ze_copy_md, md_attr_advertises_alloc_caps) {
/* ze_copy MD must be able to allocate ZE memory. */
EXPECT_NE(0u, attr.alloc_mem_types);
EXPECT_GT(attr.max_alloc, 0u);

uct_md_close(md);
EXPECT_NE(0u, md_attr().alloc_mem_types);
EXPECT_GT(md_attr().max_alloc, 0u);
}


/*
* Allocate ZE_DEVICE memory through the MD, then free it.
*/
UCS_TEST_F(test_ze_copy_md, mem_alloc_free_device) {
uct_md_h md = open_first_md();
if (md == NULL) {
UCS_TEST_SKIP_R("Could not open ze_cpy MD on this system");
}

UCS_TEST_P(test_ze_copy_md, mem_alloc_free_device) {
size_t length = 4096;
void *addr = NULL;
uct_mem_h memh = NULL;

ucs_status_t status = md->ops->mem_alloc(md, &length, &addr,
UCS_MEMORY_TYPE_ZE_DEVICE,
UCS_SYS_DEVICE_ID_UNKNOWN, 0,
"test_ze_copy", &memh);
if (status == UCS_ERR_UNSUPPORTED) {
UCS_TEST_SKIP_R("ZE_DEVICE alloc unsupported on this system");
}
ASSERT_UCS_OK(status);
ASSERT_TRUE(addr != NULL);
ASSERT_GE(length, 4096u);
alloc_device(md(), &length, &addr, &memh);
EXPECT_TRUE(addr != NULL);
EXPECT_GE(length, 4096u);

EXPECT_UCS_OK(md->ops->mem_free(md, memh));
uct_md_close(md);
EXPECT_UCS_OK(md()->ops->mem_free(md(), memh));
}


/*
* detect_memory_type must classify ZE device pointers as ZE_DEVICE.
*/
UCS_TEST_F(test_ze_copy_md, detect_memory_type_device) {
uct_md_h md = open_first_md();
if (md == NULL) {
UCS_TEST_SKIP_R("Could not open ze_cpy MD on this system");
}

size_t length = 4096;
void *addr = NULL;
uct_mem_h memh = NULL;
ucs_status_t status = md->ops->mem_alloc(md, &length, &addr,
UCS_MEMORY_TYPE_ZE_DEVICE,
UCS_SYS_DEVICE_ID_UNKNOWN, 0,
"test_ze_copy", &memh);
if (status == UCS_ERR_UNSUPPORTED) {
UCS_TEST_SKIP_R("ZE_DEVICE alloc unsupported on this system");
}
ASSERT_UCS_OK(status);
UCS_TEST_P(test_ze_copy_md, detect_memory_type_device) {
size_t length = 4096;
void *addr = NULL;
uct_mem_h memh = NULL;
alloc_device(md(), &length, &addr, &memh);

ucs_memory_type_t mem_type = UCS_MEMORY_TYPE_HOST;
EXPECT_UCS_OK(md->ops->detect_memory_type(md, addr, length, &mem_type));
EXPECT_UCS_OK(md()->ops->detect_memory_type(md(), addr, length, &mem_type));
EXPECT_EQ(UCS_MEMORY_TYPE_ZE_DEVICE, mem_type);

EXPECT_UCS_OK(md->ops->mem_free(md, memh));
uct_md_close(md);
EXPECT_UCS_OK(md()->ops->mem_free(md(), memh));
}


/*
* mem_query reports base address and allocation length back.
*/
UCS_TEST_F(test_ze_copy_md, mem_query_base_and_length) {
uct_md_h md = open_first_md();
if (md == NULL) {
UCS_TEST_SKIP_R("Could not open ze_cpy MD on this system");
}

size_t length = 8192;
void *addr = NULL;
uct_mem_h memh = NULL;
ucs_status_t status = md->ops->mem_alloc(md, &length, &addr,
UCS_MEMORY_TYPE_ZE_DEVICE,
UCS_SYS_DEVICE_ID_UNKNOWN, 0,
"test_ze_copy", &memh);
if (status == UCS_ERR_UNSUPPORTED) {
UCS_TEST_SKIP_R("ZE_DEVICE alloc unsupported on this system");
}
ASSERT_UCS_OK(status);
UCS_TEST_P(test_ze_copy_md, mem_query_base_and_length) {
size_t length = 8192;
void *addr = NULL;
uct_mem_h memh = NULL;
alloc_device(md(), &length, &addr, &memh);

uct_md_mem_attr_v2_t mem_attr = {};
mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |
UCT_MD_MEM_ATTR_V2_FIELD_BASE_ADDRESS |
UCT_MD_MEM_ATTR_V2_FIELD_ALLOC_LENGTH;
EXPECT_UCS_OK(md->ops->mem_query(md, addr, length, &mem_attr));
EXPECT_UCS_OK(md()->ops->mem_query(md(), addr, length, &mem_attr));
EXPECT_EQ(UCS_MEMORY_TYPE_ZE_DEVICE, mem_attr.mem_type);
EXPECT_TRUE(mem_attr.base_address != NULL);
EXPECT_GE(mem_attr.alloc_length, length);

EXPECT_UCS_OK(md->ops->mem_free(md, memh));
uct_md_close(md);
EXPECT_UCS_OK(md()->ops->mem_free(md(), memh));
}

_UCT_MD_INSTANTIATE_TEST_CASE(test_ze_copy_md, ze_cpy)


class test_ze_copy_memtype_cache : public test_ze_copy_md {
protected:
void init() override
{
/* Drop the cache, so that the MD open has to create it again */
ucs_memtype_cache_cleanup();
ucs_memtype_cache_global_init();
ASSERT_EQ(nullptr, ucs_memtype_cache_global_instance);

test_ze_copy_md::init();
}
};


/*
* An application allocation made after the MD is opened must be tracked in
* the memtype cache.
*/
UCS_TEST_P(test_ze_copy_memtype_cache, memtype_cache_created_at_md_open,
"MEMTYPE_CACHE=y") {
ASSERT_NE(nullptr, ucs_memtype_cache_global_instance);

const size_t length = 4096;
ucs_memory_info_t mem_info;
void *addr;

{
mem_buffer buffer(length, UCS_MEMORY_TYPE_ZE_DEVICE);

addr = buffer.ptr();
ASSERT_UCS_OK(ucs_memtype_cache_lookup(addr, length, &mem_info));
EXPECT_EQ(UCS_MEMORY_TYPE_ZE_DEVICE, mem_info.type);
}

EXPECT_EQ(UCS_ERR_NO_ELEM,
ucs_memtype_cache_lookup(addr, length, &mem_info));
}

_UCT_MD_INSTANTIATE_TEST_CASE(test_ze_copy_memtype_cache, ze_cpy)
Loading