Add gfx unit testing framework. (#1943)

* Add gfx unit testing framework.

* Fix compilation error.

* Reset gfxDebugCallback after render_test.

* Pass enabledApi flags through.

* Fix for code review suggestions.

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

8 files changed, 487 insertions, 0 deletions

diff --git a/tools/gfx-test/compute-smoke.cpp b/tools/gfx-test/compute-smoke.cpp
new file mode 100644
index 000000000..bbba72348
--- /dev/null
+++ b/tools/gfx-test/compute-smoke.cpp
@@ -0,0 +1,145 @@
+#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
+{
+    SlangResult computeSmokeTestImpl(IDevice* device, slang::UnitTestContext* context)
+    {
+        Slang::ComPtr<ITransientResourceHeap> transientHeap;
+        ITransientResourceHeap::Desc transientHeapDesc = {};
+        transientHeapDesc.constantBufferSize = 4096;
+        SLANG_RETURN_ON_FAIL(
+            device->createTransientResourceHeap(transientHeapDesc, transientHeap.writeRef()));
+
+        ComPtr<IShaderProgram> shaderProgram;
+        slang::ProgramLayout* slangReflection;
+        SLANG_RETURN_ON_FAIL(loadShaderProgram(device, shaderProgram, context->outputWriter, "compute-smoke", slangReflection));
+
+        ComputePipelineStateDesc pipelineDesc = {};
+        pipelineDesc.program = shaderProgram.get();
+        ComPtr<gfx::IPipelineState> pipelineState;
+        SLANG_RETURN_ON_FAIL(
+            device->createComputePipelineState(pipelineDesc, pipelineState.writeRef()));
+
+        const int numberCount = 4;
+        float initialData[] = { 0.0f, 1.0f, 2.0f, 3.0f };
+        IBufferResource::Desc bufferDesc = {};
+        bufferDesc.sizeInBytes = numberCount * sizeof(float);
+        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;
+        SLANG_RETURN_ON_FAIL(device->createBufferResource(
+            bufferDesc,
+            (void*)initialData,
+            numbersBuffer.writeRef()));
+
+        ComPtr<IResourceView> bufferView;
+        IResourceView::Desc viewDesc = {};
+        viewDesc.type = IResourceView::Type::UnorderedAccess;
+        viewDesc.format = Format::Unknown;
+        SLANG_RETURN_ON_FAIL(device->createBufferView(numbersBuffer, viewDesc, bufferView.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);
+
+            slang::TypeReflection* addTransformerType =
+                slangReflection->findTypeByName("AddTransformer");
+
+            // Now we can use this type to create a shader object that can be bound to the root object.
+            ComPtr<IShaderObject> transformer;
+            SLANG_RETURN_ON_FAIL(device->createShaderObject(
+                addTransformerType, ShaderObjectContainerType::None, transformer.writeRef()));
+            // Set the `c` field of the `AddTransformer`.
+            float c = 1.0f;
+            ShaderCursor(transformer).getPath("c").setData(&c, sizeof(float));
+
+            ShaderCursor entryPointCursor(
+                rootObject->getEntryPoint(0)); // get a cursor the the first entry-point.
+            // Bind buffer view to the entry point.
+            entryPointCursor.getPath("buffer").setResource(bufferView);
+
+            // Bind the previously created transformer object to root object.
+            entryPointCursor.getPath("transformer").setObject(transformer);
+
+            encoder->dispatchCompute(1, 1, 1);
+            encoder->endEncoding();
+            commandBuffer->close();
+            queue->executeCommandBuffer(commandBuffer);
+            queue->wait();
+        }
+
+        return compareComputeResult(
+            device,
+            numbersBuffer,
+            Slang::makeArray<float>(11.0f, 12.0f, 13.0f, 14.0f));
+    }
+
+    SlangResult computeSmokeTestAPI(slang::UnitTestContext* context, Slang::RenderApiFlag::Enum api)
+    {
+        if ((api & context->enabledApis) == 0)
+        {
+            return SLANG_E_NOT_AVAILABLE;
+        }
+        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:
+            return SLANG_E_NOT_AVAILABLE;
+        }
+        deviceDesc.slang.slangGlobalSession = context->slangGlobalSession;
+        const char* searchPaths[] = { "", "../../tools/gfx-test", "tools/gfx-test" };
+        deviceDesc.slang.searchPathCount = (SlangInt)SLANG_COUNT_OF(searchPaths);
+        deviceDesc.slang.searchPaths = searchPaths;
+        auto createDeviceResult = gfxCreateDevice(&deviceDesc, device.writeRef());
+        if (SLANG_FAILED(createDeviceResult))
+        {
+            return SLANG_E_NOT_AVAILABLE;
+        }
+
+        SLANG_RETURN_ON_FAIL(computeSmokeTestImpl(device, context));
+        return SLANG_OK;
+    }
+
+    SLANG_UNIT_TEST(computeSmokeD3D11)
+    {
+        return computeSmokeTestAPI(context, Slang::RenderApiFlag::D3D11);
+    }
+
+    SLANG_UNIT_TEST(computeSmokeVulkan)
+    {
+        return computeSmokeTestAPI(context, Slang::RenderApiFlag::Vulkan);
+    }
+
+}
diff --git a/tools/gfx-test/compute-smoke.slang b/tools/gfx-test/compute-smoke.slang
new file mode 100644
index 000000000..7ecdb4177
--- /dev/null
+++ b/tools/gfx-test/compute-smoke.slang
@@ -0,0 +1,68 @@
+// compute-smoke.slang
+
+// This is a copy of `shader-object.slang` in `shader-object` example
+// for use by compute-smoke 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 + 10.0f; }
+};
+
+// 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);
+}
diff --git a/tools/gfx-test/gfx-test-util.cpp b/tools/gfx-test/gfx-test-util.cpp
new file mode 100644
index 000000000..5e77879a9
--- /dev/null
+++ b/tools/gfx-test/gfx-test-util.cpp
@@ -0,0 +1,79 @@
+#include "gfx-test-util.h"
+
+#include <slang-com-ptr.h>
+
+using Slang::ComPtr;
+
+namespace gfx_test
+{
+    void diagnoseIfNeeded(ISlangWriter* diagnosticWriter, slang::IBlob* diagnosticsBlob)
+    {
+        if (diagnosticsBlob != nullptr)
+        {
+            diagnosticWriter->write((const char*)diagnosticsBlob->getBufferPointer(), diagnosticsBlob->getBufferSize());
+        }
+    }
+
+    Slang::Result loadShaderProgram(
+        gfx::IDevice* device,
+        Slang::ComPtr<gfx::IShaderProgram>& outShaderProgram,
+        ISlangWriter* diagnosticWriter,
+        const char* shaderModuleName,
+        slang::ProgramLayout*& slangReflection)
+    {
+        Slang::ComPtr<slang::ISession> slangSession;
+        SLANG_RETURN_ON_FAIL(device->getSlangSession(slangSession.writeRef()));
+        Slang::ComPtr<slang::IBlob> diagnosticsBlob;
+        slang::IModule* module = slangSession->loadModule(shaderModuleName, diagnosticsBlob.writeRef());
+        diagnoseIfNeeded(diagnosticWriter, diagnosticsBlob);
+        if (!module)
+            return SLANG_FAIL;
+
+        char const* computeEntryPointName = "computeMain";
+        ComPtr<slang::IEntryPoint> computeEntryPoint;
+        SLANG_RETURN_ON_FAIL(
+            module->findEntryPointByName(computeEntryPointName, computeEntryPoint.writeRef()));
+
+        Slang::List<slang::IComponentType*> componentTypes;
+        componentTypes.add(module);
+        componentTypes.add(computeEntryPoint);
+
+        Slang::ComPtr<slang::IComponentType> composedProgram;
+        SlangResult result = slangSession->createCompositeComponentType(
+            componentTypes.getBuffer(),
+            componentTypes.getCount(),
+            composedProgram.writeRef(),
+            diagnosticsBlob.writeRef());
+        diagnoseIfNeeded(diagnosticWriter, diagnosticsBlob);
+        SLANG_RETURN_ON_FAIL(result);
+        slangReflection = composedProgram->getLayout();
+
+        gfx::IShaderProgram::Desc programDesc = {};
+        programDesc.pipelineType = gfx::PipelineType::Compute;
+        programDesc.slangProgram = composedProgram.get();
+
+        auto shaderProgram = device->createProgram(programDesc);
+
+        outShaderProgram = shaderProgram;
+        return SLANG_OK;
+    }
+
+    Slang::Result compareComputeResult(gfx::IDevice* device, gfx::IBufferResource* buffer, uint8_t* expectedResult, size_t expectedBufferSize)
+    {
+        // Read back the results.
+        ComPtr<ISlangBlob> resultBlob;
+        SLANG_RETURN_ON_FAIL(device->readBufferResource(
+            buffer, 0, expectedBufferSize, resultBlob.writeRef()));
+        if (resultBlob->getBufferSize() < expectedBufferSize)
+            return SLANG_FAIL;
+
+        // Compare results.
+        auto result = reinterpret_cast<const uint8_t*>(resultBlob->getBufferPointer());
+        for (int i = 0; i < expectedBufferSize; i++)
+        {
+            if (expectedResult[i] != result[i])
+                return SLANG_FAIL;
+        }
+        return SLANG_OK;
+    }
+}
diff --git a/tools/gfx-test/gfx-test-util.h b/tools/gfx-test/gfx-test-util.h
new file mode 100644
index 000000000..5223ba773
--- /dev/null
+++ b/tools/gfx-test/gfx-test-util.h
@@ -0,0 +1,38 @@
+#pragma once
+
+#include "slang-gfx.h"
+#include "source/core/slang-basic.h"
+
+namespace gfx_test
+{
+        /// Helper function for print out diagnostic messages output by Slang compiler.
+    void diagnoseIfNeeded(ISlangWriter* diagnosticWriter, slang::IBlob* diagnosticsBlob);
+
+        /// Loads a compute shader module and produces a `gfx::IShaderProgram`.
+    Slang::Result loadShaderProgram(
+        gfx::IDevice* device,
+        Slang::ComPtr<gfx::IShaderProgram>& outShaderProgram,
+        ISlangWriter* diagnosticWriter,
+        const char* shaderModuleName,
+        slang::ProgramLayout*& slangReflection);
+
+        /// Reads back the content of `buffer` and compares it against `expectedResult`.
+    Slang::Result compareComputeResult(
+        gfx::IDevice* device,
+        gfx::IBufferResource* buffer,
+        uint8_t* expectedResult,
+        size_t expectedBufferSize);
+
+    template<typename T, Slang::Index count>
+    Slang::Result compareComputeResult(
+        gfx::IDevice* device,
+        gfx::IBufferResource* buffer,
+        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);
+        return compareComputeResult(device, buffer, expectedBuffer.getBuffer(), bufferSize);
+    }
+}
diff --git a/tools/render-test/render-test-main.cpp b/tools/render-test/render-test-main.cpp
index 271356cc3..7c65da5e7 100644
--- a/tools/render-test/render-test-main.cpp
+++ b/tools/render-test/render-test-main.cpp
@@ -1264,6 +1264,13 @@ static SlangResult _innerMain(Slang::StdWriters* stdWriters, SlangSession* sessi
     StdWritersDebugCallback debugCallback;
     debugCallback.writers = stdWriters;
     gfxSetDebugCallback(&debugCallback);
+    struct ResetDebugCallbackRAII
+    {
+        ~ResetDebugCallbackRAII()
+        {
+            gfxSetDebugCallback(nullptr);
+        }
+    } resetDebugCallbackRAII;
 
     // Use the profile name set on options if set
     input.profile = options.profileName.getLength() ? options.profileName : input.profile;
diff --git a/tools/slang-test/slang-test-main.cpp b/tools/slang-test/slang-test-main.cpp
index 18ac05477..9c194afd4 100644
--- a/tools/slang-test/slang-test-main.cpp
+++ b/tools/slang-test/slang-test-main.cpp
@@ -17,6 +17,8 @@
 #include "../../source/core/slang-process-util.h"
 #include "../../source/core/slang-render-api-util.h"
 
+#include "tools/unit-test/slang-unit-test.h"
+#undef SLANG_UNIT_TEST
 
 #include "directory-util.h"
 #include "test-context.h"
@@ -39,6 +41,7 @@
 #define SLANG_PRELUDE_NAMESPACE CPPPrelude
 #include "../../prelude/slang-cpp-types.h"
 
+
 using namespace Slang;
 
 // Options for a particular test
@@ -3389,6 +3392,62 @@ static void _disableCPPBackends(TestContext* context)
     }
 }
 
