Share repeated SSA bitcasts and guard emitted shader size

src/graphics/shader/recompiler/ir/passes/ConstantPropagation.cpp

@@ -1029,6 +1029,19 @@ void ConstantPropagationPass(const BlockList& blocks, uint32_t wave_size) {
 		}
 		drain();
 	} while (normalized);
+	// Repeated register reads share a bitcast only within the same SSA block.
+	std::unordered_map<const Inst*, Inst*> bitcasts;
+	for (auto* block: blocks) {
+		bitcasts.clear();
+		for (auto& inst: *block) {
+			if (inst.GetOpcode() != ValueOpcode::BitCastF32U32 &&
+			    inst.GetOpcode() != ValueOpcode::BitCastU32F32) continue;
+			const auto* source = Arg(inst, 0).TryInstruction();
+			if (source == nullptr) continue;
+			const auto [entry, inserted] = bitcasts.try_emplace(source, &inst);
+			if (!inserted) Replace(inst, Value(entry->second));
+		}
+	}
 }
 
 } // namespace Libs::Graphics::ShaderRecompiler::IR

tests/ShaderRecompilerComputeTests.cpp

@@ -1366,6 +1366,7 @@ struct TestCase {
   std::vector<std::pair<std::string, size_t>> ir_counts;
   u32 expected_mip_descriptors = 0;
   std::optional<std::vector<u32>> expected_buffer_resources;
+  size_t max_spirv_words = 0;
 };
 
 struct GraphicsCase {
@@ -1756,6 +1757,13 @@ CompiledShader CompileCase(const TestCase &test, u32 host_subgroup_size = 64) {
           "recompiler returned empty SPIR-V");
   ValidateSpirv(test.name, result.spirv);
   CheckSpirvEmission(test, result.spirv);
+  if (test.max_spirv_words != 0u) {
+    Require(test.name, "emitted shader size", result.spirv.size() <= test.max_spirv_words,
+            "SPIR-V words=" + std::to_string(result.spirv.size()) +
+                ", limit=" + std::to_string(test.max_spirv_words));
+    std::printf("[compute] %s wave=%u SPIR-V words=%zu\n", test.name,
+                result.program.wave_size, result.spirv.size());
+  }
   if (test.expected_buffer_resources) {
     std::vector<u32> resources;
     for (u32 i = 0; i < result.program.info.buffers.size(); ++i) {
@@ -29393,7 +29401,7 @@ TestCase VectorCompareClassF32() {
            O::BUFFER_STORE_DWORD, O::S_ENDPGM}};
 }
 
-TestCase VectorCompareClassF32ImmediateAndDynamicMasks() {
+TestCase VectorCompareClassF32ImmediateAndDynamicMasks(u32 wave) {
   using O = ShaderOpcode;
   // One input for each architectural class bit, in bit order.
   constexpr std::array<u32, 10> values{
@@ -29403,6 +29411,9 @@ TestCase VectorCompareClassF32ImmediateAndDynamicMasks() {
       0u, 0x3ffu, 3u, 7u, 0xfffffc00u, 0xffffffffu, 0xfffffc03u, 0xfffffc07u};
   TestCase test;
   test.name = "VectorCompareClassF32ImmediateAndDynamicMasks";
+  test.has_compute_info = true;
+  test.compute_info.wave_size = wave;
+  test.compute_info.threads_num[0] = test.compute_info.threads_num[1] = test.compute_info.threads_num[2] = 1u;
   test.initial.assign(values.begin(), values.end());
   test.initial.insert(test.initial.end(), masks.begin(), masks.end());
   test.expected = test.initial;
@@ -29426,12 +29437,17 @@ TestCase VectorCompareClassF32ImmediateAndDynamicMasks() {
         test.expected.push_back((masks[mask_index] & (1u << index)) != 0u);
       }
     }
+    AppendStoreVgpr(&code, 0, static_cast<u32>(test.expected.size()));
+    test.expected.push_back(values[index]);
   }
   AppendEnd(&code);
   test.initial.resize(test.expected.size());
   test.opcodes = {O::V_MOV_B32, O::BUFFER_LOAD_DWORD, O::V_CMP_CLASS_F32,
                   O::V_CNDMASK_B32, O::BUFFER_STORE_DWORD, O::S_ENDPGM};
   test.required_spirv = {"OpIsNan"};
+  test.ir_counts = {{"BitCastF32U32", values.size()}};
+  // Measured 56,287/117,071 words; allow only a small fixed emission margin.
+  test.max_spirv_words = wave == 32u ? 56400u : 117200u;
   return test;
 }
 
@@ -42826,7 +42842,8 @@ std::vector<TestCase> MakeCases() {
   AddCase(VectorVopcCmpxNeU64CapturedExecMask);
   AddCase(VectorVop3CmpxNeI64CapturedExecMask);
   AddCase(VectorCompareClassF32);
-  AddCase(VectorCompareClassF32ImmediateAndDynamicMasks);
+  cases.push_back(VectorCompareClassF32ImmediateAndDynamicMasks(32));
+  cases.push_back(VectorCompareClassF32ImmediateAndDynamicMasks(64));
   AddCase(VectorVopcSdwaCmpxClassF32CapturedExecMask);
   cases.push_back(VectorCmpClassF16(32));
   cases.push_back(VectorCmpClassF16(64));
@@ -48674,7 +48691,8 @@ int main(int argc, char **argv) {
       RunCase(&vulkan, VectorCompareF32DenormalModes(mode));
     }
     RunCase(&vulkan, VectorCompareClassF32());
-    RunCase(&vulkan, VectorCompareClassF32ImmediateAndDynamicMasks());
+    for (const u32 wave : {32u, 64u})
+      RunCase(&vulkan, VectorCompareClassF32ImmediateAndDynamicMasks(wave));
     RunCase(&vulkan, VectorCompareF16Ops());
     return 0;
   }

tests/shaderCfgTests.cpp

@@ -45,6 +45,7 @@
 #include <span>
 #include <sstream>
 #include <string>
+#include <tuple>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
@@ -12964,6 +12965,54 @@ void TestLateConstantPropagation() {
         "an ungrounded Phi cycle acquired a constant");
 }
 
