GFX: implement mutable shader objects. (#1963)

* GFX: implement mutable shader objects.

* Revert unnecessary changes

* Revert more changes.

* Fix clang errors.

* Fix clang/gcc errors.

* Fix clang errors.

* Remove CPU test.

* Fix after merge.

* Fix after merge.

* Remove gl test

* Code review fixes.

* Fixing all vk validation errors.

* Flush test output more often.

* Fix a crash in `specializeDynamicAssociatedTypeLookup`.

* temporarily disable std-lib-serialize test to see what happens

* Fix crashes.

* Make sure cpu gfx unit tests are properly disabled on TeamCity.

* Disable cpu test.

* Fix.

* Fix cuda.

* Disable nv-ray-tracing-motion-blur

Co-authored-by: Yong He <yhe@nvidia.com>

5 files changed, 312 insertions, 50 deletions

diff --git a/tools/gfx-unit-test/compute-smoke.cpp b/tools/gfx-unit-test/compute-smoke.cpp
index 9038cd5b6..9a8298bf8 100644
--- a/tools/gfx-unit-test/compute-smoke.cpp
+++ b/tools/gfx-unit-test/compute-smoke.cpp
@@ -96,49 +96,14 @@ namespace gfx_test
             Slang::makeArray<float>(11.0f, 12.0f, 13.0f, 14.0f));
     }
 
-    void computeSmokeTestAPI(UnitTestContext* context, Slang::RenderApiFlag::Enum api)
-    {
-        if ((api & context->enabledApis) == 0)
-        {
-            SLANG_IGNORE_TEST
-        }
-        Slang::ComPtr<IDevice> device;
-        IDevice::Desc deviceDesc = {};
-        switch (api)
-        {
-        case Slang::RenderApiFlag::D3D11:
-            deviceDesc.deviceType = gfx::DeviceType::DirectX11;
-            break;
-        case Slang::RenderApiFlag::D3D12:
-            deviceDesc.deviceType = gfx::DeviceType::DirectX12;
-            break;
-        case Slang::RenderApiFlag::Vulkan:
-            deviceDesc.deviceType = gfx::DeviceType::Vulkan;
-            break;
-        default:
-            SLANG_IGNORE_TEST
-        }
-        deviceDesc.slang.slangGlobalSession = context->slangGlobalSession;
-        const char* searchPaths[] = { "", "../../tools/gfx-unit-test", "tools/gfx-unit-test" };
-        deviceDesc.slang.searchPathCount = (SlangInt)SLANG_COUNT_OF(searchPaths);
-        deviceDesc.slang.searchPaths = searchPaths;
-        auto createDeviceResult = gfxCreateDevice(&deviceDesc, device.writeRef());
-        if (SLANG_FAILED(createDeviceResult))
-        {
-            SLANG_IGNORE_TEST
-        }
-
-        computeSmokeTestImpl(device, context);
-    }
-
     SLANG_UNIT_TEST(computeSmokeD3D11)
     {
-        computeSmokeTestAPI(unitTestContext, Slang::RenderApiFlag::D3D11);
+        runTestImpl(computeSmokeTestImpl, unitTestContext, Slang::RenderApiFlag::D3D11);
     }
 
     SLANG_UNIT_TEST(computeSmokeVulkan)
     {
-        computeSmokeTestAPI(unitTestContext, Slang::RenderApiFlag::Vulkan);
+        runTestImpl(computeSmokeTestImpl, unitTestContext, Slang::RenderApiFlag::Vulkan);
     }
 
 }
diff --git a/tools/gfx-unit-test/gfx-test-util.cpp b/tools/gfx-unit-test/gfx-test-util.cpp
index 1b23047a1..f1dd6be58 100644
--- a/tools/gfx-unit-test/gfx-test-util.cpp
+++ b/tools/gfx-unit-test/gfx-test-util.cpp
@@ -3,6 +3,14 @@
 
 #include <slang-com-ptr.h>
 
+#define GFX_ENABLE_RENDERDOC_INTEGRATION 0
+
+#if GFX_ENABLE_RENDERDOC_INTEGRATION
+#    include "external/renderdoc_app.h"
+#    define WIN32_LEAN_AND_MEAN
+#    include <Windows.h>
+#endif
+
 using Slang::ComPtr;
 
 namespace gfx_test
@@ -64,23 +72,76 @@ namespace gfx_test
         ComPtr<ISlangBlob> resultBlob;
         GFX_CHECK_CALL_ABORT(device->readBufferResource(
             buffer, 0, expectedBufferSize, resultBlob.writeRef()));
