shader: specialize bounded pointer and packed image tables

tests/ShaderRecompilerComputeTests.cpp

@@ -36505,6 +36505,113 @@ void CheckCountedRowImageSelector(VulkanHarness &vulkan) {
   }
 }
 
+void CheckPointerImageTable(VulkanHarness &vulkan) {
+  using O = ShaderOpcode;
+  constexpr u32 base = 256u, table_offset = 24364u;
+  for (const u32 variant : {0u, 1u, 2u})
+  for (const u32 wave : {32u, 64u}) {
+    const bool packed = variant == 2u;
+    const u32 count = variant * 2u;
+    TestCase test;
+    test.name = packed ? "PackedPointerImageTable" : "CountedPointerImageTable";
+    test.dispatch_x = packed ? 4u : 1u;
+    test.bda_mappings = {{0u, 0u}};
+    test.initial.resize((base + table_offset + 32u * 32u) / 4u);
+    std::fill_n(test.initial.begin(), test.dispatch_x, 0xdeadbeefu);
+    test.initial[(base + 492u) / 4u] = count;
+    for (u32 row = 0; row < 32u; ++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.initial[(base + 2756u) / 4u] = 1u;
+    test.initial[(base + 2756u + 48u) / 4u] = 31u;
+    constexpr std::array fields{31u, 0u, 7u, 1u};
+    for (u32 row = 0; row < fields.size(); ++row)
+      test.initial[(base + 5008u + row * 60u) / 4u] = 0xffe00abcu | (fields[row] << 12u);
+    test.user_data[0] = base;
+    const auto output = MakeStructuredStorageBufferData(0u, test.dispatch_x * 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 = wave;
+    test.compute_info.threads_num[0] = test.compute_info.threads_num[1] =
+        test.compute_info.threads_num[2] = 1u;
+    test.compute_info.group_id[0] = true;
+    test.compute_info.workgroup_register = 44u;
+    auto &code = test.code;
+    AppendVMovLiteral(&code, 0, 0u);
+    AppendVMovLiteral(&code, 20, 0x3f000000u);
+    AppendVMovLiteral(&code, 21, 0x3f000000u);
+    size_t loop = 0, exit = 0, skip = 0;
+    if (packed) {
+      code.insert(code.end(), {EncodeSop2(0x26, 23, 44, InlineU32(60)),
+          EncodeSop2(0x02, 23, 23, 255), 5008u,
+          EncodeSmem0(0x00, 24), EncodeSmem1(0, 23),
+          EncodeSop2(0x20, 24, 24, InlineU32(12)),
+          EncodeSop2(0x1e, 23, 24, InlineU32(5)),
+          EncodeSop2(0x0e, 23, 23, 255), 992u});
+    } else {
+      code.insert(code.end(), {EncodeSmem0(0x00, 21), EncodeSmem1(492, 125),
+                              EncodeSMovB32(20, InlineU32(0))});
+      loop = code.size();
+      code.push_back(EncodeSopc(0x0a, 20, 21));
+      exit = code.size(); code.push_back(0u);
+      code.insert(code.end(), {EncodeSop2(0x26, 23, 20, InlineU32(48)),
+          EncodeSop2(0x02, 23, 23, 255), 2756u,
+          EncodeSmem0(0x00, 24), EncodeSmem1(0, 23),
+          EncodeSopc(0x03, 24, InlineU32(0))});
+      skip = code.size(); code.push_back(0u);
+      code.push_back(EncodeSop2(0x1e, 23, 24, InlineU32(5)));
+    }
+    code.insert(code.end(), {EncodeSop2(0x02, 23, 23, 255), table_offset,
+        EncodeSmem0(0x03, 8), EncodeSmem1(0, 23),
+        EncodeMimg0(0x27, 1), EncodeMimg1(4, 20, 2, 10),
+        EncodeVop2(0x03, 0, Vgpr(4), 0)});
+    if (!packed) {
+      code[skip] = EncodeSopp(0x04, code.size() - skip - 1u);
+      code.push_back(EncodeSop2(0x02, 20, 20, InlineU32(1)));
+      code.push_back(EncodeSopp(0x02, loop - code.size() - 1u));
+      code[exit] = EncodeSopp(0x04, code.size() - exit - 1u);
+    }
+    code.push_back(EncodeSop2(0x1e, 25, 44, InlineU32(2)));
+    code.push_back(EncodeVop1(0x01, 30, 25));
+    AppendBufferStoreDword(&code, 0, 30);
+    AppendEnd(&code);
+    test.opcodes = {O::S_LOAD_DWORD, O::S_MUL_I32, O::S_ADD_I32, O::S_LSHL_B32,
+        O::S_LOAD_DWORDX8, O::IMAGE_SAMPLE, O::V_ADD_F32, O::V_MOV_B32,
+        O::BUFFER_STORE_DWORD, O::S_ENDPGM};
+    if (packed) test.opcodes.insert(test.opcodes.end(), {O::S_LSHR_B32, O::S_AND_B32});
+    else test.opcodes.insert(test.opcodes.end(), {O::S_MOV_B32, O::S_CMP_LT_U32,
+        O::S_CMP_GE_I32, O::S_CBRANCH_SCC0, O::S_BRANCH});
+    const auto compiled = CompileCase(test, vulkan.SubgroupSize());
+    Require(test.name, "bounded image candidates",
+            packed ? compiled.resources.images.size() == 32u
+                   : compiled.resources.images.size() == std::max(count, 1u),
+            "descriptor planning escaped its selector domain");
+    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);
+    auto sampler = vulkan.CreateSampler(test.name, 1);
+    vulkan.Dispatch(test, compiled, buffer, nullptr, nullptr, nullptr, nullptr, sampler, images);
+    std::vector<u32> expected(test.dispatch_x, std::bit_cast<u32>(count == 0u ? 0.0f : 34.0f));
+    if (packed) for (u32 row = 0; row < fields.size(); ++row)
+      expected[row] = std::bit_cast<u32>(static_cast<float>(fields[row] + 1u));
+    CompareWords(test, "bounded pointer selectors and packed fields", expected,
+                 vulkan.ReadBuffer(test.name, buffer, expected.size()));
+    vulkan.DestroyBuffer(&buffer);
+    vulkan.Device().destroySampler(sampler);
+    for (auto &image : images) vulkan.DestroyImage(&image);
+    std::printf("[compute] %s wave=%u count=%u ok\n", test.name, wave, count);
+  }
+}
+
 void CheckGuardedRecordImageTable(VulkanHarness &vulkan) {
   using O = ShaderOpcode;
   constexpr u32 table = 256u;
@@ -45134,6 +45241,7 @@ int main(int argc, char **argv) {
     CheckImageSamplerSpecialization();
     VulkanHarness vulkan;
     CheckCountedRowImageSelector(vulkan);
+    CheckPointerImageTable(vulkan);
     CheckGuardedRecordImageTable(vulkan);
     CheckIndirectImageKeySwitch(vulkan);
     CheckFiniteInlineSamplerPhi(vulkan);
@@ -45301,6 +45409,7 @@ int main(int argc, char **argv) {
   CheckRuntimeBufferRecords(vulkan);
   CheckComputeThreadDimensions(vulkan);
   CheckCountedRowImageSelector(vulkan);
+  CheckPointerImageTable(vulkan);
   CheckGuardedRecordImageTable(vulkan);
   CheckIndirectImageKeySwitch(vulkan);
   CheckDynamicSampleOffsets(vulkan, 32u);