MSL shader support

Co-authored-by: Isaac Marovitz <[email protected]>

Author: squidbus

XenosRecomp/CMakeLists.txt

@@ -4,9 +4,13 @@ if (WIN32)
     option(XENOS_RECOMP_DXIL "Generate DXIL shader cache" ON)
 endif()
 
+if (APPLE)
+    option(XENOS_RECOMP_AIR "Generate Metal AIR shader cache" ON)
+endif()
+
 set(SMOLV_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../thirdparty/smol-v/source")
 
-add_executable(XenosRecomp 
+add_executable(XenosRecomp
     constant_table.h
     dxc_compiler.cpp
     dxc_compiler.h
@@ -30,13 +34,6 @@ target_precompile_headers(XenosRecomp PRIVATE pch.h)
 
 if (CMAKE_CXX_COMPILER_ID STREQUAL "Clang" OR CMAKE_CXX_COMPILER_ID STREQUAL "AppleClang")
     target_compile_options(XenosRecomp PRIVATE -Wno-switch -Wno-unused-variable -Wno-null-arithmetic -fms-extensions)
-
-    include(CheckCXXSymbolExists)
-    check_cxx_symbol_exists(_LIBCPP_VERSION version LIBCPP)
-    if(LIBCPP)
-        # Allows using std::execution
-        target_compile_options(XenosRecomp PRIVATE -fexperimental-library)
-    endif()
 endif()
 
 if (WIN32)
@@ -51,3 +48,8 @@ if (XENOS_RECOMP_DXIL)
     target_compile_definitions(XenosRecomp PRIVATE XENOS_RECOMP_DXIL)
     target_link_libraries(XenosRecomp PRIVATE Microsoft::DXIL)
 endif()
+
+if (XENOS_RECOMP_AIR)
+    target_compile_definitions(XenosRecomp PRIVATE XENOS_RECOMP_AIR)
+    target_sources(XenosRecomp PRIVATE air_compiler.cpp air_compiler.h)
+endif()

XenosRecomp/air_compiler.cpp

@@ -0,0 +1,81 @@
+#include "air_compiler.h"
+
+#include <fstream>
+#include <iterator>
+#include <spawn.h>
+#include <unistd.h>
+
+struct TemporaryPath
+{
+    const std::string path;
+
+    explicit TemporaryPath(std::string_view path) : path(path) {}
+
+    ~TemporaryPath()
+    {
+        unlink(path.c_str());
+    }
+};
+
+static int executeCommand(const char** argv)
+{
+    pid_t pid;
+    if (posix_spawn(&pid, argv[0], nullptr, nullptr, const_cast<char**>(argv), nullptr) != 0)
+        return -1;
+
+    int status;
+    if (waitpid(pid, &status, 0) == -1)
+        return -1;
+
+    return status;
+}
+
+std::vector<uint8_t> AirCompiler::compile(const std::string& shaderSource)
+{
+    // Save source to a location on disk for the compiler to read.
+    char sourcePathTemplate[PATH_MAX] = "/tmp/xenos_metal_XXXXXX.metal";
+    const int sourceFd = mkstemps(sourcePathTemplate, 6);
+    if (sourceFd == -1)
+    {
+        fmt::println("Failed to create temporary file for shader source: {}", strerror(errno));
+        std::exit(1);
+    }
+
+    const TemporaryPath sourcePath(sourcePathTemplate);
+    const TemporaryPath irPath(sourcePath.path + ".ir");
+    const TemporaryPath metalLibPath(sourcePath.path + ".metallib");
+
+    const ssize_t sourceWritten = write(sourceFd, shaderSource.data(), shaderSource.size());
+    close(sourceFd);
+    if (sourceWritten < 0)
+    {
+        fmt::println("Failed to write shader source to disk: {}", strerror(errno));
+        std::exit(1);
+    }
+
+    const char* compileCommand[] = {
+        "/usr/bin/xcrun", "-sdk", "macosx", "metal", "-o", irPath.path.c_str(), "-c", sourcePath.path.c_str(), "-Wno-unused-variable", "-frecord-sources", "-gline-tables-only", "-D__air__",
+#ifdef UNLEASHED_RECOMP
+        "-DUNLEASHED_RECOMP",
+#endif
+        nullptr
+    };
+    if (const int compileStatus = executeCommand(compileCommand); compileStatus != 0)
+    {
+        fmt::println("Metal compiler exited with status: {}", compileStatus);
+        fmt::println("Generated source:\n{}", shaderSource);
+        std::exit(1);
+    }
+
+    const char* linkCommand[] = { "/usr/bin/xcrun", "-sdk", "macosx", "metallib", "-o", metalLibPath.path.c_str(), irPath.path.c_str(), nullptr };
+    if (const int linkStatus = executeCommand(linkCommand); linkStatus != 0)
+    {
+        fmt::println("Metal linker exited with status: {}", linkStatus);
+        fmt::println("Generated source:\n{}", shaderSource);
+        std::exit(1);
+    }
+
+    std::ifstream libStream(metalLibPath.path, std::ios::binary);
+    std::vector<uint8_t> data((std::istreambuf_iterator(libStream)), std::istreambuf_iterator<char>());
+    return data;
+}

XenosRecomp/air_compiler.h

@@ -0,0 +1,10 @@
+#pragma once
+
+#include <string>
+#include <vector>
+
+class AirCompiler
+{
+public:
+    [[nodiscard]] static std::vector<uint8_t> compile(const std::string& shaderSource);
+};

XenosRecomp/dxc_compiler.cpp

@@ -34,6 +34,11 @@ IDxcBlob* DxcCompiler::compile(const std::string& shaderSource, bool compilePixe
             target = L"-T vs_6_0";
     }
 
+    if (!compileLibrary)
+    {
+        args[argCount++] = L"-E shaderMain";
+    }
+
     args[argCount++] = target;
     args[argCount++] = L"-HV 2021";
     args[argCount++] = L"-all-resources-bound";

XenosRecomp/main.cpp

@@ -1,7 +1,15 @@
+#include <deque>
+#include <mutex>
+#include <thread>
+
 #include "shader.h"
 #include "shader_recompiler.h"
 #include "dxc_compiler.h"
 
+#ifdef XENOS_RECOMP_AIR
+#include "air_compiler.h"
+#endif
+
 static std::unique_ptr<uint8_t[]> readAllBytes(const char* filePath, size_t& fileSize)
 {
     FILE* file = fopen(filePath, "rb");
@@ -26,9 +34,43 @@ struct RecompiledShader
     uint8_t* data = nullptr;
     IDxcBlob* dxil = nullptr;
     std::vector<uint8_t> spirv;
+    std::vector<uint8_t> air;
     uint32_t specConstantsMask = 0;
 };
 
+void recompileShader(RecompiledShader& shader, const std::string_view include, std::atomic<uint32_t>& progress, uint32_t numShaders)
+{
+    thread_local ShaderRecompiler recompiler;
+    recompiler = {};
+    recompiler.recompile(shader.data, include);
+
+    shader.specConstantsMask = recompiler.specConstantsMask;
+
+    thread_local DxcCompiler dxcCompiler;
+
+#ifdef XENOS_RECOMP_DXIL
+    shader.dxil = dxcCompiler.compile(recompiler.out, recompiler.isPixelShader, recompiler.specConstantsMask != 0, false);
+    assert(shader.dxil != nullptr);
+    assert(*(reinterpret_cast<uint32_t *>(shader.dxil->GetBufferPointer()) + 1) != 0 && "DXIL was not signed properly!");
+#endif
+
+#ifdef XENOS_RECOMP_AIR
+    shader.air = AirCompiler::compile(recompiler.out);
+#endif
+
+    IDxcBlob* spirv = dxcCompiler.compile(recompiler.out, recompiler.isPixelShader, false, true);
+    assert(spirv != nullptr);
+
+    bool result = smolv::Encode(spirv->GetBufferPointer(), spirv->GetBufferSize(), shader.spirv, smolv::kEncodeFlagStripDebugInfo);
+    assert(result);
+
+    spirv->Release();
+
+    size_t currentProgress = ++progress;
+    if ((currentProgress % 10) == 0 || (currentProgress == numShaders - 1))
+        fmt::println("Recompiling shaders... {}%", currentProgress / float(numShaders) * 100.0f);
+}
+
 int main(int argc, char** argv)
 {
 #ifndef XENOS_RECOMP_INPUT
@@ -71,6 +113,7 @@ int main(int argc, char** argv)
     {
         std::vector<std::unique_ptr<uint8_t[]>> files;
         std::map<XXH64_hash_t, RecompiledShader> shaders;
+        std::map<XXH64_hash_t, std::string> shaderFilenames;
 
         for (auto& file : std::filesystem::recursive_directory_iterator(input))
         {
@@ -99,6 +142,7 @@ int main(int argc, char** argv)
                     {
                         shader.first->second.data = fileData.get() + i;
                         foundAny = true;
+                        shaderFilenames[hash] = file.path().string();
                     }
 
                     i += dataSize;
@@ -113,38 +157,42 @@ int main(int argc, char** argv)
                 files.emplace_back(std::move(fileData));
         }
 
-        std::atomic<uint32_t> progress = 0;
-
-        std::for_each(std::execution::par_unseq, shaders.begin(), shaders.end(), [&](auto& hashShaderPair)
-            {
-                auto& shader = hashShaderPair.second;
-
-                thread_local ShaderRecompiler recompiler;
-                recompiler = {};
-                recompiler.recompile(shader.data, include);
-
-                shader.specConstantsMask = recompiler.specConstantsMask;
-
-                thread_local DxcCompiler dxcCompiler;
-
-#ifdef XENOS_RECOMP_DXIL
-                shader.dxil = dxcCompiler.compile(recompiler.out, recompiler.isPixelShader, recompiler.specConstantsMask != 0, false);
-                assert(shader.dxil != nullptr);
-                assert(*(reinterpret_cast<uint32_t *>(shader.dxil->GetBufferPointer()) + 1) != 0 && "DXIL was not signed properly!");
-#endif
-
-                IDxcBlob* spirv = dxcCompiler.compile(recompiler.out, recompiler.isPixelShader, false, true);
-                assert(spirv != nullptr);
-
-                bool result = smolv::Encode(spirv->GetBufferPointer(), spirv->GetBufferSize(), shader.spirv, smolv::kEncodeFlagStripDebugInfo);
-                assert(result);
+        std::mutex shaderQueueMutex;
+        std::deque<XXH64_hash_t> shaderQueue;
+        for (const auto& [hash, _] : shaders)
+        {
+            shaderQueue.emplace_back(hash);
+        }
 
-                spirv->Release();
+        const uint32_t numThreads = std::max(std::thread::hardware_concurrency(), 1u);
+        fmt::println("Recompiling shaders with {} threads", numThreads);
 
-                size_t currentProgress = ++progress;
-                if ((currentProgress % 10) == 0 || (currentProgress == shaders.size() - 1))
-                    fmt::println("Recompiling shaders... {}%", currentProgress / float(shaders.size()) * 100.0f);
+        std::atomic<uint32_t> progress = 0;
+        std::vector<std::thread> threads;
+        threads.reserve(numThreads);
+        for (uint32_t i = 0; i < numThreads; i++)
+        {
+            threads.emplace_back([&]
+            {
+                while (true)
+                {
+                    XXH64_hash_t shaderHash;
+                    {
+                        std::lock_guard lock(shaderQueueMutex);
+                        if (shaderQueue.empty()) {
+                            return;
+                        }
+                        shaderHash = shaderQueue.front();
+                        shaderQueue.pop_front();
+                    }
+                    recompileShader(shaders[shaderHash], include, progress, shaders.size());
+                }
             });
+        }
+        for (auto& thread : threads)
+        {
+            thread.join();
+        }
 
         fmt::println("Creating shader cache...");
 
@@ -154,18 +202,32 @@ int main(int argc, char** argv)
 
         std::vector<uint8_t> dxil;
         std::vector<uint8_t> spirv;
+        std::vector<uint8_t> air;
 
         for (auto& [hash, shader] : shaders)
         {
-            f.println("\t{{ 0x{:X}, {}, {}, {}, {}, {} }},",
-                hash, dxil.size(), (shader.dxil != nullptr) ? shader.dxil->GetBufferSize() : 0, spirv.size(), shader.spirv.size(), shader.specConstantsMask);
+            const std::string& fullFilename = shaderFilenames[hash];
+            std::string filename = fullFilename;
+            size_t shaderPos = filename.find("shader");
+            if (shaderPos != std::string::npos) {
+                filename = filename.substr(shaderPos);
+                // Prevent bad escape sequences in Windows shader path.
+                std::replace(filename.begin(), filename.end(), '\\', '/');
+            }
+            f.println("\t{{ 0x{:X}, {}, {}, {}, {}, {}, {}, {}, \"{}\" }},",
+                hash, dxil.size(), (shader.dxil != nullptr) ? shader.dxil->GetBufferSize() : 0,
+                spirv.size(), shader.spirv.size(), air.size(), shader.air.size(), shader.specConstantsMask, filename);
 
             if (shader.dxil != nullptr)
             {
                 dxil.insert(dxil.end(), reinterpret_cast<uint8_t *>(shader.dxil->GetBufferPointer()),
                     reinterpret_cast<uint8_t *>(shader.dxil->GetBufferPointer()) + shader.dxil->GetBufferSize());
             }
-            
+
+#ifdef XENOS_RECOMP_AIR
+            air.insert(air.end(), shader.air.begin(), shader.air.end());
+#endif
+
             spirv.insert(spirv.end(), shader.spirv.begin(), shader.spirv.end());
         }
 
@@ -189,6 +251,22 @@ int main(int argc, char** argv)
         f.println("const size_t g_dxilCacheDecompressedSize = {};", dxil.size());
 #endif
 
+#ifdef XENOS_RECOMP_AIR
+        fmt::println("Compressing AIR cache...");
+
+        std::vector<uint8_t> airCompressed(ZSTD_compressBound(air.size()));
+        airCompressed.resize(ZSTD_compress(airCompressed.data(), airCompressed.size(), air.data(), air.size(), level));
+
+        f.print("const uint8_t g_compressedAirCache[] = {{");
+
+        for (auto data : airCompressed)
+            f.print("{},", data);
+
+        f.println("}};");
+        f.println("const size_t g_airCacheCompressedSize = {};", airCompressed.size());
+        f.println("const size_t g_airCacheDecompressedSize = {};", air.size());
+#endif
+
         fmt::println("Compressing SPIRV cache...");
 
         std::vector<uint8_t> spirvCompressed(ZSTD_compressBound(spirv.size()));

XenosRecomp/pch.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #ifdef _WIN32
+#define NOMINMAX
 #include <Windows.h>
 #endif