nsight Aftermath crash example (#2984)

* Small fixes and improvements around reflection tool.

* Make PrettyWriter printing a class.

* Aftermath crash demo WIP.

* Enable aftermath in test project.

* Setting failCount.

* Dumping out of source maps.

* Improve comments.
Simplify handling of compile products.

* Other small fixes to aftermath example.

* Added Emit SourceLocType.
Track sourcemap association meaning.
Improved documentation.

* Small improvements.

* Capture debug information for D3D11/D3D12/Vulkan.

* Enable debug info.

* Small improvements.

* Improve aftermath example README.md.

7 files changed, 194 insertions, 4 deletions

diff --git a/tools/gfx/d3d/d3d-swapchain.h b/tools/gfx/d3d/d3d-swapchain.h
index 0d4b3fafb..c9e0de82a 100644
--- a/tools/gfx/d3d/d3d-swapchain.h
+++ b/tools/gfx/d3d/d3d-swapchain.h
@@ -99,7 +99,15 @@ public:
     }
     virtual SLANG_NO_THROW Result SLANG_MCALL present() override
     {
-        if (SLANG_FAILED(m_swapChain->Present(m_desc.enableVSync ? 1 : 0, 0)))
+        const auto res = m_swapChain->Present(m_desc.enableVSync ? 1 : 0, 0);
+
+        // We may want to wait for crash dump completion for some kinds of debugging scenarios
+        if (res == DXGI_ERROR_DEVICE_REMOVED || res == DXGI_ERROR_DEVICE_RESET)
+        {
+            D3DUtil::waitForCrashDumpCompletion(res);
+        }
+
+        if (SLANG_FAILED(res))
         {
             return SLANG_FAIL;
         }
diff --git a/tools/gfx/d3d/d3d-util.cpp b/tools/gfx/d3d/d3d-util.cpp
index e1ffc0efc..34d615744 100644
--- a/tools/gfx/d3d/d3d-util.cpp
+++ b/tools/gfx/d3d/d3d-util.cpp
@@ -14,6 +14,14 @@
 #include "core/slang-basic.h"
 #include "core/slang-platform.h"
 
+#ifdef GFX_NV_AFTERMATH
+#   include "GFSDK_Aftermath.h"
+#   include "GFSDK_Aftermath_Defines.h"
+#   include "GFSDK_Aftermath_GpuCrashDump.h"
+
+#   include "core/slang-process.h"
+#endif
+
 namespace gfx {
 using namespace Slang;
 
@@ -870,6 +878,58 @@ Result SLANG_MCALL reportD3DLiveObjects()
     return D3DUtil::reportLiveObjects();
 }
 
+
+/* static */SlangResult D3DUtil::waitForCrashDumpCompletion(HRESULT res)
+{
+    // If it's not a device remove/reset then theres nothing to wait for
+    if (!(res == DXGI_ERROR_DEVICE_REMOVED || res == DXGI_ERROR_DEVICE_RESET))
+    {
+        return SLANG_OK;
+    }
+
+#if GFX_NV_AFTERMATH
+    {
+        GFSDK_Aftermath_CrashDump_Status status = GFSDK_Aftermath_CrashDump_Status_Unknown;
+        if (GFSDK_Aftermath_GetCrashDumpStatus(&status) != GFSDK_Aftermath_Result_Success)
+        {
+            return SLANG_FAIL;
+        }
+
+        const auto startTick = Process::getClockTick();
+        const auto frequency = Process::getClockFrequency();
+
+        float timeOutInSecs = 1.0f;
+
+        uint64_t timeOutTicks = uint64_t(frequency * timeOutInSecs) + 1;
+
+        // Loop while Aftermath crash dump data collection has not finished or
+        // the application is still processing the crash dump data.
+        while (status != GFSDK_Aftermath_CrashDump_Status_CollectingDataFailed &&
+            status != GFSDK_Aftermath_CrashDump_Status_Finished &&
+            Process::getClockTick() - startTick < timeOutTicks)
+        {
+            // Sleep a couple of milliseconds and poll the status again.
+            Process::sleepCurrentThread(50);
+            if (GFSDK_Aftermath_GetCrashDumpStatus(&status) != GFSDK_Aftermath_Result_Success)
+            {
+                return SLANG_FAIL;
+            }
+        }
+
+        if (status == GFSDK_Aftermath_CrashDump_Status_Finished)
+        {
+            return SLANG_OK;
+        }
+        else
+        {
+            return SLANG_E_TIME_OUT;
+        }
+    }
+#endif
+
+    return SLANG_OK;
+}
+
 /* static */SlangResult D3DUtil::findAdapters(DeviceCheckFlags flags, const AdapterLUID* adapterLUID, IDXGIFactory* dxgiFactory, List<ComPtr<IDXGIAdapter>>& outDxgiAdapters)
 {
     outDxgiAdapters.clear();
diff --git a/tools/gfx/d3d/d3d-util.h b/tools/gfx/d3d/d3d-util.h
index a35928f47..ce40ec722 100644
--- a/tools/gfx/d3d/d3d-util.h
+++ b/tools/gfx/d3d/d3d-util.h
@@ -121,6 +121,10 @@ class D3DUtil
     static D3D12_RESOURCE_STATES getResourceState(ResourceState state);
 
     static SlangResult reportLiveObjects();
+
+        /// Call after a DXGI_ERROR_DEVICE_REMOVED/DXGI_ERROR_DEVICE_RESET on present, to wait for 
+        /// dumping to complete. Will return SLANG_OK if wait happened successfully
+    static SlangResult waitForCrashDumpCompletion(HRESULT res);
 };
 
 #if SLANG_GFX_HAS_DXR_SUPPORT
diff --git a/tools/gfx/d3d11/d3d11-device.cpp b/tools/gfx/d3d11/d3d11-device.cpp
index c10a608dc..96a5043fb 100644
--- a/tools/gfx/d3d11/d3d11-device.cpp
+++ b/tools/gfx/d3d11/d3d11-device.cpp
@@ -16,6 +16,12 @@
 
 #include "d3d11-helper-functions.h"
 
+#ifdef GFX_NV_AFTERMATH
+#   include "GFSDK_Aftermath.h"
+#   include "GFSDK_Aftermath_Defines.h"
+#   include "GFSDK_Aftermath_GpuCrashDump.h"
+#endif
+
 namespace gfx
 {
 
@@ -23,6 +29,7 @@ using namespace Slang;
 
 namespace d3d11
 {
+
 SlangResult DeviceImpl::initialize(const Desc& desc)
 {
     SLANG_RETURN_ON_FAIL(slangContext.initialize(
@@ -148,11 +155,42 @@ SlangResult DeviceImpl::initialize(const Desc& desc)
                 m_device.writeRef(),
                 &featureLevel,
                 m_immediateContext.writeRef());
+
+#ifdef GFX_NV_AFTERMATH
+            if (SLANG_SUCCEEDED(res))
+            {
+                if (deviceCheckFlags & DeviceCheckFlag::UseDebug) 
+                {
+                    // Initialize Nsight Aftermath for this device.
+                    // This combination of flags is not necessarily appropriate for real world usage 
+                    const uint32_t aftermathFlags =
+                        GFSDK_Aftermath_FeatureFlags_EnableMarkers |             // Enable event marker tracking.
+                        GFSDK_Aftermath_FeatureFlags_CallStackCapturing |        // Enable automatic call stack event markers.
+                        GFSDK_Aftermath_FeatureFlags_EnableResourceTracking |    // Enable tracking of resources.
+                        GFSDK_Aftermath_FeatureFlags_GenerateShaderDebugInfo |   // Generate debug information for shaders.
+                        GFSDK_Aftermath_FeatureFlags_EnableShaderErrorReporting; // Enable additional runtime shader error reporting.
+
+                    auto initResult = GFSDK_Aftermath_DX11_Initialize(
+                        GFSDK_Aftermath_Version_API,
+                        aftermathFlags,
+                        m_device);
+
+                    if (initResult != GFSDK_Aftermath_Result_Success)
+                    {
+                        SLANG_ASSERT_FAILURE("Unable to initialize aftermath");
+                        // Unable to initialize aftermath
+                        return SLANG_FAIL;
+                    }
+                }
+            }
+#endif
+
             // Check if successfully constructed - if so we are done.
             if (SLANG_SUCCEEDED(res))
             {
                 break;
             }
+
         }
         // If res is failure, means all styles have have failed, and so initialization fails.
         if (SLANG_FAILED(res))
diff --git a/tools/gfx/d3d12/d3d12-device.cpp b/tools/gfx/d3d12/d3d12-device.cpp
index 863326a94..68078b445 100644
--- a/tools/gfx/d3d12/d3d12-device.cpp
+++ b/tools/gfx/d3d12/d3d12-device.cpp
@@ -28,6 +28,12 @@
 #    include "../nvapi/nvapi-include.h"
 #endif
 
+#ifdef GFX_NV_AFTERMATH
+#   include "GFSDK_Aftermath.h"
+#   include "GFSDK_Aftermath_Defines.h"
+#   include "GFSDK_Aftermath_GpuCrashDump.h"
+#endif
+
 namespace gfx
 {
 namespace d3d12
@@ -37,6 +43,13 @@ using namespace Slang;
 
 static const uint32_t D3D_FEATURE_LEVEL_12_2 = 0xc200;
 
+
+#if GFX_NV_AFTERMATH
+/* static */const bool DeviceImpl::g_isAftermathEnabled = true;
+#else
+/* static */const bool DeviceImpl::g_isAftermathEnabled = false;
+#endif
+
 struct ShaderModelInfo
 {
     D3D_SHADER_MODEL shaderModel;
@@ -286,7 +299,7 @@ Result DeviceImpl::_createDevice(
     D3D_FEATURE_LEVEL featureLevel,
     D3D12DeviceInfo& outDeviceInfo)
 {
-    if (m_dxDebug && (deviceCheckFlags & DeviceCheckFlag::UseDebug))
+    if (m_dxDebug && (deviceCheckFlags & DeviceCheckFlag::UseDebug) && !g_isAftermathEnabled)
     {
         m_dxDebug->EnableDebugLayer();
     }
@@ -319,7 +332,7 @@ Result DeviceImpl::_createDevice(
         return SLANG_FAIL;
     }
 
-    if (m_dxDebug && (deviceCheckFlags & DeviceCheckFlag::UseDebug))
+    if (m_dxDebug && (deviceCheckFlags & DeviceCheckFlag::UseDebug) && !g_isAftermathEnabled)
     {
         ComPtr<ID3D12InfoQueue> infoQueue;
         if (SLANG_SUCCEEDED(device->QueryInterface(infoQueue.writeRef())))
@@ -373,6 +386,35 @@ Result DeviceImpl::_createDevice(
         }
     }
 
+
+#ifdef GFX_NV_AFTERMATH
+    {
+        if ((deviceCheckFlags & DeviceCheckFlag::UseDebug) && g_isAftermathEnabled)
+        {
+            // Initialize Nsight Aftermath for this device.
+            // This combination of flags is not necessarily appropraite for real world usage 
+            const uint32_t aftermathFlags =  
+                GFSDK_Aftermath_FeatureFlags_EnableMarkers |                // Enable event marker tracking.
+                GFSDK_Aftermath_FeatureFlags_CallStackCapturing |           // Enable automatic call stack event markers.
+                GFSDK_Aftermath_FeatureFlags_EnableResourceTracking |       // Enable tracking of resources.
+                GFSDK_Aftermath_FeatureFlags_GenerateShaderDebugInfo |      // Generate debug information for shaders.
+                GFSDK_Aftermath_FeatureFlags_EnableShaderErrorReporting;    // Enable additional runtime shader error reporting.
+            
+            auto initResult = GFSDK_Aftermath_DX12_Initialize(
+                GFSDK_Aftermath_Version_API,
+                aftermathFlags,
+                device);
+            
+            if ( initResult != GFSDK_Aftermath_Result_Success)
+            {
+                SLANG_ASSERT_FAILURE("Unable to initialize aftermath");
+                // Unable to initialize 
+                return SLANG_FAIL;
+            }
+        }
+    }
+#endif
+
     // Get the descs
     {
         adapter->GetDesc(&outDeviceInfo.m_desc);
@@ -474,7 +516,9 @@ Result DeviceImpl::initialize(const Desc& desc)
     }
 #endif
 
-    if (ENABLE_DEBUG_LAYER || isGfxDebugLayerEnabled())
+
+    // If Aftermath is enabled, we can't enable the D3D12 debug layer as well
+    if (ENABLE_DEBUG_LAYER || isGfxDebugLayerEnabled() && !g_isAftermathEnabled)
     {
         m_D3D12GetDebugInterface =
             (PFN_D3D12_GET_DEBUG_INTERFACE)loadProc(d3dModule, "D3D12GetDebugInterface");
diff --git a/tools/gfx/d3d12/d3d12-device.h b/tools/gfx/d3d12/d3d12-device.h
index 975ba419b..6bbdde9d0 100644
--- a/tools/gfx/d3d12/d3d12-device.h
+++ b/tools/gfx/d3d12/d3d12-device.h
@@ -56,6 +56,8 @@ public:
 
     ComPtr<ID3D12Debug> m_dxDebug;
 
+    static const bool g_isAftermathEnabled;
+
     D3D12DeviceInfo m_deviceInfo;
     ID3D12Device* m_device = nullptr;
     ID3D12Device5* m_device5 = nullptr;
diff --git a/tools/gfx/vulkan/vk-device.cpp b/tools/gfx/vulkan/vk-device.cpp
index 654353d63..8f7a88886 100644
--- a/tools/gfx/vulkan/vk-device.cpp
+++ b/tools/gfx/vulkan/vk-device.cpp
@@ -18,6 +18,12 @@
 
 #include "vk-helper-functions.h"
 
+#ifdef GFX_NV_AFTERMATH
+#   include "GFSDK_Aftermath.h"
+#   include "GFSDK_Aftermath_Defines.h"
+#   include "GFSDK_Aftermath_GpuCrashDump.h"
+#endif
+
 namespace gfx
 {
 
@@ -658,6 +664,34 @@ Result DeviceImpl::initVulkanInstanceAndDevice(
     m_queueFamilyIndex = m_api.findQueue(VK_QUEUE_GRAPHICS_BIT | VK_QUEUE_COMPUTE_BIT);
     assert(m_queueFamilyIndex >= 0);
 
+#if defined(GFX_NV_AFTERMATH)
+    VkDeviceDiagnosticsConfigCreateInfoNV aftermathInfo = {};
+
+    {
+        // Enable NV_device_diagnostic_checkpoints extension to be able to
+        // use Aftermath event markers.
+        deviceExtensions.add(VK_NV_DEVICE_DIAGNOSTIC_CHECKPOINTS_EXTENSION_NAME);
+
+        // Enable NV_device_diagnostics_config extension to configure Aftermath
+        // features.
+        deviceExtensions.add(VK_NV_DEVICE_DIAGNOSTICS_CONFIG_EXTENSION_NAME);
+
+        // Set up device creation info for Aftermath feature flag configuration.
+        VkDeviceDiagnosticsConfigFlagsNV aftermathFlags =
+            VK_DEVICE_DIAGNOSTICS_CONFIG_ENABLE_AUTOMATIC_CHECKPOINTS_BIT_NV |  // Enable automatic call stack checkpoints.
+            VK_DEVICE_DIAGNOSTICS_CONFIG_ENABLE_RESOURCE_TRACKING_BIT_NV |      // Enable tracking of resources.
+            VK_DEVICE_DIAGNOSTICS_CONFIG_ENABLE_SHADER_DEBUG_INFO_BIT_NV;       // Generate debug information for shaders.
+            // Not available on the version of Vulkan currently building with.
+            //VK_DEVICE_DIAGNOSTICS_CONFIG_ENABLE_SHADER_ERROR_REPORTING_BIT_NV;  // Enable additional runtime shader error reporting.
+
+        aftermathInfo.sType = VK_STRUCTURE_TYPE_DEVICE_DIAGNOSTICS_CONFIG_CREATE_INFO_NV;
+        aftermathInfo.flags = aftermathFlags;
+
+        aftermathInfo.pNext = deviceCreateInfo.pNext;
+        deviceCreateInfo.pNext = &aftermathInfo;
+    }
+#endif
+
     if (handles[2].handleValue == 0)
     {
         float queuePriority = 0.0f;