shader: reuse dominating descriptor loads across image samples

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

@@ -1596,7 +1596,7 @@ private:
 		return false;
 	}
 
-	bool InitializedBefore(const Inst& initial, const Inst& use, Value witness) const {
+	bool AvailableBefore(const Inst& initial, const Inst& use, Value witness = {}) const {
 		if (initial.Parent() == use.Parent()) return BeforeInBlock(initial, use);
 		std::vector<const Block*> pending {use.Parent()};
 		for (size_t index = 0; index < pending.size(); ++index) {
@@ -1604,10 +1604,12 @@ private:
 			if (current == initial.Parent()) continue;
 			if (current->ImmPredecessors().empty()) return false;
 			for (const auto* previous: current->ImmPredecessors()) {
+				if (!witness.IsEmpty()) {
 					const auto edge = ConditionalEdge(previous, current);
 					// The proven local-zero invocation makes an all-inactive bypass impossible.
 					if (edge && edge->lanes == LaneQuantifier::All && !edge->positive &&
 					    Implies(witness, edge->condition)) continue;
+				}
 				if (std::ranges::find(pending, previous) == pending.end()) pending.push_back(previous);
 			}
 		}
@@ -1817,8 +1819,8 @@ private:
 			uint32_t initialized = 0;
 			for (const auto& initial: initializers) {
 				const auto witness = initial.inst->Arg(initial.inst->NumArgs() - 1u);
-				if (!InitializedBefore(*initial.inst, *append.atomic, witness) ||
+				if (!AvailableBefore(*initial.inst, *append.atomic, witness) ||
-				    !InitializedBefore(*initial.inst, *counter_read, witness)) return {};
+				    !AvailableBefore(*initial.inst, *counter_read, witness)) return {};
 				if (append.counter >= initial.address &&
 				    uint64_t {append.counter} < uint64_t {initial.address} + initial.words * 4u) ++initialized;
 			}
@@ -2382,7 +2384,7 @@ private:
 			    current_handle->GetOpcode() != (memory->kind == ResourceKind::ScalarAddress
 			                                      ? ValueOpcode::GetAddressResource
 			                                      : ValueOpcode::GetBufferResource) ||
-			    (memory->kind == ResourceKind::ScalarAddress && read->Parent() != handle.Parent()) ||
+			    (memory->kind == ResourceKind::ScalarAddress && !AvailableBefore(*read, handle)) ||
 			    (table_handle != nullptr &&
 			     !EquivalentValue(m_program, Value(table_handle), Value(current_handle))) ||
 			    (!indexed && !MatchTableOffset(read->Arg(1), current_key, offset, stride))) {

tests/ResourceTrackingTests.cpp

@@ -566,12 +566,14 @@ void TestGuardedDirectImageTable() {
     Bounded, BoundMissing, BoundUnrelated, BoundBypass, BoundAnyLane,
     BoundDiamond, BoundBranchSpecific, BoundMismatched,
     BoundExec, BoundExecUnguarded, BoundExecMismatched, BoundExecAnd, BoundExecUnprovenActive,
-    Finite, FiniteZero, FiniteMissing, FiniteWrong, FiniteSigned, FiniteAdd, FiniteXor
+    Finite, FiniteZero, FiniteMissing, FiniteWrong, FiniteSigned, FiniteAdd, FiniteXor,
+    FiniteShared, FiniteSharedBypass, FiniteLate
   };
   const auto make_plan = [](Guard guard) {
     const bool finite = guard >= Guard::Finite;
     const bool bounded = guard >= Guard::Bounded && !finite;
     const bool exec_bound = bounded && guard >= Guard::BoundExec;
+    const bool shared = guard == Guard::FiniteShared || guard == Guard::FiniteSharedBypass;
     const bool bypass = guard == Guard::Bypass || guard == Guard::BoundBypass;
     const uint32_t table_offset = finite ? 3344u : exec_bound ? 672u : bounded ? 3696u : 344u;
     const uint32_t table_stride = exec_bound ? 784u : bounded ? 272u : 32u;
@@ -582,6 +584,8 @@ void TestGuardedDirectImageTable() {
     auto *before_sample = fixture.AddBlock();
     auto *sample = fixture.AddBlock();
     auto *exit = fixture.AddBlock();
+    auto *descriptor_block = guard == Guard::FiniteSharedBypass ? fixture.AddBlock()
+                              : shared ? before_sample : sample;
     const auto varying = bounded || finite
         ? fixture.Emit(ValueOpcode::GetBuiltin,
             {Value(static_cast<uint32_t>(StageInputKind::LocalInvocationId)), Value(0u)})
@@ -689,6 +693,11 @@ void TestGuardedDirectImageTable() {
           .true_block = sample, .false_block = latch};
       sample->terminator = {.kind = CFG::TerminatorKind::Branch, .true_block = latch};
       latch->terminator = {.kind = CFG::TerminatorKind::Branch, .true_block = middle};
+      if (guard == Guard::FiniteSharedBypass) {
+        before_sample->terminator.true_block = descriptor_block;
+        before_sample->terminator.false_block = sample;
+        descriptor_block->terminator = {.kind = CFG::TerminatorKind::Branch, .true_block = sample};
+      }
     }
     if (guard == Guard::BoundDiamond || exec_bound) {
       middle->condition = fixture.Emit(ValueOpcode::IEqual32,
@@ -729,7 +738,7 @@ void TestGuardedDirectImageTable() {
           ValueOpcode::LoadAddressU32, {srt, Value(0u), Value(0u), Value(true)},
           fixture.AddMemory(memory, 0x20));
     }
-    fixture.block = sample;
+    fixture.block = descriptor_block;
     if (!bounded && !finite) key = fixture.Emit(ValueOpcode::FindILsb32, {mask});
     const auto offset = fixture.Emit(ValueOpcode::IAdd32,
         {bounded ? fixture.Emit(ValueOpcode::IMul32, {key, Value(table_stride)})
@@ -752,12 +761,26 @@ void TestGuardedDirectImageTable() {
       }
     }
     const auto image = fixture.Image(words, 0x128);
+    if (guard == Guard::FiniteLate) {
+      auto &instructions = fixture.block->Instructions();
+      const auto at = [&](Value value) {
+        return std::ranges::find_if(instructions,
+            [&](const Inst &inst) { return &inst == value.Instruction(); });
+      };
+      instructions.splice(at(words[0]), instructions, at(image));
+    }
     const auto sampler = fixture.Sampler({Value(0u), Value(0u), Value(0u), Value(0u)});
     MemoryInfo memory;
     memory.kind = ResourceKind::Image;
     memory.image_dimension = Decoder::ImageDimension::Dim2D;
     fixture.Emit(ValueOpcode::ImageSampleRaw, {image, sampler, fixture.ImageAddress(), active},
                  fixture.AddMemory(memory, 0x128));
+    if (shared) {
+      fixture.block = sample;
+      fixture.Emit(ValueOpcode::ImageSampleRaw,
+          {fixture.Image(words, 0x12c), sampler, fixture.ImageAddress(), active},
+          fixture.AddMemory(memory, 0x12c));
+    }
     if (guard == Guard::BoundExecUnguarded || guard == Guard::BoundExecMismatched) {
       const auto other_active = guard == Guard::BoundExecUnguarded ? Value(true)
           : fixture.Emit(ValueOpcode::ULessThan32, {key, fixture.UserData(2)});
@@ -812,7 +835,8 @@ void TestGuardedDirectImageTable() {
                            Guard::BoundBranchSpecific, Guard::BoundMismatched,
                            Guard::BoundExecUnguarded, Guard::BoundExecMismatched,
                            Guard::BoundExecAnd, Guard::BoundExecUnprovenActive,
-                           Guard::FiniteMissing, Guard::FiniteWrong, Guard::FiniteSigned}) {
+                           Guard::FiniteMissing, Guard::FiniteWrong, Guard::FiniteSigned,
+                           Guard::FiniteSharedBypass, Guard::FiniteLate}) {
     CheckFatal([&] { make_plan(guard); }, "not a valid runtime value",
                "direct table accepted a selector without a dominating range guard");
   }
@@ -939,12 +963,14 @@ void TestGuardedDirectImageTable() {
               records.descriptor_reads == 0u,
           "unrepresentable table span was read before validating its count");
   }
-  for (const auto guard : {Guard::Finite, Guard::FiniteZero, Guard::FiniteAdd, Guard::FiniteXor}) {
+  for (const auto guard : {Guard::Finite, Guard::FiniteZero, Guard::FiniteAdd,
+                           Guard::FiniteXor, Guard::FiniteShared}) {
     const auto finite = make_plan(guard);
     Check(finite.capture_specialization_reads && finite.info.images[1].written,
           "finite bit table with image writes omitted specialization read capture");
     LinearTestMemory records;
-    const uint32_t count = guard == Guard::FiniteZero ? 0u : guard == Guard::Finite ? 12u : 32u;
+    const uint32_t count = guard == Guard::FiniteZero ? 0u
+        : guard == Guard::FiniteAdd || guard == Guard::FiniteXor ? 32u : 12u;
     const uint64_t first = records.base + 3344u;
     fill_table(records, first);
     // The GPU may select no bits or a sparse subset. Null rows remain valid candidates.

tests/ShaderRecompilerComputeTests.cpp

@@ -36568,7 +36568,14 @@ void CheckGuardedRecordImageTable(VulkanHarness &vulkan) {
     code.insert(code.end(), {EncodeSmem0(0x03, 8), EncodeSmem1(0, 23),
                             EncodeMimg0(0x27, 1), EncodeMimg1(4, 20, 2, 10),
                             EncodeVop2(0x03, 0, Vgpr(4), 0)});
-    if (!bitset) {
+    if (bitset) {
+      // An optional second sample reuses the dominating descriptor load.
+      code.push_back(EncodeSopc(0x0d, 22, InlineU32(0)));
+      const auto skip_second = code.size(); code.push_back(0u);
+      code.insert(code.end(), {EncodeMimg0(0x27, 1), EncodeMimg1(4, 20, 2, 10),
+                              EncodeVop2(0x03, 0, Vgpr(4), 0)});
+      code[skip_second] = EncodeSopp(0x05, code.size() - skip_second - 1u);
+    } else {
       code[skip_first] = EncodeSopp(0x05, code.size() - skip_first - 1u);
       // Both sides of the optional first sample reach the second descriptor.
       code.insert(code.end(), {EncodeSop2(0x02, 23, 23, InlineU32(32)),
@@ -36593,7 +36600,8 @@ void CheckGuardedRecordImageTable(VulkanHarness &vulkan) {
         O::S_MUL_I32, O::S_ADD_I32, O::S_LOAD_DWORDX8, O::IMAGE_SAMPLE,
         O::V_ADD_F32, O::S_LSHL_B32, O::S_ANDN2_B32, O::S_BRANCH,
         O::IMAGE_STORE, O::S_ENDPGM};
-    if (bitset) test.opcodes.insert(test.opcodes.end(), {O::S_AND_B32, O::S_BFE_U32});
+    if (bitset) test.opcodes.insert(test.opcodes.end(), {O::S_AND_B32, O::S_BFE_U32,
+        O::S_BITCMP1_B32, O::S_CBRANCH_SCC1});
     else test.opcodes.insert(test.opcodes.end(), {O::S_CMP_GE_U32, O::S_CBRANCH_SCC1,
                                                 O::S_BITCMP1_B32});
     if (masked) test.opcodes.insert(test.opcodes.end(), {O::S_CSELECT_B64,
@@ -36624,10 +36632,10 @@ void CheckGuardedRecordImageTable(VulkanHarness &vulkan) {
     vulkan.Dispatch(test, compiled, buffer, nullptr, nullptr, &storage, nullptr, sampler, images);
     std::vector<u32> expected(4u * lanes, 0u);
     if (count != 0u) {
-      expected[0] = std::bit_cast<u32>(bitset ? (count == 3 ? 4.0f : 1.0f)
+      expected[0] = std::bit_cast<u32>(bitset ? (count == 3 ? 8.0f : 2.0f)
                                                     : (count == 3 ? 44.0f : 11.0f));
       expected[lanes] = std::bit_cast<u32>(bitset ? 2.0f : 20.0f);
-      expected[2u * lanes] = std::bit_cast<u32>(bitset ? 1.0f : 11.0f);
+      expected[2u * lanes] = std::bit_cast<u32>(bitset ? 2.0f : 11.0f);
     }
     CompareWords(test, "GPU worklist, guarded count and image writes", expected,
                  vulkan.ReadImage(test.name, &storage));