AIR Compilation, Compression, and Packing

Signed-off-by: Isaac Marovitz <isaacryu@icloud.com>

XenosRecomp/CMakeLists.txt

@@ -10,8 +10,10 @@ endif()
 
 set(SMOLV_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../thirdparty/smol-v/source")
 
-add_executable(XenosRecomp 
+add_executable(XenosRecomp
     constant_table.h
+    air_compiler.cpp
+    air_compiler.h
     dxc_compiler.cpp
     dxc_compiler.h
     main.cpp

XenosRecomp/air_compiler.cpp

@@ -0,0 +1,109 @@
+#include "air_compiler.h"
+
+#include <iostream>
+#include <spawn.h>
+#include <unistd.h>
+#include <sys/wait.h>
+#include <cstdio>
+
+std::vector<uint8_t> AirCompiler::compile(const std::string& shaderSource) {
+    // First, generate AIR from shader source
+    std::string inputFile = ".metal";
+    int tmpFD = makeTemporaryFile(inputFile);
+    write(tmpFD, shaderSource.data(), shaderSource.size());
+    close(tmpFD);
+
+    std::string irFile = ".ir";
+    tmpFD = makeTemporaryFile(irFile);
+    close(tmpFD);
+
+    pid_t pid;
+    char* airArgv[] = { "xcrun", "-sdk", "macosx", "metal", "-o", irFile.data(), "-c", inputFile.data(), "-D__air__", "-DUNLEASHED_RECOMP", "-Wno-unused-variable", nullptr };
+
+    if (posix_spawn(&pid, "/usr/bin/xcrun", nullptr, nullptr, airArgv, nullptr) != 0) {
+        unlink(inputFile.data());
+        unlink(irFile.data());
+        fmt::println("Failed to spawn AIR xcrun process");
+        exit(1);
+    }
+
+    int status;
+
+    if (waitpid(pid, &status, 0) == -1) {
+        unlink(inputFile.data());
+        unlink(irFile.data());
+        fmt::println("Failed to wait AIR xcrun process");
+        exit(1);
+    }
+
+    if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
+        unlink(inputFile.data());
+        unlink(irFile.data());
+        fmt::println("AIR xcrun exited with code {}", WEXITSTATUS(status));
+        fmt::println("{}", shaderSource);
+        exit(1);
+    }
+
+    unlink(inputFile.data());
+
+    // Now we need to turn the AIR into a .metallib
+    std::string libFile = ".metallib";
+    tmpFD = makeTemporaryFile(libFile);
+    close(tmpFD);
+
+    char* libArgv[] = { "xcrun", "-sdk", "macosx", "metallib", "-o", libFile.data(), irFile.data(), nullptr };
+
+    if (posix_spawn(&pid, "/usr/bin/xcrun", nullptr, nullptr, libArgv, nullptr) != 0) {
+        unlink(irFile.data());
+        unlink(libFile.data());
+        fmt::println("Failed to spawn .metallib xcrun process");
+        exit(1);
+    }
+
+    if (waitpid(pid, &status, 0) == -1) {
+        unlink(irFile.data());
+        unlink(libFile.data());
+        fmt::println("Failed to wait .metallib xcrun process");
+        exit(1);
+    }
+
+    if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
+        unlink(irFile.data());
+        unlink(libFile.data());
+        fmt::println(".metallib exited with code {}", WEXITSTATUS(status));
+        exit(1);
+    }
+
+    // Read from built .metallib
+    FILE* file = fopen(libFile.data(), "rb");
+    fseek(file, 0, SEEK_END);
+    size_t fileSize = ftell(file);
+    fseek(file, 0, SEEK_SET);
+    auto data = std::vector<uint8_t>(fileSize);
+    fread(data.data(), 1, fileSize, file);
+    fclose(file);
+
+    // Cleanup temporary files
+    unlink(irFile.data());
+    unlink(libFile.data());
+
+    return data;
+}
+
+int AirCompiler::makeTemporaryFile(std::string &extension) {
+    const std::string path = "/tmp/xenos_metal_XXXXXX";
+
+    size_t size = path.size() + extension.size() + 1;
+    char fullTemplate[size];
+    snprintf(fullTemplate, size, "%s%s", path.c_str(), extension.c_str());
+
+    int tmpFD = mkstemps(fullTemplate, extension.size());
+    if (tmpFD == -1) {
+        fmt::println("Failed to open temporary file, \"{}\"!", std::string(fullTemplate));
+        unlink(fullTemplate);
+        exit(1);
+    }
+
+    extension = fullTemplate;
+    return tmpFD;
+}

XenosRecomp/air_compiler.h

@@ -0,0 +1,9 @@
+#pragma once
+
+struct AirCompiler
+{
+    static std::vector<uint8_t> compile(const std::string& shaderSource);
+
+private:
+    static int makeTemporaryFile(std::string& extension);
+};

XenosRecomp/main.cpp

@@ -1,6 +1,7 @@
 #include "shader.h"
 #include "shader_recompiler.h"
 #include "dxc_compiler.h"
+#include "air_compiler.h"
 
 static std::unique_ptr<uint8_t[]> readAllBytes(const char* filePath, size_t& fileSize)
 {
@@ -26,6 +27,7 @@ struct RecompiledShader
     uint8_t* data = nullptr;
     IDxcBlob* dxil = nullptr;
     std::vector<uint8_t> spirv;
+    std::vector<uint8_t> air;
     uint32_t specConstantsMask = 0;
 };
 
@@ -133,6 +135,10 @@ int main(int argc, char** argv)
                 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);
 
@@ -154,18 +160,24 @@ 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);
+            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);
 
             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 +201,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()));