D3D12 and Vulkan to CUDA Texture Sharing (#2038)

5 files changed, 284 insertions, 9 deletions

diff --git a/tools/gfx-unit-test/gfx-test-util.cpp b/tools/gfx-unit-test/gfx-test-util.cpp
index aafa3fc59..952638fd8 100644
--- a/tools/gfx-unit-test/gfx-test-util.cpp
+++ b/tools/gfx-unit-test/gfx-test-util.cpp
@@ -66,6 +66,20 @@ namespace gfx_test
         return SLANG_OK;
     }
 
+    void compareComputeResult(gfx::IDevice* device, gfx::ITextureResource* texture, gfx::ResourceState state, float* expectedResult, size_t expectedBufferSize)
+    {
+        // Read back the results.
+        ComPtr<ISlangBlob> resultBlob;
+        size_t rowPitch = 0;
+        size_t pixelSize = 0;
+        GFX_CHECK_CALL_ABORT(device->readTextureResource(
+            texture, state, resultBlob.writeRef(), &rowPitch, &pixelSize));
+        SLANG_CHECK(resultBlob->getBufferSize() == expectedBufferSize);
+        auto result = (float*)resultBlob->getBufferPointer();
+        // Compare results.
+        SLANG_CHECK(memcmp(resultBlob->getBufferPointer(), expectedResult, expectedBufferSize) == 0);
+    }
+
     void compareComputeResult(gfx::IDevice* device, gfx::IBufferResource* buffer, uint8_t* expectedResult, size_t expectedBufferSize)
     {
         // Read back the results.
diff --git a/tools/gfx-unit-test/gfx-test-util.h b/tools/gfx-unit-test/gfx-test-util.h
index 01a367915..07513c39a 100644
--- a/tools/gfx-unit-test/gfx-test-util.h
+++ b/tools/gfx-unit-test/gfx-test-util.h
@@ -25,6 +25,14 @@ namespace gfx_test
         uint8_t* expectedResult,
         size_t expectedBufferSize);
 
+    /// Reads back the content of `buffer` and compares it against `expectedResult`.
+    void compareComputeResult(
+        gfx::IDevice* device,
+        gfx::ITextureResource* texture,
+        gfx::ResourceState state,
+        float* expectedResult,
+        size_t expectedBufferSize);
+
         /// Reads back the content of `buffer` and compares it against `expectedResult` with a set tolerance.
     void compareComputeResultFuzzy(
         gfx::IDevice* device,
@@ -45,6 +53,22 @@ namespace gfx_test
         if (std::is_same<T, float>::value) return compareComputeResultFuzzy(device, buffer, (float*)expectedBuffer.getBuffer(), bufferSize);
         return compareComputeResult(device, buffer, expectedBuffer.getBuffer(), bufferSize);
     }
+
+//     TODO: Implement compareComputeResultFuzzy() and keep or just directly use compareComputeResult() above and add a second overload for uint/int? 
+//     template<typename T, Slang::Index count>
+//     void compareComputeResult(
+//         gfx::IDevice* device,
+//         gfx::ITextureResource* texture,
+//         gfx::ResourceState state,
+//         Slang::Array<T, count> expectedResult)
+//     {
+//         Slang::List<uint8_t> expectedBuffer;
+//         size_t bufferSize = sizeof(T) * count;
+//         expectedBuffer.setCount(bufferSize);
+//         memcpy(expectedBuffer.getBuffer(), expectedResult.begin(), bufferSize);
+//         if (std::is_same<T, float>::value) return compareComputeResultFuzzy(device, buffer, (float*)expectedBuffer.getBuffer(), bufferSize);
+//         return compareComputeResult(device, texture, state, expectedBuffer.getBuffer(), bufferSize);
+//     }
     
     Slang::ComPtr<gfx::IDevice> createTestingDevice(UnitTestContext* context, Slang::RenderApiFlag::Enum api);
 
diff --git a/tools/gfx-unit-test/get-shared-handle.cpp b/tools/gfx-unit-test/shared-buffers-tests.cpp
index 8b991f15c..fe7757083 100644
--- a/tools/gfx-unit-test/get-shared-handle.cpp
+++ b/tools/gfx-unit-test/shared-buffers-tests.cpp
@@ -9,7 +9,7 @@ using namespace gfx;
 
 namespace gfx_test
 {
-    void sharedHandleTestImpl(IDevice* srcDevice, IDevice* dstDevice, UnitTestContext* context)
+    void sharedBufferTestImpl(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.
@@ -98,7 +98,7 @@ namespace gfx_test
             Slang::makeArray<float>(1.0f, 2.0f, 3.0f, 4.0f));
     }
 
