Inline common IR operands and share storage with Phi pairs

src/graphics/shader/recompiler/ir/Value.cpp

@@ -2,7 +2,7 @@
 
 #include <algorithm>
 #include <cstring>
-#include <limits>
+#include <memory>
 
 namespace Libs::Graphics::ShaderRecompiler::IR {
 
@@ -135,10 +135,16 @@ bool Value::operator==(const Value& other) const {
 }
 
 Inst::Inst(ValueOpcode value_opcode, uint64_t value_flags)
-    : opcode(value_opcode), flags(value_flags) {
-	const auto count = NumArgsOf(opcode);
-	if (count != std::numeric_limits<size_t>::max()) {
-		args.resize(count);
+    : opcode(value_opcode),
+      num_args(value_opcode == ValueOpcode::Phi ? PhiArity
+                                               : static_cast<uint8_t>(NumArgsOf(value_opcode))),
+      flags(value_flags) {
+	if (num_args == PhiArity) {
+		std::destroy_at(&fixed_args);
+		std::construct_at(&phi_args);
+	} else if (num_args > InlineArity) {
+		std::destroy_at(&fixed_args);
+		std::construct_at(&large_args, num_args);
 	}
 }
 
@@ -154,8 +160,8 @@ Type Inst::GetType() const {
 	if (opcode == ValueOpcode::Phi) {
 		return static_cast<Type>(flags);
 	}
-	if (opcode == ValueOpcode::Identity && !args.empty()) {
-		return args.front().GetType();
+	if (opcode == ValueOpcode::Identity && num_args != 0) {
+		return Arg(0).GetType();
 	}
 	return TypeOf(opcode);
 }
@@ -173,21 +179,25 @@ size_t Inst::UseCount() const {
 }
 
 size_t Inst::NumArgs() const {
-	return args.size();
+	return num_args == PhiArity ? phi_args.size() : num_args;
 }
 
 size_t Inst::NumPhiBlocks() const {
-	return phi_blocks.size();
+	return num_args == PhiArity ? phi_args.size() : 0;
 }
 
 Value Inst::Arg(size_t index) const {
-	EXIT_IF(index >= args.size());
-	return args[index];
+	if (num_args <= InlineArity) {
+		EXIT_IF(index >= num_args);
+		return fixed_args[index];
+	}
+	EXIT_IF(index >= NumArgs());
+	return num_args == PhiArity ? phi_args[index].second : large_args[index];
 }
 
 Block* Inst::PhiBlock(size_t index) const {
-	EXIT_IF(opcode != ValueOpcode::Phi || index >= phi_blocks.size());
-	return phi_blocks[index];
+	EXIT_IF(opcode != ValueOpcode::Phi || index >= phi_args.size());
+	return phi_args[index].first;
 }
 
 Block* Inst::Parent() const {
@@ -203,15 +213,17 @@ void Inst::SetParent(Block* block) {
 }
 
 void Inst::SetArg(size_t index, Value value) {
-	if (index >= args.size()) {
-		EXIT_IF(NumArgsOf(opcode) != std::numeric_limits<size_t>::max());
-		args.resize(index + 1);
-	}
-	const auto old = args[index];
+	const auto old = Arg(index);
 	if (auto* old_inst = old.TryInstruction(); old_inst != nullptr) {
 		RemoveUse(old_inst, index);
 	}
-	args[index] = value;
+	if (num_args <= InlineArity) {
+		fixed_args[index] = value;
+	} else if (num_args == PhiArity) {
+		phi_args[index].second = value;
+	} else {
+		large_args[index] = value;
+	}
 	if (auto* new_inst = value.TryInstruction(); new_inst != nullptr) {
 		AddUse(new_inst, index);
 	}
@@ -219,9 +231,8 @@ void Inst::SetArg(size_t index, Value value) {
 
 void Inst::AddPhiOperand(Block* predecessor, Value value) {
 	EXIT_IF(opcode != ValueOpcode::Phi);
-	const auto index = args.size();
-	args.push_back(value);
-	phi_blocks.push_back(predecessor);
+	const auto index = phi_args.size();
+	phi_args.emplace_back(predecessor, value);
 	if (auto* value_inst = value.TryInstruction(); value_inst != nullptr) {
 		AddUse(value_inst, index);
 	}
@@ -234,21 +245,12 @@ void Inst::ReplaceUsesWith(Value replacement, bool preserve) {
 	}
 	Invalidate();
 	if (preserve) {
-		ReplaceOpcode(ValueOpcode::Identity);
-		args.resize(1);
+		opcode = ValueOpcode::Identity;
+		num_args = 1;
 		SetArg(0, replacement);
 	}
 }
 
-void Inst::ReplaceOpcode(ValueOpcode value_opcode) {
-	opcode           = value_opcode;
-	const auto count = NumArgsOf(opcode);
-	if (count != std::numeric_limits<size_t>::max()) {
-		EXIT_IF(!args.empty() && args.size() != count);
-		args.resize(count);
-	}
-}
-
 void Inst::Invalidate() {
 	ClearArgs();
 	opcode = ValueOpcode::Void;
@@ -269,13 +271,21 @@ void Inst::RemoveUse(Inst* used, size_t operand) {
 }
 
 void Inst::ClearArgs() {
-	for (size_t index = 0; index < args.size(); index++) {
-		if (auto* value_inst = args[index].TryInstruction(); value_inst != nullptr) {
+	for (size_t index = 0; index < NumArgs(); index++) {
+		if (auto* value_inst = Arg(index).TryInstruction(); value_inst != nullptr) {
 			RemoveUse(value_inst, index);
 		}
 	}
-	args.clear();
-	phi_blocks.clear();
+	if (num_args == PhiArity) {
+		std::destroy_at(&phi_args);
+		std::construct_at(&fixed_args);
+	} else if (num_args > InlineArity) {
+		std::destroy_at(&large_args);
+		std::construct_at(&fixed_args);
+	} else {
+		fixed_args.fill(Value {});
+	}
+	num_args = 0;
 }
 
 } // namespace Libs::Graphics::ShaderRecompiler::IR

src/graphics/shader/recompiler/ir/Value.h

@@ -4,6 +4,7 @@
 #include "graphics/shader/recompiler/ir/Reg.h"
 #include "graphics/shader/recompiler/ir/opcodes/ValueOpcodes.h"
 
+#include <array>
 #include <bit>
 #include <cstdint>
 #include <cstring>
@@ -133,7 +134,6 @@ public:
 	void SetArg(size_t index, Value value);
 	void AddPhiOperand(Block* predecessor, Value value);
 	void ReplaceUsesWith(Value replacement, bool preserve = true);
-	void ReplaceOpcode(ValueOpcode opcode);
 	void Invalidate();
 
 	template <typename T>
@@ -156,13 +156,22 @@ private:
 	void RemoveUse(Inst* used, size_t operand);
 	void ClearArgs();
 
+	static constexpr uint8_t InlineArity = 4;
+	static constexpr uint8_t PhiArity = UINT8_MAX;
+
 	ValueOpcode         opcode;
+	uint8_t             num_args;
 	mutable uint32_t    evaluation_index = UINT32_MAX;
 	uint64_t            flags;
 	Block*              parent = nullptr;
-	std::vector<Value>  args;
-	std::vector<Block*> phi_blocks;
+	union {
+		std::array<Value, InlineArity> fixed_args {};
+		std::vector<Value> large_args;
+		std::vector<std::pair<Block*, Value>> phi_args;
+	};
 	std::vector<Use>    uses;
 };
 
+static_assert(sizeof(Inst) <= 112, "Inst operand storage unintentionally increased");
+
 } // namespace Libs::Graphics::ShaderRecompiler::IR

src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.h

@@ -9,7 +9,7 @@
 
 namespace Libs::Graphics::ShaderRecompiler::IR {
 
-enum class ValueOpcode {
+enum class ValueOpcode : uint16_t {
 #define VALUE_OPCODE(name, ...) name,
 #include "graphics/shader/recompiler/ir/opcodes/ValueOpcodes.inc"
 #undef VALUE_OPCODE

tests/ScalarProvenanceTests.cpp

@@ -281,6 +281,61 @@ void TestInvariantAndDivergentPhi() {
         "divergent phi was accepted");
 }
 
+void TestOperandMutationLifecycle() {
+  Fixture fixture(6);
+  const auto definition = fixture.Emit(ValueOpcode::IAdd32, {Value(0x55u), Value(0u)});
+  auto *source = definition.TryInstruction();
+  const auto inserted = fixture.Emit(ValueOpcode::BitFieldInsert,
+      {definition, definition, Value(4u), Value(4u)});
+  Check(source->UseCount() == 2, "inline operands lost distinct reverse-use indices");
+  inserted.TryInstruction()->SetArg(1, Value(3u));
+  Check(source->UseCount() == 1, "inline operand mutation did not detach its old use");
+  source->ReplaceUsesWith(Value(0xaau));
+  uint32_t result = 0;
+  Check(source->UseCount() == 0 &&
+            SrtWalker(fixture.program, {}).Evaluate(inserted, result) && result == 0x3au,
+        "inline operand replacement changed the evaluated value");
+
+  for (const auto opcode : {ValueOpcode::GetImageResource, ValueOpcode::MakeImageAddress}) {
+    auto &large = fixture.program.value_storage.emplace_back(opcode);
+    auto &replacement = fixture.program.value_storage.emplace_back(opcode);
+    const auto count = large.NumArgs();
+    for (size_t index = 0; index < count; ++index) {
+      large.SetArg(index, definition);
+      replacement.SetArg(index, Value(static_cast<uint32_t>(index)));
+    }
+    const auto user = fixture.Emit(ValueOpcode::Identity, {Value(&large)}, 0, 1);
+    Check(source->UseCount() == count, "large operands lost reverse-use indices");
+    large.ReplaceUsesWith(Value(&replacement));
+    Check(source->UseCount() == 0 && large.GetOpcode() == ValueOpcode::Identity &&
+              large.NumArgs() == 1 && large.Arg(0) == Value(&replacement) &&
+              user.Resolve() == Value(&replacement) && replacement.UseCount() == 2,
+          "large operand replacement did not preserve users or detach old operands");
+    large.Invalidate();
+    large.Invalidate();
+    Check(large.NumArgs() == 0 && replacement.UseCount() == 1,
+          "repeated invalidation detached a surviving large-value user");
+  }
+
+  auto &phi = fixture.BlockAt(5).AppendNewInst(ValueOpcode::Phi);
+  phi.SetFlags(Type::U32);
+  for (size_t index = 0; index < 5; ++index) {
+    phi.AddPhiOperand(&fixture.BlockAt(index), definition);
+  }
+  phi.SetArg(3, Value(7u));
+  Check(phi.NumArgs() == 5 && phi.NumPhiBlocks() == 5 && source->UseCount() == 4,
+        "Phi mutation lost its incoming values or reverse uses");
+  for (size_t index = 0; index < 5; ++index) {
+    Check(phi.PhiBlock(index) == &fixture.BlockAt(index),
+          "Phi mutation changed an incoming predecessor");
+  }
+  const auto sum = fixture.Emit(ValueOpcode::IAdd32, {Value(&phi), Value(5u)}, 0, 5);
+  phi.ReplaceUsesWith(Value(6u));
+  Check(source->UseCount() == 0 && phi.NumPhiBlocks() == 0 && phi.NumArgs() == 1 &&
+            SrtWalker(fixture.program, {}).Evaluate(sum, result) && result == 11,
+        "Phi replacement retained incoming uses or changed the consumer value");
+}
+
 void TestControlDependentStandaloneLoadStaysTyped() {
   Fixture fixture(3);
   const auto memory = fixture.AddMemory(ResourceKind::ScalarAddress);
@@ -634,6 +689,7 @@ int main() {
     TestShaderBaseAndUserData();
     TestCarryAndBitFields();
     TestInvariantAndDivergentPhi();
+    TestOperandMutationLifecycle();
     TestControlDependentStandaloneLoadStaysTyped();
     TestRuntime64BitDescriptorOps();
     TestUniformFirstLaneSamplerLod();