shader: track material keys selected from active lanes

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

@@ -967,6 +967,33 @@ private:
 		return edge && edge->positive ? edge->condition : Value {};
 	}
 
+	bool HasActiveLane(Value mask, const Block* use) const {
+		mask = mask.Resolve();
+		const auto incoming_is_nonempty = [&](const Block* from, const Block* to,
+		                                     Value incoming, const Block* header) {
+			for (size_t depth = 0; depth < m_program.blocks.size(); ++depth) {
+				const auto edge = ConditionalEdge(from, to);
+				if (edge && edge->positive && Implies(edge->condition, incoming)) return true;
+				if (from == header || from->ImmSuccessors().size() != 1u ||
+				    from->ImmPredecessors().size() != 1u) return false;
+				to = from;
+				from = from->ImmPredecessors()[0];
+			}
+			return false;
+		};
+		const auto* phi = mask.TryInstruction();
+		if (phi != nullptr && phi->GetOpcode() == ValueOpcode::Phi &&
+		    phi->GetType() == Type::U1 && phi->NumArgs() != 0u) {
+			for (size_t arm = 0; arm < phi->NumArgs(); ++arm) {
+				if (!incoming_is_nonempty(phi->PhiBlock(arm), phi->Parent(),
+				                          phi->Arg(arm), phi->Parent())) return false;
+			}
+			return true;
+		}
+		return use != nullptr && use->ImmPredecessors().size() == 1u &&
+		       incoming_is_nonempty(use->ImmPredecessors()[0], use, mask, nullptr);
+	}
+
 	Value SimplifyGuard(Value guard) const {
 		const auto invariant = ResolveInvariantPhi(m_program, guard);
 		return (invariant.IsEmpty() ? guard : invariant).Resolve();
@@ -1074,10 +1101,14 @@ private:
 		        EquivalentValue(m_program, inst->Arg(1), index));
 	}
 
