From fb3ba62b3a47ad645b007d5031ed9f8aaa7cb5c0 Mon Sep 17 00:00:00 2001
From: ReinUsesLisp <reinuseslisp@airmail.cc>
Date: Mon, 17 May 2021 19:24:09 -0300
Subject: [PATCH] glasm: Fix aliased bitcasts ref counting

---
 .../glasm/emit_glasm_bitwise_conversion.cpp   | 10 ++---
 .../backend/glasm/reg_alloc.cpp               | 39 +++++++++++++++----
 .../backend/glasm/reg_alloc.h                 |  6 +++
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/src/shader_recompiler/backend/glasm/emit_glasm_bitwise_conversion.cpp b/src/shader_recompiler/backend/glasm/emit_glasm_bitwise_conversion.cpp
index a6c66b826..cdbf6e93e 100644
--- a/src/shader_recompiler/backend/glasm/emit_glasm_bitwise_conversion.cpp
+++ b/src/shader_recompiler/backend/glasm/emit_glasm_bitwise_conversion.cpp
@@ -12,12 +12,10 @@ static void Alias(IR::Inst& inst, const IR::Value& value) {
     if (value.IsImmediate()) {
         return;
     }
-    IR::Inst* const value_inst{value.InstRecursive()};
-    if (inst.GetOpcode() == IR::Opcode::Identity) {
-        value_inst->DestructiveAddUsage(inst.UseCount());
-        value_inst->DestructiveRemoveUsage();
-    }
-    inst.SetDefinition(value_inst->Definition<Id>());
+    IR::Inst& value_inst{RegAlloc::AliasInst(*value.Inst())};
+    value_inst.DestructiveAddUsage(inst.UseCount());
+    value_inst.DestructiveRemoveUsage();
+    inst.SetDefinition(value_inst.Definition<Id>());
 }
 
 void EmitIdentity(EmitContext&, IR::Inst& inst, const IR::Value& value) {
diff --git a/src/shader_recompiler/backend/glasm/reg_alloc.cpp b/src/shader_recompiler/backend/glasm/reg_alloc.cpp
index 0e38f467f..707b22247 100644
--- a/src/shader_recompiler/backend/glasm/reg_alloc.cpp
+++ b/src/shader_recompiler/backend/glasm/reg_alloc.cpp
@@ -30,9 +30,10 @@ Value RegAlloc::Consume(const IR::Value& value) {
 }
 
 void RegAlloc::Unref(IR::Inst& inst) {
-    inst.DestructiveRemoveUsage();
-    if (!inst.HasUses()) {
-        Free(inst.Definition<Id>());
+    IR::Inst& value_inst{AliasInst(inst)};
+    value_inst.DestructiveRemoveUsage();
+    if (!value_inst.HasUses()) {
+        Free(value_inst.Definition<Id>());
     }
 }
 
@@ -99,10 +100,7 @@ Value RegAlloc::PeekInst(IR::Inst& inst) {
 }
 
 Value RegAlloc::ConsumeInst(IR::Inst& inst) {
-    inst.DestructiveRemoveUsage();
-    if (!inst.HasUses()) {
-        Free(inst.Definition<Id>());
-    }
+    Unref(inst);
     return PeekInst(inst);
 }
 
@@ -138,4 +136,31 @@ void RegAlloc::Free(Id id) {
     }
 }
 
+/*static*/ bool RegAlloc::IsAliased(const IR::Inst& inst) {
+    switch (inst.GetOpcode()) {
+    case IR::Opcode::Identity:
+    case IR::Opcode::BitCastU16F16:
+    case IR::Opcode::BitCastU32F32:
+    case IR::Opcode::BitCastU64F64:
+    case IR::Opcode::BitCastF16U16:
+    case IR::Opcode::BitCastF32U32:
+    case IR::Opcode::BitCastF64U64:
+        return true;
+    default:
+        return false;
+    }
+}
+
+/*static*/ IR::Inst& RegAlloc::AliasInst(IR::Inst& inst) {
+    IR::Inst* it{&inst};
+    while (IsAliased(*it)) {
+        const IR::Value arg{it->Arg(0)};
+        if (arg.IsImmediate()) {
+            break;
+        }
+        it = arg.InstRecursive();
+    }
+    return *it;
+}
+
 } // namespace Shader::Backend::GLASM
diff --git a/src/shader_recompiler/backend/glasm/reg_alloc.h b/src/shader_recompiler/backend/glasm/reg_alloc.h
index ede6edd1f..41b7c92be 100644
--- a/src/shader_recompiler/backend/glasm/reg_alloc.h
+++ b/src/shader_recompiler/backend/glasm/reg_alloc.h
@@ -126,6 +126,12 @@ public:
         return num_used_long_registers;
     }
 
+    /// Returns true if the instruction is expected to be aliased to another
+    static bool IsAliased(const IR::Inst& inst);
+
+    /// Returns the underlying value out of an alias sequence
+    static IR::Inst& AliasInst(IR::Inst& inst);
+
 private:
     static constexpr size_t NUM_REGS = 4096;
     static constexpr size_t NUM_ELEMENTS = 4;