perf(recompiler): scan image-table key buffers in bulk

core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp

@@ -1240,6 +1240,43 @@ void verifyImageTableBulkReads() {
     require(cutByRun->snapshot.tables == cutByWord->snapshot.tables && cutWordReads.wordReads == 12u && cutRunReads.wordReads == 0u && cutRunReads.bulkWords == 12u, "image table bulk reads: a record cut by an unmapped page read other words in runs than word by word");
 }
 
+void verifyImageTableBulkKeyScan() {
+    using namespace ShaderRecompiler;
+    namespace Abi = ImageTableAbi;
+    static TableTexture textures[2];
+    alignas(4096) static TableGuest guest;
+    std::fill(guest.heap.begin(), guest.heap.end(), std::array<std::uint32_t, 8>{});
+    guest.heap[0] = TableT(textures[0].bytes.data(), TableFormat8888UNorm, TableType2D);
+    guest.heap[1] = TableT(textures[1].bytes.data(), TableFormat32UInt, TableType2D);
+    guest.heap[3] = guest.heap[0];
+    guest.materials = {{{0u, 1u, 0u, 0u}, {0u, 0x08000003u, 0u, 0u}, {0u, 1u, 0u, 0u}}};
+    FillTableSrt(guest, guest.heap.data(), 32u, 4u);
+    const auto userData = TableUserData(guest);
+    const auto request = TableRequest(TableMaterialCode, userData);
+    const auto capture = [&](CountingMemory& memory, bool bulk) {
+        SrtRuntime runtime;
+        runtime.userData = request.context.userData;
+        runtime.shaderBase = request.shader.codeAddress;
+        runtime.userContext = &memory;
+        runtime.readMemory = &CountingRead;
+        runtime.readSpecializationMemory = &CountingRead;
+        runtime.accessible = &CountingAccessible;
+        if (bulk) runtime.readWords = &CountingReadWords;
+        return CaptureResources(request, runtime);
+    };
+    const auto begin = reinterpret_cast<std::uintptr_t>(guest.materials.data());
+    CountingMemory byWordReads{begin, begin + sizeof(guest.materials)};
+    CountingMemory byRunReads = byWordReads;
+    const auto byWord = capture(byWordReads, false);
+    const auto byRun = capture(byRunReads, true);
+    require(byWordReads.wordReads == guest.materials.size(), "image table bulk key scan: without a bulk reader the keys were not read one word each");
+    require(byRunReads.wordReads == 0u, "image table bulk key scan: the keys were read word by word although a bulk reader exists");
+    require(byRun->snapshot.tables == byWord->snapshot.tables, "image table bulk key scan: scanning the keys in bulk changed the snapshot");
+    const auto& snapshot = byRun->snapshot.tables.tables.at(0);
+    const auto outside = Abi::PoisonCode(static_cast<std::uint32_t>(Abi::PoisonReason::OutsideDomain));
+    require(snapshot.keys == 4u && snapshot.codes[0] == snapshot.codes[3] && snapshot.codes[0] != snapshot.codes[1] && snapshot.codes[2] == outside, "image table bulk key scan: the narrowed codes are wrong");
+}
+
 void verifyImageTables() {
     using namespace ShaderRecompiler;
     namespace Abi = ImageTableAbi;
@@ -5212,6 +5249,7 @@ int main(int argc, char** argv) {
             verifyImageTableKeys();
             verifyImageTables();
             verifyImageTableBulkReads();
+            verifyImageTableBulkKeyScan();
             verifyImageTableSrtKey();
             verifyImageTableMaskedOffset();
             verifyImageTableKeyUnion();

core/shader/recompiler/Optimization/src/ResourceMaterializer.cpp

@@ -505,27 +505,42 @@ class TableSnapshotter {
                 skip();
                 return;
             }
-            for (std::uint64_t index = 0; index < positions; index++) {
-                const auto address = base + first + index * domain.stride;
-                if (!PageAccessible(address, base, base + size)) {
-                    skip();
-                    return;
+            if (runtime.readWords != nullptr && positions != 0u && domain.stride % 4u == 0u && domain.stride <= 64u) {
+                const auto words = static_cast<std::size_t>(((positions - 1u) * domain.stride) / 4u + 1u);
+                for (std::uint64_t page = (base + first) & ~(PageBytes - 1u); page < base + first + words * 4u; page += PageBytes) {
+                    if (!PageAccessible(std::max(page, base + first), base, base + size)) {
+                        skip();
+                        return;
+                    }
+                }
+                if (keyWords.size() < words) keyWords.resize(words);
+                if (!runtime.readWords(runtime.userContext, base + first, keyWords.data(), words)) throw std::runtime_error("image table: failed to read guest memory at 0x" + hexText(base + first));
+                const auto step = static_cast<std::size_t>(domain.stride / 4u);
+                values.reserve(values.size() + positions);
+                for (std::size_t index = 0; index < words; index += step) values.push_back(keyWords[index]);
+            } else {
+                for (std::uint64_t index = 0; index < positions; index++) {
+                    const auto address = base + first + index * domain.stride;
+                    if (!PageAccessible(address, base, base + size)) {
+                        skip();
+                        return;
+                    }
+                    values.push_back(Read(address));
                 }
-                values.push_back(Read(address));
             }
             if (positions != 0u) scanned.emplace_back(base, size);
         }
         records.states.assign(records.keys, RecordState::OutsideDomain);
         records.outside = false;
+        const std::uint32_t unwrapped = std::numeric_limits<std::uint32_t>::max() / column.stride;
         for (const auto value : values) {
             const std::uint32_t relative = value * column.stride;
-            const bool aligned = relative % column.stride == 0u;
-            const auto record = relative / column.stride;
-            if (aligned && record < records.keys) {
+            const auto record = value <= unwrapped ? value : relative % column.stride == 0u ? relative / column.stride : records.keys;
+            if (record < records.keys) {
                 records.states[record] = RecordState::Read;
                 continue;
             }
-            if (column.address || ImageTableAbi::ScalarBufferDword(relative + column.addend, column.offset, records.size).has_value()) records.outside = true;
+            if (!records.outside && (column.address || ImageTableAbi::ScalarBufferDword(relative + column.addend, column.offset, records.size).has_value())) records.outside = true;
         }
         records.narrowed = true;
         for (const auto& [base, size] : scanned) tables.ranges.emplace_back(base, size);
@@ -774,6 +789,7 @@ class TableSnapshotter {
     std::unordered_map<WordsHash::Key, std::uint32_t, WordsHash> wordsIndex;
     std::unordered_map<std::uint64_t, bool> pages;
     std::unordered_map<std::uint64_t, bool> wholePages;
+    std::vector<std::uint32_t> keyWords;
 };
 
 void materializeSnapshot(const IrResourcePlan& plan, const SrtRuntime& runtime, SrtWalker& walker, ResourceSnapshot& snapshot, std::vector<std::uint8_t>& activeSources) {