1# Tencent is pleased to support the open source community by making ncnn available.
2#
3# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
4#
5# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
6# in compliance with the License. You may obtain a copy of the License at
7#
8# https://opensource.org/licenses/BSD-3-Clause
9#
10# Unless required by applicable law or agreed to in writing, software distributed
11# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
12# CONDITIONS OF ANY KIND, either express or implied. See the License for the
13# specific language governing permissions and limitations under the License.
14
15import pytest
16
17import ncnn
18
19
20def test_vk_blob_allocator():
21    if not hasattr(ncnn, "get_gpu_count"):
22        return
23
24    vkdev = ncnn.get_gpu_device(0)
25    assert vkdev is not None
26    allocator = ncnn.VkBlobAllocator(vkdev)
27    assert allocator.buffer_memory_type_index >= 0
28    assert allocator.image_memory_type_index >= 0
29
30    mappable = allocator.mappable
31    allocator.mappable = not mappable
32    assert allocator.mappable == (not mappable)
33
34    coherent = allocator.coherent
35    allocator.coherent = not coherent
36    assert allocator.coherent == (not coherent)
37
38    bufmem = allocator.fastMalloc(10 * 1024)
39    assert bufmem is not None
40    allocator.fastFree(bufmem)
41
42    imgmem = allocator.fastMalloc(4, 4, 3, 4, 1)
43    assert imgmem is not None
44    allocator.fastFree(imgmem)
45
46
47def test_vk_weight_allocator():
48    if not hasattr(ncnn, "get_gpu_count"):
49        return
50
51    vkdev = ncnn.get_gpu_device(0)
52    assert vkdev is not None
53    allocator = ncnn.VkWeightAllocator(vkdev)
54    assert allocator.buffer_memory_type_index >= 0
55    assert allocator.image_memory_type_index >= 0
56
57    mappable = allocator.mappable
58    allocator.mappable = not mappable
59    assert allocator.mappable == (not mappable)
60
61    coherent = allocator.coherent
62    allocator.coherent = not coherent
63    assert allocator.coherent == (not coherent)
64
65    bufmem = allocator.fastMalloc(10 * 1024)
66    assert bufmem is not None
67    allocator.fastFree(bufmem)
68
69    imgmem = allocator.fastMalloc(4, 4, 3, 4, 1)
70    assert imgmem is not None
71    allocator.fastFree(imgmem)
72
73
74def test_vk_staging_allocator():
75    if not hasattr(ncnn, "get_gpu_count"):
76        return
77
78    vkdev = ncnn.get_gpu_device(0)
79    assert vkdev is not None
80    allocator = ncnn.VkStagingAllocator(vkdev)
81    assert allocator.buffer_memory_type_index >= 0
82    assert allocator.image_memory_type_index >= 0
83
84    mappable = allocator.mappable
85    allocator.mappable = not mappable
86    assert allocator.mappable == (not mappable)
87
88    coherent = allocator.coherent
89    allocator.coherent = not coherent
90    assert allocator.coherent == (not coherent)
91
92    bufmem = allocator.fastMalloc(10 * 1024)
93    assert bufmem is not None
94    allocator.fastFree(bufmem)
95
96    imgmem = allocator.fastMalloc(4, 4, 3, 4, 1)
97    assert imgmem is not None
98    allocator.fastFree(imgmem)
99
100
101def test_vk_weight_staging_allocator():
102    if not hasattr(ncnn, "get_gpu_count"):
103        return
104
105    vkdev = ncnn.get_gpu_device(0)
106    assert vkdev is not None
107    allocator = ncnn.VkWeightStagingAllocator(vkdev)
108    assert allocator.buffer_memory_type_index >= 0
109    assert allocator.image_memory_type_index >= 0
110
111    mappable = allocator.mappable
112    allocator.mappable = not mappable
113    assert allocator.mappable == (not mappable)
114
115    coherent = allocator.coherent
116    allocator.coherent = not coherent
117    assert allocator.coherent == (not coherent)
118
119    bufmem = allocator.fastMalloc(10 * 1024)
120    assert bufmem is not None
121    allocator.fastFree(bufmem)
122
123    imgmem = allocator.fastMalloc(4, 4, 3, 4, 1)
124    assert imgmem is None
125