main.c - main - redgpu/vkcuda

main.c
712 строк26 KB
Mary Guillemard
Use a fat binary with all major generation supported
14 окт 2025, 22:35
14 окт 2025, 22:35caa44be
#define GLFW_INCLUDE_VULKAN

#include "volk.h"
#include <GLFW/glfw3.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

const uint8_t triangle_vert[] =
{
#embed "triangle.vert.spv"
};

const uint8_t triangle_frag[] =
{
#embed "triangle.frag.spv"
};

const uint8_t fatbin_data[] =
{
#embed "frog.fatbin"
};

uint32_t findMemoryType(VkPhysicalDevice physDev, uint32_t typeFilter, VkMemoryPropertyFlags props) {
    VkPhysicalDeviceMemoryProperties memProps;
    vkGetPhysicalDeviceMemoryProperties(physDev, &memProps);
    for (uint32_t i = 0; i < memProps.memoryTypeCount; i++) {
        if ((typeFilter & (1 << i)) && (memProps.memoryTypes[i].propertyFlags & props) == props) {
            return i;
        }
    }
    return 0;
}

#define WIDTH (1280)
#define HEIGHT (720)

#define USE_CUBIN (1)

#define vk_check(x) do { vr = x; if (vr != VK_SUCCESS) { fprintf(stderr, #x " failed! (VkResult %d 0x%x)\n", vr, vr); abort(); } } while(0)