+    /// Loads a DLL containing unit test functions and run them one by one.
+static SlangResult runUnitTestModule(TestContext* context, TestOptions& testOptions, const char* moduleName)
+{
+    SharedLibrary::Handle moduleHandle;
+    SLANG_RETURN_ON_FAIL(SharedLibrary::load(
+        Path::combine(context->exeDirectoryPath, moduleName).getBuffer(),
+        moduleHandle));
+    slang::UnitTestGetModuleFunc getModuleFunc =
+        (slang::UnitTestGetModuleFunc) SharedLibrary::findSymbolAddressByName(
+            moduleHandle, "slangUnitTestGetModule");
+    if (!getModuleFunc)
+        return SLANG_FAIL;
+
+    slang::IUnitTestModule* testModule = getModuleFunc();
+    if (!testModule)
+        return SLANG_FAIL;
+
+    slang::UnitTestContext unitTestContext;
+    unitTestContext.slangGlobalSession = context->getSession();
+    unitTestContext.workDirectory = "";
+    unitTestContext.enabledApis = context->options.enabledApis;
+    auto testCount = testModule->getTestCount();
+    for (SlangInt i = 0; i < testCount; i++)
+    {
+        StringBuilder sb;
+        StringWriter messageWriter(&sb, WriterFlag::Enum::AutoFlush);
+        auto testFunc = testModule->getTestFunc(i);
+        auto testName = testModule->getTestName(i);
+
+        StringBuilder filePath;
+        filePath << moduleName << "/" << testName << ".internal";
+
+        testOptions.command = filePath;
+
+        if (shouldRunTest(context, testOptions.command))
+        {
+            if (testPassesCategoryMask(context, testOptions))
+            {
+                unitTestContext.outputWriter = &messageWriter;
+                TestReporter::get()->startTest(testOptions.command);
+                auto result = testFunc(&unitTestContext);
+                if (sb.getLength())
+                    TestReporter::get()->message(TestMessageType::Info, sb.ProduceString());
+
+                if (result == SLANG_E_NOT_AVAILABLE)
+                    TestReporter::get()->addResult(TestResult::Ignored);
+                else if (SLANG_FAILED(result))
+                    TestReporter::get()->addResult(TestResult::Fail);
+                else
+                    TestReporter::get()->addResult(TestResult::Pass);
+                TestReporter::get()->endTest();
+            }
+        }
+    }
+    return SLANG_OK;
+}
 
 SlangResult innerMain(int argc, char** argv)
 {
@@ -3594,6 +3653,11 @@ SlangResult innerMain(int argc, char** argv)
                 cur = cur->m_next;
             }
 