-        if (resultBlob->getBufferSize() < expectedBufferSize)
+        SLANG_CHECK(resultBlob->getBufferSize() == expectedBufferSize);
+        // Compare results.
+        SLANG_CHECK(memcmp(resultBlob->getBufferPointer(), expectedResult, expectedBufferSize) == 0);
+    }
+
+    Slang::ComPtr<gfx::IDevice> createTestingDevice(UnitTestContext* context, Slang::RenderApiFlag::Enum api)
+    {
+        Slang::ComPtr<gfx::IDevice> device;
+        gfx::IDevice::Desc deviceDesc = {};
+        switch (api)
         {
-            getTestReporter()->addResult(TestResult::Fail);
-            return;
+        case Slang::RenderApiFlag::D3D11:
+            deviceDesc.deviceType = gfx::DeviceType::DirectX11;
+            break;
+        case Slang::RenderApiFlag::D3D12:
+            deviceDesc.deviceType = gfx::DeviceType::DirectX12;
+            break;
+        case Slang::RenderApiFlag::Vulkan:
+            deviceDesc.deviceType = gfx::DeviceType::Vulkan;
+            break;
+        case Slang::RenderApiFlag::CPU:
+            deviceDesc.deviceType = gfx::DeviceType::CPU;
+            break;
+        case Slang::RenderApiFlag::CUDA:
+            deviceDesc.deviceType = gfx::DeviceType::CUDA;
+            break;
+        case Slang::RenderApiFlag::OpenGl:
+            deviceDesc.deviceType = gfx::DeviceType::OpenGl;
+            break;
+        default:
+            SLANG_IGNORE_TEST
         }
-
-        // Compare results.
-        auto result = reinterpret_cast<const uint8_t*>(resultBlob->getBufferPointer());
-        for (int i = 0; i < expectedBufferSize; i++)
+        deviceDesc.slang.slangGlobalSession = context->slangGlobalSession;
+        const char* searchPaths[] = { "", "../../tools/gfx-unit-test", "tools/gfx-unit-test" };
+        deviceDesc.slang.searchPathCount = (SlangInt)SLANG_COUNT_OF(searchPaths);
+        deviceDesc.slang.searchPaths = searchPaths;
+        auto createDeviceResult = gfxCreateDevice(&deviceDesc, device.writeRef());
+        if (SLANG_FAILED(createDeviceResult))
         {
-            if (expectedResult[i] != result[i])
-            {
-                getTestReporter()->addResult(TestResult::Fail);
-                return;
-            }
+            SLANG_IGNORE_TEST
+        }
+        return device;
+    }
 
+#if GFX_ENABLE_RENDERDOC_INTEGRATION
+    RENDERDOC_API_1_1_2* rdoc_api = NULL;
+    void initializeRenderDoc()
+    {
+        if (HMODULE mod = GetModuleHandleA("renderdoc.dll"))
+        {
+            pRENDERDOC_GetAPI RENDERDOC_GetAPI =
+                (pRENDERDOC_GetAPI)GetProcAddress(mod, "RENDERDOC_GetAPI");
+            int ret = RENDERDOC_GetAPI(eRENDERDOC_API_Version_1_1_2, (void**)&rdoc_api);
+            assert(ret == 1);
         }
-        getTestReporter()->addResult(TestResult::Pass);
     }
+    void renderDocBeginFrame()
+    {
+        if (!rdoc_api) initializeRenderDoc();
+        if (rdoc_api) rdoc_api->StartFrameCapture(nullptr, nullptr);
+    }
+    void renderDocEndFrame()
+    {
+        if (rdoc_api)
+            rdoc_api->EndFrameCapture(nullptr, nullptr);
+        _fgetchar();
+    }
+#else
+    void initializeRenderDoc() {}
+    void renderDocBeginFrame() {}
+    void renderDocEndFrame() {}
+#endif
 }
diff --git a/tools/gfx-unit-test/gfx-test-util.h b/tools/gfx-unit-test/gfx-test-util.h
index c7db19394..b357a310a 100644
--- a/tools/gfx-unit-test/gfx-test-util.h
+++ b/tools/gfx-unit-test/gfx-test-util.h
@@ -2,6 +2,8 @@
 
 #include "slang-gfx.h"
 #include "source/core/slang-basic.h"
+#include "source/core/slang-render-api-util.h"
+#include "tools/unit-test/slang-unit-test.h"
 
 namespace gfx_test
 {
@@ -35,6 +37,37 @@ namespace gfx_test
         memcpy(expectedBuffer.getBuffer(), expectedResult.begin(), bufferSize);
         return compareComputeResult(device, buffer, expectedBuffer.getBuffer(), bufferSize);
     }