-    void sharedHandleTestAPI(UnitTestContext* context, Slang::RenderApiFlag::Enum srcApi, Slang::RenderApiFlag::Enum dstApi)
+    void sharedBufferTestAPI(UnitTestContext* context, Slang::RenderApiFlag::Enum srcApi, Slang::RenderApiFlag::Enum dstApi)
     {
         auto srcDevice = createTestingDevice(context, srcApi);
         auto dstDevice = createTestingDevice(context, dstApi);
@@ -107,19 +107,17 @@ namespace gfx_test
             SLANG_IGNORE_TEST;
         }
 
-        sharedHandleTestImpl(srcDevice, dstDevice, context);
+        sharedBufferTestImpl(srcDevice, dstDevice, context);
     }
 #if SLANG_WIN64
-    SLANG_UNIT_TEST(sharedHandleD3D12ToCUDA)
+    SLANG_UNIT_TEST(sharedBufferD3D12ToCUDA)
     {
-        sharedHandleTestAPI(unitTestContext, Slang::RenderApiFlag::D3D12, Slang::RenderApiFlag::CUDA);
+        sharedBufferTestAPI(unitTestContext, Slang::RenderApiFlag::D3D12, Slang::RenderApiFlag::CUDA);
     }
 
-#if SLANG_WINDOWS_FAMILY // TODO: Remove when Linux support is added
-    SLANG_UNIT_TEST(sharedHandleVulkanToCUDA)
+    SLANG_UNIT_TEST(sharedBufferVulkanToCUDA)
     {
-        sharedHandleTestAPI(unitTestContext, Slang::RenderApiFlag::Vulkan, Slang::RenderApiFlag::CUDA);
+        sharedBufferTestAPI(unitTestContext, Slang::RenderApiFlag::Vulkan, Slang::RenderApiFlag::CUDA);
     }
 #endif
-#endif
 }
