video_core: Implement an arm64 shader-jit backend (#7002)

* externals: Add oaksim submodule

Used for emitting ARM64 assembly

* common: Implement aarch64 ABI

Utilize oaknut to implement a stack frame.

* tests: Allow shader-jit tests for x64 and a64

Run the shader-jit tests for both x86_64 and arm64 targets

* video_core: Initialize arm64 shader-jit backend

Passes all current unit tests!

* shader_jit_a64: protect/unprotect memory when jit-ing

Required on MacOS. Memory needs to be fully unprotected and then
re-protected when writing or there will be memory access errors on
MacOS.

* shader_jit_a64: Fix ARM64-Imm overflow

These conditionals were throwing exceptions since the immediate values
were overflowing the available space in the `EOR` instructions. Instead
they are generated from `MOV` and then `EOR`-ed after.

* shader_jit_a64: Fix Geometry shader conditional

* shader_jit_a64: Replace `ADRL` with `MOVP2R`

Fixes some immediate-generation exceptions.

* common/aarch64: Fix CallFarFunction

* shader_jit_a64: Optimize `SantitizedMul`

Co-authored-by: merryhime <merryhime@users.noreply.github.com>

* shader_jit_a64: Fix address register offset behavior

Based on https://github.com/citra-emu/citra/pull/6942
Passes unit tests.

* shader_jit_a64: Fix `RET` address offset

A64 stack is 16-byte aligned rather than 8. So a direct port of the x64
code won't work. Fixes weird branches into invalid memory for any
shaders with subroutines.

* shader_jit_a64: Increase max program size

Tuned for A64 program size.

* shader_jit_a64: Use `UBFX` for extracting loop-state

Co-authored-by: JosJuice <JosJuice@users.noreply.github.com>

* shader_jit_a64: Optimize `SUB+CMP` to `SUBS`

Co-authored-by: JosJuice <JosJuice@users.noreply.github.com>

* shader_jit_a64: Optimize `CMP+B` to `CBNZ`

Co-authored-by: JosJuice <JosJuice@users.noreply.github.com>

* shader_jit_a64: Use `FMOV` for `ONE` vector

Co-authored-by: JosJuice <JosJuice@users.noreply.github.com>

* shader_jit_a64: Remove x86-specific documentation

* shader_jit_a64: Use `UBFX` to extract exponent

Co-authored-by: JosJuice <JosJuice@users.noreply.github.com>

* shader_jit_a64: Remove redundant MIN/MAX `SRC2`-NaN check

Special handling only needs to check SRC1 for NaN, not SRC2.
It would work as follows in the four possible cases:

No NaN: No special handling needed.
Only SRC1 is NaN: The special handling is triggered because SRC1 is NaN, and SRC2 is picked.
Only SRC2 is NaN: FMAX automatically picks SRC2 because it always picks the NaN if there is one.
Both SRC1 and SRC2 are NaN: The special handling is triggered because SRC1 is NaN, and SRC2 is picked.

Co-authored-by: JosJuice <JosJuice@users.noreply.github.com>

* shader_jit/tests:: Add catch-stringifier for vec2f/vec3f

* shader_jit/tests: Add Dest Mask unit test

* shader_jit_a64: Fix Dest-Mask `BSL` operand order

Passes the dest-mask unit tests now.

* shader_jit_a64: Use `MOVI` for DestEnable mask

Accelerate certain cases of masking with MOVI as well

Co-authored-by: JosJuice <JosJuice@users.noreply.github.com>

* shader_jit/tests: Add source-swizzle unit test

This is not expansive. Generating all `4^4` cases seems to make Catch2
crash. So I've added some component-masking(non-reordering) tests based
on the Dest-Mask unit-test and some additional ones to test
broadcasts/splats and component re-ordering.

* shader_jit_a64: Fix swizzle index generation

This was still generating `SHUFPS` indices and not the ones that we wanted for the `TBL` instruction. Passes all unit tests now.

* shader_jit/tests: Add `ShaderSetup` constructor to `ShaderTest`

Rather than using the direct output of `CompileShaderSetup` allow a
`ShaderSetup` object to be passed in directly.  This enabled the ability
emit assembly that is not directly supported by nihstro.

* shader_jit/tests: Add `CALL` unit-test

Tests nested `CALL` instructions to eventually reach an `EX2`
instruction.

EX2 is picked in particular since it is implemented as an even deeper
dispatch and ensures subroutines are properly implemented between `CALL`
instructions and implementation-calls.

* shader_jit_a64: Fix nested `BL` subroutines

`lr` was getting writen over by nested calls to `BL`, causing undefined
behavior with mixtures of `CALL`, `EX2`, and `LG2` instructions.

Each usage of `BL` is now protected with a stach push/pop to preserve
and restore teh `lr` register to allow nested subroutines to work
properly.

* shader_jit/tests: Allocate generated tests on heap

Each of these generated shader-test objects were causing the stack to
overflow.  Allocate each of the generated tests on the heap and use
unique_ptr so they only exist within the life-time of the `REQUIRE`
statement.

* shader_jit_a64: Preserve `lr` register from external function calls

`EMIT` makes an external function call, and should be preserving `lr`

* shader_jit/tests: Add `MAD` unit-test

The Inline Asm version requires an upstream fix:
https://github.com/neobrain/nihstro/issues/68

Instead, the program code is manually configured and added.

* shader_jit/tests: Fix uninitialized instructions

These `union`-type instruction-types were uninitialized, causing tests
to indeterminantly fail at times.

* shader_jit_a64: Remove unneeded `MOV`

Residue from the direct-port of x64 code.

* shader_jit_a64: Use `std::array` for `instr_table`

Add some type-safety and const-correctness around this type as well.

* shader_jit_a64: Avoid c-style offset casting

Add some more const-correctness to this function as well.

* video_core: Add arch preprocessor comments

* common/aarch64: Use X16 as the veneer register

https://developer.arm.com/documentation/102374/0101/Procedure-Call-Standard

* shader_jit/tests: Add uniform reading unit-test

Particularly to ensure that addresses are being properly truncated

* common/aarch64: Use `X0` as `ABI_RETURN`

`X8` is used as the indirect return result value in the case that the
result is bigger than 128-bits. Principally `X0` is the general-case
return register though.

* common/aarch64: Add veneer register note

`LR` is generally overwritten by `BLR` anyways, and would also be a safe
veneer to utilize for far-calls.

* shader_jit_a64: Remove unneeded scratch register from `SanitizedMul`

* shader_jit_a64: Fix CALLU condition

Should be `EQ` not `NE`. Fixes the regression on Kid Icarus.
No known regressions anymore!

---------

Co-authored-by: merryhime <merryhime@users.noreply.github.com>
Co-authored-by: JosJuice <JosJuice@users.noreply.github.com>

src/video_core/CMakeLists.txt

@@ -149,6 +149,10 @@ add_library(video_core STATIC
     shader/shader.h
     shader/shader_interpreter.cpp
     shader/shader_interpreter.h
+    shader/shader_jit_a64.cpp
+    shader/shader_jit_a64_compiler.cpp
+    shader/shader_jit_a64.h
+    shader/shader_jit_a64_compiler.h
     shader/shader_jit_x64.cpp
     shader/shader_jit_x64_compiler.cpp
     shader/shader_jit_x64.h
@@ -177,6 +181,10 @@ if ("x86_64" IN_LIST ARCHITECTURE)
     target_link_libraries(video_core PUBLIC xbyak)
 endif()
 
+if ("arm64" IN_LIST ARCHITECTURE)
+    target_link_libraries(video_core PUBLIC oaknut)
+endif()
+
 if (CITRA_USE_PRECOMPILED_HEADERS)
     target_precompile_headers(video_core PRIVATE precompiled_headers.h)
 endif()

src/video_core/shader/shader.cpp

@@ -15,7 +15,9 @@
 #include "video_core/shader/shader_interpreter.h"
 #if CITRA_ARCH(x86_64)
 #include "video_core/shader/shader_jit_x64.h"
-#endif // CITRA_ARCH(x86_64)
+#elif CITRA_ARCH(arm64)
+#include "video_core/shader/shader_jit_a64.h"
+#endif
 #include "video_core/video_core.h"
 
 namespace Pica::Shader {
@@ -141,27 +143,29 @@ MICROPROFILE_DEFINE(GPU_Shader, "GPU", "Shader", MP_RGB(50, 50, 240));
 
 #if CITRA_ARCH(x86_64)
 static std::unique_ptr<JitX64Engine> jit_engine;
-#endif // CITRA_ARCH(x86_64)
+#elif CITRA_ARCH(arm64)
+static std::unique_ptr<JitA64Engine> jit_engine;
+#endif
 static InterpreterEngine interpreter_engine;
 
 ShaderEngine* GetEngine() {
-#if CITRA_ARCH(x86_64)
+#if CITRA_ARCH(x86_64) || CITRA_ARCH(arm64)
     // TODO(yuriks): Re-initialize on each change rather than being persistent
     if (VideoCore::g_shader_jit_enabled) {
         if (jit_engine == nullptr) {
-            jit_engine = std::make_unique<JitX64Engine>();
+            jit_engine = std::make_unique<decltype(jit_engine)::element_type>();
         }
         return jit_engine.get();
     }
-#endif // CITRA_ARCH(x86_64)
+#endif // CITRA_ARCH(x86_64) || CITRA_ARCH(arm64)
 
     return &interpreter_engine;
 }
 
 void Shutdown() {
-#if CITRA_ARCH(x86_64)
+#if CITRA_ARCH(x86_64) || CITRA_ARCH(arm64)
     jit_engine = nullptr;
-#endif // CITRA_ARCH(x86_64)
+#endif // CITRA_ARCH(x86_64) || CITRA_ARCH(arm64)
 }
 
 } // namespace Pica::Shader

src/video_core/shader/shader_jit_a64.cpp

@@ -0,0 +1,51 @@
+// Copyright 2023 Citra Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#include "common/arch.h"
+#if CITRA_ARCH(arm64)
+
+#include "common/assert.h"
+#include "common/microprofile.h"
+#include "video_core/shader/shader.h"
+#include "video_core/shader/shader_jit_a64.h"
+#include "video_core/shader/shader_jit_a64_compiler.h"
+
+namespace Pica::Shader {
+
+JitA64Engine::JitA64Engine() = default;
+JitA64Engine::~JitA64Engine() = default;
+
+void JitA64Engine::SetupBatch(ShaderSetup& setup, unsigned int entry_point) {
+    ASSERT(entry_point < MAX_PROGRAM_CODE_LENGTH);
+    setup.engine_data.entry_point = entry_point;
+
+    u64 code_hash = setup.GetProgramCodeHash();
+    u64 swizzle_hash = setup.GetSwizzleDataHash();
+
+    u64 cache_key = code_hash ^ swizzle_hash;
+    auto iter = cache.find(cache_key);
+    if (iter != cache.end()) {
+        setup.engine_data.cached_shader = iter->second.get();
+    } else {
+        auto shader = std::make_unique<JitShader>();
+        shader->Compile(&setup.program_code, &setup.swizzle_data);
+        setup.engine_data.cached_shader = shader.get();
+        cache.emplace_hint(iter, cache_key, std::move(shader));
+    }
+}
+
+MICROPROFILE_DECLARE(GPU_Shader);
+
+void JitA64Engine::Run(const ShaderSetup& setup, UnitState& state) const {
+    ASSERT(setup.engine_data.cached_shader != nullptr);
+
+    MICROPROFILE_SCOPE(GPU_Shader);
+
+    const JitShader* shader = static_cast<const JitShader*>(setup.engine_data.cached_shader);
+    shader->Run(setup, state, setup.engine_data.entry_point);
+}
+
+} // namespace Pica::Shader
+
+#endif // CITRA_ARCH(arm64)

src/video_core/shader/shader_jit_a64.h

@@ -0,0 +1,33 @@
+// Copyright 2023 Citra Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include "common/arch.h"
+#if CITRA_ARCH(arm64)
+
+#include <memory>
+#include <unordered_map>
+#include "common/common_types.h"
+#include "video_core/shader/shader.h"
+
+namespace Pica::Shader {
+
+class JitShader;
+
+class JitA64Engine final : public ShaderEngine {
+public:
+    JitA64Engine();
+    ~JitA64Engine() override;
+
+    void SetupBatch(ShaderSetup& setup, unsigned int entry_point) override;
+    void Run(const ShaderSetup& setup, UnitState& state) const override;
+
+private:
+    std::unordered_map<u64, std::unique_ptr<JitShader>> cache;
+};
+
+} // namespace Pica::Shader
+
+#endif // CITRA_ARCH(arm64)

src/video_core/shader/shader_jit_a64_compiler.h

@@ -0,0 +1,146 @@
+// Copyright 2023 Citra Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include "common/arch.h"
+#if CITRA_ARCH(arm64)
+
+#include <array>
+#include <bitset>
+#include <cstddef>
+#include <optional>
+#include <utility>
+#include <vector>
+#include <nihstro/shader_bytecode.h>
+#include <oaknut/code_block.hpp>
+#include <oaknut/oaknut.hpp>
+#include "common/common_types.h"
+#include "video_core/shader/shader.h"
+
+using nihstro::Instruction;
+using nihstro::OpCode;
+using nihstro::SourceRegister;
+using nihstro::SwizzlePattern;
+
+namespace Pica::Shader {
+
+/// Memory allocated for each compiled shader
+constexpr std::size_t MAX_SHADER_SIZE = MAX_PROGRAM_CODE_LENGTH * 256;
+
+/**
+ * This class implements the shader JIT compiler. It recompiles a Pica shader program into x86_64
+ * code that can be executed on the host machine directly.
+ */
+class JitShader : private oaknut::CodeBlock, public oaknut::CodeGenerator {
+public:
+    JitShader();
+
+    void Run(const ShaderSetup& setup, UnitState& state, unsigned offset) const {
+        program(&setup.uniforms, &state, instruction_labels[offset].ptr<const std::byte*>());
+    }
+
+    void Compile(const std::array<u32, MAX_PROGRAM_CODE_LENGTH>* program_code,
+                 const std::array<u32, MAX_SWIZZLE_DATA_LENGTH>* swizzle_data);
+
+    void Compile_ADD(Instruction instr);
+    void Compile_DP3(Instruction instr);
+    void Compile_DP4(Instruction instr);
+    void Compile_DPH(Instruction instr);
+    void Compile_EX2(Instruction instr);
+    void Compile_LG2(Instruction instr);
+    void Compile_MUL(Instruction instr);
+    void Compile_SGE(Instruction instr);
+    void Compile_SLT(Instruction instr);
+    void Compile_FLR(Instruction instr);
+    void Compile_MAX(Instruction instr);
+    void Compile_MIN(Instruction instr);
+    void Compile_RCP(Instruction instr);
+    void Compile_RSQ(Instruction instr);
+    void Compile_MOVA(Instruction instr);
+    void Compile_MOV(Instruction instr);
+    void Compile_NOP(Instruction instr);
+    void Compile_END(Instruction instr);
+    void Compile_BREAKC(Instruction instr);
+    void Compile_CALL(Instruction instr);
+    void Compile_CALLC(Instruction instr);
+    void Compile_CALLU(Instruction instr);
+    void Compile_IF(Instruction instr);
+    void Compile_LOOP(Instruction instr);
+    void Compile_JMP(Instruction instr);
+    void Compile_CMP(Instruction instr);
+    void Compile_MAD(Instruction instr);
+    void Compile_EMIT(Instruction instr);
+    void Compile_SETE(Instruction instr);
+
+private:
+    void Compile_Block(unsigned end);
+    void Compile_NextInstr();
+
+    void Compile_SwizzleSrc(Instruction instr, unsigned src_num, SourceRegister src_reg,
+                            oaknut::QReg dest);
+    void Compile_DestEnable(Instruction instr, oaknut::QReg dest);
+
+    /**
+     * Compiles a `MUL src1, src2` operation, properly handling the PICA semantics when multiplying
+     * zero by inf. Clobbers `src2` and `scratch`.
+     */
+    void Compile_SanitizedMul(oaknut::QReg src1, oaknut::QReg src2, oaknut::QReg scratch0);
+
+    void Compile_EvaluateCondition(Instruction instr);
+    void Compile_UniformCondition(Instruction instr);
+
+    /**
+     * Emits the code to conditionally return from a subroutine envoked by the `CALL` instruction.
+     */
+    void Compile_Return();
+
+    std::bitset<64> PersistentCallerSavedRegs();
+
+    /**
+     * Assertion evaluated at compile-time, but only triggered if executed at runtime.
+     * @param condition Condition to be evaluated.
+     * @param msg       Message to be logged if the assertion fails.
+     */
+    void Compile_Assert(bool condition, const char* msg);
+
+    /**
+     * Analyzes the entire shader program for `CALL` instructions before emitting any code,
+     * identifying the locations where a return needs to be inserted.
+     */
+    void FindReturnOffsets();
+
+    /**
+     * Emits data and code for utility functions.
+     */
+    void CompilePrelude();
+    oaknut::Label CompilePrelude_Log2();
+    oaknut::Label CompilePrelude_Exp2();
+
+    const std::array<u32, MAX_PROGRAM_CODE_LENGTH>* program_code = nullptr;
+    const std::array<u32, MAX_SWIZZLE_DATA_LENGTH>* swizzle_data = nullptr;
+
+    /// Mapping of Pica VS instructions to pointers in the emitted code
+    std::array<oaknut::Label, MAX_PROGRAM_CODE_LENGTH> instruction_labels;
+
+    /// Labels pointing to the end of each nested LOOP block. Used by the BREAKC instruction to
+    /// break out of a loop.
+    std::vector<oaknut::Label> loop_break_labels;
+
+    /// Offsets in code where a return needs to be inserted
+    std::vector<unsigned> return_offsets;
+
+    unsigned program_counter = 0; ///< Offset of the next instruction to decode
+    u8 loop_depth = 0;            ///< Depth of the (nested) loops currently compiled
+
+    using CompiledShader = void(const void* setup, void* state, const std::byte* start_addr);
+    CompiledShader* program = nullptr;
+
+    oaknut::Label log2_subroutine;
+    oaknut::Label exp2_subroutine;
+};
+
+} // namespace Pica::Shader
+
+#endif