Fix issue with slang-embed & include ordering (#5680)

* Fix issue with slang-embed & include ordering

* Update CMakeLists.txt

3 files changed, 22 insertions, 4 deletions

diff --git a/prelude/CMakeLists.txt b/prelude/CMakeLists.txt
index 67dc89d99..3b0e2cf46 100644
--- a/prelude/CMakeLists.txt
+++ b/prelude/CMakeLists.txt
@@ -11,7 +11,9 @@ foreach(input ${prelude_headers})
     set(output "${CMAKE_CURRENT_BINARY_DIR}/${input_name}.cpp")
     add_custom_command(
         OUTPUT ${output}
-        COMMAND slang-embed "${input}" ${output}
+        COMMAND
+            slang-embed "${input}" "${CMAKE_CURRENT_LIST_DIR}/../include"
+            ${output}
         DEPENDS ${input} slang-embed
         VERBATIM
     )
diff --git a/prelude/slang-torch-prelude.h b/prelude/slang-torch-prelude.h
index d303c1045..8ece877b6 100644
--- a/prelude/slang-torch-prelude.h
+++ b/prelude/slang-torch-prelude.h
@@ -1,10 +1,13 @@
 // Prelude for PyTorch cpp binding.
 
+// clang-format off
+#include <torch/extension.h>
+// clang-format on
+
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/cuda/CUDAUtils.h>
 #include <stdexcept>
 #include <string>
-#include <torch/extension.h>
 #include <vector>
 
 #ifdef SLANG_LLVM
diff --git a/tools/slang-embed/slang-embed.cpp b/tools/slang-embed/slang-embed.cpp
index eb493be4b..8cba887f3 100644
--- a/tools/slang-embed/slang-embed.cpp
+++ b/tools/slang-embed/slang-embed.cpp
@@ -63,6 +63,7 @@ struct App
 {
     char const* appName = "slang-embed";
     char const* inputPath = nullptr;
+    char const* includeDir = nullptr;
     char const* outputPath = nullptr;
     Slang::HashSet<Slang::String> includedFiles;
 
@@ -85,13 +86,19 @@ struct App
 
         if (argc > 0)
         {
+            includeDir = *argv++;
+            argc--;
+        }
+
+        if (argc > 0)
+        {
             outputPath = *argv++;
             argc--;
         }
 
         if (!inputPath || (argc != 0))
         {
-            fprintf(stderr, "usage: %s inputPath [outputPath]\n", appName);
+            fprintf(stderr, "usage: %s inputPath includeDir [outputPath]\n", appName);
             exit(1);
         }
     }
@@ -137,7 +144,13 @@ struct App
                 auto path =
                     Slang::Path::combine(Slang::Path::getParentDirectory(inputPath), fileName);
                 if (!Slang::File::exists(path))
-                    goto normalProcess;
+                {
+                    // Try looking in the include directory.
+                    path = Slang::Path::combine(includeDir, fileName);
+
+                    if (!Slang::File::exists(path))
+                        goto normalProcess;
+                }
                 processInputFile(outputFile, path.getUnownedSlice());
                 continue;
             }