Program cache: one head reader for both readers of a record (the prediction finds its programs again)

gpu/vk/vk_shader.cpp

@@ -1104,6 +1104,35 @@ struct Reader {
     }
 };
 
+/// (hookshot) A program record's head, what comes before its stages: written and read only here, so that program_cache_read and
+/// program_cache_each read the same fields. program_cache_each read only the first ones and took the mesh fallback's five words
+/// for the stage count: no kept program reached the prediction (vk_predict.cpp, "predicted 0 pipelines")
+void put_head(std::vector<uint8_t>& o, const Program& p) {
+    put(o, (uint32_t)((p.compute ? 1 : 0) | (p.mesh ? 2 : 0) | (p.fallback ? 4 : 0) | (p.fallback_task ? 8 : 0)));
+    put_list(o, p.answers);
+    put_list(o, p.questions);
+    put(o, p.workgroup);
+    put(o, p.shared_memory);
+    put(o, (uint32_t)p.fallback_topology);
+    put(o, p.fallback_vertices_per_workgroup);
+    put(o, p.fallback_record_words);
+    put(o, p.fallback_payload_words);
+    put(o, p.fallback_task_threads);
+}
+void get_head(Reader& r, Program& p) {
+    const uint32_t kind = r.get<uint32_t>();
+    p.compute = (kind & 1) != 0, p.mesh = (kind & 2) != 0, p.fallback = (kind & 4) != 0, p.fallback_task = (kind & 8) != 0;
+    r.list(p.answers);
+    r.list(p.questions);
+    p.workgroup = r.get<std::array<uint32_t, 3>>();
+    p.shared_memory = r.get<uint32_t>();
+    p.fallback_topology = (VkPrimitiveTopology)r.get<uint32_t>();
+    p.fallback_vertices_per_workgroup = r.get<uint32_t>();
+    p.fallback_record_words = r.get<uint32_t>();
+    p.fallback_payload_words = r.get<uint32_t>();
+    p.fallback_task_threads = r.get<uint32_t>();
+}
+
 /// Shader::Info is plain data but for these members (the boost containers are heap vectors here, static_vector too);
 /// they go on their own, the rest as it lies
 void put_info(std::vector<uint8_t>& o, const Shader::Info& i) {
@@ -1154,17 +1183,7 @@ void program_cache_save(const Program& p, const std::vector<std::vector<uint32_t
     if (!disk_on || spirv.size() != p.stages.size()) return;
     std::vector<uint8_t> o;
     o.reserve(256 << 10);
-    put(o, (uint32_t)((p.compute ? 1 : 0) | (p.mesh ? 2 : 0) | (p.fallback ? 4 : 0) |
-                      (p.fallback_task ? 8 : 0)));
-    put_list(o, p.answers);
-    put_list(o, p.questions);
-    put(o, p.workgroup);
-    put(o, p.shared_memory);
-    put(o, (uint32_t)p.fallback_topology);
-    put(o, p.fallback_vertices_per_workgroup);
-    put(o, p.fallback_record_words);
-    put(o, p.fallback_payload_words);
-    put(o, p.fallback_task_threads);
+    put_head(o, p);
     put(o, (uint64_t)p.stages.size());
     for (size_t k = 0; k < p.stages.size(); k++) {
         const Stage& st = p.stages[k];
@@ -1214,19 +1233,8 @@ void program_cache_read(FILE* in, std::mutex* mu, uint64_t key, const std::vecto
         Reader r{data.data() + 16, data.data() + data.size()};
         auto p = std::make_unique<Program>();
         p->key = key;
-        const uint32_t kind = r.get<uint32_t>();
-        p->compute = (kind & 1) != 0, p->mesh = (kind & 2) != 0, p->fallback = (kind & 4) != 0,
-        p->fallback_task = (kind & 8) != 0;
-        r.list(p->answers);
-        r.list(p->questions);
+        get_head(r, *p);
         p->variant = answers_hash(*p);
-        p->workgroup = r.get<std::array<uint32_t, 3>>();
-        p->shared_memory = r.get<uint32_t>();
-        p->fallback_topology = (VkPrimitiveTopology)r.get<uint32_t>();
-        p->fallback_vertices_per_workgroup = r.get<uint32_t>();
-        p->fallback_record_words = r.get<uint32_t>();
-        p->fallback_payload_words = r.get<uint32_t>();
-        p->fallback_task_threads = r.get<uint32_t>();
         const uint64_t n = r.count(1);
         bool ok = !r.bad && n <= 8;   // (a program has at most a few stages)
         for (uint64_t k = 0; ok && k < n; k++) {
@@ -2049,12 +2057,8 @@ void program_cache_each(const std::function<void(const KeptProgram&)>& fn) {
         Reader r{data.data(), data.data() + data.size()};
         Program p;
         p.key = key;
-        const uint32_t kind = r.get<uint32_t>();
-        r.list(p.answers);
-        r.list(p.questions);   // (the state list's, vk_draw.cpp)
-        k.key = key, k.variant = answers_hash(p), k.mesh = (kind & 2) != 0;
-        r.get<std::array<uint32_t, 3>>();
-        r.get<uint32_t>();   // (shared memory)
+        get_head(r, p);   // (the whole head: reading past only part of it took the fallback's fields for the stage count)
+        k.key = key, k.variant = answers_hash(p), k.mesh = p.mesh;
         const uint64_t n = r.count(1);
         k.stages.resize(n <= 8 ? n : 0);
         for (auto& [stage, spirv] : k.stages) {

tests/run_all.py

@@ -39,6 +39,7 @@
     "tests/test_cpp_preflight.py": ["--cxx", "g++"],
     "tests/test_invalidate_vcore.py": ["--cxx", "g++"],
     "tests/test_seed_cs_pointers.py": ["--cxx", "g++"],
+    "tests/test_program_record_head.py": ["--cxx", "g++"],
     "tests/test_lower_atomg_cas.py": ["--cxx", "g++"],
     "nx2/test_cbuf_table_extent.py": ["--cxx", "g++"],
     "nx2/test_dump_ssbo_read.py": ["--cxx", "g++"],

tests/test_program_record_head.py

@@ -0,0 +1,117 @@
+"""A program record's head is read whole by both readers of hookshot_programs.bin.
+
+program_cache_save writes, before a program's stages, its kind, texture answers and questions, workgroup, shared memory and the
+mesh fallback's five words. program_cache_read read all of them; program_cache_each (the prediction's, vk_predict.cpp) stopped
+after shared memory and took the fallback's first two words for the stage count: 0 for every program that is not a fallback one,
+so no kept program reached the prediction ("predicted 0 pipelines ... for 0 program variants"). The head is now written by one
+function (put_head) and read by one (get_head). This builds put/put_list/Reader/put_head/get_head as they are in the source, with
+a stand-in Program, and checks that a fallback and an ordinary program each come back whole with the stage count after them; and
+that program_cache_save, program_cache_read and program_cache_each use those two functions.
+python -I -B tests/test_program_record_head.py [--cxx g++]
+"""
+import argparse
+import subprocess
+import tempfile
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+SHADER = ROOT / "gpu" / "vk" / "vk_shader.cpp"
+
+
+def block(text, marker, end="}"):
+    start = text.index(marker)
+    at = text.index("{", start) + 1
+    depth = 1
+    while depth:
+        depth += (text[at] == "{") - (text[at] == "}")
+        at += 1
+    return text[start:at] + (";" if end == "};" else "")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--cxx", default="g++")
+    a = ap.parse_args()
+    src = SHADER.read_text(encoding="utf-8")
+    parts = [
+        block(src, "void put_bytes("),
+        src[src.index("template <class T> void put("):src.index("\n", src.index("template <class T> void put("))],
+        block(src, "template <class V> void put_list("),
+        block(src, "struct Reader {", "};"),
+        block(src, "void put_head("),
+        block(src, "void get_head("),
+    ]
+    cpp = r"""
+#include <array>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <vector>
+enum VkPrimitiveTopology { VK_PRIMITIVE_TOPOLOGY_POINT_LIST = 0, VK_PRIMITIVE_TOPOLOGY_LINE_LIST = 1, VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST = 3 };
+struct Program {   // (the fields of vk_shader.h's Program that a record's head holds)
+    bool compute = false, mesh = false, fallback = false, fallback_task = false;
+    std::array<uint32_t, 3> workgroup{1, 1, 1};
+    uint32_t shared_memory = 0;
+    VkPrimitiveTopology fallback_topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
+    uint32_t fallback_vertices_per_workgroup = 1, fallback_record_words = 1, fallback_payload_words = 0, fallback_task_threads = 1;
+    struct Answer {
+        uint16_t stage, desc;
+        uint8_t type;
+        bool image = false;
+    };
+    std::vector<Answer> answers;
+    std::vector<uint8_t> questions;
+};
+""" + "\n".join(parts) + r"""
+int fails = 0;
+#define WANT(c) do { if (!(c)) std::printf("%s: %s\n", name, #c), fails++; } while (0)
+void check(const char* name, const Program& p, uint64_t stages) {
+    std::vector<uint8_t> o;
+    put_head(o, p);
+    put(o, stages);   // (as program_cache_save: the stage count right after the head)
+    Reader r{o.data(), o.data() + o.size()};
+    Program q;
+    get_head(r, q);
+    WANT(r.get<uint64_t>() == stages);
+    WANT(!r.bad && r.p == r.end);
+    WANT(q.compute == p.compute && q.mesh == p.mesh && q.fallback == p.fallback && q.fallback_task == p.fallback_task);
+    WANT(q.workgroup == p.workgroup && q.shared_memory == p.shared_memory);
+    WANT(q.fallback_topology == p.fallback_topology && q.fallback_vertices_per_workgroup == p.fallback_vertices_per_workgroup);
+    WANT(q.fallback_record_words == p.fallback_record_words && q.fallback_payload_words == p.fallback_payload_words);
+    WANT(q.fallback_task_threads == p.fallback_task_threads);
+    WANT(q.answers.size() == p.answers.size() && q.questions == p.questions);
+    for (size_t i = 0; i < q.answers.size() && i < p.answers.size(); i++)
+        WANT(q.answers[i].stage == p.answers[i].stage && q.answers[i].desc == p.answers[i].desc && q.answers[i].type == p.answers[i].type);
+}
+int main() {
+    Program plain;
+    plain.answers = {{1, 3, 0x81, false}};
+    plain.questions = {1, 0, 2};
+    plain.shared_memory = 256;
+    check("graphics", plain, 2);
+    Program fb;
+    fb.mesh = false, fb.fallback = true, fb.fallback_task = true;
+    fb.workgroup = {32, 1, 1};
+    fb.fallback_topology = VK_PRIMITIVE_TOPOLOGY_LINE_LIST;
+    fb.fallback_vertices_per_workgroup = 384, fb.fallback_record_words = 9, fb.fallback_payload_words = 4, fb.fallback_task_threads = 8;
+    check("mesh fallback", fb, 4);
+    if (fails) return 1;
+    std::puts("ok");
+}
+"""
+    with tempfile.TemporaryDirectory(prefix="program-head-") as tmp:
+        c, exe = Path(tmp) / "fixture.cpp", Path(tmp) / "fixture"
+        c.write_text(cpp, encoding="utf-8")
+        subprocess.run([a.cxx, "-std=c++20", "-Wall", str(c), "-o", str(exe)], check=True)
+        got = subprocess.run([str(exe)], capture_output=True, text=True)
+        if got.returncode or got.stdout.strip() != "ok":
+            raise SystemExit("a program record's head:\n" + got.stdout + got.stderr)
+    for marker, want in (("void program_cache_save(", "put_head(o, p);"), ("void program_cache_read(", "get_head(r, *p);"),
+                         ("void program_cache_each(", "get_head(r, p);")):
+        if want not in block(src, marker):
+            raise SystemExit("%s does not use %s" % (marker.rstrip("("), want))
+    print("PASS: a program record's head comes back whole, fallback or not, in both readers")
+
+
+if __name__ == "__main__":
+    main()