2 files changed, 43 insertions, 10 deletions

diff --git a/source/core/slang-test-tool-util.cpp b/source/core/slang-test-tool-util.cpp
index f8e163758..5e31dc4d7 100644
--- a/source/core/slang-test-tool-util.cpp
+++ b/source/core/slang-test-tool-util.cpp
@@ -75,29 +75,54 @@ static SlangResult _addCUDAPrelude(const String& parentPath, slang::IGlobalSessi
     return SLANG_OK;
 }
 
-/* static */SlangResult TestToolUtil::setSessionDefaultPrelude(const char* exePath, slang::IGlobalSession* session)
+/* static */SlangResult TestToolUtil::setSessionDefaultPrelude(const PreludeInfo& info, slang::IGlobalSession* session)
 {
     // Set the prelude to a path
-    String canonicalPath;
-    if (SLANG_SUCCEEDED(Path::getCanonical(exePath, canonicalPath)))
+    if (info.exePath)
     {
-        // Get the directory
-        String parentPath = Path::getParentDirectory(canonicalPath);
+        String exePath(info.exePath);
 
-        if (SLANG_FAILED(_addCPPPrelude(parentPath, session)))
+        String canonicalPath;
+        if (SLANG_SUCCEEDED(Path::getCanonical(exePath, canonicalPath)))
         {
-            SLANG_ASSERT(!"Couldn't find the C++ prelude relative to the executable");
+            // Get the directory
+            String parentPath = Path::getParentDirectory(canonicalPath);
+
+            if (SLANG_FAILED(_addCPPPrelude(parentPath, session)))
+            {
+                SLANG_ASSERT(!"Couldn't find the C++ prelude relative to the executable");
+            }
+
+            if (SLANG_FAILED(_addCUDAPrelude(parentPath, session)))
+            {
+                SLANG_ASSERT(!"Couldn't find the CUDA prelude relative to the executable");
+            }
         }
-
-        if (SLANG_FAILED(_addCUDAPrelude(parentPath, session)))
+    }
+    // If the nvAPI path is set, and we find nvHLSLExtns.h, put that in the HLSL prelude
+    if (info.nvAPIPath)
+    {
+        String includePath;
+        if (SLANG_SUCCEEDED(_calcIncludePath(info.nvAPIPath, "nvHLSLExtns.h", includePath)))
         {
-            SLANG_ASSERT(!"Couldn't find the CUDA prelude relative to the executable");
+            StringBuilder buf;
+
+            buf << "#include \"" << includePath << "\"\n";
+
+            session->setLanguagePrelude(SLANG_SOURCE_LANGUAGE_HLSL, buf.getBuffer());
+            return SLANG_OK;
         }
     }
 
     return SLANG_OK;
 }
 
+/* static */SlangResult TestToolUtil::setSessionDefaultPrelude(const char* exePath, slang::IGlobalSession* session)
+{
+    PreludeInfo info;
+    info.exePath = exePath;
+    return setSessionDefaultPrelude(info, session);
+}
 
 }
 
diff --git a/source/core/slang-test-tool-util.h b/source/core/slang-test-tool-util.h
index 9df2a6d6a..43bd49c9d 100644
--- a/source/core/slang-test-tool-util.h
+++ b/source/core/slang-test-tool-util.h
@@ -36,6 +36,12 @@ enum class ToolReturnCodeSpan
 /* Utility functions for 'test tools' */
 struct TestToolUtil
 {
+    struct PreludeInfo
+    {
+        const char* exePath = nullptr;
+        const char* nvAPIPath = nullptr;
+    };
+
     typedef SlangResult(*InnerMainFunc)(Slang::StdWriters* stdWriters, SlangSession* session, int argc, const char*const* argv);
 
         /// If the test failed to run or was ignored then we are done
@@ -48,6 +54,8 @@ struct TestToolUtil
     static ToolReturnCode getReturnCode(SlangResult res);
 
         /// Sets the default preludes on the session based on the executable path
+    static SlangResult setSessionDefaultPrelude(const PreludeInfo& preludeInfo, slang::IGlobalSession* session);
+
     static SlangResult setSessionDefaultPrelude(const char* exePath, slang::IGlobalSession* session);
 };