Add support for buffer sharing from Vulkan to CUDA (#2008)

* Added getSharedHandle() and additional code to handle shareable buffer creation to Buffer::init() and initVulkanInstanceAndDevice() for Vulkan; Modified createBufferFromSharedHandle() in CUDA to assign externalMemoryHandleDesc.type based on the type of handle being provided; Added an additional test case to get-shared-handle.cpp testing Vulkan to CUDA

* Added createBufferFromNativeHandle() to Vulkan and enabled corresponding test

* disable cuda

* Fixed getSharedHandle() for D3D12 buffers assigning Win32 as the handle's source

* Removed a dangling comment inside Buffer::init()

* Added a missing override; Added code to check that a physical device supports the necessary external memory extensions before adding them to the deviceExtensions list; Added #if SLANG_WINDOWS_FAMILY guards around all Windows-specific code and sharedHandleVulkanToCUDA test (which uses said platform-specific code)

* Added Windows check around vkGetMemoryWin32HandleKHR in vk-api.h

* Added missing Windows check around BufferResourceImpl destructor

* Added a temporary hack to ensure synchronization between devices, which solves an issue with buffer sharing resulting in incorrect values being read back; Added #if SLANG_WIN64 around all CUDA tests as the backend currently only supports running CUDA on 64-bit (despite devices being created successfully in a 32-bit config)

5 files changed, 163 insertions, 12 deletions

diff --git a/tools/gfx/cuda/render-cuda.cpp b/tools/gfx/cuda/render-cuda.cpp
index 2dc1bc1e1..9f52c4c62 100644
--- a/tools/gfx/cuda/render-cuda.cpp
+++ b/tools/gfx/cuda/render-cuda.cpp
@@ -1842,8 +1842,17 @@ public:
         // memory association, we first need to fill in a descriptor struct.
         cudaExternalMemoryHandleDesc externalMemoryHandleDesc;
         memset(&externalMemoryHandleDesc, 0, sizeof(externalMemoryHandleDesc));
-        // TODO: Change according to the type of handle being passed in
-        externalMemoryHandleDesc.type = cudaExternalMemoryHandleTypeD3D12Resource;
+        switch (handle.api)
+        {
+        case InteropHandleAPI::D3D12:
+            externalMemoryHandleDesc.type = cudaExternalMemoryHandleTypeD3D12Resource;
+            break;
+        case InteropHandleAPI::Vulkan:
+            externalMemoryHandleDesc.type = cudaExternalMemoryHandleTypeOpaqueWin32;
+            break;
+        default:
+            return SLANG_FAIL;
+        }
         externalMemoryHandleDesc.handle.win32.handle = (void*)handle.handleValue;
         externalMemoryHandleDesc.size = desc.sizeInBytes;
         externalMemoryHandleDesc.flags = cudaExternalMemoryDedicated;
diff --git a/tools/gfx/d3d12/render-d3d12.cpp b/tools/gfx/d3d12/render-d3d12.cpp
index d14d72b84..accd9297c 100644
--- a/tools/gfx/d3d12/render-d3d12.cpp
+++ b/tools/gfx/d3d12/render-d3d12.cpp
@@ -263,7 +263,7 @@ public:
             auto pResource = m_resource.getResource();
             pResource->GetDevice(IID_PPV_ARGS(pDevice.writeRef()));
             SLANG_RETURN_ON_FAIL(pDevice->CreateSharedHandle(pResource, NULL, GENERIC_ALL, nullptr, (HANDLE*)&outHandle->handleValue));
-            outHandle->api = InteropHandleAPI::Win32;
+            outHandle->api = InteropHandleAPI::D3D12;
             sharedHandle = *outHandle;
             return SLANG_OK;
         }
diff --git a/tools/gfx/immediate-renderer-base.cpp b/tools/gfx/immediate-renderer-base.cpp
index 33b545f26..774162eb3 100644
--- a/tools/gfx/immediate-renderer-base.cpp
+++ b/tools/gfx/immediate-renderer-base.cpp
@@ -369,7 +369,7 @@ public:
     virtual SLANG_NO_THROW void SLANG_MCALL close() override { }
 
     virtual SLANG_NO_THROW Result SLANG_MCALL
