From 9edf237eefdf46c286bedac85b02d2e6421ef72c Mon Sep 17 00:00:00 2001 From: Isaac Marovitz Date: Mon, 24 Mar 2025 20:24:39 -0400 Subject: [PATCH] AIR Compilation, Compression, and Packing Signed-off-by: Isaac Marovitz --- XenosRecomp/CMakeLists.txt | 4 +- XenosRecomp/air_compiler.cpp | 109 +++++++++++++++++++++++++++++++++++ XenosRecomp/air_compiler.h | 9 +++ XenosRecomp/main.cpp | 34 ++++++++++- 4 files changed, 152 insertions(+), 4 deletions(-) create mode 100644 XenosRecomp/air_compiler.cpp create mode 100644 XenosRecomp/air_compiler.h diff --git a/XenosRecomp/CMakeLists.txt b/XenosRecomp/CMakeLists.txt index fe2b1d5..818733a 100644 --- a/XenosRecomp/CMakeLists.txt +++ b/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 diff --git a/XenosRecomp/air_compiler.cpp b/XenosRecomp/air_compiler.cpp new file mode 100644 index 0000000..ddc2ad6 --- /dev/null +++ b/XenosRecomp/air_compiler.cpp @@ -0,0 +1,109 @@ +#include "air_compiler.h" + +#include +#include +#include +#include +#include + +std::vector 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(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; +} diff --git a/XenosRecomp/air_compiler.h b/XenosRecomp/air_compiler.h new file mode 100644 index 0000000..917cfd8 --- /dev/null +++ b/XenosRecomp/air_compiler.h @@ -0,0 +1,9 @@ +#pragma once + +struct AirCompiler +{ + static std::vector compile(const std::string& shaderSource); + +private: + static int makeTemporaryFile(std::string& extension); +}; \ No newline at end of file diff --git a/XenosRecomp/main.cpp b/XenosRecomp/main.cpp index d015145..c6e2260 100644 --- a/XenosRecomp/main.cpp +++ b/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 readAllBytes(const char* filePath, size_t& fileSize) { @@ -26,6 +27,7 @@ struct RecompiledShader uint8_t* data = nullptr; IDxcBlob* dxil = nullptr; std::vector spirv; + std::vector air; uint32_t specConstantsMask = 0; }; @@ -133,6 +135,10 @@ int main(int argc, char** argv) assert(*(reinterpret_cast(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 dxil; std::vector spirv; + std::vector 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(shader.dxil->GetBufferPointer()), reinterpret_cast(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 airCompressed(ZSTD_compressBound(air.size())); + airCompressed.reserve(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 spirvCompressed(ZSTD_compressBound(spirv.size()));