+    
+    Slang::ComPtr<gfx::IDevice> createTestingDevice(UnitTestContext* context, Slang::RenderApiFlag::Enum api);
+
+    void initializeRenderDoc();
+    void renderDocBeginFrame();
+    void renderDocEndFrame();
+
+    template<typename ImplFunc>
+    void runTestImpl(const ImplFunc& f, UnitTestContext* context, Slang::RenderApiFlag::Enum api)
+    {
+        if ((api & context->enabledApis) == 0)
+        {
+            SLANG_IGNORE_TEST
+        }
+        auto device = createTestingDevice(context, api);
+        if (!device)
+        {
+            SLANG_IGNORE_TEST
+        }
+        try
+        {
+            renderDocBeginFrame();
+            f(device, context);
+        }
+        catch (AbortTestException& e)
+        {
+            renderDocEndFrame();
+            throw e;
+        }
+        renderDocEndFrame();
+    }
 
 #define GFX_CHECK_CALL(x) SLANG_CHECK(!SLANG_FAILED(x))
 #define GFX_CHECK_CALL_ABORT(x) SLANG_CHECK_ABORT(!SLANG_FAILED(x))
diff --git a/tools/gfx-unit-test/mutable-shader-object.cpp b/tools/gfx-unit-test/mutable-shader-object.cpp
new file mode 100644
index 000000000..efaf95761
--- /dev/null
+++ b/tools/gfx-unit-test/mutable-shader-object.cpp
@@ -0,0 +1,135 @@
+#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 mutableShaderObjectTestImpl(IDevice* device, UnitTestContext* context)
+    {
+        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, "mutable-shader-object", "computeMain", slangReflection));
+
+        ComputePipelineStateDesc pipelineDesc = {};
+        pipelineDesc.program = shaderProgram.get();
+        ComPtr<gfx::IPipelineState> pipelineState;
+        GFX_CHECK_CALL_ABORT(
+            device->createComputePipelineState(pipelineDesc, pipelineState.writeRef()));
+
+        float initialData[] = { 0.0f, 1.0f, 2.0f, 3.0f };
+        const int numberCount = SLANG_COUNT_OF(initialData);
+        IBufferResource::Desc bufferDesc = {};
+        bufferDesc.sizeInBytes = sizeof(initialData);
+        bufferDesc.format = gfx::Format::Unknown;
+        bufferDesc.elementSize = sizeof(float);
+        bufferDesc.allowedStates = ResourceStateSet(
+            ResourceState::ShaderResource,
+            ResourceState::UnorderedAccess,
+            ResourceState::CopyDestination,
+            ResourceState::CopySource);
+        bufferDesc.defaultState = ResourceState::UnorderedAccess;
+        bufferDesc.cpuAccessFlags = AccessFlag::Write | AccessFlag::Read;
+
+        ComPtr<IBufferResource> numbersBuffer;
+        GFX_CHECK_CALL_ABORT(device->createBufferResource(
+            bufferDesc,
+            (void*)initialData,
+            numbersBuffer.writeRef()));
+
+        ComPtr<IResourceView> bufferView;
+        IResourceView::Desc viewDesc = {};
+        viewDesc.type = IResourceView::Type::UnorderedAccess;
+        viewDesc.format = Format::Unknown;
+        GFX_CHECK_CALL_ABORT(device->createBufferView(numbersBuffer, viewDesc, bufferView.writeRef()));
+
+        {
+            slang::TypeReflection* addTransformerType =
+                slangReflection->findTypeByName("AddTransformer");
+
+            ComPtr<IShaderObject> transformer;
+            GFX_CHECK_CALL_ABORT(device->createMutableShaderObject(
+                addTransformerType, ShaderObjectContainerType::None, transformer.writeRef()));
+            // Set the `c` field of the `AddTransformer`.
+            float c = 1.0f;
+            ShaderCursor(transformer).getPath("c").setData(&c, sizeof(float));
+
+            ICommandQueue::Desc queueDesc = { ICommandQueue::QueueType::Graphics };
+            auto queue = device->createCommandQueue(queueDesc);
+
+            auto commandBuffer = transientHeap->createCommandBuffer();
+            auto encoder = commandBuffer->encodeComputeCommands();
+
+            auto rootObject = encoder->bindPipeline(pipelineState);
+
+            auto entryPointCursor = ShaderCursor(rootObject->getEntryPoint(0));
+
+            entryPointCursor.getPath("buffer").setResource(bufferView);
+
+            // Bind the previously created transformer object to root object.
+            ComPtr<IShaderObject> transformerVersion;
+            transformer->getCurrentVersion(transientHeap, transformerVersion.writeRef());
+            entryPointCursor.getPath("transformer").setObject(transformerVersion);
+
+            encoder->dispatchCompute(1, 1, 1);
+            encoder->endEncoding();
+
+            auto barrierEncoder = commandBuffer->encodeResourceCommands();
+            barrierEncoder->bufferBarrier(1, numbersBuffer.readRef(), ResourceState::UnorderedAccess, ResourceState::UnorderedAccess);
+            barrierEncoder->endEncoding();
+
+            encoder = commandBuffer->encodeComputeCommands();
+
+            rootObject = encoder->bindPipeline(pipelineState);
+            entryPointCursor = ShaderCursor(rootObject->getEntryPoint(0));
+
+            // Mutate `transformer` object and run again.
+            c = 2.0f;
+            ShaderCursor(transformer).getPath("c").setData(&c, sizeof(float));
+            transformer->getCurrentVersion(transientHeap, transformerVersion.writeRef());
+            entryPointCursor.getPath("buffer").setResource(bufferView);
+            entryPointCursor.getPath("transformer").setObject(transformerVersion);
+            encoder->dispatchCompute(1, 1, 1);
+            encoder->endEncoding();
+
+            commandBuffer->close();
+            queue->executeCommandBuffer(commandBuffer);
+            queue->wait();
+        }
+
+        compareComputeResult(
+            device,
+            numbersBuffer,
+            Slang::makeArray<float>(3.0f, 4.0f, 5.0f, 6.0f));
+    }
+
+    //SLANG_UNIT_TEST(mutableShaderObjectCPU)
+    //{
+    //    runTestImpl(mutableShaderObjectTestImpl, unitTestContext, Slang::RenderApiFlag::CPU);
+    //}
+
+    SLANG_UNIT_TEST(mutableShaderObjectD3D11)
+    {
+        runTestImpl(mutableShaderObjectTestImpl, unitTestContext, Slang::RenderApiFlag::D3D11);
+    }
+
+    SLANG_UNIT_TEST(mutableShaderObjectD3D12)
+    {
+        runTestImpl(mutableShaderObjectTestImpl, unitTestContext, Slang::RenderApiFlag::D3D12);
+    }
+
+    SLANG_UNIT_TEST(mutableShaderObjectVulkan)
+    {
+        runTestImpl(mutableShaderObjectTestImpl, unitTestContext, Slang::RenderApiFlag::Vulkan);
+    }
+}
diff --git a/tools/gfx-unit-test/mutable-shader-object.slang b/tools/gfx-unit-test/mutable-shader-object.slang
new file mode 100644
index 000000000..42d520f13
--- /dev/null
+++ b/tools/gfx-unit-test/mutable-shader-object.slang
@@ -0,0 +1,68 @@
+// mutable-shader-object.slang
+
+// This is a copy of `shader-object.slang` in `shader-object` example
+// for use by mutable-shader-object gfx unit test.
+
+// This file implements a simple compute shader that transforms
+// input floating point numbers stored in a `RWStructuredBuffer`.
+// Specifically, for each number x from input buffer, compute
+// f(x) and store the result back in the same buffer.
+
+// The compute shader supports multiple transformation functions,
+// such add(x, c) which returns x+c, or mul(x, c) which returns x*c.
+// This functions are implemented as types that conforms to the
+// `ITransformer` interface.
+
+// The main entry point function takes a parameter of `ITransformer`
+// type, and applies the transformation to numbers in the input
+// buffer. By defining the shader parameter using interfaces,
+// we enable the flexiblity to generate either specialized compute
+// kernels that performs specific transformation or a general
+// kernel that can perform any transformations encoded by the
+// parameter at run-time, without changing any shader code or
+// host-application logic for setting and preparing shader parameters.
+
+// Defines the transformer interface, which implements a single
+// `transform` operation.
+interface ITransformer
+{
+    float transform(float x);
+}
+
+// Represents a transform function f(x) = x + c.
+struct AddTransformer : ITransformer
+{
+    float c;
+    float transform(float x) { return x + c; }
+};
+
+// Represents a transform function f(x) = x * c.
+struct MulTransformer : ITransformer
+{
+    float c;
+    float transform(float x) { return x * c; }
+};
+
+// Represents a composite function f(x) = f0(f1(x));
+struct CompositeTransformer : ITransformer
+{
+    ITransformer func0;
+    ITransformer func1;
+    float transform(float x)
+    {
+        return func0.transform(func1.transform(x));
+    }
+};
+
+// Main entry-point. Applies the transformation encoded by `transformer`
+// to all elements in `buffer`.
+[shader("compute")]
+[numthreads(4,1,1)]
+void computeMain(
+    uint3 sv_dispatchThreadID : SV_DispatchThreadID,
+    uniform RWStructuredBuffer<float> buffer,
+    uniform ITransformer transformer)
+{
+    var input = buffer[sv_dispatchThreadID.x];
+    buffer[sv_dispatchThreadID.x] = transformer.transform(input);
+}