-        getNativeHandle(NativeHandle* outHandle)
+        getNativeHandle(NativeHandle* outHandle) override
     {
         *outHandle = 0;
         return SLANG_OK;
diff --git a/tools/gfx/vulkan/render-vk.cpp b/tools/gfx/vulkan/render-vk.cpp
index b68812af6..a3f5e078b 100644
--- a/tools/gfx/vulkan/render-vk.cpp
+++ b/tools/gfx/vulkan/render-vk.cpp
@@ -40,6 +40,10 @@
 #undef None
 #endif
 
+#if SLANG_WINDOWS_FAMILY
+#include <dxgi1_2.h>
+#endif
+
 namespace gfx {
 using namespace Slang;
 
@@ -78,6 +82,10 @@ public:
         const IBufferResource::Desc& desc,
         const void* initData,
         IBufferResource** outResource) override;
+    virtual SLANG_NO_THROW Result SLANG_MCALL createBufferFromNativeHandle(
+        InteropHandle handle,
+        const IBufferResource::Desc& srcDesc,
+        IBufferResource** outResource) override;
     virtual SLANG_NO_THROW Result SLANG_MCALL
         createSamplerState(ISamplerState::Desc const& desc, ISamplerState** outSampler) override;
 
@@ -149,7 +157,13 @@ public:
     {
         public:
             /// Initialize a buffer with specified size, and memory props
-        Result init(const VulkanApi& api, size_t bufferSize, VkBufferUsageFlags usage, VkMemoryPropertyFlags reqMemoryProperties);
+        Result init(
+            const VulkanApi& api,
+            size_t bufferSize,
+            VkBufferUsageFlags usage,
+            VkMemoryPropertyFlags reqMemoryProperties,
+            bool isShared = false,
+            VkExternalMemoryHandleTypeFlagsKHR extMemHandleType = 0);
 
             /// Returns true if has been initialized
         bool isInitialized() const { return m_api != nullptr; }
@@ -193,6 +207,16 @@ public:
             assert(renderer);
         }
 
+        ~BufferResourceImpl()
+        {
+            if (sharedHandle.handleValue != 0)
+            {
+#if SLANG_WINDOWS_FAMILY
+                CloseHandle((HANDLE)sharedHandle.handleValue);
+#endif
+            }
+        }
+
         RefPtr<VKDevice> m_renderer;
         Buffer m_buffer;
         Buffer m_uploadBuffer;
@@ -217,8 +241,31 @@ public:
 
         virtual SLANG_NO_THROW Result SLANG_MCALL getSharedHandle(InteropHandle* outHandle) override
         {
+            // Check if a shared handle already exists for this resource.
+            if (sharedHandle.handleValue != 0)
+            {
+                *outHandle = sharedHandle;
+                return SLANG_OK;
+            }
+            
+            // If a shared handle doesn't exist, create one and store it.
+#if SLANG_WINDOWS_FAMILY
+            VkMemoryGetWin32HandleInfoKHR info = {};
+            info.sType = VK_STRUCTURE_TYPE_MEMORY_GET_WIN32_HANDLE_INFO_KHR;
+            info.pNext = nullptr;
+            info.memory = m_buffer.m_memory;
+            info.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT;
+
+            auto api = m_buffer.m_api;
+            PFN_vkGetMemoryWin32HandleKHR vkCreateSharedHandle;
+            vkCreateSharedHandle = api->vkGetMemoryWin32HandleKHR;
+            if (!vkCreateSharedHandle)
+            {
+                return SLANG_FAIL;
+            }
+            SLANG_RETURN_ON_FAIL(vkCreateSharedHandle(api->m_device, &info, (HANDLE*)&outHandle->handleValue) != VK_SUCCESS);
+#endif
             outHandle->api = InteropHandleAPI::Vulkan;
-            outHandle->handleValue = 0;
             return SLANG_OK;
         }
     };
