1#include "taichi/rhi/vulkan/vulkan_common.h"
2#include "taichi/rhi/device.h"
3
4#include "taichi/rhi/vulkan/vulkan_loader.h"
5
6#ifdef __APPLE__
7// For `runtime_lib_dir()`
8#include "taichi/util/lang_util.h"
9// For `glfwInitVulkanLoader`
10#include "GLFW/glfw3.h"
11#endif
12
13namespace taichi::lang {
14namespace vulkan {
15
16VulkanLoader::VulkanLoader() {
17}
18
19bool VulkanLoader::check_vulkan_device() {
20#ifdef __APPLE__
21 glfwInitVulkanLoader(vkGetInstanceProcAddr);
22#endif
23
24 bool found_device_with_compute = false;
25
26 // We create an temporary Vulkan instance to probe the Vulkan devices.
27 // Otherwise, in the case of a CPU only VM with Vulkan installed, Vulkan will
28 // not run as there is no GPU available, but the fallback will not happen
29 // because Vulkan API is available.
30
31 VkApplicationInfo app_info{};
32 app_info.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
33 app_info.pApplicationName = "Checking Vulkan Device";
34 app_info.applicationVersion = VK_MAKE_VERSION(1, 0, 0);
35 app_info.pEngineName = "No Engine";
36 app_info.engineVersion = VK_MAKE_VERSION(1, 0, 0);
37 app_info.apiVersion = VK_API_VERSION_1_0;
38
39 VkInstanceCreateInfo create_info{};
40 create_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
41 create_info.pApplicationInfo = &app_info;
42
43 VkInstance instance{VK_NULL_HANDLE};
44 VkResult res = vkCreateInstance(&create_info, kNoVkAllocCallbacks, &instance);
45
46 do {
47 if (res != VK_SUCCESS) {
48 RHI_LOG_ERROR("Can not create Vulkan instance");
49 break;
50 }
51
52 load_instance(instance);
53
54 uint32_t device_count = 0;
55 vkEnumeratePhysicalDevices(instance, &device_count, nullptr);
56
57 if (device_count == 0) {
58 RHI_LOG_ERROR("Can not find Vulkan capable devices");
59 break;
60 }
61
62 std::vector<VkPhysicalDevice> devices(device_count);
63 vkEnumeratePhysicalDevices(instance, &device_count, devices.data());
64
65 for (int i = 0; i < devices.size(); i++) {
66 const auto &physical_device = devices[i];
67
68 uint32_t queue_family_count = 0;
69 vkGetPhysicalDeviceQueueFamilyProperties(physical_device,
70 &queue_family_count, nullptr);
71 if (queue_family_count > 0) {
72 std::vector<VkQueueFamilyProperties> queue_families(queue_family_count);
73 vkGetPhysicalDeviceQueueFamilyProperties(
74 physical_device, &queue_family_count, queue_families.data());
75
76 for (auto &queue : queue_families) {
77 if (queue.queueFlags & VK_QUEUE_COMPUTE_BIT) {
78 found_device_with_compute = true;
79 }
80 }
81 }
82 }
83 } while (false);
84
85 if (instance) {
86 vkDestroyInstance(instance, kNoVkAllocCallbacks);
87 }
88
89 return found_device_with_compute;
90}
91
92bool VulkanLoader::init(PFN_vkGetInstanceProcAddr get_proc_addr) {
93 std::call_once(init_flag_, [&]() {
94 if (initialized_) {
95 return;
96 }
97 // (penguinliong) So that MoltenVK instances can be imported.
98 if (get_proc_addr != nullptr) {
99 volkInitializeCustom(get_proc_addr);
100 initialized_ = true;
101 return;
102 }
103#if defined(__APPLE__)
104 vulkan_rt_ = std::make_unique<DynamicLoader>(runtime_lib_dir() +
105 "/libMoltenVK.dylib");
106 PFN_vkGetInstanceProcAddr get_proc_addr =
107 (PFN_vkGetInstanceProcAddr)vulkan_rt_->load_function(
108 "vkGetInstanceProcAddr");
109
110 volkInitializeCustom(get_proc_addr);
111 initialized_ = true;
112#else
113 VkResult result = volkInitialize();
114 initialized_ = result == VK_SUCCESS;
115#endif
116 initialized_ = initialized_ && check_vulkan_device();
117 const char *id = std::getenv("TI_VISIBLE_DEVICE");
118 if (id) {
119 set_vulkan_visible_device(id);
120 }
121 });
122 return initialized_;
123}
124
125void VulkanLoader::load_instance(VkInstance instance) {
126 vulkan_instance_ = instance;
127 volkLoadInstance(instance);
128}
129void VulkanLoader::load_device(VkDevice device) {
130 vulkan_device_ = device;
131 volkLoadDevice(device);
132}
133
134PFN_vkVoidFunction VulkanLoader::load_function(const char *name) {
135 auto result =
136 vkGetInstanceProcAddr(VulkanLoader::instance().vulkan_instance_, name);
137 if (result == nullptr) {
138 char msg_buf[256];
139 snprintf(msg_buf, sizeof(msg_buf), "Failed to load vulkan function %s",
140 name);
141 RHI_LOG_ERROR(msg_buf);
142 }
143 return result;
144}
145
146bool is_vulkan_api_available() {
147 return VulkanLoader::instance().init();
148}
149
150void set_vulkan_visible_device(std::string id) {
151 VulkanLoader::instance().visible_device_id = id;
152}
153
154} // namespace vulkan
155} // namespace taichi::lang
156