+            {
+                TestOptions testOptions;
+                testOptions.categories.add(unitTestCategory);
+                runUnitTestModule(&context, testOptions, "gfx-test-tool");
+            }
             TestReporter::set(nullptr);
         }
 
diff --git a/tools/unit-test/slang-unit-test.cpp b/tools/unit-test/slang-unit-test.cpp
new file mode 100644
index 000000000..a31614c05
--- /dev/null
+++ b/tools/unit-test/slang-unit-test.cpp
@@ -0,0 +1,48 @@
+#include "slang-unit-test.h"
+#include "slang.h"
+#include "source/core/slang-basic.h"
+
+struct SlangUnitTest
+{
+    const char* name;
+    slang::UnitTestFunc func;
+};
+
+class SlangUnitTestModule : public slang::IUnitTestModule
+{
+public:
+    Slang::List<SlangUnitTest> tests;
+
+    virtual SlangInt getTestCount() override
+    {
+        return tests.getCount();
+    }
+    virtual const char* getTestName(SlangInt index) override
+    {
+        return tests[index].name;
+    }
+
+    virtual slang::UnitTestFunc getTestFunc(SlangInt index) override
+    {
+        return tests[index].func;
+    }
+};
+
+SlangUnitTestModule* _getTestModule()
+{
+    static SlangUnitTestModule testModule;
+    return &testModule;
+}
+
+extern "C"
+{
+SLANG_DLL_EXPORT slang::IUnitTestModule* slangUnitTestGetModule()
+{
+    return _getTestModule();
+}
+}
+
+slang::UnitTestRegisterHelper::UnitTestRegisterHelper(const char* name, UnitTestFunc testFunc)
+{
+    _getTestModule()->tests.add(SlangUnitTest{ name, testFunc });
+}
diff --git a/tools/unit-test/slang-unit-test.h b/tools/unit-test/slang-unit-test.h
new file mode 100644
index 000000000..7651e6b46
--- /dev/null
+++ b/tools/unit-test/slang-unit-test.h
@@ -0,0 +1,38 @@
+#pragma once
+
+#include "slang.h"
+#include "source/core/slang-render-api-util.h"
+
+namespace slang
+{
+    struct UnitTestContext
+    {
+        slang::IGlobalSession* slangGlobalSession;
+        const char* workDirectory;
+        ISlangWriter* outputWriter;
+        Slang::RenderApiFlags enabledApis;
+    };
+
+    typedef SlangResult (*UnitTestFunc)(UnitTestContext*);
+
+    class IUnitTestModule
+    {
+    public:
+        virtual SlangInt getTestCount() = 0;
+        virtual const char* getTestName(SlangInt index) = 0;
+        virtual UnitTestFunc getTestFunc(SlangInt index) = 0;
+    };
+
+    class UnitTestRegisterHelper
+    {
+    public:
+        UnitTestRegisterHelper(const char* name, UnitTestFunc testFunc);
+    };
+
+    typedef slang::IUnitTestModule* (*UnitTestGetModuleFunc)();
+
+#define SLANG_UNIT_TEST(name) \
+    SlangResult name(slang::UnitTestContext* context); \
+    slang::UnitTestRegisterHelper _##name##RegisterHelper(#name, name); \
+    SlangResult name(slang::UnitTestContext* context)
+}