-	bool ClearsFirstSetBit(Value value, Value mask) const {
+	bool MaskOnlyLosesBits(Value value, Value mask) const {
 		const auto* update = value.Resolve().TryInstruction();
-		if (update == nullptr || update->GetOpcode() != ValueOpcode::BitwiseXor32 ||
+		if (update == nullptr || update->NumArgs() != 2u) return false;
-		    update->NumArgs() != 2u) return false;
+		if (update->GetOpcode() == ValueOpcode::BitwiseAnd32) {
+			return EquivalentValue(m_program, update->Arg(0), mask) ||
+			       EquivalentValue(m_program, update->Arg(1), mask);
+		}
+		if (update->GetOpcode() != ValueOpcode::BitwiseXor32) return false;
 		Value bit;
 		if (EquivalentValue(m_program, update->Arg(0), mask)) bit = update->Arg(1);
 		else if (EquivalentValue(m_program, update->Arg(1), mask)) bit = update->Arg(0);
@@ -1102,13 +1133,13 @@ private:
 		       first->NumArgs() == 1u && EquivalentValue(m_program, first->Arg(0), mask);
 	}
 
-	Value InitialClearingMask(Value value, const Block* update_block) const {
+	Value InitialCandidateMask(Value value, const Block* update_block) const {
 		const auto* phi = value.Resolve().TryInstruction();
 		if (phi == nullptr || phi->GetOpcode() != ValueOpcode::Phi ||
 		    phi->NumArgs() != 2u || phi->GetType() != Type::U32) return {};
 		for (uint32_t back = 0; back < 2u; ++back) {
 			if (phi->PhiBlock(back) != update_block ||
-			    !ClearsFirstSetBit(phi->Arg(back), value)) continue;
+			    !MaskOnlyLosesBits(phi->Arg(back), value)) continue;
 			const auto initial = phi->Arg(back ^ 1u).Resolve();
 			return ValidateRuntimeValue(m_program, initial, RuntimeValueType::Integer)
 			           ? initial : Value {};
@@ -1139,7 +1170,7 @@ private:
 			if (first == nullptr || first->GetOpcode() != ValueOpcode::FindILsb32 ||
 			    first->NumArgs() != 1u || !ImpliesNonzero(bit_guard, first->Arg(0)))
 				continue;
-			const auto initial_mask = InitialClearingMask(first->Arg(0), phi->PhiBlock(bit_arm));
+			const auto initial_mask = InitialCandidateMask(first->Arg(0), phi->PhiBlock(bit_arm));
 			if (initial_mask.IsEmpty()) continue;
 			for (uint32_t sentinel_arm = 0; sentinel_arm < 3u; ++sentinel_arm) {
 				if (sentinel_arm == bit_arm) continue;
@@ -1183,10 +1214,32 @@ private:
 
 	bool MatchUniformizedMaterialKey(Value key, const Inst& image,
 	                                DescriptorSource::IndirectImage& indirect,
-	                                DescriptorSource& material_source, uint32_t pc) {
+	                                DescriptorSource& material_source) {
-		const auto guard = PositiveLaneWitness(image.Parent());
+		Value guard;
-		if (guard.IsEmpty()) return false;
+		Value local;
-		const auto local = EqualLocalKey(guard, key);
+		const auto* first = key.Resolve().TryInstruction();
+		if (first != nullptr && first->GetOpcode() == ValueOpcode::ReadFirstLane &&
+		    first->NumArgs() == 2u) {
+			guard = first->Arg(1).Resolve();
+			// Empty EXEC selects lane zero, which may not have loaded a material key.
+			if (!HasActiveLane(guard, first->Parent())) return false;
+			local = first->Arg(0).Resolve();
+			const auto* phi = guard.TryInstruction();
+			if (phi != nullptr && phi->GetOpcode() == ValueOpcode::Phi &&
+			    phi->GetType() == Type::U1 && phi->NumArgs() == 2u) {
+				for (uint32_t initial = 0; initial < 2u; ++initial) {
+					if (Implies(phi->Arg(initial ^ 1u), guard)) {
+						// The backedge only removes lanes from the initial mask.
+						guard = phi->Arg(initial).Resolve();
+						break;
+					}
+				}
+			}
+		} else {
+			guard = PositiveLaneWitness(image.Parent());
+			if (guard.IsEmpty()) return false;
+			local = EqualLocalKey(guard, key);
+		}
 		const auto* selected = local.Resolve().TryInstruction();
 		if (selected == nullptr || selected->GetOpcode() != ValueOpcode::SelectU32 ||
 		    selected->NumArgs() != 3u || !Implies(guard, selected->Arg(0))) return false;
@@ -1272,7 +1325,7 @@ private:
 		return {};
 	}
 
-	bool TryMakeIndirectImage(Inst& handle, uint32_t pc, IndirectImagePlan& plan) {
+	bool TryMakeIndirectImage(Inst& handle, IndirectImagePlan& plan) {
 		if (handle.GetOpcode() != ValueOpcode::GetImageResource || handle.NumArgs() != 8u) {
 			return false;
 		}
@@ -1339,7 +1392,7 @@ private:
 				indirect.key_count = BoundedLoopCount(key, handle.Parent());
 			}
 			if (indirect.key_count.IsEmpty() &&
-			    !MatchUniformizedMaterialKey(key, handle, indirect, material_source, pc)) {
+			    !MatchUniformizedMaterialKey(key, handle, indirect, material_source)) {
 				return false;
 			}
 			if ((table_offset & 3u) != 0u ||
@@ -1418,7 +1471,7 @@ private:
 					continue;
 				}
 				IndirectImagePlan plan;
-				if (TryMakeIndirectImage(*handle, inst.Flags<MemoryFlags>().pc, plan)) {
+				if (TryMakeIndirectImage(*handle, plan)) {
 					m_indirect_images.push_back(std::move(plan));
 				}
 			}

tests/ResourceTrackingTests.cpp

@@ -1031,11 +1031,16 @@ void TestBoundedComputeImageLoop() {
 void TestUniformizedMaterialImageKeys() {
   namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
   enum class Variant {
-    Valid, Plain, WrongUpdate, WrongEquality, WrongExit, WrongCarry, WrongBackedge
+    Valid, Plain, WrongUpdate, WrongEquality, WrongExit, WrongCarry, WrongBackedge,
+    AndNot, Subset, FirstLane, FirstLaneDirect, FirstLaneAndNot, EmptyFirstEntry, WrongFirstBackedge,
+    WideningFirstEntry, WideningFirstBackedge, WrongFirstLoadMask
   };
   const auto make_plan = [](Variant variant) {
     Fixture fixture;
     fixture.program.wave_size = 64u;
+    const bool first_lane = variant >= Variant::FirstLane;
+    const bool first_lane_loop = first_lane && variant != Variant::FirstLaneDirect;
+    const uint32_t sample_entry_id = first_lane_loop ? 10u : 9u;
     const auto branch = [&](Value predicate, CFG::BranchCondition kind, Block *block) {
       return variant == Variant::Plain ? predicate : fixture.Emit(
           ValueOpcode::ConditionRef, {predicate}, kind, block);
@@ -1049,6 +1054,8 @@ void TestUniformizedMaterialImageKeys() {
     auto *choose = fixture.AddBlock();
     auto *sample = fixture.AddBlock();
     auto *done = fixture.AddBlock();
+    auto *sample_header = first_lane_loop ? fixture.AddBlock() : nullptr;
+    auto *sample_entry = first_lane ? fixture.AddBlock() : nullptr;
     entry->AddBranch(header);
     header->AddBranch(inactive);
     inactive->AddBranch(merge);
@@ -1058,8 +1065,20 @@ void TestUniformizedMaterialImageKeys() {
     bit->AddBranch(header);
     bit->AddBranch(merge);
     merge->AddBranch(choose);
-    choose->AddBranch(sample);
+    choose->AddBranch(first_lane ? sample_entry : sample);
     choose->AddBranch(done);
+    if (first_lane) {
+      sample_entry->AddBranch(first_lane_loop ? sample_header : sample);
+      fixture.program.block_info[sample_entry_id].terminator = {
+          .kind = CFG::TerminatorKind::Branch, .true_block = first_lane_loop ? 9u : 7u};
+    }
+    if (first_lane_loop) {
+      sample_header->AddBranch(sample);
+      sample->AddBranch(sample_header);
+      fixture.program.block_info[9].terminator = {
+          .kind = CFG::TerminatorKind::Branch, .true_block = 7u,
+          .merge_block = 8u, .continue_block = 7u, .loop_header = true};
+    }
     sample->AddBranch(done);
     fixture.program.block_info[0].terminator = {
         .kind = CFG::TerminatorKind::Branch, .true_block = 1u};
@@ -1102,9 +1121,15 @@ void TestUniformizedMaterialImageKeys() {
         ValueOpcode::BitwiseAnd32, {first, Value(31u)}, 0, bit);
     const auto one_bit = fixture.Emit(
         ValueOpcode::ShiftLeftLogical32, {Value(1u), position}, 0, bit);
+    const bool and_not = variant == Variant::AndNot || variant == Variant::FirstLaneAndNot;
+    const auto removed = and_not
+        ? fixture.Emit(ValueOpcode::BitwiseNot32, {one_bit}, 0, bit)
+        : variant == Variant::Subset ? fixture.UserData(7u) : one_bit;
     const auto cleared = fixture.Emit(
-        variant == Variant::WrongUpdate ? ValueOpcode::BitwiseOr32 : ValueOpcode::BitwiseXor32,
+        variant == Variant::WrongUpdate ? ValueOpcode::BitwiseOr32
-        {mask, one_bit}, 0, bit);
+        : and_not || variant == Variant::Subset ? ValueOpcode::BitwiseAnd32
+                                              : ValueOpcode::BitwiseXor32,
+        {mask, removed}, 0, bit);
     const auto continuation = fixture.Emit(
         ValueOpcode::LogicalAnd, {bit_guard, active_on_entry}, 0, bit);
     fixture.program.block_info[4].condition = branch(
@@ -1155,25 +1180,54 @@ void TestUniformizedMaterialImageKeys() {
     const auto base = fixture.Address(fixture.UserData(0u), fixture.UserData(1u));
     MemoryInfo material_memory;
     material_memory.kind = ResourceKind::Global;
+    const auto load_mask = variant == Variant::WrongFirstLoadMask
+        ? fixture.Emit(ValueOpcode::LogicalNot, {material_guard}, 0, choose)
+        : material_guard;
     const auto loaded = fixture.Emit(
         ValueOpcode::LoadAddressU32,
-        {base, material_offset, Value(0u), material_guard},
+        {base, material_offset, Value(0u), load_mask},
         fixture.AddMemory(material_memory, 0x1a88u), choose);
     const auto local_key = fixture.Emit(
         ValueOpcode::SelectU32, {material_guard, loaded, arbitrary}, 0, choose);
-    const auto key = fixture.Emit(
+    Inst *sample_active_phi = nullptr;
-        ValueOpcode::ReadLane, {local_key, Value(0u)}, 0, choose);
+    auto sample_active = material_guard;
+    if (first_lane) {
+      const auto initial = variant == Variant::EmptyFirstEntry ? Value(false)
+          : variant == Variant::WideningFirstEntry ? enabled
+          : fixture.Emit(ValueOpcode::LogicalAnd,
+              {material_guard, fixture.Emit(ValueOpcode::SLessThanEqual32,
+                  {Value(0u), local_key}, 0, choose)}, 0, choose);
+      sample_active = initial;
+      if (first_lane_loop) {
+        sample_active_phi = &sample_header->AppendNewInst(
+            ValueOpcode::Phi, {}, static_cast<uint64_t>(Type::U1));
+        sample_active_phi->AddPhiOperand(sample_entry, initial);
+        sample_active = Value(sample_active_phi);
+      }
+      fixture.program.block_info[6].condition = branch(
+          variant == Variant::EmptyFirstEntry ? enabled : initial,
+          CFG::BranchCondition::ExecNonZero, choose);
+      fixture.program.block_info[6].terminator = {
+          .kind = CFG::TerminatorKind::ConditionalBranch,
+          .true_block = sample_entry_id, .false_block = 8u};
+    }
+    auto *key_block = first_lane ? sample : choose;
+    const auto key = first_lane
+        ? fixture.Emit(ValueOpcode::ReadFirstLane, {local_key, sample_active}, 0, key_block)
+        : fixture.Emit(ValueOpcode::ReadLane, {local_key, Value(0u)}, 0, key_block);
     const auto compared = fixture.Emit(
         ValueOpcode::IEqual32,
-        {key, variant == Variant::WrongEquality ? arbitrary : local_key}, 0, choose);
+        {key, variant == Variant::WrongEquality ? arbitrary : local_key}, 0, key_block);
     const auto sample_guard = fixture.Emit(
-        ValueOpcode::LogicalAnd, {material_guard, compared}, 0, choose);
+        ValueOpcode::LogicalAnd, {sample_active, compared}, 0, key_block);
-    fixture.program.block_info[6].condition = branch(
+    if (!first_lane) {
-        fixture.Emit(ValueOpcode::LogicalNot, {sample_guard}, 0, choose),
+      fixture.program.block_info[6].condition = branch(
-        CFG::BranchCondition::ExecZero, choose);
+          fixture.Emit(ValueOpcode::LogicalNot, {sample_guard}, 0, choose),
-    fixture.program.block_info[6].terminator = {
+          CFG::BranchCondition::ExecZero, choose);
-        .kind = CFG::TerminatorKind::ConditionalBranch,
+      fixture.program.block_info[6].terminator = {
-        .true_block = 8u, .false_block = 7u};
+          .kind = CFG::TerminatorKind::ConditionalBranch,
+          .true_block = 8u, .false_block = 7u};
+    }
     const auto table_offset = fixture.Emit(
         ValueOpcode::IAdd32,
         {fixture.Emit(ValueOpcode::ShiftLeftLogical32,
@@ -1203,10 +1257,33 @@ void TestUniformizedMaterialImageKeys() {
     MemoryInfo image_memory;
     image_memory.kind = ResourceKind::Image;
     image_memory.image_dimension = Decoder::ImageDimension::Dim2D;
-    fixture.Emit(ValueOpcode::ImageSampleRaw, {image, sampler, address},
+    const auto sampled = fixture.Emit(ValueOpcode::ImageSampleRaw, {image, sampler, address},
-                 fixture.AddMemory(image_memory, 0x1aecu), sample);
+        fixture.AddMemory(image_memory, 0x1aecu), sample);
-    fixture.program.block_info[7].terminator = {
+    if (first_lane) {
-        .kind = CFG::TerminatorKind::Branch, .true_block = 8u};
+      const auto component = fixture.Emit(
+          ValueOpcode::CompositeExtractU32x4, {sampled, Value(0u)}, 0, sample);
+      fixture.Emit(ValueOpcode::ReferenceU32,
+          {fixture.Emit(ValueOpcode::SelectU32,
+              {variant == Variant::EmptyFirstEntry ? material_guard : sample_guard,
+               component, Value(0u)}, 0, sample)}, 0, sample);
+    }
+    if (first_lane_loop) {
+      auto remaining = fixture.Emit(ValueOpcode::LogicalAnd,
+          {sample_active, fixture.Emit(ValueOpcode::LogicalNot,
+              {sample_guard}, 0, sample)}, 0, sample);
+      if (variant == Variant::WideningFirstBackedge)
+        remaining = fixture.Emit(ValueOpcode::LogicalOr, {remaining, enabled}, 0, sample);
+      sample_active_phi->AddPhiOperand(sample, remaining);
+      fixture.program.block_info[7].condition = branch(
+          variant == Variant::WrongFirstBackedge ? sample_guard : remaining,
+          CFG::BranchCondition::ExecNonZero, sample);
+      fixture.program.block_info[7].terminator = {
+          .kind = CFG::TerminatorKind::ConditionalBranch,
+          .true_block = 9u, .false_block = 8u};
+    } else {
+      fixture.program.block_info[7].terminator = {
+          .kind = CFG::TerminatorKind::Branch, .true_block = 8u};
+    }
     fixture.program.block_info[8].terminator.kind = CFG::TerminatorKind::Return;
     const auto output = fixture.Emit(
         ValueOpcode::GetBufferResource,
@@ -1232,6 +1309,11 @@ void TestUniformizedMaterialImageKeys() {
   };
   auto plan = make_plan(Variant::Valid);
   make_plan(Variant::Plain);
+  make_plan(Variant::AndNot);
+  make_plan(Variant::Subset);
+  make_plan(Variant::FirstLane);
+  make_plan(Variant::FirstLaneDirect);
+  plan = make_plan(Variant::FirstLaneAndNot);
   Check(plan.requires_specialization_memory &&
             plan.descriptor_sources[plan.info.images[0].source]
                 .indirect_image->selector_mask.Resolve().TryInstruction() != nullptr,
@@ -1295,6 +1377,12 @@ void TestUniformizedMaterialImageKeys() {
     CheckFatal([&] { make_plan(variant); }, "not a valid runtime value",
                "inactive material lane did not preserve its selected index across iterations");
   }
+  for (const auto variant : {Variant::EmptyFirstEntry, Variant::WrongFirstBackedge,
+                            Variant::WideningFirstEntry, Variant::WideningFirstBackedge,
+                            Variant::WrongFirstLoadMask}) {
+    CheckFatal([&] { make_plan(variant); }, "not a valid runtime value",
+               "first-lane material key escaped its active load mask");
+  }
 }
 
 void TestImageDescriptorFields() {