shader: bound workgroup selector prefixes from runtime starting indices

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

@@ -577,8 +577,9 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 			    (material.Type() != 0u || material.SwizzleEnabled() || material.AddTid() ||
 			     material.OutOfBounds() != 0u || uint64_t {selector->offset} + 4u > step ||
 			     uint64_t {first} + count > material.NumRecords())) return false;
-			const auto end = count != 0u
+			const auto last = count != 0u ? uint64_t {first} + count - 1u : 0u;
-			    ? (uint64_t {first} + count - 1u) * step + selector->offset + 4u : 0u;
+			if (last > UINT32_MAX) return false;
+			const auto end = count != 0u ? last * step + selector->offset + 4u : 0u;
 			if (end > uint64_t {UINT32_MAX} + 1u ||
 			    (indirect.workgroup_axis == UINT32_MAX && end > size))
 				return false;

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

@@ -2204,6 +2204,7 @@ private:
 		const auto* memory = ScalarReadMemory(*read, memory_index);
 		Value index;
 		Value count;
+		Value first(0u);
 		uint32_t workgroup_axis = UINT32_MAX;
 		DescriptorSource count_source;
 		uint32_t row_stride = 0;
@@ -2214,6 +2215,20 @@ private:
 		    ((offset | stride | memory->offset) & 3u) == 0u &&
 		    uint64_t {offset} + memory->offset <= UINT32_MAX) {
 			workgroup_axis = WorkgroupAxis(index);
+			if (workgroup_axis == UINT32_MAX) {
+				const auto* sum = index.Resolve().TryInstruction();
+				if (sum != nullptr && sum->GetOpcode() == ValueOpcode::IAdd32) {
+					for (uint32_t arm = 0; arm < 2u; ++arm) {
+						const auto axis = WorkgroupAxis(sum->Arg(arm));
+						const auto base = sum->Arg(arm ^ 1u);
+						if (axis == UINT32_MAX ||
+						    !ValidateRuntimeValue(m_program, base, RuntimeValueType::Integer)) continue;
+						workgroup_axis = axis;
+						first = base;
+						break;
+					}
+				}
+			}
 			const auto* bound = BoundedLoop(index, read->Parent());
 			if (bound != nullptr && bound->GetOpcode() == ValueOpcode::ULessThan32) {
 				count = bound->Arg(1);
@@ -2269,7 +2284,7 @@ private:
 		if (rows) indirect.selector->rows.emplace(
 		    DescriptorSource::IndirectDescriptor::SelectorRead::CountedRows {
 		        .source = InternSource(count_source), .stride = row_stride});
-		if (!count.IsEmpty() || workgroup_axis != UINT32_MAX) indirect.selector_first = Value(0u);
+		if (!count.IsEmpty() || workgroup_axis != UINT32_MAX) indirect.selector_first = first;
 		return true;
 	}
 

tests/ResourceTrackingTests.cpp

@@ -5236,8 +5236,9 @@ void TestWorkgroupSelectedBufferRecords() {
 
 void TestGpuSelectedFormattedRecordStore() {
   namespace Prospero = Libs::Graphics::Prospero;
+  enum class Prefix { None, Forward, Reversed, Lane, TwoAxes };
   const auto make = [](uint32_t stride = 592u, uint32_t shift = 16u, bool split = false,
-                       bool wide = false) {
+                       bool wide = false, Prefix prefix = Prefix::None) {
     auto fixture = std::make_unique<Fixture>();
     auto& f = *fixture;
     std::array<Value, 4> table_words, selector_words;
@@ -5247,8 +5248,20 @@ void TestGpuSelectedFormattedRecordStore() {
     }
     const auto table = wide ? f.Address(table_words[0], table_words[1]) : f.Buffer(table_words);
     const auto lane = f.Emit(ValueOpcode::LaneId);
+    auto selector_index = lane;
+    if (prefix != Prefix::None) {
+      const auto group = f.Emit(ValueOpcode::GetBuiltin,
+          {Value(static_cast<uint32_t>(StageInputKind::WorkgroupId)), Value(0u)});
+      const auto first = prefix == Prefix::Lane ? lane : prefix == Prefix::TwoAxes
+          ? f.Emit(ValueOpcode::GetBuiltin,
+              {Value(static_cast<uint32_t>(StageInputKind::WorkgroupId)), Value(1u)})
+          : f.Emit(ValueOpcode::LoadAddressU32, {table, Value(0u), Value(0u), Value(true)},
+              f.AddMemory({.kind = ResourceKind::ScalarAddress, .offset = 0x234u}, 0x1088u));
+      selector_index = f.Emit(ValueOpcode::IAdd32,
+          {prefix == Prefix::Reversed ? group : first, prefix == Prefix::Reversed ? first : group});
+    }
     const auto packed = wide ? f.Emit(ValueOpcode::ReadConstBuffer,
-        {f.Buffer(selector_words), f.Emit(ValueOpcode::IMul32, {lane, Value(8u)})},
+        {f.Buffer(selector_words), f.Emit(ValueOpcode::IMul32, {selector_index, Value(8u)})},
         f.AddMemory({.kind = ResourceKind::ScalarBuffer, .offset = 4u}, 0x1c4cu)) :
         f.Emit(ValueOpcode::LoadBufferU32,
         {f.Buffer(selector_words), lane, Value(0u), Value(0u), Value(true)},
@@ -5514,6 +5527,97 @@ void TestGpuSelectedFormattedRecordStore() {
   for (const auto [shift, split] : {std::pair{0u, false}, std::pair{16u, true}})
     CheckFatal([&] { make(16u, shift, split, true); }, "not a valid runtime value",
                "finite RGBA16 store accepted an unbounded key or independently wrapping DWORDs");
+
+  struct PrefixMemory : LinearTestMemory { uint32_t read_attempts = 0u; } prefix_memory;
+  const auto read_prefix = [](void* userdata, uint64_t address, std::span<uint32_t> words) {
+    auto* memory = static_cast<PrefixMemory*>(userdata);
+    ++memory->read_attempts;
+    return ReadLinearTestMemory(memory, address, words);
+  };
+  std::ranges::fill(prefix_memory.words, 0xdeadbeefu);
+  constexpr uint32_t selector_offset = 0x1000u;
+  constexpr uint32_t first_offset = table_offset - 48u + 0x234u;
+  for (uint32_t row = 0; row < 2u; ++row) {
+    const std::array descriptor{0x100000u + row * 256u, 8u << 16u, 20u, rgba};
+    std::copy(descriptor.begin(), descriptor.end(),
+              prefix_memory.words.begin() + table_offset / 4u + row * 4u);
+  }
+  const std::array numeric_descriptor{0xbf8d214eu, 0xbf78904bu, 0x3e930d39u, 0x33d6bf95u};
+  std::copy(numeric_descriptor.begin(), numeric_descriptor.end(),
+            prefix_memory.words.begin() + table_offset / 4u + 4u * 4u);
+  for (uint32_t index = 0; index < 12u; ++index)
+    prefix_memory.words[selector_offset / 4u + index * 2u + 1u] =
+        ((index < 4u ? index & 1u : index < 6u ? 1u : 4u) << 16u) | 0xbeefu;
+  std::array<uint32_t, 8> prefix_data{uint32_t(prefix_memory.base) + table_offset - 48u, 0u, 0u, 0u,
+      uint32_t(prefix_memory.base) + selector_offset, 8u << 16u, 12u, 0x5204u};
+  std::array counts{2u, 1u, 1u};
+  SrtRuntime prefix_runtime{.user_data = prefix_data, .read_memory = read_prefix,
+      .userdata = &prefix_memory, .read_specialization_memory = read_prefix,
+      .workgroup_counts = counts};
+  prefix_memory.watched_address = prefix_memory.base + selector_offset + 4u;
+  for (const auto prefix : {Prefix::Forward, Prefix::Reversed}) {
+    const auto selected = make(16u, 16u, false, true, prefix);
+    auto selected_plan = ExtractResourcePlan(selected->program);
+    auto& indirect = *selected_plan.descriptor_sources[
+        selected_plan.info.buffers[wide_root].source].indirect_descriptor;
+    Check(indirect.workgroup_axis == 0u && !indirect.selector_first.IsImmediate() &&
+              indirect.key_count.IsEmpty() && indirect.selector && indirect.selector->stride == 8u &&
+              indirect.selector->offset == 4u && indirect.table_immediate == 48u,
+          "runtime-first selector lost its workgroup prefix or captured table geometry");
+    const auto refresh = [&] {
+      prefix_memory.watched_reads = 0u;
+      prefix_memory.read_attempts = 0u;
+      return MaterializeResources(selected_plan, prefix_runtime, snapshot, specialization);
+    };
+    for (const uint32_t first : {0u, 2u, 4u, 0u}) {
+      prefix_memory.words[first_offset / 4u] = first;
+      Check(refresh() && snapshot.buffers.size() == selected_plan.info.buffers.size() + (first == 4u ? 0u : 1u) &&
+                snapshot.buffers[wide_root].dwords[0] == 0x100000u + (first == 4u ? 256u : 0u),
+            "runtime-first refresh froze a prior selector range or enumerated unrelated numeric rows");
+      for (const auto read : {std::pair{prefix_memory.base + first_offset, uint64_t{4u}},
+                             std::pair{prefix_memory.base + selector_offset + first * 8u + 4u, uint64_t{12u}}})
+        Check(std::ranges::find(snapshot.specialization_reads, read) != snapshot.specialization_reads.end(),
+              "runtime first or selected prefix escaped specialization read capture");
+    }
+    prefix_memory.fail_address = prefix_memory.base + first_offset;
+    Check(!refresh() && prefix_memory.read_attempts == 1u,
+          "unreadable runtime first reached the selector payload");
+    prefix_memory.fail_address = UINT64_MAX;
+    prefix_memory.words[first_offset / 4u] = 6u;
+    Check(!refresh(), "an actually selected invalid numeric descriptor was silently omitted");
+    for (const uint32_t first : {UINT32_MAX, 0x20000000u}) {
+      prefix_memory.words[first_offset / 4u] = first;
+      Check(!refresh() && prefix_memory.read_attempts == 1u,
+            "wrapped first-plus-group or scaled prefix read selector payload");
+    }
+    prefix_memory.words[first_offset / 4u] = UINT32_MAX;
+    counts[0] = 7u;
+    indirect.selector->stride = 0xfffffffcu;
+    Check(!refresh() && prefix_memory.read_attempts == 1u,
+          "wrapped first-plus-group overflowed the host prefix product before rejection");
+    indirect.selector->stride = 8u;
+    counts[0] = 1u;
+    prefix_memory.words[first_offset / 4u] = 0x1fffffffu;
+    Check(refresh() && prefix_memory.watched_reads == 0u && snapshot.buffers[wide_root].dwords[0] == 0x100000u,
+          "last nonwrapping selector DWORD lost scalar OOB key zero");
+    prefix_memory.words[first_offset / 4u] = UINT32_MAX;
+    for (const auto empty : {std::array{0u, 1u, 1u}, std::array{2u, 0u, 1u}, std::array{2u, 1u, 0u}}) {
+      counts = empty;
+      Check(refresh() && prefix_memory.watched_reads == 0u &&
+                snapshot.buffers.size() == selected_plan.info.buffers.size() &&
+                snapshot.buffers[wide_root].dwords[0] == 0u,
+            "empty dispatch evaluated an unreachable wrapped selector prefix");
+    }
+    counts = {2u, 1u, 1u};
+  }
+  for (const auto prefix : {Prefix::Lane, Prefix::TwoAxes}) {
+    const auto rejected = make(16u, 16u, false, true, prefix);
+    const auto& indirect = rejected->program.descriptor_sources[
+        rejected->program.info.buffers[wide_root].source].indirect_descriptor;
+    Check(indirect && indirect->workgroup_axis == UINT32_MAX && !indirect->selector &&
+              indirect->key_count == Value(32u),
+          "lane-varying or multiple-axis first incorrectly narrowed the finite descriptor domain");
+  }
 }
 
 void TestFiniteBufferSelection() {

tests/ShaderRecompilerComputeTests.cpp

@@ -10968,14 +10968,25 @@ public:
       constexpr auto record_keys = base + 0xd1000;
       constexpr auto record_table = base + 0xd2000;
       constexpr auto record_output = base + 0xd7000;
-      std::vector<u32> record_shader{
+      for (const bool based : {false, true}) {
-          EncodeSop2(0x26, 16, 13, InlineU32(12)),
+        std::vector<u32> record_shader;
-          EncodeSmem0(0x08, 16, 0), EncodeSmem1(0, 16), EncodeSopp(0x0c, 0),
+        if (based) {
+          // The later captured shader starts its prefix at a clean SRT DWORD.
+          record_shader = {EncodeSmem0(0x00, 18, 2), EncodeSmem1(0x234, 125), EncodeSopp(0x0c, 0),
+              EncodeSop2(0x02, 18, 18, 13),
+              EncodeSop2(0x1e, 16, 18, InlineU32(3)),
+              EncodeSop2(0x02, 16, 16, InlineU32(4))};
+        } else record_shader = {EncodeSop2(0x26, 16, 13, InlineU32(12))};
+        record_shader.insert(record_shader.end(), {
+            EncodeSmem0(0x08, 16, 0), EncodeSmem1(0, 16), EncodeSopp(0x0c, 0)});
+        if (based) record_shader.insert(record_shader.end(), {
+            EncodeSop2(0x27, 16, 16, 255), (5u << 16u) | 16u});
+        record_shader.insert(record_shader.end(), {
             EncodeSop2(0x26, 17, 16, 255), 192u,
-          EncodeSmem0(0x0a, 20, 2), EncodeSmem1(160, 17),
+            EncodeSmem0(based ? 0x02 : 0x0a, 20, 2), EncodeSmem1(160, 17),
-          EncodeSmem0(0x0a, 24, 2), EncodeSmem1(176, 17), EncodeSopp(0x0c, 0),
+            EncodeSmem0(based ? 0x02 : 0x0a, 24, 2), EncodeSmem1(176, 17), EncodeSopp(0x0c, 0),
             EncodeMubuf0(0x83, 0, true, false), EncodeMubuf1(4, 5, 0),
-          EncodeSopp(0x0c, 0)};
+            EncodeSopp(0x0c, 0)});
         AppendVop3(&record_shader, 0x346, 8, 13, 12, Vgpr(0));
         record_shader.insert(record_shader.end(), {
             EncodeMubuf0(0x1d, 0, true, false), EncodeMubuf1(4, 2, 8)});
@@ -11017,28 +11028,32 @@ public:
           processor.DispatchDirect(1, 1, 1, 0x41u);
         };
         for (const u32 wave : {32u, 64u}) {
+          u32 first = based ? 2u : 0u;
           for (const auto dimensions : {std::array{2u, 1u, 1u}, std::array{0u, 1u, 1u},
                                        std::array{2u, 0u, 1u}, std::array{2u, 1u, 0u},
                                        std::array{1u, 1u, 1u}, std::array{2u, 1u, 1u}}) {
             const bool empty = std::ranges::find(dimensions, 0u) != dimensions.end();
             const std::array<u32, 3> stale_args{};
-          std::array<u32, 7> stale_keys;
+            std::array<u32, 10> stale_keys;
             stale_keys.fill(sentinel);
             Require(name, "record fixture reset", context.InvalidateMemory(record_args, 0x8000),
                     "record fixture escaped its mapped allocation");
             LibKernel::Memory::WriteBacking(record_args, stale_args.data(), sizeof(stale_args));
             LibKernel::Memory::WriteBacking(record_keys, stale_keys.data(), sizeof(stale_keys));
+            LibKernel::Memory::WriteBacking(record_table + 0x234, &first, sizeof(first));
             std::vector<u32> expected_output(wave * 4u, sentinel);
             LibKernel::Memory::WriteBacking(record_output, expected_output.data(), expected_output.size() * 4u);
             const std::array<u32, 4> untouched{sentinel, sentinel, sentinel, sentinel};
             for (u32 row = 0; row < 2u; ++row)
               LibKernel::Memory::WriteBacking(base + 0xd5000 + row * 0x1000, untouched.data(), sizeof(untouched));
             produce(record_args, 4u, dimensions);
-          produce(record_keys, 12u, {1u, 0u, 22369622u});
+            produce(record_keys + (based ? first * 8u + 4u : 0u), based ? 8u : 12u,
+                    based ? std::array{0x1beefu, 0xbeefu, 0x1fbeefu}
+                          : std::array{1u, 0u, 22369622u});
             std::array<u32, 3> cpu_args{};
-          std::array<u32, 7> cpu_keys{};
+            std::array<u32, 10> cpu_keys{};
             Require(name, "GPU record dependencies remain pending",
-                  cache.HasGpuDirtyBytes(record_args, 12u) && cache.HasGpuDirtyBytes(record_keys, 16u) &&
+                    cache.HasGpuDirtyBytes(record_args, 12u) && cache.HasGpuDirtyBytes(record_keys, sizeof(stale_keys)) &&
                         LibKernel::Memory::TryReadBacking(record_args, cpu_args.data(), sizeof(cpu_args)) &&
                         LibKernel::Memory::TryReadBacking(record_keys, cpu_keys.data(), sizeof(cpu_keys)) &&
                         cpu_args == stale_args && cpu_keys == stale_keys,
@@ -11046,8 +11061,9 @@ public:
             shaders.SetCsShader({.data_addr = reinterpret_cast<uint64_t>(record_shader.data()),
                                  .num_thread_x = wave, .num_thread_y = 1, .num_thread_z = 1,
                                  .user_sgpr = 13, .tgid_x_en = true});
-          set_buffer(0, record_keys, 3u, 12u);
+            set_buffer(0, record_keys, based ? 5u : 3u, based ? 8u : 12u);
             set_buffer(4, record_table, 2u, 192u);
+            if (based) shaders.SetCsUserSgpr(5, record_table >> 32u, HW::UserSgprType::Unknown);
             set_buffer(8, record_output, wave * 2u, 8u);
             shaders.SetCsUserSgpr(12, std::countr_zero(wave), HW::UserSgprType::Unknown);
             processor.DispatchIndirect(record_args, 0x41u |
@@ -11081,6 +11097,8 @@ public:
                           target, actual.data(), sizeof(actual)) && actual == expected,
                       "masked/OOB record stores changed an inactive row or adjacent bytes");
             }
+            if (based) first ^= 2u;
+          }
         }
       }