backend/arm64/reg_alloc: Support multiple locks on a location

src/dynarmic/backend/arm64/emit_arm64.cpp

@@ -152,6 +152,8 @@ EmittedBlockInfo EmitArm64(oaknut::CodeGenerator& code, IR::Block block, const E
             ASSERT_FALSE("Invalid opcode: {}", inst->GetOpcode());
             break;
         }
+
+        reg_alloc.AssertAllUnlocked();
     }
 
     reg_alloc.AssertNoMoreUses();

src/dynarmic/backend/arm64/reg_alloc.cpp

@@ -91,15 +91,16 @@ bool HostLocInfo::Contains(const IR::Inst* value) const {
 }
 
 void HostLocInfo::SetupScratchLocation() {
-    ASSERT(values.empty());
+    ASSERT(IsCompletelyEmpty());
-    locked = true;
+    locked++;
     realized = true;
 }
 
 void HostLocInfo::SetupLocation(const IR::Inst* value) {
+    ASSERT(IsCompletelyEmpty());
     values.clear();
     values.emplace_back(value);
-    locked = true;
+    locked++;
     realized = true;
     uses_this_inst = 0;
     accumulated_uses = 0;
@@ -123,10 +124,9 @@ void HostLocInfo::UpdateUses() {
     uses_this_inst = 0;
 
     if (accumulated_uses == expected_uses) {
-        *this = {};
+        values.clear();
-    } else {
+        accumulated_uses = 0;
-        realized = false;
+        expected_uses = 0;
-        locked = false;
     }
 }
 
@@ -187,7 +187,16 @@ void RegAlloc::DefineAsRegister(IR::Inst* inst, oaknut::Reg reg) {
     ASSERT(!ValueLocation(inst));
     auto& info = reg.is_vector() ? fprs[reg.index()] : gprs[reg.index()];
     ASSERT(info.IsCompletelyEmpty());
-    info.SetupLocation(inst);
+    info.values.emplace_back(inst);
+    info.expected_uses += inst->UseCount();
+}
+
+void RegAlloc::AssertAllUnlocked() const {
+    const auto is_unlocked = [](const auto& i) { return !i.locked && !i.realized; };
+    ASSERT(std::all_of(gprs.begin(), gprs.end(), is_unlocked));
+    ASSERT(std::all_of(fprs.begin(), fprs.end(), is_unlocked));
+    ASSERT(is_unlocked(flags));
+    ASSERT(std::all_of(spills.begin(), spills.end(), is_unlocked));
 }
 
 void RegAlloc::AssertNoMoreUses() const {
@@ -329,12 +338,6 @@ template int RegAlloc::RealizeWriteImpl<HostLoc::Kind::Gpr>(const IR::Inst* valu
 template int RegAlloc::RealizeWriteImpl<HostLoc::Kind::Fpr>(const IR::Inst* value);
 template int RegAlloc::RealizeWriteImpl<HostLoc::Kind::Flags>(const IR::Inst* value);
 
-void RegAlloc::Unlock(HostLoc host_loc) {
-    HostLocInfo& info = ValueInfo(host_loc);
-    ASSERT(info.locked && info.realized);
-    info.UpdateUses();
-}
-
 int RegAlloc::AllocateRegister(const std::array<HostLocInfo, 32>& regs, const std::vector<int>& order) const {
     const auto empty = std::find_if(order.begin(), order.end(), [&](int i) { return regs[i].IsImmediatelyAllocatable(); });
     if (empty != order.end()) {
@@ -392,6 +395,8 @@ void RegAlloc::ReadWriteFlags(Argument& read, IR::Inst* write) {
         ASSERT_FALSE("Invalid current location for flags");
     }
     flags.SetupLocation(write);
+    flags.locked--;
+    flags.realized = false;
 }
 
 void RegAlloc::SpillFlags() {

src/dynarmic/backend/arm64/reg_alloc.h

@@ -102,8 +102,7 @@ public:
 
 private:
     friend class RegAlloc;
-    explicit RAReg(RegAlloc& reg_alloc, bool write, const IR::Value& value)
+    explicit RAReg(RegAlloc& reg_alloc, bool write, const IR::Value& value);
-            : reg_alloc{reg_alloc}, write{write}, value{value} {}
 
     RAReg(const RAReg&) = delete;
     RAReg& operator=(const RAReg&) = delete;
@@ -120,7 +119,7 @@ private:
 
 struct HostLocInfo {
     std::vector<const IR::Inst*> values;
-    bool locked = false;
+    size_t locked = 0;
     bool realized = false;
     size_t uses_this_inst = 0;
     size_t accumulated_uses = 0;
@@ -145,16 +144,16 @@ public:
     ArgumentInfo GetArgumentInfo(IR::Inst* inst);
     bool IsValueLive(IR::Inst* inst) const;
 
-    auto ReadX(Argument& arg) { return RAReg<oaknut::XReg>{*this, false, PreReadImpl(arg.value)}; }
+    auto ReadX(Argument& arg) { return RAReg<oaknut::XReg>{*this, false, arg.value}; }
-    auto ReadW(Argument& arg) { return RAReg<oaknut::WReg>{*this, false, PreReadImpl(arg.value)}; }
+    auto ReadW(Argument& arg) { return RAReg<oaknut::WReg>{*this, false, arg.value}; }
 
-    auto ReadQ(Argument& arg) { return RAReg<oaknut::QReg>{*this, false, PreReadImpl(arg.value)}; }
+    auto ReadQ(Argument& arg) { return RAReg<oaknut::QReg>{*this, false, arg.value}; }
-    auto ReadD(Argument& arg) { return RAReg<oaknut::DReg>{*this, false, PreReadImpl(arg.value)}; }
+    auto ReadD(Argument& arg) { return RAReg<oaknut::DReg>{*this, false, arg.value}; }
-    auto ReadS(Argument& arg) { return RAReg<oaknut::SReg>{*this, false, PreReadImpl(arg.value)}; }
+    auto ReadS(Argument& arg) { return RAReg<oaknut::SReg>{*this, false, arg.value}; }
-    auto ReadH(Argument& arg) { return RAReg<oaknut::HReg>{*this, false, PreReadImpl(arg.value)}; }
+    auto ReadH(Argument& arg) { return RAReg<oaknut::HReg>{*this, false, arg.value}; }
-    auto ReadB(Argument& arg) { return RAReg<oaknut::BReg>{*this, false, PreReadImpl(arg.value)}; }
+    auto ReadB(Argument& arg) { return RAReg<oaknut::BReg>{*this, false, arg.value}; }
 
-    auto ReadFlags(Argument& arg) { return RAReg<FlagsTag>{*this, false, PreReadImpl(arg.value)}; }
+    auto ReadFlags(Argument& arg) { return RAReg<FlagsTag>{*this, false, arg.value}; }
 
     template<size_t size>
     auto ReadReg(Argument& arg) {
@@ -242,6 +241,7 @@ public:
         (rs.Realize(), ...);
     }
 
+    void AssertAllUnlocked() const;
     void AssertNoMoreUses() const;
 
 private:
@@ -249,20 +249,12 @@ private:
     template<typename>
     friend struct RAReg;
 
-    const IR::Value& PreReadImpl(const IR::Value& value) {
-        if (!value.IsImmediate()) {
-            ValueInfo(value.GetInst()).locked = true;
-        }
-        return value;
-    }
-
     template<HostLoc::Kind kind>
     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);
 
     int AllocateRegister(const std::array<HostLocInfo, 32>& regs, const std::vector<int>& order) const;
     void SpillGpr(int index);
@@ -287,10 +279,29 @@ private:
     mutable std::mt19937 rand_gen;
 };
 
+template<typename T>
+RAReg<T>::RAReg(RegAlloc& reg_alloc, bool write, const IR::Value& value)
+        : reg_alloc{reg_alloc}, write{write}, value{value} {
+    if (!write && !value.IsImmediate()) {
+        reg_alloc.ValueInfo(value.GetInst()).locked++;
+    }
+}
+
 template<typename T>
 RAReg<T>::~RAReg() {
+    if (value.IsImmediate()) {
         if (reg) {
-        reg_alloc.Unlock(HostLoc{kind, reg->index()});
+            // Immediate in scratch register
+            HostLocInfo& info = reg_alloc.ValueInfo(HostLoc{kind, reg->index()});
+            info.locked--;
+            info.realized = false;
+        }
+    } else {
+        HostLocInfo& info = reg_alloc.ValueInfo(value.GetInst());
+        info.locked--;
+        if (reg) {
+            info.realized = false;
+        }
     }
 }