cuda-python/benchmarks/cuda_bindings/pytest-legacy/test_pointer_attributes.py at a368a48e07969e911f62d35589c2bc8328da4b1e · NVIDIA/cuda-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
# SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE

import random

import pytest

from conftest import ASSERT_DRV
from cuda.bindings import driver as cuda

random.seed(0)

idx = 0


def query_attribute(attribute, ptrs):
    global idx
    ptr = ptrs[idx]
    idx = (idx + 1) % len(ptrs)

    cuda.cuPointerGetAttribute(attribute, ptr)


def query_attributes(attributes, ptrs):
    global idx
    ptr = ptrs[idx]
    idx = (idx + 1) % len(ptrs)

    cuda.cuPointerGetAttributes(len(attributes), attributes, ptr)


@pytest.mark.benchmark(group="pointer-attributes")
# Measure cuPointerGetAttribute in the same way as C benchmarks
def test_pointer_get_attribute(benchmark, init_cuda):
    _ = init_cuda

    ptrs = []
    for _ in range(500):
        err, ptr = cuda.cuMemAlloc(1 << 18)
        ASSERT_DRV(err)
        ptrs.append(ptr)

    random.shuffle(ptrs)

    benchmark(query_attribute, cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_MEMORY_TYPE, ptrs)

    for p in ptrs:
        (err,) = cuda.cuMemFree(p)
        ASSERT_DRV(err)


@pytest.mark.benchmark(group="pointer-attributes")
# Measure cuPointerGetAttributes with all attributes
def test_pointer_get_attributes_all(benchmark, init_cuda):
    _ = init_cuda

    ptrs = []
    for _ in range(500):
        err, ptr = cuda.cuMemAlloc(1 << 18)
        ASSERT_DRV(err)
        ptrs.append(ptr)

    random.shuffle(ptrs)

    attributes = [
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_CONTEXT,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_HOST_POINTER,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_P2P_TOKENS,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_SYNC_MEMOPS,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_BUFFER_ID,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_IS_MANAGED,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_IS_LEGACY_CUDA_IPC_CAPABLE,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_RANGE_SIZE,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_MAPPED,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_IS_GPU_DIRECT_RDMA_CAPABLE,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_ACCESS_FLAGS,
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_MEMPOOL_HANDLE,
    ]

    benchmark(query_attributes, attributes, ptrs)

    for p in ptrs:
        (err,) = cuda.cuMemFree(p)
        ASSERT_DRV(err)


@pytest.mark.benchmark(group="pointer-attributes")
# Measure cuPointerGetAttributes with a single attribute
def test_pointer_get_attributes_single(benchmark, init_cuda):
    _ = init_cuda

    ptrs = []
    for _ in range(500):
        err, ptr = cuda.cuMemAlloc(1 << 18)
        ASSERT_DRV(err)
        ptrs.append(ptr)

    random.shuffle(ptrs)

    attributes = [
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
    ]

    benchmark(query_attributes, attributes, ptrs)

    for p in ptrs:
        (err,) = cuda.cuMemFree(p)
        ASSERT_DRV(err)