shader: bound waterfall image table keys converted from clamped float arithmetic

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

@@ -7,6 +7,7 @@
 
 #include <algorithm>
 #include <bit>
+#include <cmath>
 #include <fmt/format.h>
 #include <map>
 #include <numeric>
@@ -3014,6 +3015,18 @@ private:
 		return true;
 	}
 
+	// The EXEC predicate of the image access through use; empty for any other use.
+	Value AccessPredicate(const Use& use, bool sampler) const {
+		const auto& access = *use.user;
+		if (use.operand != (sampler ? 1u : 0u) ||
+		    ImageOpcodeInfoOf(access.GetOpcode()).access == ImageAccess::None) return {};
+		if (sampler && (access.GetOpcode() != ValueOpcode::ImageSampleRaw ||
+		    (m_program.memory_info[access.Flags<MemoryFlags>().index].image_sample_flags &
+		     Decoder::ImageSampleFlagAdjust) != 0u)) return {};
+		const auto mask = access.Arg(access.NumArgs() - 1u);
+		return mask.GetType() == Type::U1 ? mask : Value {};
+	}
+
 	bool TryMakeTableDescriptor(Inst& handle, const DescriptorSource& descriptor,
 	                            IndirectDescriptorPlan& plan) {
 		const bool sampler = handle.GetOpcode() == ValueOpcode::GetSamplerResource;
@@ -3102,7 +3115,19 @@ private:
 			    !MatchUniformizedMaterialKey(key, handle, indirect, material_source, false) &&
 			    !MatchUniformizedMaterialKey(key, handle, indirect, material_source, true) &&
 			    !MatchReversedImageRange(key, indirect) && !MatchWorkgroupKey(key, indirect)) {
-				const auto maximum = SelectorMaximum(key);
+				auto maximum = SelectorMaximum(key);
+				if (maximum >= MaxIndirectDescriptorProbes) {
+					// Each image access bounds the key on the lanes it executes.
+					maximum = 0u;
+					for (const auto& use: handle.Uses()) {
+						const auto active = AccessPredicate(use, false);
+						if (active.IsEmpty()) {
+							maximum = UINT32_MAX;
+							break;
+						}
+						maximum = std::max(maximum, SelectorMaximum(key, {}, 0, active));
+					}
+				}
 				if (maximum >= MaxIndirectDescriptorProbes) return false;
 				indirect.key_count = Value(maximum + 1u);
 			}
@@ -3111,26 +3136,16 @@ private:
 			// Keep the GPU selector and descriptor reads, and preserve the proven scalar
 			// offset domain within the bounded image table.
 			if (sampler || plan.split_offsets) return false;