int main() {
    volkInitialize();

    glfwInit();
    glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
    glfwWindowHint(GLFW_RESIZABLE, GLFW_FALSE);
    GLFWwindow* window = glfwCreateWindow(WIDTH, HEIGHT, "VK_NVX_binary_import example", NULL, NULL);
    
    VkResult vr;

    VkApplicationInfo appInfo = {
        .sType = VK_STRUCTURE_TYPE_APPLICATION_INFO,
        .apiVersion = VK_API_VERSION_1_4
    };
    
    uint32_t glfwExtCount;
    const char** glfwExt = glfwGetRequiredInstanceExtensions(&glfwExtCount);
    
    VkInstanceCreateInfo instInfo = {
        .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
        .pApplicationInfo = &appInfo,
        .enabledExtensionCount = glfwExtCount,
        .ppEnabledExtensionNames = glfwExt
    };
    
    VkInstance instance;
    vk_check(vkCreateInstance(&instInfo, NULL, &instance));
    volkLoadInstance(instance);
    
    VkSurfaceKHR surface;
    vk_check(glfwCreateWindowSurface(instance, window, NULL, &surface));
    
    uint32_t devCount;
    vk_check(vkEnumeratePhysicalDevices(instance, &devCount, NULL));
    VkPhysicalDevice* devices = malloc(devCount * sizeof(VkPhysicalDevice));
    vk_check(vkEnumeratePhysicalDevices(instance, &devCount, devices));
    VkPhysicalDevice physDev = devices[0];
    free(devices);
    
    uint32_t queueFamCount;
    vkGetPhysicalDeviceQueueFamilyProperties(physDev, &queueFamCount, NULL);
    VkQueueFamilyProperties* queueFams = malloc(queueFamCount * sizeof(VkQueueFamilyProperties));
    vkGetPhysicalDeviceQueueFamilyProperties(physDev, &queueFamCount, queueFams);
    
    uint32_t gfxQueueFam = 0;
    for (uint32_t i = 0; i < queueFamCount; i++) {
        VkBool32 presentSupport;
        vk_check(vkGetPhysicalDeviceSurfaceSupportKHR(physDev, i, surface, &presentSupport));
        if ((queueFams[i].queueFlags & VK_QUEUE_GRAPHICS_BIT) && presentSupport) {
            gfxQueueFam = i;
            break;
        }
    }
    free(queueFams);
    
    float queuePriority = 1.0f;
    VkDeviceQueueCreateInfo queueInfo = {
        .sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
        .queueFamilyIndex = gfxQueueFam,
        .queueCount = 1,
        .pQueuePriorities = &queuePriority
    };

    VkPhysicalDeviceVulkan14Features vulkan14features = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_4_FEATURES,
        .maintenance5 = VK_TRUE,
    };
    
    VkPhysicalDeviceFeatures2 features2 = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2,
        .pNext = &vulkan14features,
    };
    
    const char* devExt[] = { VK_KHR_SWAPCHAIN_EXTENSION_NAME, VK_NVX_BINARY_IMPORT_EXTENSION_NAME, VK_NVX_IMAGE_VIEW_HANDLE_EXTENSION_NAME };
    VkDeviceCreateInfo devInfo = {
        .sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
        .pNext = &features2,
        .queueCreateInfoCount = 1,
        .pQueueCreateInfos = &queueInfo,
        .enabledExtensionCount = sizeof( devExt ) / sizeof( devExt[0] ),
        .ppEnabledExtensionNames = devExt
    };
    
    VkDevice device;
    vk_check(vkCreateDevice(physDev, &devInfo, NULL, &device));
    volkLoadDevice(device);
    
    VkQueue queue;
    vkGetDeviceQueue(device, gfxQueueFam, 0, &queue);
    
    VkSurfaceCapabilitiesKHR surfCaps;
    vk_check(vkGetPhysicalDeviceSurfaceCapabilitiesKHR(physDev, surface, &surfCaps));
    
    uint32_t fmtCount;
    vk_check(vkGetPhysicalDeviceSurfaceFormatsKHR(physDev, surface, &fmtCount, NULL));
    VkSurfaceFormatKHR* surfFmts = malloc(fmtCount * sizeof(VkSurfaceFormatKHR));
    vk_check(vkGetPhysicalDeviceSurfaceFormatsKHR(physDev, surface, &fmtCount, surfFmts));
    VkSurfaceFormatKHR surfFmt = surfFmts[0];
    free(surfFmts);

    VkExtent2D currentExtent = surfCaps.currentExtent;
    if (currentExtent.width == UINT32_MAX) {
        currentExtent.width = WIDTH;
        currentExtent.height = HEIGHT;
    }
    
    VkSwapchainCreateInfoKHR swapInfo = {
        .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
        .surface = surface,
        .minImageCount = surfCaps.minImageCount + 1,
        .imageFormat = surfFmt.format,
        .imageColorSpace = surfFmt.colorSpace,
        .imageExtent = currentExtent,
        .imageArrayLayers = 1,
        .imageUsage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_STORAGE_BIT,
        .imageSharingMode = VK_SHARING_MODE_EXCLUSIVE,
        .preTransform = surfCaps.currentTransform,
        .compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
        .presentMode = VK_PRESENT_MODE_FIFO_KHR,
        .clipped = VK_TRUE
    };
    
    VkSwapchainKHR swapchain;
    vk_check(vkCreateSwapchainKHR(device, &swapInfo, NULL, &swapchain));
    
    uint32_t imgCount;
    vk_check(vkGetSwapchainImagesKHR(device, swapchain, &imgCount, NULL));
    VkImage* swapImgs = malloc(imgCount * sizeof(VkImage));
    vk_check(vkGetSwapchainImagesKHR(device, swapchain, &imgCount, swapImgs));
    
    VkImage renderImg;
    VkImageView renderImgView;
    {
        // Create temporary render target image
        VkImageCreateInfo imgInfo = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
            .imageType = VK_IMAGE_TYPE_2D,
            .format = surfFmt.format,
            .extent = {currentExtent.width, currentExtent.height, 1},
            .mipLevels = 1,
            .arrayLayers = 1,
            .samples = VK_SAMPLE_COUNT_1_BIT,
            .tiling = VK_IMAGE_TILING_OPTIMAL,
            .usage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
            .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
            .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED
        };
        
        vk_check(vkCreateImage(device, &imgInfo, NULL, &renderImg));
        
        VkMemoryRequirements memReqs;
        vkGetImageMemoryRequirements(device, renderImg, &memReqs);
        
        VkMemoryAllocateInfo allocInfo = {
            .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
            .allocationSize = memReqs.size,
            .memoryTypeIndex = findMemoryType(physDev, memReqs.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
        };
        
        VkDeviceMemory renderImgMem;
        vk_check(vkAllocateMemory(device, &allocInfo, NULL, &renderImgMem));
        vk_check(vkBindImageMemory(device, renderImg, renderImgMem, 0));
        
        VkImageViewCreateInfo viewInfo = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
            .image = renderImg,
            .viewType = VK_IMAGE_VIEW_TYPE_2D,
            .format = surfFmt.format,
            .components = {VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY, 
                        VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY},
            .subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}
        };
        
        vk_check(vkCreateImageView(device, &viewInfo, NULL, &renderImgView));
    }

    VkImage cudaOutImg;
    VkImageView cudaOutImgView;
    {
        VkFormat cudaFormat = surfFmt.format;
        //cudaFormat = VK_FORMAT_B8G8R8A8_UNORM;

        // Create temporary cuda output image
        VkImageCreateInfo imgInfo = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
            .imageType = VK_IMAGE_TYPE_2D,
            .format = cudaFormat,
            .extent = {currentExtent.width, currentExtent.height, 1},
            .mipLevels = 1,
            .arrayLayers = 1,
            .samples = VK_SAMPLE_COUNT_1_BIT,
            .tiling = VK_IMAGE_TILING_OPTIMAL,
            .usage = VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT,
            .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
            .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED
        };
        
        vk_check(vkCreateImage(device, &imgInfo, NULL, &cudaOutImg));
        
        VkMemoryRequirements memReqs;
        vkGetImageMemoryRequirements(device, cudaOutImg, &memReqs);
        
        VkMemoryAllocateInfo allocInfo = {
            .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
            .allocationSize = memReqs.size,
            .memoryTypeIndex = findMemoryType(physDev, memReqs.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
        };
        
        VkDeviceMemory cudaOutImgMem;
        vk_check(vkAllocateMemory(device, &allocInfo, NULL, &cudaOutImgMem));
        vk_check(vkBindImageMemory(device, cudaOutImg, cudaOutImgMem, 0));
        
        VkImageViewCreateInfo viewInfo = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
            .image = cudaOutImg,
            .viewType = VK_IMAGE_VIEW_TYPE_2D,
            .format = cudaFormat,
            .components = {VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY, 
                        VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY},
            .subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}
        };
        
        vk_check(vkCreateImageView(device, &viewInfo, NULL, &cudaOutImgView));
    }
    
    VkAttachmentDescription colorAttach = {
        .format = surfFmt.format,
        .samples = VK_SAMPLE_COUNT_1_BIT,
        .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
        .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
        .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
        .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
        .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
        .finalLayout = VK_IMAGE_LAYOUT_GENERAL,
    };
    
    VkAttachmentReference colorAttachRef = {
        .attachment = 0,
        .layout = VK_IMAGE_LAYOUT_GENERAL
    };
    
    VkSubpassDescription subpass = {
        .pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS,
        .colorAttachmentCount = 1,
        .pColorAttachments = &colorAttachRef
    };
    
    VkSubpassDependency dep = {
        .srcSubpass = VK_SUBPASS_EXTERNAL,
        .dstSubpass = 0,
        .srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
        .srcAccessMask = 0,
        .dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
        .dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT
    };
    
    VkRenderPassCreateInfo rpInfo = {
        .sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO,
        .attachmentCount = 1,
        .pAttachments = &colorAttach,
        .subpassCount = 1,
        .pSubpasses = &subpass,
        .dependencyCount = 1,
        .pDependencies = &dep
    };
    
    VkRenderPass renderPass;
    vk_check(vkCreateRenderPass(device, &rpInfo, NULL, &renderPass));
    
    VkShaderModuleCreateInfo vertShaderInfo = {
        .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
        .codeSize = sizeof(triangle_vert),
        .pCode = (uint32_t*)triangle_vert,
    };

    VkPipelineShaderStageCreateInfo vertStage = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
        .pNext = &vertShaderInfo,
        .stage = VK_SHADER_STAGE_VERTEX_BIT,
        .module = NULL,
        .pName = "main"
    };
    
    VkShaderModuleCreateInfo fragShaderInfo = {
        .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
        .codeSize = sizeof(triangle_frag),
        .pCode = (uint32_t*)triangle_frag,
    };

    VkPipelineShaderStageCreateInfo fragStage = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
        .pNext = &fragShaderInfo,
        .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
        .module = NULL,
        .pName = "main"
    };
    
    VkPipelineShaderStageCreateInfo stages[] = {vertStage, fragStage};
    
    VkPipelineVertexInputStateCreateInfo vertInput = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
        .vertexBindingDescriptionCount = 0,
        .vertexAttributeDescriptionCount = 0
    };
    
    VkPipelineInputAssemblyStateCreateInfo inputAsm = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
        .topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
        .primitiveRestartEnable = VK_FALSE
    };
    
    VkViewport viewport = {0, 0, currentExtent.width, currentExtent.height, 0, 1};
    VkRect2D scissor = {0, 0, currentExtent.width, currentExtent.height};
    
    VkPipelineViewportStateCreateInfo viewportState = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
        .viewportCount = 1,
        .pViewports = &viewport,
        .scissorCount = 1,
        .pScissors = &scissor
    };
    
    VkPipelineRasterizationStateCreateInfo raster = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO,
        .depthClampEnable = VK_FALSE,
        .rasterizerDiscardEnable = VK_FALSE,
        .polygonMode = VK_POLYGON_MODE_FILL,
        .lineWidth = 1.0f,
        .cullMode = VK_CULL_MODE_BACK_BIT,
        .frontFace = VK_FRONT_FACE_CLOCKWISE,
        .depthBiasEnable = VK_FALSE
    };
    
    VkPipelineMultisampleStateCreateInfo multisample = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO,
        .sampleShadingEnable = VK_FALSE,
        .rasterizationSamples = VK_SAMPLE_COUNT_1_BIT
    };
    
    VkPipelineColorBlendAttachmentState colorBlendAttach = {
        .colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | 
                         VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT,
        .blendEnable = VK_FALSE
    };
    
    VkPipelineColorBlendStateCreateInfo colorBlend = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
        .logicOpEnable = VK_FALSE,
        .attachmentCount = 1,
        .pAttachments = &colorBlendAttach
    };
    
    VkPipelineLayoutCreateInfo layoutInfo = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO
    };
    
    VkPipelineLayout pipelineLayout;
    vk_check(vkCreatePipelineLayout(device, &layoutInfo, NULL, &pipelineLayout));
    
    VkGraphicsPipelineCreateInfo pipelineInfo = {
        .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
        .stageCount = 2,
        .pStages = stages,
        .pVertexInputState = &vertInput,
        .pInputAssemblyState = &inputAsm,
        .pViewportState = &viewportState,
        .pRasterizationState = &raster,
        .pMultisampleState = &multisample,
        .pColorBlendState = &colorBlend,
        .layout = pipelineLayout,
        .renderPass = renderPass,
        .subpass = 0
    };
    
    VkPipeline pipeline;
    vk_check(vkCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, &pipelineInfo, NULL, &pipeline));
    
    VkFramebufferCreateInfo fbInfo = {
        .sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO,
        .renderPass = renderPass,
        .attachmentCount = 1,
        .pAttachments = &renderImgView,
        .width = currentExtent.width,
        .height = currentExtent.height,
        .layers = 1
    };
    
    VkFramebuffer fb;
    vk_check(vkCreateFramebuffer(device, &fbInfo, NULL, &fb));
    
    VkCommandPoolCreateInfo poolInfo = {
        .sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
        .flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
        .queueFamilyIndex = gfxQueueFam
    };
    
    VkCommandPool cmdPool;
    vk_check(vkCreateCommandPool(device, &poolInfo, NULL, &cmdPool));
    
    VkCommandBufferAllocateInfo cbAllocInfo = {
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
        .commandPool = cmdPool,
        .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
        .commandBufferCount = 1
    };
    
    VkCommandBuffer cmdBuf;
    vk_check(vkAllocateCommandBuffers(device, &cbAllocInfo, &cmdBuf));
    
    VkSemaphoreCreateInfo semInfo = {.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO};
    VkFenceCreateInfo fenceInfo = {
        .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
        .flags = VK_FENCE_CREATE_SIGNALED_BIT
    };
    
    VkSemaphore imgAvail[8], renderDone[8];
    VkFence inFlight;
    for ( int i = 0; i < 8; i++ ) {
        vk_check(vkCreateSemaphore(device, &semInfo, NULL, &imgAvail[i]));
        vk_check(vkCreateSemaphore(device, &semInfo, NULL, &renderDone[i]));
    }
    vk_check(vkCreateFence(device, &fenceInfo, NULL, &inFlight));

    VkSampler sampler;
    VkSamplerCreateInfo samplerInfo = {
        .sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
        .magFilter = VK_FILTER_LINEAR,
        .minFilter = VK_FILTER_LINEAR,
        .mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR,
        .addressModeU = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeV = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .mipLodBias = 0.0f,
        .anisotropyEnable = false,
        .maxAnisotropy = 0.0f,
        .compareEnable = false,
        .minLod = -16.0f,
        .maxLod = 16.0f,
    };
    vk_check(vkCreateSampler(device, &samplerInfo, NULL, &sampler));

    // Cubin bits and bobs
    VkCuModuleNVX cuModule = VK_NULL_HANDLE;
    VkCuFunctionNVX cuFunction = VK_NULL_HANDLE;
    uint64_t cuInputImageHandle = 0;
    uint64_t cuOutputImageHandle = 0;

    VkCuModuleTexturingModeCreateInfoNVX cuModuleTexturingModeInfo = {
        .sType = VK_STRUCTURE_TYPE_CU_MODULE_TEXTURING_MODE_CREATE_INFO_NVX,
        .use64bitTexturing = VK_TRUE,
    };

    VkCuModuleCreateInfoNVX cuModuleInfo = {
        .sType = VK_STRUCTURE_TYPE_CU_MODULE_CREATE_INFO_NVX,
        .pNext = &cuModuleTexturingModeInfo,
        .dataSize = sizeof(fatbin_data),
        .pData = fatbin_data,
    };
    vk_check(vkCreateCuModuleNVX(device, &cuModuleInfo, NULL, &cuModule));
    assert(cuModule != VK_NULL_HANDLE);

    VkCuFunctionCreateInfoNVX cuFunctionInfo = {
        .sType = VK_STRUCTURE_TYPE_CU_FUNCTION_CREATE_INFO_NVX,
        .module = cuModule,
        .pName = "add_frog_to_image",
    };
    vk_check(vkCreateCuFunctionNVX(device, &cuFunctionInfo, NULL, &cuFunction));
    assert(cuFunction != VK_NULL_HANDLE);

    VkImageViewHandleInfoNVX cuInputImageHandleInfo = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_HANDLE_INFO_NVX,
        .imageView = renderImgView,
        .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
        .sampler = sampler,
    };
    cuInputImageHandle = vkGetImageViewHandle64NVX(device, &cuInputImageHandleInfo);
    assert(cuInputImageHandle != 0);

    VkImageViewHandleInfoNVX cuOutputImageHandleInfo = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_HANDLE_INFO_NVX,
        .imageView = cudaOutImgView,
        .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
    };
    cuOutputImageHandle = vkGetImageViewHandle64NVX(device, &cuOutputImageHandleInfo);
    assert(cuOutputImageHandle != 0);

    // We don't use this, I don't know how one would use this.
    VkImageViewAddressPropertiesNVX cuOutputImageAddressProperties = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_ADDRESS_PROPERTIES_NVX
    };
    vk_check(vkGetImageViewAddressNVX(device, cudaOutImgView, &cuOutputImageAddressProperties));
    
    uint32_t frame_idx = 0;
    while (!glfwWindowShouldClose(window)) {
        glfwPollEvents();
        
        vk_check(vkWaitForFences(device, 1, &inFlight, VK_TRUE, UINT64_MAX));
        vk_check(vkResetFences(device, 1, &inFlight));
        
        uint32_t imgIdx;
        vk_check(vkAcquireNextImageKHR(device, swapchain, UINT64_MAX, imgAvail[frame_idx], VK_NULL_HANDLE, &imgIdx));
        
        vk_check(vkResetCommandBuffer(cmdBuf, 0));
        
        VkCommandBufferBeginInfo beginInfo = {
            .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO
        };
        vk_check(vkBeginCommandBuffer(cmdBuf, &beginInfo));
        
        // Render to temporary image
        VkClearValue clearColor = {{{0.0f, 0.0f, 0.0f, 1.0f}}};
        VkRenderPassBeginInfo rpBeginInfo = {
            .sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
            .renderPass = renderPass,
            .framebuffer = fb,
            .renderArea = {{0, 0}, currentExtent},
            .clearValueCount = 1,
            .pClearValues = &clearColor
        };
        
        vkCmdBeginRenderPass(cmdBuf, &rpBeginInfo, VK_SUBPASS_CONTENTS_INLINE);
        vkCmdBindPipeline(cmdBuf, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
        vkCmdDraw(cmdBuf, 3, 1, 0, 0);
        vkCmdEndRenderPass(cmdBuf);

        VkImage stagingImage = renderImg;
        if (USE_CUBIN) {
            // Transition cuda image from undefined to general
            VkImageMemoryBarrier barrier1 = {
                .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
                .oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
                .newLayout = VK_IMAGE_LAYOUT_GENERAL,
                .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
                .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
                .image = cudaOutImg,
                .subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1},
                .srcAccessMask = 0,
                .dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT
            };
            
            vkCmdPipelineBarrier(cmdBuf, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, // XXX?
                                0, 0, NULL, 0, NULL, 1, &barrier1);

            // I have no idea what sync is needed here.
            VkMemoryBarrier fatBarrier = {};
            fatBarrier.srcAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT;
            fatBarrier.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT;
            // fat barrier
            vkCmdPipelineBarrier(cmdBuf, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 1, &fatBarrier, 0, NULL, 0, NULL);

            union cuda_param_t
            {
                uint32_t u32;
                uint64_t u64;
                VkDeviceAddress addr;
            } params[4];

            params[0].u64 = cuOutputImageHandle;
            params[1].u64 = cuInputImageHandle;
            params[2].u32 = WIDTH;
            params[3].u32 = HEIGHT;

            const void *const pparams[4] =
            {
                &params[0],
                &params[1],
                &params[2],
                &params[3],
            };

            VkCuLaunchInfoNVX cuLaunchInfo = {
                .sType = VK_STRUCTURE_TYPE_CU_LAUNCH_INFO_NVX,
                .function = cuFunction,
                .gridDimX = (WIDTH + 7) / 8,
                .gridDimY = (HEIGHT + 7) / 8,
                .gridDimZ = 1,
                .blockDimX = 8,
                .blockDimY = 8,
                .blockDimZ = 1,
                .sharedMemBytes = 0,
                .paramCount = sizeof( params ) / sizeof( params[0] ),
                .pParams = pparams,
                .extraCount = 0,
                .pExtras = NULL,
            };
            vkCmdCuLaunchKernelNVX(cmdBuf, &cuLaunchInfo);

            // fat barrier
            vkCmdPipelineBarrier(cmdBuf, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 1, &fatBarrier, 0, NULL, 0, NULL);

            stagingImage = cudaOutImg;
        }

        // Transition swapchain image to general
        VkImageMemoryBarrier barrier1 = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
            .oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
            .newLayout = VK_IMAGE_LAYOUT_GENERAL,
            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .image = swapImgs[imgIdx],
            .subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1},
            .srcAccessMask = 0,
            .dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT
        };
        
        vkCmdPipelineBarrier(cmdBuf, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT,
                            0, 0, NULL, 0, NULL, 1, &barrier1);
        
        // Copy render image to swapchain image
        VkImageCopy copyRegion = {
            .srcSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1},
            .srcOffset = {0, 0, 0},
            .dstSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1},
            .dstOffset = {0, 0, 0},
            .extent = {currentExtent.width, currentExtent.height, 1}
        };
        
        vkCmdCopyImage(cmdBuf, stagingImage, VK_IMAGE_LAYOUT_GENERAL,
                    swapImgs[imgIdx], VK_IMAGE_LAYOUT_GENERAL, 1, &copyRegion);
        
        // Transition swapchain image to present
        VkImageMemoryBarrier barrier2 = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
            .oldLayout = VK_IMAGE_LAYOUT_GENERAL,
            .newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .image = swapImgs[imgIdx],
            .subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1},
            .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
            .dstAccessMask = 0
        };
        
        vkCmdPipelineBarrier(cmdBuf, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT,
                                0, 0, NULL, 0, NULL, 1, &barrier2);
        
        vk_check(vkEndCommandBuffer(cmdBuf));
        
        VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
        VkSubmitInfo submitInfo = {
            .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
            .waitSemaphoreCount = 1,
            .pWaitSemaphores = &imgAvail[frame_idx],
            .pWaitDstStageMask = waitStages,
            .commandBufferCount = 1,
            .pCommandBuffers = &cmdBuf,
            .signalSemaphoreCount = 1,
            .pSignalSemaphores = &renderDone[frame_idx]
        };
        
        vk_check(vkQueueSubmit(queue, 1, &submitInfo, inFlight));
        
        VkPresentInfoKHR presentInfo = {
            .sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
            .waitSemaphoreCount = 1,
            .pWaitSemaphores = &renderDone[frame_idx],
            .swapchainCount = 1,
            .pSwapchains = &swapchain,
            .pImageIndices = &imgIdx
        };
        
        vk_check(vkQueuePresentKHR(queue, &presentInfo));

        frame_idx = (frame_idx + 1) % 8;
    }

    return 0;
}