@@ -5368,7 +5415,13 @@ Result VKDevice::PipelineCommandEncoder::bindRootShaderObjectImpl(
 
 /* !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! VKDevice::Buffer !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! */
 
-Result VKDevice::Buffer::init(const VulkanApi& api, size_t bufferSize, VkBufferUsageFlags usage, VkMemoryPropertyFlags reqMemoryProperties)
+Result VKDevice::Buffer::init(
+    const VulkanApi& api,
+    size_t bufferSize,
+    VkBufferUsageFlags usage,
+    VkMemoryPropertyFlags reqMemoryProperties,
+    bool isShared,
+    VkExternalMemoryHandleTypeFlagsKHR extMemHandleType)
 {
     assert(!isInitialized());
 
@@ -5379,8 +5432,17 @@ Result VKDevice::Buffer::init(const VulkanApi& api, size_t bufferSize, VkBufferU
     VkBufferCreateInfo bufferCreateInfo = { VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO };
     bufferCreateInfo.size = bufferSize;
     bufferCreateInfo.usage = usage;
-    SLANG_VK_CHECK(api.vkCreateBuffer(api.m_device, &bufferCreateInfo, nullptr, &m_buffer));
+    bufferCreateInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
+
+    VkExternalMemoryBufferCreateInfo externalMemoryBufferCreateInfo = { VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_BUFFER_CREATE_INFO };
+    if (isShared)
+    {
+        externalMemoryBufferCreateInfo.handleTypes = extMemHandleType;
+        bufferCreateInfo.pNext = &externalMemoryBufferCreateInfo;
+    }
 
+    SLANG_VK_CHECK(api.vkCreateBuffer(api.m_device, &bufferCreateInfo, nullptr, &m_buffer));
+    
     VkMemoryRequirements memoryReqs = {};
     api.vkGetBufferMemoryRequirements(api.m_device, m_buffer, &memoryReqs);
 
@@ -5388,15 +5450,34 @@ Result VKDevice::Buffer::init(const VulkanApi& api, size_t bufferSize, VkBufferU
     assert(memoryTypeIndex >= 0);
 
     VkMemoryPropertyFlags actualMemoryProperites = api.m_deviceMemoryProperties.memoryTypes[memoryTypeIndex].propertyFlags;
-
     VkMemoryAllocateInfo allocateInfo = { VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO };
     allocateInfo.allocationSize = memoryReqs.size;
     allocateInfo.memoryTypeIndex = memoryTypeIndex;
+#if SLANG_WINDOWS_FAMILY
+    VkExportMemoryWin32HandleInfoKHR exportMemoryWin32HandleInfo = { VK_STRUCTURE_TYPE_EXPORT_MEMORY_WIN32_HANDLE_INFO_KHR };
+    VkExportMemoryAllocateInfoKHR exportMemoryAllocateInfo = { VK_STRUCTURE_TYPE_EXPORT_MEMORY_ALLOCATE_INFO_KHR };
+    if (isShared)
+    {
+        exportMemoryWin32HandleInfo.pNext = nullptr;
+        exportMemoryWin32HandleInfo.pAttributes = nullptr;
+        exportMemoryWin32HandleInfo.dwAccess = DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE;
+        exportMemoryWin32HandleInfo.name = NULL;
+
+        exportMemoryAllocateInfo.pNext =
+            extMemHandleType & VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT_KHR
+            ? &exportMemoryWin32HandleInfo
+            : NULL;
+        exportMemoryAllocateInfo.handleTypes = extMemHandleType;
+        allocateInfo.pNext = &exportMemoryAllocateInfo;
+    }
+#endif
     VkMemoryAllocateFlagsInfo flagInfo = {VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO};
     if (usage & VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT)
     {
         flagInfo.deviceMask = 1;
         flagInfo.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT;
+
+        flagInfo.pNext = allocateInfo.pNext;
         allocateInfo.pNext = &flagInfo;
     }
  
@@ -5546,9 +5627,10 @@ Result VKDevice::initVulkanInstanceAndDevice(const InteropHandle* handles, bool
         applicationInfo.engineVersion = 1;
         applicationInfo.applicationVersion = 1;
 
-        Array<const char*, 4> instanceExtensions;
+        Array<const char*, 5> instanceExtensions;
 
         instanceExtensions.add(VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME);
+        instanceExtensions.add(VK_KHR_EXTERNAL_MEMORY_CAPABILITIES_EXTENSION_NAME);
 
         // Software (swiftshader) implementation currently does not support surface extension,
         // so only use it with a hardware implementation.
@@ -5789,7 +5871,7 @@ Result VKDevice::initVulkanInstanceAndDevice(const InteropHandle* handles, bool
         deviceFeatures2.pNext = &extendedFeatures.atomicFloatFeatures;
 
         m_api.vkGetPhysicalDeviceFeatures2(m_api.m_physicalDevice, &deviceFeatures2);
-
+        
         if (deviceFeatures2.features.shaderResourceMinLod)
         {
             m_features.add("shader-resource-min-lod");
@@ -5912,6 +5994,30 @@ Result VKDevice::initVulkanInstanceAndDevice(const InteropHandle* handles, bool
             deviceExtensions.add(VK_KHR_BUFFER_DEVICE_ADDRESS_EXTENSION_NAME);
             m_features.add("inline-uniform-block");
         }
+
+        uint32_t extensionCount = 0;
+        m_api.vkEnumerateDeviceExtensionProperties(m_api.m_physicalDevice, NULL, &extensionCount, NULL);
+        Slang::List<VkExtensionProperties> extensions;
+        extensions.setCount(extensionCount);
+        m_api.vkEnumerateDeviceExtensionProperties(m_api.m_physicalDevice, NULL, &extensionCount, extensions.getBuffer());
+
+        HashSet<String> extensionNames;
+        for (const auto& e : extensions)
+        {
+            extensionNames.Add(e.extensionName);
+        }
+
+        if (extensionNames.Contains("VK_KHR_external_memory"))
+        {
+            deviceExtensions.add(VK_KHR_EXTERNAL_MEMORY_EXTENSION_NAME);
+#if SLANG_WINDOWS_FAMILY
+            if (extensionNames.Contains("VK_KHR_external_memory_win32"))
+            {
+                deviceExtensions.add(VK_KHR_EXTERNAL_MEMORY_WIN32_EXTENSION_NAME);
+            }
+#endif
+            m_features.add("external-memory");
+        }
     }
     if (m_api.m_module->isSoftware())
     {
@@ -6770,7 +6876,14 @@ Result VKDevice::createBufferResource(const IBufferResource::Desc& descIn, const
     }
 
     RefPtr<BufferResourceImpl> buffer(new BufferResourceImpl(desc, this));
-    SLANG_RETURN_ON_FAIL(buffer->m_buffer.init(m_api, desc.sizeInBytes, usage, reqMemoryProperties));
+    if (desc.isShared)
+    {
+        SLANG_RETURN_ON_FAIL(buffer->m_buffer.init(m_api, desc.sizeInBytes, usage, reqMemoryProperties, desc.isShared, VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT));
+    }
+    else
+    {
+        SLANG_RETURN_ON_FAIL(buffer->m_buffer.init(m_api, desc.sizeInBytes, usage, reqMemoryProperties));
+    }
 
     if ((desc.cpuAccessFlags & AccessFlag::Write) || initData)
     {
@@ -6801,6 +6914,23 @@ Result VKDevice::createBufferResource(const IBufferResource::Desc& descIn, const
     return SLANG_OK;
 }
 
+Result VKDevice::createBufferFromNativeHandle(InteropHandle handle, const IBufferResource::Desc& srcDesc, IBufferResource** outResource)
+{
+    RefPtr<BufferResourceImpl> buffer(new BufferResourceImpl(srcDesc, this));
+
+    if (handle.api == InteropHandleAPI::Vulkan)
+    {
+        buffer->m_buffer.m_buffer = (VkBuffer)handle.handleValue;
+    }
+    else
+    {
+        return SLANG_FAIL;
+    }
+
+    returnComPtr(outResource, buffer);
+    return SLANG_OK;
+}
+
 VkFilter translateFilterMode(TextureFilteringMode mode)
 {
     switch (mode)
diff --git a/tools/gfx/vulkan/vk-api.h b/tools/gfx/vulkan/vk-api.h
index 0bb3339fb..543ad90d1 100644
--- a/tools/gfx/vulkan/vk-api.h
+++ b/tools/gfx/vulkan/vk-api.h
@@ -9,6 +9,7 @@ namespace gfx {
     x(vkGetInstanceProcAddr) \
     x(vkCreateInstance) \
     x(vkEnumerateInstanceLayerProperties) \
+    x(vkEnumerateDeviceExtensionProperties) \
     x(vkDestroyInstance) \
     /* */
 
@@ -154,7 +155,18 @@ namespace gfx {
     x(vkAcquireNextImageKHR) \
     /* */
 
+#if SLANG_WINDOWS_FAMILY
+#   define VK_API_DEVICE_PLATFORM_OPT_PROCS(x) \
+    x(vkGetMemoryWin32HandleKHR) \
+    /* */
+#else
+#   define VK_API_DEVICE_PLATFORM_OPT_PROCS(x) \
+    x(vkGetMemoryFdKHR) \
+    /* */
+#endif
+
 #define VK_API_DEVICE_OPT_PROCS(x) \
+    VK_API_DEVICE_PLATFORM_OPT_PROCS(x) \
     x(vkCmdSetPrimitiveTopologyEXT) \
     x(vkGetBufferDeviceAddress) \
     x(vkGetBufferDeviceAddressKHR) \