gfx: support d3d12 root parameters (#2122)

* Various fixes to gfx.

* Fix.

* Fixes.

* Fix.

* gfx: support root parameter via user-defined attribute.

* Fix.

* Fix.

* Skip d3d12 tests on win x86.

* Fixes.

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

4 files changed, 182 insertions, 0 deletions

diff --git a/tools/gfx-unit-test/gfx-test-util.cpp b/tools/gfx-unit-test/gfx-test-util.cpp
index 1133e3842..ecd19f179 100644
--- a/tools/gfx-unit-test/gfx-test-util.cpp
+++ b/tools/gfx-unit-test/gfx-test-util.cpp
@@ -202,6 +202,14 @@ namespace gfx_test
         const char* searchPaths[] = { "", "../../tools/gfx-unit-test", "tools/gfx-unit-test" };
         deviceDesc.slang.searchPathCount = (SlangInt)SLANG_COUNT_OF(searchPaths);
         deviceDesc.slang.searchPaths = searchPaths;
+
+        gfx::D3D12DeviceExtendedDesc extDesc = {};
+        extDesc.rootParameterShaderAttributeName = "root";
+
+        deviceDesc.extendedDescCount = 1;
+        void* extDescPtr = &extDesc;
+        deviceDesc.extendedDescs = &extDescPtr;
+
         auto createDeviceResult = gfxCreateDevice(&deviceDesc, device.writeRef());
         if (SLANG_FAILED(createDeviceResult))
         {
diff --git a/tools/gfx-unit-test/gfx-test-util.h b/tools/gfx-unit-test/gfx-test-util.h
index 67660d4f5..2977ac599 100644
--- a/tools/gfx-unit-test/gfx-test-util.h
+++ b/tools/gfx-unit-test/gfx-test-util.h
@@ -87,6 +87,13 @@ namespace gfx_test
         {
             SLANG_IGNORE_TEST
         }
+#if SLANG_WIN32
+        // Skip d3d12 tests on x86 now since dxc doesn't function correctly there on Windows 11.
+        if (api == Slang::RenderApiFlag::D3D12)
+        {
+            SLANG_IGNORE_TEST
+        }
+#endif
         try
         {
             renderDocBeginFrame();
diff --git a/tools/gfx-unit-test/root-shader-parameter.cpp b/tools/gfx-unit-test/root-shader-parameter.cpp
new file mode 100644
index 000000000..a7b92843d
--- /dev/null
+++ b/tools/gfx-unit-test/root-shader-parameter.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
+{
+    ComPtr<IBufferResource> createBuffer(IDevice* device, uint32_t content)
+    {
+        ComPtr<IBufferResource> buffer;
+        IBufferResource::Desc bufferDesc = {};
+        bufferDesc.sizeInBytes = sizeof(uint32_t);
+        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.memoryType = MemoryType::DeviceLocal;
+
+        ComPtr<IBufferResource> numbersBuffer;
+        GFX_CHECK_CALL_ABORT(
+            device->createBufferResource(bufferDesc, (void*)&content, buffer.writeRef()));
+
+        return buffer;
+    }
+    void rootShaderParameterTestImpl(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, "root-shader-parameter", "computeMain", slangReflection));
+
+        ComputePipelineStateDesc pipelineDesc = {};
+        pipelineDesc.program = shaderProgram.get();
+        ComPtr<gfx::IPipelineState> pipelineState;
+        GFX_CHECK_CALL_ABORT(
+            device->createComputePipelineState(pipelineDesc, pipelineState.writeRef()));
+
+        Slang::List<ComPtr<IBufferResource>> buffers;
+        Slang::List<ComPtr<IResourceView>> srvs, uavs;
+
+        for (uint32_t i = 0; i < 9; i++)
+        {
+            buffers.add(createBuffer(device, i == 0 ? 10 : i));
+
+            ComPtr<IResourceView> bufferView;
+            IResourceView::Desc viewDesc = {};
+            viewDesc.type = IResourceView::Type::UnorderedAccess;
+            viewDesc.format = Format::Unknown;
+            GFX_CHECK_CALL_ABORT(
+                device->createBufferView(buffers[i], nullptr, viewDesc, bufferView.writeRef()));
+            uavs.add(bufferView);
+
+            viewDesc.type = IResourceView::Type::ShaderResource;
+            viewDesc.format = Format::Unknown;
+            GFX_CHECK_CALL_ABORT(
+                device->createBufferView(buffers[i], nullptr, viewDesc, bufferView.writeRef()));
+            srvs.add(bufferView);
+        }
+
+        ComPtr<IShaderObject> rootObject;
+        device->createMutableRootShaderObject(shaderProgram, rootObject.writeRef());
+
+        ComPtr<IShaderObject> g, s1, s2;
+        device->createMutableShaderObject(
+            slangReflection->findTypeByName("S0"), ShaderObjectContainerType::None, g.writeRef());
+        device->createMutableShaderObject(
+            slangReflection->findTypeByName("S1"), ShaderObjectContainerType::None, s1.writeRef());
+        device->createMutableShaderObject(
+            slangReflection->findTypeByName("S1"), ShaderObjectContainerType::None, s2.writeRef());
+
+        {
+            auto cursor = ShaderCursor(s1);
+            cursor["c0"].setResource(srvs[2]);
+            cursor["c1"].setResource(uavs[3]);
+            cursor["c2"].setResource(srvs[4]);
+        }
+        {
+            auto cursor = ShaderCursor(s2);
+            cursor["c0"].setResource(srvs[5]);
+            cursor["c1"].setResource(uavs[6]);
+            cursor["c2"].setResource(srvs[7]);
+        }
+        {
+            auto cursor = ShaderCursor(g);
+            cursor["b0"].setResource(srvs[0]);
+            cursor["b1"].setResource(srvs[1]);
+            cursor["s1"].setObject(s1);
+            cursor["s2"].setObject(s2);
+        }
+        {
+            auto cursor = ShaderCursor(rootObject);
+            cursor["g"].setObject(g);
+            cursor["buffer"].setResource(uavs[8]);
+        }
+
+        {
+            ICommandQueue::Desc queueDesc = { ICommandQueue::QueueType::Graphics };
+            auto queue = device->createCommandQueue(queueDesc);
+
+            auto commandBuffer = transientHeap->createCommandBuffer();
+            {
+                auto encoder = commandBuffer->encodeComputeCommands();
+                auto root = encoder->bindPipeline(pipelineState);
+                root->copyFrom(rootObject, transientHeap);
+                encoder->dispatchCompute(1, 1, 1);
+                encoder->endEncoding();
+            }
+
+            commandBuffer->close();
+            queue->executeCommandBuffer(commandBuffer);
+            queue->waitOnHost();
+        }
+
+        compareComputeResult(
+            device, buffers[8], Slang::makeArray<uint32_t>(10 - 1 + 2 - 3 + 4 + 5 - 6 + 7));
+    }
+
+    SLANG_UNIT_TEST(rootShaderParameterD3D12)
+    {
+        runTestImpl(rootShaderParameterTestImpl, unitTestContext, Slang::RenderApiFlag::D3D12);
+    }
+}
diff --git a/tools/gfx-unit-test/root-shader-parameter.slang b/tools/gfx-unit-test/root-shader-parameter.slang
new file mode 100644
index 000000000..bf7af1851
--- /dev/null
+++ b/tools/gfx-unit-test/root-shader-parameter.slang
@@ -0,0 +1,32 @@
+// root-shader-parameter.slang
+
+// Test use of root shader parameters.
+[__AttributeUsage(_AttributeTargets.Var)]
+struct rootAttribute {};
+
+struct S1
+{
+    StructuredBuffer<uint> c0;
+    [root] RWStructuredBuffer<uint> c1;
+    StructuredBuffer<uint> c2;
+}
+
+struct S0
+{
+    StructuredBuffer<uint> b0;
+    [root] StructuredBuffer<uint> b1;
+    ParameterBlock<S1> s1;
+    ConstantBuffer<S1> s2;
+}
+
+ParameterBlock<S0> g;
+[root] RWStructuredBuffer<uint> buffer;
+
+[shader("compute")]
+[numthreads(1,1,1)]
+void computeMain(
+    uint3 sv_dispatchThreadID : SV_DispatchThreadID)
+{
+    buffer[0] = g.b0[0] - g.b1[0] + g.s1.c0[0] - g.s1.c1[0] + g.s1.c2[0] + g.s2.c0[0] - g.s2.c1[0] + g.s2.c2[0];
+    // 10-1+2-3+4+5-6+7
+}