diff --git a/tools/gfx-unit-test/shared-textures-tests.cpp b/tools/gfx-unit-test/shared-textures-tests.cpp
new file mode 100644
index 000000000..896d6d6e6
--- /dev/null
+++ b/tools/gfx-unit-test/shared-textures-tests.cpp
@@ -0,0 +1,221 @@
+#include "tools/unit-test/slang-unit-test.h"
+
+#include "slang-gfx.h"
+#include "gfx-test-util.h"
+#include "tools/gfx-util/shader-cursor.h"
+#include "source/core/slang-basic.h"
+
+using namespace gfx;
+
+namespace gfx_test
+{
+    void setUpAndRunShader(
+        IDevice* device,
+        ComPtr<ITextureResource> tex,
+        ComPtr<IResourceView> texView,
+        ComPtr<IResourceView> bufferView,
+        const char* entryPoint,
+        ComPtr<ISamplerState> sampler = nullptr)
+    {
+        Slang::ComPtr<ITransientResourceHeap> transientHeap;
+        ITransientResourceHeap::Desc transientHeapDesc = {};
+        transientHeapDesc.constantBufferSize = 4096;
+        GFX_CHECK_CALL_ABORT(
+            device->createTransientResourceHeap(transientHeapDesc, transientHeap.writeRef()));
+
+        ComPtr<IShaderProgram> shaderProgram;
+        slang::ProgramLayout* slangReflection;
+        GFX_CHECK_CALL_ABORT(loadComputeProgram(device, shaderProgram, "trivial-copy", entryPoint, slangReflection));
+
+        ComputePipelineStateDesc pipelineDesc = {};
+        pipelineDesc.program = shaderProgram.get();
+        ComPtr<gfx::IPipelineState> pipelineState;
+        GFX_CHECK_CALL_ABORT(
+            device->createComputePipelineState(pipelineDesc, pipelineState.writeRef()));
+
+        // We have done all the set up work, now it is time to start recording a command buffer for
+        // GPU execution.
+        {
+            ICommandQueue::Desc queueDesc = { ICommandQueue::QueueType::Graphics };
+            auto queue = device->createCommandQueue(queueDesc);
+
+            auto commandBuffer = transientHeap->createCommandBuffer();
+            auto encoder = commandBuffer->encodeComputeCommands();
+
+            auto rootObject = encoder->bindPipeline(pipelineState);
+
+            ShaderCursor entryPointCursor(
+                rootObject->getEntryPoint(0)); // get a cursor the the first entry-point.
+
+            auto& desc = *tex->getDesc();
+            entryPointCursor["width"].setData(desc.size.width);
+            entryPointCursor["height"].setData(desc.size.height);
+
+            // Bind texture view to the entry point
+            entryPointCursor["tex"].setResource(texView);
+
+            if (sampler) entryPointCursor["sampler"].setSampler(sampler);
+
+            // Bind buffer view to the entry point.
+            entryPointCursor["buffer"].setResource(bufferView);
+
+            encoder->dispatchCompute(1, 1, 1);
+            encoder->endEncoding();
+            commandBuffer->close();
+            queue->executeCommandBuffer(commandBuffer);
+            queue->waitOnHost();
+        }
+    }
+
+    ComPtr<ITextureResource> createTexture(IDevice* device, ITextureResource::Size extents, gfx::Format format, ITextureResource::SubresourceData* initialData)
+    {
+        ITextureResource::Desc texDesc = {};
+        texDesc.type = IResource::Type::Texture2D;
+        texDesc.numMipLevels = 1;
+        texDesc.arraySize = 1;
+        texDesc.size = extents;
+        texDesc.defaultState = ResourceState::UnorderedAccess;
+        texDesc.allowedStates = ResourceStateSet(
+            ResourceState::ShaderResource,
+            ResourceState::UnorderedAccess,
+            ResourceState::CopyDestination,
+            ResourceState::CopySource);
+        texDesc.format = format;
+        texDesc.isShared = true;
+
+        ComPtr<ITextureResource> inTex;
+        GFX_CHECK_CALL_ABORT(device->createTextureResource(
+            texDesc,
+            initialData,
+            inTex.writeRef()));
+        return inTex;
+    }
+
+    ComPtr<IResourceView> createTexView(IDevice* device, ComPtr<ITextureResource> inTexture)
+    {
+        ComPtr<IResourceView> texView;
+        IResourceView::Desc texViewDesc = {};
+        texViewDesc.type = IResourceView::Type::UnorderedAccess;
+        texViewDesc.format = inTexture->getDesc()->format; // TODO: Handle typeless formats - gfxIsTypelessFormat(format) ? convertTypelessFormat(format) : format;
+        GFX_CHECK_CALL_ABORT(device->createTextureView(inTexture, texViewDesc, texView.writeRef()));
+        return texView;
+    }
+
+    template <typename T>
+    ComPtr<IBufferResource> createBuffer(IDevice* device, int size, void* initialData)
+    {
+        IBufferResource::Desc bufferDesc = {};
+        bufferDesc.sizeInBytes = size * sizeof(T);
+        bufferDesc.format = gfx::Format::Unknown;
+        bufferDesc.elementSize = sizeof(T);
+        bufferDesc.allowedStates = ResourceStateSet(
+            ResourceState::ShaderResource,
+            ResourceState::UnorderedAccess,
+            ResourceState::CopyDestination,
+            ResourceState::CopySource);
+        bufferDesc.defaultState = ResourceState::UnorderedAccess;
+        bufferDesc.cpuAccessFlags = AccessFlag::Write | AccessFlag::Read;
+
+        ComPtr<IBufferResource> outBuffer;
+        GFX_CHECK_CALL_ABORT(device->createBufferResource(
+            bufferDesc,
+            initialData,
+            outBuffer.writeRef()));
+        return outBuffer;
+    }
+
+    ComPtr<IResourceView> createOutBufferView(IDevice* device, ComPtr<IBufferResource> outBuffer)
+    {
+        ComPtr<IResourceView> bufferView;
+        IResourceView::Desc viewDesc = {};
+        viewDesc.type = IResourceView::Type::UnorderedAccess;
+        viewDesc.format = Format::Unknown;
+        GFX_CHECK_CALL_ABORT(device->createBufferView(outBuffer, viewDesc, bufferView.writeRef()));
+        return bufferView;
+    }
+
+    void sharedTextureTestImpl(IDevice* srcDevice, IDevice* dstDevice, UnitTestContext* context)
+    {
+        ISamplerState::Desc samplerDesc;
+        auto sampler = dstDevice->createSamplerState(samplerDesc);
+
+        float initFloatData[16] = { 0.0f };
+        auto floatResults = createBuffer<float>(dstDevice, 16, initFloatData);
+        auto floatBufferView = createOutBufferView(dstDevice, floatResults);
+
+        uint32_t initUintData[16] = { 0u };
+        auto uintResults = createBuffer<uint32_t>(dstDevice, 16, initUintData);
+        auto uintBufferView = createOutBufferView(dstDevice, uintResults);
+
+        int32_t initIntData[16] = { 0 };
+        auto intResults = createBuffer<uint32_t>(dstDevice, 16, initIntData);
+        auto intBufferView = createOutBufferView(dstDevice, intResults);
+
+        ITextureResource::Size size = {};
+        size.width = 2;
+        size.height = 2;
+        size.depth = 1;
+
+        ITextureResource::Size bcSize = {};
+        bcSize.width = 4;
+        bcSize.height = 4;
+        bcSize.depth = 1;
+
+        {
+            float texData[] = { 1.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f, 0.0f, 1.0f,
+                                0.0f, 0.0f, 1.0f, 1.0f, 0.5f, 0.5f, 0.5f, 1.0f };
+            ITextureResource::SubresourceData subData = { (void*)texData, 32, 0 };
+
+            // Create a shareable texture using srcDevice, get its handle, then create a texture using the handle using
+            // dstDevice. Read back the texture and check that its contents are correct.
+            auto srcTexture = createTexture(srcDevice, size, gfx::Format::R32G32B32A32_FLOAT, &subData);
+
+            InteropHandle sharedHandle;
+            GFX_CHECK_CALL_ABORT(srcTexture->getSharedHandle(&sharedHandle));
+            ComPtr<ITextureResource> dstTexture;
+            size_t sizeInBytes = 0;
+            size_t alignment = 0;
+            GFX_CHECK_CALL_ABORT(srcDevice->getTextureAllocationInfo(*(srcTexture->getDesc()), &sizeInBytes, &alignment));
+            GFX_CHECK_CALL_ABORT(dstDevice->createTextureFromSharedHandle(sharedHandle, *(srcTexture->getDesc()), sizeInBytes, dstTexture.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(
+                dstDevice,
+                dstTexture,
+                ResourceState::ShaderResource,
+                texData,
+                sizeof(texData));
+
+            auto texView = createTexView(dstDevice, dstTexture);
+            setUpAndRunShader(dstDevice, dstTexture, texView, floatBufferView, "copyTexFloat4");
+            compareComputeResult(
+                dstDevice,
+                floatResults,
+                Slang::makeArray<float>(1.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f, 0.0f, 1.0f,
+                                        0.0f, 0.0f, 1.0f, 1.0f, 0.5f, 0.5f, 0.5f, 1.0f));
+        }
+    }
+
+    void sharedTextureTestAPI(UnitTestContext* context, Slang::RenderApiFlag::Enum srcApi, Slang::RenderApiFlag::Enum dstApi)
+    {
+        auto srcDevice = createTestingDevice(context, srcApi);
+        auto dstDevice = createTestingDevice(context, dstApi);
+        if (!srcDevice || !dstDevice)
+        {
+            SLANG_IGNORE_TEST;
+        }
+
+        sharedTextureTestImpl(srcDevice, dstDevice, context);
+    }
+#if SLANG_WIN64
+    SLANG_UNIT_TEST(sharedTextureD3D12ToCUDA)
+    {
+        sharedTextureTestAPI(unitTestContext, Slang::RenderApiFlag::D3D12, Slang::RenderApiFlag::CUDA);
+    }
+
+    SLANG_UNIT_TEST(sharedTextureVulkanToCUDA)
+    {
+        sharedTextureTestAPI(unitTestContext, Slang::RenderApiFlag::Vulkan, Slang::RenderApiFlag::CUDA);
+    }
+#endif
+}
diff --git a/tools/gfx-unit-test/trivial-copy.slang b/tools/gfx-unit-test/trivial-copy.slang
new file mode 100644
index 000000000..acf091ffa
--- /dev/null
+++ b/tools/gfx-unit-test/trivial-copy.slang
@@ -0,0 +1,18 @@
+// trivial-copy.slang
+
+// Copy the contents of "tex" into "buffer". These are for textures containing FLOAT data.
+[shader("compute")]
+[numthreads(4,1,1)]
+void copyTexFloat4(
+    uint3 sv_dispatchThreadID : SV_DispatchThreadID,
+    uniform uint width,
+    uniform uint height,
+    uniform RWTexture2D<float4> tex,
+    uniform RWStructuredBuffer<float> buffer)
+{
+    float4 result = tex[uint2(sv_dispatchThreadID.x % width, sv_dispatchThreadID.x / width)];
+    buffer[sv_dispatchThreadID.x * 4] = result.r;
+    buffer[sv_dispatchThreadID.x * 4 + 1] = result.g;
+    buffer[sv_dispatchThreadID.x * 4 + 2] = result.b;
+    buffer[sv_dispatchThreadID.x * 4 + 3] = result.a;
+}