From 02cfbb8b0b1e61a601b67a68dfe0e0d5679a0b37 Mon Sep 17 00:00:00 2001
From: Merry <git@mary.rs>
Date: Sun, 24 Jul 2022 18:53:20 +0100
Subject: [PATCH] backend/arm64/reg_alloc: Generate immediates when required

---
 src/dynarmic/backend/arm64/reg_alloc.cpp | 52 +++++++++++++++++++++---
 src/dynarmic/backend/arm64/reg_alloc.h   | 36 ++++++++--------
 2 files changed, 67 insertions(+), 21 deletions(-)

diff --git a/src/dynarmic/backend/arm64/reg_alloc.cpp b/src/dynarmic/backend/arm64/reg_alloc.cpp
index 611e498d..dbdd72d8 100644
--- a/src/dynarmic/backend/arm64/reg_alloc.cpp
+++ b/src/dynarmic/backend/arm64/reg_alloc.cpp
@@ -90,6 +90,12 @@ bool HostLocInfo::Contains(const IR::Inst* value) const {
     return std::find(values.begin(), values.end(), value) != values.end();
 }
 
+void HostLocInfo::SetupScratchLocation() {
+    ASSERT(values.empty());
+    locked = true;
+    realized = true;
+}
+
 void HostLocInfo::SetupLocation(const IR::Inst* value) {
     values.clear();
     values.emplace_back(value);
@@ -192,9 +198,45 @@ void RegAlloc::AssertNoMoreUses() const {
     ASSERT(std::all_of(spills.begin(), spills.end(), is_empty));
 }
 
+template<HostLoc::Kind kind>
+int RegAlloc::GenerateImmediate(const IR::Value& value) {
+    if constexpr (kind == HostLoc::Kind::Gpr) {
+        const int new_location_index = AllocateRegister(gprs, gpr_order);
+        SpillGpr(new_location_index);
+        gprs[new_location_index].SetupScratchLocation();
+
+        code.MOV(oaknut::XReg{new_location_index}, value.GetImmediateAsU64());
+
+        return new_location_index;
+    } else if constexpr (kind == HostLoc::Kind::Fpr) {
+        const int new_location_index = AllocateRegister(fprs, fpr_order);
+        SpillFpr(new_location_index);
+        fprs[new_location_index].SetupScratchLocation();
+
+        code.MOV(Xscratch0, value.GetImmediateAsU64());
+        code.FMOV(oaknut::DReg{new_location_index}, Xscratch0);
+
+        return new_location_index;
+    } else if constexpr (kind == HostLoc::Kind::Flags) {
+        SpillFlags();
+        flags.SetupScratchLocation();
+
+        code.MOV(Xscratch0, value.GetImmediateAsU64());
+        code.MSR(oaknut::SystemReg::NZCV, Xscratch0);
+
+        return 0;
+    } else {
+        static_assert(kind == HostLoc::Kind::Fpr || kind == HostLoc::Kind::Gpr || kind == HostLoc::Kind::Flags);
+    }
+}
+
 template<HostLoc::Kind required_kind>
-int RegAlloc::RealizeReadImpl(const IR::Inst* value) {
-    const auto current_location = ValueLocation(value);
+int RegAlloc::RealizeReadImpl(const IR::Value& value) {
+    if (value.IsImmediate()) {
+        return GenerateImmediate<required_kind>(value);
+    }
+
+    const auto current_location = ValueLocation(value.GetInst());
     ASSERT(current_location);
 
     if (current_location->kind == required_kind) {
@@ -280,9 +322,9 @@ int RegAlloc::RealizeWriteImpl(const IR::Inst* value) {
     }
 }
 
-template int RegAlloc::RealizeReadImpl<HostLoc::Kind::Gpr>(const IR::Inst* value);
-template int RegAlloc::RealizeReadImpl<HostLoc::Kind::Fpr>(const IR::Inst* value);
-template int RegAlloc::RealizeReadImpl<HostLoc::Kind::Flags>(const IR::Inst* value);
+template int RegAlloc::RealizeReadImpl<HostLoc::Kind::Gpr>(const IR::Value& value);
+template int RegAlloc::RealizeReadImpl<HostLoc::Kind::Fpr>(const IR::Value& value);
+template int RegAlloc::RealizeReadImpl<HostLoc::Kind::Flags>(const IR::Value& value);
 template int RegAlloc::RealizeWriteImpl<HostLoc::Kind::Gpr>(const IR::Inst* value);
 template int RegAlloc::RealizeWriteImpl<HostLoc::Kind::Fpr>(const IR::Inst* value);
 template int RegAlloc::RealizeWriteImpl<HostLoc::Kind::Flags>(const IR::Inst* value);
diff --git a/src/dynarmic/backend/arm64/reg_alloc.h b/src/dynarmic/backend/arm64/reg_alloc.h
index c483256d..a3154c16 100644
--- a/src/dynarmic/backend/arm64/reg_alloc.h
+++ b/src/dynarmic/backend/arm64/reg_alloc.h
@@ -102,7 +102,7 @@ public:
 
 private:
     friend class RegAlloc;
-    explicit RAReg(RegAlloc& reg_alloc, bool write, const IR::Inst* value)
+    explicit RAReg(RegAlloc& reg_alloc, bool write, const IR::Value& value)
             : reg_alloc{reg_alloc}, write{write}, value{value} {}
 
     RAReg(const RAReg&) = delete;
@@ -114,7 +114,7 @@ private:
 
     RegAlloc& reg_alloc;
     bool write;
-    const IR::Inst* value;
+    const IR::Value value;
     std::optional<T> reg;
 };
 
@@ -127,6 +127,7 @@ struct HostLocInfo {
     size_t expected_uses = 0;
 
     bool Contains(const IR::Inst*) const;
+    void SetupScratchLocation();
     void SetupLocation(const IR::Inst*);
     bool IsCompletelyEmpty() const;
     bool IsImmediatelyAllocatable() const;
@@ -183,16 +184,16 @@ public:
         }
     }
 
-    auto WriteX(IR::Inst* inst) { return RAReg<oaknut::XReg>{*this, true, inst}; }
-    auto WriteW(IR::Inst* inst) { return RAReg<oaknut::WReg>{*this, true, inst}; }
+    auto WriteX(IR::Inst* inst) { return RAReg<oaknut::XReg>{*this, true, IR::Value{inst}}; }
+    auto WriteW(IR::Inst* inst) { return RAReg<oaknut::WReg>{*this, true, IR::Value{inst}}; }
 
-    auto WriteQ(IR::Inst* inst) { return RAReg<oaknut::QReg>{*this, true, inst}; }
-    auto WriteD(IR::Inst* inst) { return RAReg<oaknut::DReg>{*this, true, inst}; }
-    auto WriteS(IR::Inst* inst) { return RAReg<oaknut::SReg>{*this, true, inst}; }
-    auto WriteH(IR::Inst* inst) { return RAReg<oaknut::HReg>{*this, true, inst}; }
-    auto WriteB(IR::Inst* inst) { return RAReg<oaknut::BReg>{*this, true, inst}; }
+    auto WriteQ(IR::Inst* inst) { return RAReg<oaknut::QReg>{*this, true, IR::Value{inst}}; }
+    auto WriteD(IR::Inst* inst) { return RAReg<oaknut::DReg>{*this, true, IR::Value{inst}}; }
+    auto WriteS(IR::Inst* inst) { return RAReg<oaknut::SReg>{*this, true, IR::Value{inst}}; }
+    auto WriteH(IR::Inst* inst) { return RAReg<oaknut::HReg>{*this, true, IR::Value{inst}}; }
+    auto WriteB(IR::Inst* inst) { return RAReg<oaknut::BReg>{*this, true, IR::Value{inst}}; }
 
-    auto WriteFlags(IR::Inst* inst) { return RAReg<FlagsTag>{*this, true, inst}; }
+    auto WriteFlags(IR::Inst* inst) { return RAReg<FlagsTag>{*this, true, IR::Value{inst}}; }
 
     template<size_t size>
     auto WriteReg(IR::Inst* inst) {
@@ -248,14 +249,17 @@ private:
     template<typename>
     friend struct RAReg;
 
-    const IR::Inst* PreReadImpl(const IR::Value& value) {
-        const IR::Inst* inst = value.GetInst();
-        ValueInfo(inst).locked = true;
-        return inst;
+    const IR::Value& PreReadImpl(const IR::Value& value) {
+        if (!value.IsImmediate()) {
+            ValueInfo(value.GetInst()).locked = true;
+        }
+        return value;
     }
 
     template<HostLoc::Kind kind>
-    int RealizeReadImpl(const IR::Inst* value);
+    int GenerateImmediate(const IR::Value& value);
+    template<HostLoc::Kind kind>
+    int RealizeReadImpl(const IR::Value& value);
     template<HostLoc::Kind kind>
     int RealizeWriteImpl(const IR::Inst* value);
     void Unlock(HostLoc host_loc);
@@ -292,7 +296,7 @@ RAReg<T>::~RAReg() {
 
 template<typename T>
 void RAReg<T>::Realize() {
-    reg = T{write ? reg_alloc.RealizeWriteImpl<kind>(value) : reg_alloc.RealizeReadImpl<kind>(value)};
+    reg = T{write ? reg_alloc.RealizeWriteImpl<kind>(value.GetInst()) : reg_alloc.RealizeReadImpl<kind>(value)};
 }
 
 }  // namespace Dynarmic::Backend::Arm64