1 files changed, 61 insertions, 2 deletions

diff --git a/tools/gfx-unit-test/get-shared-handle.cpp b/tools/gfx-unit-test/get-shared-handle.cpp
index 7c091abc8..06114f139 100644
--- a/tools/gfx-unit-test/get-shared-handle.cpp
+++ b/tools/gfx-unit-test/get-shared-handle.cpp
@@ -11,6 +11,8 @@ namespace gfx_test
 {
     void sharedHandleTestImpl(IDevice* srcDevice, IDevice* dstDevice, UnitTestContext* context)
     {
+        // Create a shareable buffer using srcDevice, get its handle, then create a buffer using the handle using
+        // dstDevice. Read back the buffer and check that its contents are correct.
         const int numberCount = 4;
         float initialData[] = { 0.0f, 1.0f, 2.0f, 3.0f };
         IBufferResource::Desc bufferDesc = {};
@@ -36,6 +38,9 @@ namespace gfx_test
         GFX_CHECK_CALL_ABORT(srcBuffer->getSharedHandle(&sharedHandle));
         ComPtr<IBufferResource> dstBuffer;
         GFX_CHECK_CALL_ABORT(dstDevice->createBufferFromSharedHandle(sharedHandle, bufferDesc, dstBuffer.writeRef()));
+        // Reading back the buffer from srcDevice to make sure it's been filled in before reading anything back from dstDevice
+        // TODO: Implement actual synchronization (and not this hacky solution)
+        compareComputeResult(srcDevice, srcBuffer, Slang::makeArray<float>(0.0f, 1.0f, 2.0f, 3.0f));
 
         InteropHandle testHandle;
         GFX_CHECK_CALL_ABORT(dstBuffer->getNativeResourceHandle(&testHandle));
@@ -43,6 +48,54 @@ namespace gfx_test
         SLANG_CHECK(testDesc->elementSize == sizeof(float));
         SLANG_CHECK(testDesc->sizeInBytes == numberCount * sizeof(float));
         compareComputeResult(dstDevice, dstBuffer, Slang::makeArray<float>(0.0f, 1.0f, 2.0f, 3.0f));
+
+        // Check that dstBuffer can be successfully used in a compute dispatch using dstDevice.
+        Slang::ComPtr<ITransientResourceHeap> transientHeap;
+        ITransientResourceHeap::Desc transientHeapDesc = {};
+        transientHeapDesc.constantBufferSize = 4096;
+        GFX_CHECK_CALL_ABORT(
+            dstDevice->createTransientResourceHeap(transientHeapDesc, transientHeap.writeRef()));
+
+        ComPtr<IShaderProgram> shaderProgram;
+        slang::ProgramLayout* slangReflection;
+        GFX_CHECK_CALL_ABORT(loadComputeProgram(dstDevice, shaderProgram, "compute-trivial", "computeMain", slangReflection));
+
+        ComputePipelineStateDesc pipelineDesc = {};
+        pipelineDesc.program = shaderProgram.get();
+        ComPtr<gfx::IPipelineState> pipelineState;
+        GFX_CHECK_CALL_ABORT(
+            dstDevice->createComputePipelineState(pipelineDesc, pipelineState.writeRef()));
+
+        ComPtr<IResourceView> bufferView;
+        IResourceView::Desc viewDesc = {};
+        viewDesc.type = IResourceView::Type::UnorderedAccess;
+        viewDesc.format = Format::Unknown;
+        GFX_CHECK_CALL_ABORT(dstDevice->createBufferView(dstBuffer, viewDesc, bufferView.writeRef()));
+
+        {
+            ICommandQueue::Desc queueDesc = { ICommandQueue::QueueType::Graphics };
+            auto queue = dstDevice->createCommandQueue(queueDesc);
+
+            auto commandBuffer = transientHeap->createCommandBuffer();
+            auto encoder = commandBuffer->encodeComputeCommands();
+
+            auto rootObject = encoder->bindPipeline(pipelineState);
+
+            ShaderCursor rootCursor(rootObject);
+            // Bind buffer view to the entry point.
+            rootCursor.getPath("buffer").setResource(bufferView);
+
+            encoder->dispatchCompute(1, 1, 1);
+            encoder->endEncoding();
+            commandBuffer->close();
+            queue->executeCommandBuffer(commandBuffer);
+            queue->wait();
+        }
+
+        compareComputeResult(
+            dstDevice,
+            dstBuffer,
+            Slang::makeArray<float>(1.0f, 2.0f, 3.0f, 4.0f));
     }
 
     void sharedHandleTestAPI(UnitTestContext* context, Slang::RenderApiFlag::Enum srcApi, Slang::RenderApiFlag::Enum dstApi)
@@ -56,11 +109,17 @@ namespace gfx_test
 
         sharedHandleTestImpl(srcDevice, dstDevice, context);
     }
-#if 0
-    // Temporarily disabled due to inconsistent test results on TC
+#if SLANG_WIN64
     SLANG_UNIT_TEST(sharedHandleD3D12ToCUDA)
     {
         sharedHandleTestAPI(unitTestContext, Slang::RenderApiFlag::D3D12, Slang::RenderApiFlag::CUDA);
     }
+
+#if SLANG_WINDOWS_FAMILY // TODO: Remove when Linux support is added
+    SLANG_UNIT_TEST(sharedHandleVulkanToCUDA)
+    {
+        sharedHandleTestAPI(unitTestContext, Slang::RenderApiFlag::Vulkan, Slang::RenderApiFlag::CUDA);
+    }
+#endif
 #endif
 }