+void TestRepeatedBitcastPropagation() {
+  using namespace ShaderRecompiler::IR;
+  Block entry, left, right;
+  IREmitter e(&entry), l(&left), r(&right);
+  const auto raw = e.GetUserData(static_cast<ScalarReg>(0));
+  const auto other = e.GetUserData(static_cast<ScalarReg>(1));
+  const auto fp = e.Emit(ValueOpcode::ConvertF32U32, {raw});
+  for (const auto [opcode, source, distinct] : {
+       std::tuple{ValueOpcode::BitCastF32U32, Value(raw), Value(other)},
+       std::tuple{ValueOpcode::BitCastU32F32, fp,
+                  e.Emit(ValueOpcode::ConvertF32U32, {other})}}) {
+    const auto first = e.Emit(opcode, {source});
+    const auto duplicate = e.Emit(opcode, {source});
+    const auto separate = e.Emit(opcode, {distinct});
+    const auto left_cast = l.Emit(opcode, {source});
+    const auto left_duplicate = l.Emit(opcode, {source});
+    const auto right_cast = r.Emit(opcode, {source});
+    const auto use = opcode == ValueOpcode::BitCastF32U32
+        ? e.Emit(ValueOpcode::FPIsNan32, {duplicate})
+        : e.Emit(ValueOpcode::IEqual32, {duplicate, Value(0u)});
+    ConstantPropagationPass({&entry, &left, &right});
+    Check(duplicate.Resolve() == first && left_duplicate.Resolve() == left_cast &&
+              use.Instruction()->Arg(0).Resolve() == first,
+          "repeated same-block casts were not shared with their earlier definition");
+    Check(separate.Resolve() == separate && left_cast.Resolve() == left_cast &&
+              right_cast.Resolve() == right_cast && left_cast != first && right_cast != left_cast,
+          "cast reuse merged distinct values or crossed block boundaries");
+  }
+  const auto as_float = e.Emit(ValueOpcode::BitCastF32U32, {raw});
+  const auto raw_roundtrip = e.Emit(ValueOpcode::BitCastU32F32, {as_float});
+  const auto as_bits = e.Emit(ValueOpcode::BitCastU32F32, {fp});
+  const auto float_roundtrip = e.Emit(ValueOpcode::BitCastF32U32, {as_bits});
+  ConstantPropagationPass({&entry});
+  Check(raw_roundtrip.Resolve() == raw && float_roundtrip.Resolve() == fp,
+        "sharing casts prevented inverse casts from folding");
+  for (const uint32_t bits : {0u, 0x80000000u, 1u, 0x80000001u, 0x7f800000u,
+                             0xff800000u, 0x7fc12345u, 0x7f812345u}) {
+    const auto constant_float = e.Emit(ValueOpcode::BitCastF32U32, {Value(bits)});
+    const auto constant_bits = e.Emit(ValueOpcode::BitCastU32F32,
+                                      {Value::F32(std::bit_cast<float>(bits))});
+    ConstantPropagationPass({&entry});
+    Check(constant_float.Resolve().IsImmediate() &&
+              std::bit_cast<uint32_t>(constant_float.Resolve().F32Value()) == bits &&
+              constant_bits.Resolve() == Value(bits),
+          "constant bitcast folding changed signed zero, subnormal, infinity or NaN bits");
+  }
+}
+
 void TestMaskedShiftConstantPropagation() {
   using namespace ShaderRecompiler::IR;
   for (const uint32_t count : {0u, 1u, 8u, 31u, 32u, 33u, 40u, 63u, 64u, UINT32_MAX}) {
@@ -16722,6 +16771,7 @@ int main() {
   TestValuePhiValidation();
   TestLaneMaskProjection();
   TestLateConstantPropagation();
+  TestRepeatedBitcastPropagation();
   TestMaskedShiftConstantPropagation();
   TestU64ShiftConstantPropagation();
   TestU64RegisterPairRoundTrip();