-			const auto active = [&](const Use& use, bool sampler) -> Value {
-				const auto& access = *use.user;
-				if (use.operand != (sampler ? 1u : 0u) ||
-				    ImageOpcodeInfoOf(access.GetOpcode()).access == ImageAccess::None) return {};
-				if (sampler && (access.GetOpcode() != ValueOpcode::ImageSampleRaw ||
-				    (m_program.memory_info[access.Flags<MemoryFlags>().index].image_sample_flags &
-				     Decoder::ImageSampleFlagAdjust) != 0u)) return {};
-				const auto mask = access.Arg(access.NumArgs() - 1u);
-				return mask.GetType() == Type::U1 ? mask : Value {};
-			};
 			bool shared_selector = false;
 			if (!plan.table_indexed) {
 				for (const auto& owner: m_indirect_descriptors) {
 					const auto& selected = *m_sources[owner.source].indirect_descriptor;
 					if (!selected.selector || !EquivalentValue(m_program, key, owner.key) ||
 					    !std::ranges::all_of(handle.Uses(), [&](const Use& target) {
-						    const auto target_active = active(target, false);
+						    const auto target_active = AccessPredicate(target, false);
 						    return !target_active.IsEmpty() &&
 						        std::ranges::any_of(owner.handle->Uses(), [&](const Use& witness) {
-							        const auto owner_active = active(witness,
+							        const auto owner_active = AccessPredicate(witness,
 							            owner.handle->GetOpcode() == ValueOpcode::GetSamplerResource);
 							        return !owner_active.IsEmpty() && AvailableBefore(*witness.user, *target.user) &&
 							            Implies(target_active, owner_active);
@@ -3451,6 +3466,42 @@ private:
 		return result;
 	}
 
+	// Bounds of an F32 value on the lanes where guard holds; NaN when unknown.
+	std::pair<float, float> FloatRange(Value value, Value guard, uint32_t depth = 0) const {
+		constexpr std::pair unknown {NAN, NAN};
+		value = ActiveValue(value, guard);
+		if (value.IsImmediate() && value.GetType() == Type::F32)
+			return {value.F32Value(), value.F32Value()};
+		const auto* inst = value.TryInstruction();
+		// Each multiply-add evaluates three operands; keep the search shallow.
+		if (depth > 12u || inst == nullptr) return unknown;
+		const auto range = [&](uint32_t arg) {
+			return FloatRange(inst->Arg(arg), guard, depth + 1u);
+		};
+		switch (inst->GetOpcode()) {
+			case ValueOpcode::BitCastF32U32:
+			case ValueOpcode::BitCastU32F32: return range(0);
+			// Clamping maps NaN to zero.
+			case ValueOpcode::FPSaturate32: return {0.0f, 1.0f};
+			case ValueOpcode::FPFloor32: {
+				const auto [low, high] = range(0);
+				return {std::floor(low), std::floor(high)};
+			}
+			case ValueOpcode::FPMad32: {
+				// The rounded multiply and add are monotonic, so the bounds lie at the corners.
+				const auto a = range(0), b = range(1), c = range(2);
+				if (!std::isfinite(a.first) || !std::isfinite(a.second) ||
+				    !std::isfinite(b.first) || !std::isfinite(b.second)) return unknown;
+				const auto [low, high] = std::minmax({a.first * b.first, a.first * b.second,
+				                                      a.second * b.first, a.second * b.second});
+				const std::pair bounds {low + c.first, high + c.second};
+				if (!std::isfinite(bounds.first) || !std::isfinite(bounds.second)) return unknown;
+				return bounds;
+			}
+			default: return unknown;
+		}
+	}
+
 	uint32_t SelectorMaximum(Value value, SelectorBounds bounds = {}, uint32_t depth = 0,
 	                         Value guard = {}) const {
 		if (depth > 32u) return UINT32_MAX;
@@ -3466,9 +3517,13 @@ private:
 		switch (inst->GetOpcode()) {
 			case ValueOpcode::ReadFirstLane:
 				return SelectedSelectorDomain(*inst, guard).maximum;
-			case ValueOpcode::ReadLane:
+			case ValueOpcode::ReadLane: {
-				// ReadLane ignores EXEC; only context-free bounds on every source lane apply.
+				// ReadLane ignores EXEC; only context-free bounds on every source lane apply,
-				return PossibleU32Bits(inst->Arg(0));
+				// unless guard proves a waterfall key equal to the executing lanes' local value.
+				const auto local = EqualLocalKey(guard, value);
+				return local.IsEmpty() ? PossibleU32Bits(inst->Arg(0))
+				                       : SelectorMaximum(local, bounds, depth + 1u, guard);
+			}
 			case ValueOpcode::Phi: {
 				const auto initial = InitialCandidateMask(value);
 				if (!initial.IsEmpty()) return SelectorMaximum(initial, bounds, depth + 1u, guard);
@@ -3495,6 +3550,18 @@ private:
 				               NonzeroOnEntry(inst->Arg(0), inst->Parent()) ? 31u : UINT32_MAX;
 			case ValueOpcode::UMin32:
 			case ValueOpcode::BitwiseAnd32: return std::min(maximum(0), maximum(1));
+			case ValueOpcode::SMin32: {
+				// Signed and unsigned minimums agree on nonnegative operands.
+				const auto left = maximum(0), right = maximum(1);
+				if (std::max(left, right) <= INT32_MAX) return std::min(left, right);
+				break;
+			}
+			case ValueOpcode::ConvertS32F32: {
+				// The conversion truncates toward zero.
+				const auto [low, high] = FloatRange(inst->Arg(0), guard);
+				if (low >= 0.0f && high < 0x1p31f) return static_cast<uint32_t>(high);
+				break;
+			}
 			case ValueOpcode::SelectU32:
 				if (!guard.IsEmpty() && Implies(guard, inst->Arg(0))) return maximum(1);
 				return SelectorPredicateFalse(inst->Arg(0), bounds, depth + 1u, guard)

tests/ShaderRecompilerComputeTests.cpp

@@ -41907,6 +41907,89 @@ void CheckWorkgroupDividedImageTable(VulkanHarness &vulkan) {
   std::printf("[compute] %s ok\n", test.name);
 }
 
+void CheckWaterfallFloorImageTable(VulkanHarness &vulkan) {
+  using O = ShaderOpcode;
+  constexpr u32 base = 256u, table_offset = 0x100u, lanes = 64u, active = 60u;
+  constexpr float scale = 1.0f / 56.0f;
+  TestCase test;
+  test.name = "WaterfallFloorImageTable";
+  test.bda_mappings = {{0u, 0u}};
+  test.initial.resize((base + table_offset + 8u * 32u) / 4u);
+  for (u32 row = 0; row < 8u; ++row) {
+    auto image = MakeSampledTextureData(Prospero::BufferFormat::k32Float);
+    image[0] = row + 1u;
+    image[3] = DstSel(4, 0, 0, 1) | (static_cast<u32>(Prospero::ImageType::kColor2D) << 28u);
+    std::copy_n(image.begin(), 8, test.initial.begin() + (base + table_offset + row * 32u) / 4u);
+  }
+  test.user_data[0] = base;
+  const auto output = MakeStructuredStorageBufferData(0u, lanes * 4u);
+  std::copy_n(output.begin(), 4, test.user_data.begin() + 48);
+  test.has_user_data = test.has_compute_info = true;
+  test.compute_info.wave_size = 64u;
+  test.compute_info.threads_num[0] = lanes;
+  test.compute_info.thread_ids_num = 1u;
+  // Lanes below 60 compute row floor(7 - 7 * saturate(lane / 56)); a waterfall loop reads that
+  // row and row min(row + 1, 7).
+  auto &code = test.code;
+  AppendVMovLiteral(&code, 20, 0u);
+  AppendVMovLiteral(&code, 21, 0u);
+  code.insert(code.end(), {EncodeVopc(0xc4, InlineU32(active), 0), EncodeSop1(0x24, 38, 106)});
+  code.push_back(EncodeVop1(0x06, 1, Vgpr(0)));
+  AppendSMovLiteral(&code, 24, std::bit_cast<u32>(scale));
+  AppendVop3(&code, 0x108, 1, 24, Vgpr(1), 0, 0, 0, true);
+  AppendSMovLiteral(&code, 25, std::bit_cast<u32>(-7.0f));
+  code.insert(code.end(), {EncodeVop2(0x21, 2, 25, 1), std::bit_cast<u32>(7.0f),
+      EncodeVop1(0x0d, 3, Vgpr(2)), EncodeSop1(0x04, 34, 126)});
+  const auto loop = code.size();
+  code.push_back(EncodeSop1(0x14, 4, 34));
+  AppendVop3(&code, 0x360, 5, Vgpr(3), 4);
+  AppendVop3(&code, 0xc2, 36, 5, Vgpr(3));
+  code.insert(code.end(), {EncodeSop1(0x24, 30, 36),
+      EncodeSop2(0x1e, 6, 5, InlineU32(5)), EncodeSop2(0x02, 6, 6, 255), table_offset,
+      EncodeSmem0(0x03, 8), EncodeSmem1(0, 6),
+      EncodeMimg0(0x00, 1), EncodeMimg1(4, 20, 2, 10),
+      EncodeSop2(0x02, 7, 5, InlineU32(1)), EncodeSop2(0x06, 7, 7, InlineU32(7)),
+      EncodeSop2(0x1e, 7, 7, InlineU32(5)), EncodeSop2(0x02, 7, 7, 255), table_offset,
+      EncodeSmem0(0x03, 16), EncodeSmem1(0, 7),
+      EncodeMimg0(0x00, 1), EncodeMimg1(5, 20, 4, 10),
+      EncodeSop2(0x15, 34, 34, 36), EncodeSop1(0x04, 126, 30)});
+  code.push_back(EncodeSopp(0x05, loop - code.size() - 1u));
+  code.insert(code.end(), {EncodeVop2(0x20, 6, Vgpr(5), 4), std::bit_cast<u32>(16.0f),
+      EncodeVop2(0x1a, 30, InlineU32(2), 0)});
+  AppendBufferStoreDword(&code, 6, 30);
+  AppendEnd(&code);
+  test.opcodes = {O::V_MOV_B32, O::V_CMP_GT_U32, O::S_AND_SAVEEXEC_B64, O::V_CVT_F32_U32,
+                  O::S_MOV_B32, O::V_MUL_F32, O::V_MADAK_F32, O::V_CVT_FLR_I32_F32, O::S_MOV_B64,
+                  O::S_FF1_I32_B64, O::V_READLANE_B32, O::V_CMP_EQ_U32, O::S_LSHL_B32,
+                  O::S_ADD_I32, O::S_LOAD_DWORDX8, O::IMAGE_LOAD, O::S_MIN_I32, O::S_ANDN2_B64,
+                  O::S_CBRANCH_SCC1, O::V_MADMK_F32, O::V_LSHLREV_B32, O::BUFFER_STORE_DWORD,
+                  O::S_ENDPGM};
+  const auto compiled = CompileCase(test, vulkan.SubgroupSize());
+  Require(test.name, "table rows", compiled.resources.images.size() == 8u,
+          "descriptor planning did not bound the waterfall rows to the table");
+  std::vector<VulkanHarness::Image> images;
+  for (const auto &descriptor : compiled.resources.images)
+    images.push_back(vulkan.CreateImage2D(test.name, 1, 1, vk::Format::eR32Sfloat,
+        vk::ImageUsageFlagBits::eSampled,
+        {std::bit_cast<u32>(static_cast<float>(descriptor.dwords[0]))}, 1,
+        vk::ImageLayout::eShaderReadOnlyOptimal));
+  auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size(), true);
+  vulkan.Dispatch(test, compiled, buffer, nullptr, nullptr, nullptr, nullptr, {}, images);
+  std::vector<u32> expected;
+  for (u32 lane = 0; lane < active; ++lane) {
+    const float product = -7.0f * std::clamp(static_cast<float>(lane) * scale, 0.0f, 1.0f);
+    const auto row = static_cast<u32>(std::floor(product + 7.0f));
+    const auto next = std::min(row + 1u, 7u);
+    expected.push_back(std::bit_cast<u32>(static_cast<float>(row + 1u + 16u * (next + 1u))));
+  }
+  expected.resize(lanes);
+  CompareWords(test, "waterfall rows", expected,
+               vulkan.ReadBuffer(test.name, buffer, expected.size()));
+  vulkan.DestroyBuffer(&buffer);
+  for (auto &image : images) vulkan.DestroyImage(&image);
+  std::printf("[compute] %s ok\n", test.name);
+}
+
 void CheckNestedPointerImageTable(VulkanHarness &vulkan) {
   using O = ShaderOpcode;
   constexpr u32 list_address = 160;
@@ -52371,6 +52454,7 @@ int main(int argc, char **argv) {
     CheckCountedRowImageSelector(vulkan);
     CheckPointerImageTable(vulkan);
     CheckWorkgroupDividedImageTable(vulkan);
+    CheckWaterfallFloorImageTable(vulkan);
     CheckNestedPointerImageTable(vulkan);
     CheckGuardedRecordImageTable(vulkan);
     CheckIndirectImageKeySwitch(vulkan);
@@ -52550,6 +52634,7 @@ int main(int argc, char **argv) {
   CheckCountedRowImageSelector(vulkan);
   CheckPointerImageTable(vulkan);
   CheckWorkgroupDividedImageTable(vulkan);
+  CheckWaterfallFloorImageTable(vulkan);
   CheckNestedPointerImageTable(vulkan);
   CheckGuardedRecordImageTable(vulkan);
   CheckIndirectImageKeySwitch(vulkan);