graphics: support dynamic storage mip writes

src/graphics/host_gpu/renderer/image/image.h

@@ -55,6 +55,7 @@ struct ImageBinding {
 	bool is_target     = false;
 	bool needs_rebind  = false;
 	bool force_general = false;
+	bool shader_write  = false;
 };
 
 class Image final {

src/graphics/host_gpu/renderer/image/imageView.h

@@ -120,7 +120,11 @@ IsSupportedStorageImageResource(const ShaderRecompiler::IR::ImageResource& resou
 	        resource.dimension == ShaderRecompiler::Decoder::ImageDimension::Dim2D ||
 	        resource.dimension == ShaderRecompiler::Decoder::ImageDimension::Dim3D ||
 	        resource.dimension == ShaderRecompiler::Decoder::ImageDimension::Dim2DArray) &&
-	       resource.mip_mode == ShaderRecompiler::IR::ImageMipMode::None && resource.written &&
+	       ((resource.mip_mode == ShaderRecompiler::IR::ImageMipMode::None &&
+	         resource.mip_count == 1u) ||
+	        (resource.mip_mode == ShaderRecompiler::IR::ImageMipMode::DynamicStorage &&
+	         resource.mip_count != 0u)) &&
+	       resource.written &&
 	       (!resource.atomic ||
 	        (resource.kind == ShaderRecompiler::IR::ResourceKind::StorageImageUint &&
 	         resource.read)) &&

src/graphics/host_gpu/renderer/pipeline/descriptorCache.cpp

@@ -94,10 +94,14 @@ vk::DescriptorBufferInfo BufferInfo(const BufferView& view) {
 
 } // namespace
 
-vk::DescriptorImageInfo DescriptorCache::MakeImageInfo(const TextureBinding& texture) {
+vk::DescriptorImageInfo DescriptorCache::MakeImageInfo(const TextureBinding& texture,
-	EXIT_IF(!texture.image_id || texture.image_view == nullptr ||
+                                                       uint32_t              element) {
-	        texture.layout == vk::ImageLayout::eUndefined);
+	const auto view =
-	return {nullptr, texture.image_view, texture.layout};
+	    texture.mip_views.empty()
+	        ? (element == 0u ? texture.image_view : vk::ImageView {})
+	        : (element < texture.mip_views.size() ? texture.mip_views[element] : vk::ImageView {});
+	EXIT_IF(!texture.image_id || view == nullptr || texture.layout == vk::ImageLayout::eUndefined);
+	return {nullptr, view, texture.layout};
 }
 
 DescriptorCache::~DescriptorCache() {
@@ -229,14 +233,17 @@ VulkanDescriptorSet& DescriptorCache::GetDescriptor(Stage
 	auto* set = Allocate(stage, program);
 	EXIT_NOT_IMPLEMENTED(set == nullptr);
 
-	const auto descriptor_count = program.info.buffers.size() + program.info.images.size() +
+	size_t descriptor_count = 0;
-	                              program.info.samplers.size() + program.info.addresses.size() + 3u;
+	for (const auto& binding: program.bindings.descriptors) {
+		descriptor_count += DescriptorCount(binding);
+	}
 	std::vector<vk::DescriptorBufferInfo> buffer_infos;
 	std::vector<vk::DescriptorImageInfo>  image_infos;
 	std::vector<vk::WriteDescriptorSet>   writes;
 	buffer_infos.reserve(descriptor_count);
 	image_infos.reserve(descriptor_count);
 	writes.reserve(program.bindings.descriptors.size());
+	std::vector<uint32_t> image_occurrences(data.images.size());
 
 	for (const auto& binding: program.bindings.descriptors) {
 		vk::WriteDescriptorSet write {};
@@ -273,7 +280,7 @@ VulkanDescriptorSet& DescriptorCache::GetDescriptor(Stage
 			default: {
 				for (const auto resource: binding.resources) {
 					const auto& texture = data.images.at(resource);
-					image_infos.push_back(MakeImageInfo(texture));
+					image_infos.push_back(MakeImageInfo(texture, image_occurrences.at(resource)++));
 				}
 				break;
 			}
@@ -286,6 +293,12 @@ VulkanDescriptorSet& DescriptorCache::GetDescriptor(Stage
 		}
 		writes.push_back(write);
 	}
+	for (uint32_t i = 0; i < data.images.size(); i++) {
+		const auto expected = data.images[i].mip_views.empty()
+		                          ? 1u
+		                          : static_cast<uint32_t>(data.images[i].mip_views.size());
+		EXIT_IF(image_occurrences[i] != expected);
+	}
 	m_graphics.device.updateDescriptorSets(static_cast<uint32_t>(writes.size()), writes.data(), 0,
 	                                       nullptr);
 	return *set;

src/graphics/host_gpu/renderer/pipeline/descriptorCache.h

@@ -47,10 +47,11 @@ public:
 	enum class Stage { Unknown, Vertex, Pixel, Compute };
 
 	struct TextureBinding {
-		ImageId                 image_id;
+		ImageId                    image_id;
-		vk::ImageView           image_view = nullptr;
+		vk::ImageView              image_view = nullptr;
-		TextureCache::ImageDesc desc;
+		TextureCache::ImageDesc    desc;
-		vk::ImageLayout         layout = vk::ImageLayout::eUndefined;
+		vk::ImageLayout            layout = vk::ImageLayout::eUndefined;
+		std::vector<vk::ImageView> mip_views;
 	};
 
 	struct NativeDescriptors {
@@ -94,7 +95,8 @@ private:
 		int                next_free_pool = -1;
 	};
 
-	static vk::DescriptorImageInfo MakeImageInfo(const TextureBinding& texture);
+	static vk::DescriptorImageInfo MakeImageInfo(const TextureBinding& texture,
+	                                             uint32_t              element = 0);
 	void                           CreatePool();
 	VulkanDescriptorSet* Allocate(Stage stage, const ShaderRecompiler::IR::Program& program);
 	vk::DescriptorSetLayout

src/graphics/host_gpu/renderer/pipeline/descriptors.cpp

@@ -786,9 +786,10 @@ void RenderExecutor::BindImage(ImageId id, bool storage) {
 		return;
 	}
 	if (image.binding.is_bound) {
-		image.binding.force_general |= storage;
+		image.binding.force_general |= image.binding.shader_write != storage;
 	}
 	image.binding.is_bound = true;
+	image.binding.shader_write |= storage;
 	TrackImageBinding(id);
 }
 
@@ -904,8 +905,27 @@ void RenderExecutor::RebindImages(CommandBuffer&                     buffer,
 			BindImage(images[i].image_id,
 			          images[i].desc.type == TextureCache::BindingType::Storage);
 		}
-		auto& binding      = images[i];
+		auto& binding = images[i];
-		binding.image_view = texture_cache.FindTexture(binding.image_id, binding.desc);
+		binding.mip_views.clear();
+		const auto& resource = program.info.images[i];
+		if (resource.mip_mode == ShaderRecompiler::IR::ImageMipMode::DynamicStorage) {
+			EXIT_IF(resource.mip_count == 0u ||
+			        resource.mip_count != binding.desc.view_info.level_count);
+			binding.mip_views.reserve(resource.mip_count);
+			for (uint32_t mip = 0; mip < resource.mip_count; mip++) {
+				auto desc = binding.desc;
+				desc.view_info.base_level += mip;
+				desc.view_info.level_count = 1;
+				binding.mip_views.push_back(texture_cache.FindTexture(binding.image_id, desc));
+			}
+			binding.image_view = binding.mip_views.front();
+		} else {
+			auto desc = binding.desc;
+			if (desc.type == TextureCache::BindingType::Storage) {
+				desc.view_info.level_count = 1;
+			}
+			binding.image_view = texture_cache.FindTexture(binding.image_id, desc);
+		}
 		auto&      image   = texture_cache.GetImage(binding.image_id);
 		const bool storage = binding.desc.type == TextureCache::BindingType::Storage;
 		image.usage.storage |= storage;
@@ -994,8 +1014,11 @@ void RenderExecutor::CommitBindings(CommandBuffer&                     buffer,
 			              range, vk_buffer);
 		} else if ((image.binding.force_general || image.binding.is_target) &&
 		           !image.info.IsDepth()) {
+			const vk::AccessFlags2 storage_access = image.binding.shader_write
+			                                            ? vk::AccessFlagBits2::eShaderWrite
+			                                            : vk::AccessFlags2 {};
 			image.Transit(vk::ImageLayout::eGeneral,
-			              vk::AccessFlagBits2::eShaderRead |
+			              vk::AccessFlagBits2::eShaderRead | storage_access |
 			                  vk::AccessFlagBits2::eColorAttachmentRead |
 			                  vk::AccessFlagBits2::eColorAttachmentWrite,
 			              {}, vk_buffer);

src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp

@@ -90,7 +90,13 @@ bool ValidateNativeProgram(const IR::Program& program, std::string* error) {
 			return Fail(error, "native shader plan has an invalid image class");
 		}
 		present[static_cast<size_t>(kind)] = true;
-		expected[static_cast<size_t>(kind)].push_back(i);
+		const auto dynamic = program.info.images[i].mip_mode == IR::ImageMipMode::DynamicStorage;
+		const auto count   = dynamic ? program.info.images[i].mip_count : 1u;
+		if (count == 0u || (!dynamic && program.info.images[i].mip_count != 1u)) {
+			return Fail(error, "native shader plan has an invalid image mip descriptor count");
+		}
+		expected[static_cast<size_t>(kind)].insert(expected[static_cast<size_t>(kind)].end(), count,
+		                                           i);
 	}
 	if (!program.info.samplers.empty()) {
 		Expect(Kind::Samplers, Dense(program.info.samplers.size()));

src/graphics/shader/recompiler/backend/spirv/spirvEmitterAnalysis.cpp

@@ -469,6 +469,48 @@ uint32_t LoadStorageImageDescriptor(EmitterState& state, uint32_t resource, bool
 	return image;
 }
 
+void EmitStorageImageWrite(EmitterState& state, uint32_t resource, bool uint_image,
+                           ImageViewKind view, uint32_t mip_lod, uint32_t coord, uint32_t texel) {
+	const auto  kind    = StorageBindingKind(uint_image, view);
+	const auto  binding = ResourceForDescriptor(state, kind, resource);
+	const auto& image   = state.program.info.images.at(resource);
+	const auto  LoadAt  = [&](uint32_t array_index) {
+		const auto pointer = DescriptorElementPointer(
+		    state, StorageImagePointerType(state, uint_image, view),
+		    StorageImageVariable(state, uint_image, view), array_index, kind, resource,
+		    "storage image descriptor array was not emitted");
+		const auto descriptor = state.builder.AllocateId();
+		state.builder.AddFunction(
+		    {OpLoad, StorageImageType(state, uint_image, view), descriptor, pointer});
+		return descriptor;
+	};
+	if (image.mip_mode != IR::ImageMipMode::DynamicStorage) {
+		state.builder.AddFunction({OpImageWrite, LoadAt(binding.array_index), coord, texel});
+		return;
+	}
+	if (image.mip_count == 0u) {
+		ExitDescriptorBindingFailure(state, kind, resource,
+		                             "dynamic storage image has no mip descriptors");
+	}
+
+	const auto            merge_label = state.builder.AllocateId();
+	std::vector<uint32_t> labels(image.mip_count);
+	std::vector<uint32_t> words {OpSwitch, mip_lod, merge_label};
+	for (uint32_t mip = 0; mip < image.mip_count; mip++) {
+		labels[mip] = state.builder.AllocateId();
+		words.push_back(mip);
+		words.push_back(labels[mip]);
+	}
+	state.builder.AddFunction({OpSelectionMerge, merge_label, SelectionControlNone});
+	state.builder.AddFunction(words);
+	for (uint32_t mip = 0; mip < image.mip_count; mip++) {
+		state.builder.AddFunction({OpLabel, labels[mip]});
+		state.builder.AddFunction({OpImageWrite, LoadAt(binding.array_index + mip), coord, texel});
+		state.builder.AddFunction({OpBranch, merge_label});
+	}
+	state.builder.AddFunction({OpLabel, merge_label});
+}
+
 uint32_t ExecutionModelForStage(ShaderType stage) {
 	switch (stage) {
 		case ShaderType::Vertex: return ExecutionModelVertex;

src/graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h

@@ -731,6 +731,9 @@ uint32_t StorageImageDescriptorPointer(EmitterState& state, uint32_t resource, b
 uint32_t LoadStorageImageDescriptor(EmitterState& state, uint32_t resource, bool uint_image,
                                     uint32_t use_pc, ImageViewKind view = ImageViewKind::Dim2D);
 
+void EmitStorageImageWrite(EmitterState& state, uint32_t resource, bool uint_image,
+                           ImageViewKind view, uint32_t mip_lod, uint32_t coord, uint32_t texel);
+
 uint32_t ExecutionModelForStage(ShaderType stage);
 
 uint32_t ConstantU32(EmitterState& state, uint32_t value);

src/graphics/shader/recompiler/backend/spirv/spirvEmitterValueImage.cpp

@@ -408,11 +408,13 @@ bool EmitValueImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 		const auto uint_image = mem.kind == IR::ResourceKind::StorageImageUint;
 		const auto view       = StorageImageViewKind(state, mem, uint_image, pc);
 		EmitIfCondition(state, ctx.Arg(inst, 3), [&]() {
-			state.builder.AddFunction(
+			const auto mip_lod =
-			    {OpImageWrite,
+			    state.program.info.images[mem.resource].mip_mode == IR::ImageMipMode::DynamicStorage
-			     LoadStorageImageDescriptor(state, mem.resource, uint_image, pc, view),
+			        ? LodU32(ctx, mem, *address, view)
-			     CoordU32(ctx, mem, *address, view),
+			        : 0u;
-			     StoreTexel(ctx, mem, ctx.Arg(inst, 2), uint_image)});
+			const auto coord = CoordU32(ctx, mem, *address, view);
+			const auto texel = StoreTexel(ctx, mem, ctx.Arg(inst, 2), uint_image);
+			EmitStorageImageWrite(state, mem.resource, uint_image, view, mip_lod, coord, texel);
 		});
 		return true;
 	}

src/graphics/shader/recompiler/ir/ShaderIR.h

@@ -194,6 +194,7 @@ struct ImageResource {
 	ResourceKind            kind            = ResourceKind::None;
 	Decoder::ImageDimension dimension       = Decoder::ImageDimension::Unknown;
 	ImageMipMode            mip_mode        = ImageMipMode::None;
+	uint32_t                mip_count       = 1;
 	uint32_t                storage_swizzle = StorageImageIdentitySwizzle;
 	bool                    read            = false;
 	bool                    written         = false;

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

@@ -196,7 +196,16 @@ bool AllocateBindings(Program& program, const BindingLayoutOptions& options, std
 			}
 			return false;
 		}
-		image_groups[static_cast<size_t>(group - ImageBindingKinds.begin())].push_back(i);
+		auto&      resources = image_groups[static_cast<size_t>(group - ImageBindingKinds.begin())];
+		const auto dynamic   = program.info.images[i].mip_mode == ImageMipMode::DynamicStorage;
+		const auto count     = dynamic ? program.info.images[i].mip_count : 1u;
+		if (count == 0u || (!dynamic && program.info.images[i].mip_count != 1u)) {
+			if (error != nullptr) {
+				*error = "image has an invalid specialized mip descriptor count";
+			}
+			return false;
+		}
+		resources.insert(resources.end(), count, i);
 	}
 	for (uint32_t i = 0; i < image_groups.size(); i++) {
 		if (!image_groups[i].empty()) {

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

@@ -72,6 +72,15 @@ bool DescriptorIsCube(const DescriptorValue& descriptor) {
 	       Prospero::ImageType::kCube;
 }
 
+uint32_t StorageMipCount(const ImageResource& image, const DescriptorValue& descriptor) {
+	if (image.mip_mode != ImageMipMode::DynamicStorage || NullImageDescriptor(descriptor)) {
+		return 1;
+	}
+	const auto base = (descriptor.dwords[3] >> 12u) & 0xfu;
+	const auto last = (descriptor.dwords[3] >> 16u) & 0xfu;
+	return base <= last ? last - base + 1u : 0u;
+}
+
 bool DecodeBufferDescriptor(const DescriptorValue& descriptor, ShaderBufferResource& result) {
 	if (descriptor.dword_count != std::size(result.fields)) {
 		return false;
@@ -201,6 +210,13 @@ bool ValidateResourceSpecialization(const Program& program, const ResourceSnapsh
 	for (uint32_t i = 0; i < program.info.images.size(); i++) {
 		const auto& image      = program.info.images[i];
 		const auto& descriptor = snapshot.images[i];
+		const auto  mip_count  = StorageMipCount(image, descriptor);
+		if (mip_count == 0u || mip_count != image.mip_count) {
+			if (error != nullptr) {
+				*error = fmt::format("storage image descriptor {} changed mip range", i);
+			}
+			return false;
+		}
 		if (NullImageDescriptor(descriptor)) {
 			bool canonical_kind =
 			    image.kind == ResourceKind::Image || image.kind == ResourceKind::StorageImage;
@@ -404,6 +420,13 @@ bool SpecializeResources(Program& program, const ResourceSnapshot& snapshot, std
 	for (uint32_t i = 0; i < next.images.size(); i++) {
 		const auto& descriptor = snapshot.images[i];
 		auto&       image      = next.images[i];
+		image.mip_count        = StorageMipCount(image, descriptor);
+		if (image.mip_count == 0u) {
+			if (error != nullptr) {
+				*error = fmt::format("storage image descriptor {} has an invalid mip range", i);
+			}
+			return false;
+		}
 		if (NullImageDescriptor(descriptor)) {
 			image.dimension = Decoder::ImageDimension::Dim2D;
 			image.cube      = false;

src/graphics/shader/shader.cpp

@@ -1022,9 +1022,10 @@ static std::string ShaderDescribeSpecialization(const ShaderRecompiler::IR::Prog
 	}
 	for (uint32_t i = 0; i < program.info.images.size(); i++) {
 		const auto& image = program.info.images[i];
-		ret += fmt::format(" i{}[kind={} dim={} swizzle=0x{:03x}]", i,
+		ret += fmt::format(
-		                   static_cast<uint32_t>(image.kind),
+		    " i{}[kind={} dim={} mip={}/{} swizzle=0x{:03x}]", i, static_cast<uint32_t>(image.kind),
-		                   static_cast<uint32_t>(image.dimension), image.storage_swizzle);
+		    static_cast<uint32_t>(image.dimension), static_cast<uint32_t>(image.mip_mode),
+		    image.mip_count, image.storage_swizzle);
 	}
 	for (uint32_t i = 0; i < program.info.addresses.size(); i++) {
 		ret += fmt::format(" a{}[base=0x{:x}]", i, program.info.addresses[i].specialized_base);
@@ -1062,6 +1063,7 @@ static void ShaderAppendNativeSpecialization(std::vector<uint32_t>&
 		ids.push_back(static_cast<uint32_t>(image.kind));
 		ids.push_back(static_cast<uint32_t>(image.dimension));
 		ids.push_back(static_cast<uint32_t>(image.mip_mode));
+		ids.push_back(image.mip_count);
 		ids.push_back(image.storage_swizzle);
 	}
 	ids.push_back(static_cast<uint32_t>(program.info.addresses.size()));

tests/ResourceTrackingTests.cpp

@@ -329,6 +329,112 @@ void TestImagesSamplersAndAliases() {
         "buffer/image descriptor alias was not linked");
 }
 
+void TestDynamicStorageMipTracking() {
+  Fixture fixture;
+  std::array<Value, 8> image_words;
+  for (uint32_t index = 0; index < image_words.size(); index++) {
+    image_words[index] = fixture.UserData(index);
+  }
+  const auto data = fixture.Emit(ValueOpcode::CompositeConstructU32x4,
+                                 {Value(1u), Value(2u), Value(3u), Value(4u)});
+  const auto AddStore = [&](uint32_t pc, bool has_mip, Value lod) {
+    const auto handle = fixture.Image(image_words, pc);
+    const auto address = fixture.Emit(
+        ValueOpcode::MakeImageAddress,
+        {Value(0u), Value(0u), lod, Value(0u), Value(0u), Value(0u), Value(0u),
+         Value(0u), Value(0u), Value(0u), Value(0u), Value(0u), Value(0u)});
+    MemoryInfo memory;
+    memory.kind = ResourceKind::StorageImage;
+    memory.image_dimension = Decoder::ImageDimension::Dim2D;
+    memory.image_address_components = has_mip ? 3u : 2u;
+    memory.image_has_mip = has_mip;
+    const auto flags = fixture.AddMemory(memory, pc);
+    fixture.Emit(ValueOpcode::ImageWrite, {handle, address, data, Value(true)},
+                 flags);
+    return std::pair{handle, flags.index};
+  };
+
+  const auto plain = AddStore(4, false, Value(0u));
+  const auto mip1 = AddStore(8, true, Value(1u));
+  const auto mip2 = AddStore(12, true, Value(2u));
+  const auto dynamic = AddStore(16, true, fixture.UserData(8));
+  fixture.PlanAndTrack();
+
+  const auto &images = fixture.program.info.images;
+  Check(images.size() == 2 && images[0].mip_mode == ImageMipMode::None &&
+            images[0].mip_count == 1 &&
+            images[1].mip_mode == ImageMipMode::DynamicStorage &&
+            images[1].mip_count == 1,
+        "storage mip writes did not share one dynamic logical resource");
+  Check(plain.first.Instruction()->Flags<uint32_t>() == 0 &&
+            mip1.first.Instruction()->Flags<uint32_t>() == 1 &&
+            mip2.first.Instruction()->Flags<uint32_t>() == 1 &&
+            dynamic.first.Instruction()->Flags<uint32_t>() == 1 &&
+            fixture.program.values->memory_info[plain.second].resource == 0 &&
+            fixture.program.values->memory_info[mip1.second].resource == 1 &&
+            fixture.program.values->memory_info[mip2.second].resource == 1 &&
+            fixture.program.values->memory_info[dynamic.second].resource == 1,
+        "dynamic storage mip handles and memory metadata were not patched");
+
+  DescriptorValue descriptor{};
+  descriptor.dwords[0] = 0x1000u;
+  descriptor.dwords[1] =
+      static_cast<uint32_t>(
+          Libs::Graphics::Prospero::BufferFormat::k32_32_32_32Float)
+      << 20u;
+  descriptor.dwords[2] = 3u | (3u << 14u);
+  descriptor.dwords[3] =
+      Libs::Graphics::DstSel(4, 5, 6, 7) | (1u << 12u) | (3u << 16u) |
+      (static_cast<uint32_t>(Libs::Graphics::Prospero::ImageType::kColor2D)
+       << 28u);
+  descriptor.dwords[5] = 3u << 4u;
+  descriptor.dword_count = 8;
+  ResourceSnapshot snapshot;
+  snapshot.images.assign(images.size(), descriptor);
+  std::string error;
+  Check(SpecializeResources(fixture.program, snapshot, &error) &&
+            fixture.program.info.images[1].mip_count == 3 &&
+            ValidateResourceSpecialization(fixture.program, snapshot, &error),
+        "base-1 through last-3 dynamic storage range was not specialized");
+  ShaderComputeInputInfo compute{};
+  Check(CollectShaderInfo(fixture.program, {.compute = &compute}, &error) &&
+            AllocateBindings(fixture.program, {}, &error),
+        "dynamic storage mip bindings were not allocated");
+  const auto *storage_binding = FindBinding(
+      fixture.program.bindings, DescriptorBindingKind::Storage2D);
+  Check(storage_binding != nullptr &&
+            storage_binding->resources == std::vector<uint32_t>({0, 1, 1, 1}),
+        "dynamic storage mip descriptors were not expanded consecutively");
+
+  Program null_program;
+  null_program.values = std::make_shared<ValueProgram>();
+  null_program.resource_tracking_complete = true;
+  ImageResource null_image;
+  null_image.kind = ResourceKind::StorageImage;
+  null_image.dimension = Decoder::ImageDimension::Dim2D;
+  null_image.mip_mode = ImageMipMode::DynamicStorage;
+  null_image.written = true;
+  null_program.info.images.push_back(null_image);
+  ResourceSnapshot null_snapshot;
+  DescriptorValue null_descriptor{};
+  null_descriptor.dword_count = 8;
+  null_snapshot.images.push_back(null_descriptor);
+  Check(SpecializeResources(null_program, null_snapshot, &error) &&
+            null_program.info.images[0].mip_count == 1 &&
+            ValidateResourceSpecialization(null_program, null_snapshot, &error),
+        "canonical null dynamic storage image did not retain one descriptor");
+
+  snapshot.images[1].dwords[3] =
+      (snapshot.images[1].dwords[3] & ~(0xfu << 16u)) | (2u << 16u);
+  Check(!ValidateResourceSpecialization(fixture.program, snapshot, &error),
+        "a changed dynamic storage mip count reused the specialization");
+  snapshot.images[1].dwords[3] =
+      (snapshot.images[1].dwords[3] & ~((0xfu << 12u) | (0xfu << 16u))) |
+      (4u << 12u) | (3u << 16u);
+  Check(!ValidateResourceSpecialization(fixture.program, snapshot, &error),
+        "an inverted dynamic storage mip range was accepted");
+}
+
 void TestSrtFlatteningAndRuntimeMemoization() {
   Fixture fixture;
   const auto base =
@@ -697,6 +803,7 @@ int main() {
     Run("scalar/vector alias", TestScalarAndVectorBufferAlias);
     Run("runtime unsigned min", TestRuntimeUnsignedMinDescriptor);
     Run("images and samplers", TestImagesSamplersAndAliases);
+    Run("dynamic storage mips", TestDynamicStorageMipTracking);
     Run("SRT runtime", TestSrtFlatteningAndRuntimeMemoization);
     Run("dynamic SRT", TestDynamicSrtReadRemainsExplicit);
     Run("phi validation", TestPhiValidation);

tests/ShaderRecompilerComputeTests.cpp

@@ -392,8 +392,9 @@ struct RenderExecutorTestAccess {
 
 struct DescriptorCacheTestAccess {
   static vk::DescriptorImageInfo
-  MakeImageInfo(const DescriptorCache::TextureBinding &binding) {
+  MakeImageInfo(const DescriptorCache::TextureBinding &binding,
-    return DescriptorCache::MakeImageInfo(binding);
+                uint32_t element = 0) {
+    return DescriptorCache::MakeImageInfo(binding, element);
   }
 };
 
@@ -862,11 +863,7 @@ struct TestCase {
   std::vector<u32> expected_gds;
   std::vector<std::pair<std::string, size_t>> decoded_counts;
   std::vector<std::pair<std::string, size_t>> ir_counts;
-};
+  u32 expected_storage_mip_descriptors = 0;
-
-struct SkippedCase {
-  const char *name = "";
-  const char *reason = "";
 };
 
 struct GraphicsCase {
@@ -1090,6 +1087,29 @@ CompiledShader CompileCase(const TestCase &test) {
             text + " count=" + std::to_string(actual) +
                 ", expected=" + std::to_string(expected));
   }
+  if (test.expected_storage_mip_descriptors != 0) {
+    const auto image = std::ranges::find_if(
+        result.program.info.images, [](const auto &resource) {
+          return resource.mip_mode ==
+                 ShaderRecompiler::IR::ImageMipMode::DynamicStorage;
+        });
+    Require(test.name, "dynamic storage mip specialization",
+            image != result.program.info.images.end() &&
+                image->mip_count == test.expected_storage_mip_descriptors,
+            "runtime descriptor range did not specialize the mip count");
+    const auto resource = static_cast<u32>(
+        image - result.program.info.images.begin());
+    const auto binding = std::ranges::find_if(
+        result.program.bindings.descriptors, [&](const auto &group) {
+          return std::ranges::find(group.resources, resource) !=
+                 group.resources.end();
+        });
+    Require(test.name, "dynamic storage mip binding",
+            binding != result.program.bindings.descriptors.end() &&
+                std::ranges::count(binding->resources, resource) ==
+                    test.expected_storage_mip_descriptors,
+            "dynamic storage image did not receive one descriptor per mip");
+  }
   if (test.force_shader_data_storage) {
     result.program.bindings = {};
     result.program.binding_layout_complete = false;
@@ -6522,6 +6542,79 @@ public:
       std::copy(std::begin(storage.fields), std::end(storage.fields),
                 storage_descriptor.dwords.begin());
       storage_descriptor.dword_count = 8;
+      auto mipped_storage = storage;
+      constexpr uint64_t mipped_storage_address = base + 0x100000;
+      const auto encoded_mipped_storage_address = mipped_storage_address >> 8u;
+      mipped_storage.fields[0] =
+          static_cast<uint32_t>(encoded_mipped_storage_address);
+      mipped_storage.fields[1] =
+          static_cast<uint32_t>(encoded_mipped_storage_address >> 32u) |
+          (static_cast<uint32_t>(stencil_format) << 20u) | (3u << 30u);
+      mipped_storage.fields[2] = 1u | (7u << 14u);
+      mipped_storage.fields[3] =
+          DstSel(4, 5, 6, 7) | (1u << 12u) | (3u << 16u) |
+          (static_cast<uint32_t>(linear) << 20u) |
+          (static_cast<uint32_t>(Prospero::ImageType::kColor2D) << 28u);
+      mipped_storage.fields[5] = 0x00700000u | (3u << 4u);
+      ShaderRecompiler::IR::DescriptorValue mipped_storage_descriptor{};
+      std::copy(std::begin(mipped_storage.fields),
+                std::end(mipped_storage.fields),
+                mipped_storage_descriptor.dwords.begin());
+      mipped_storage_descriptor.dword_count = 8;
+      auto mipped_storage_resource = storage_resource;
+      mipped_storage_resource.mip_mode =
+          ShaderRecompiler::IR::ImageMipMode::DynamicStorage;
+      mipped_storage_resource.mip_count = 3;
+      auto plain_mipped_storage_binding =
+          RenderExecutorTestAccess::ResolveTexture(
+              executor, storage_resource, mipped_storage_descriptor);
+      auto mipped_storage_binding =
+          RenderExecutorTestAccess::ResolveTexture(
+              executor, mipped_storage_resource, mipped_storage_descriptor);
+      DescriptorCache::PreparedBindings mipped_prepared{};
+      auto mipped_program =
+          std::make_shared<ShaderRecompiler::IR::Program>();
+      mipped_program->info.images.push_back(storage_resource);
+      mipped_program->info.images.push_back(mipped_storage_resource);
+      auto mipped_snapshot =
+          std::make_shared<ShaderRecompiler::IR::ResourceSnapshot>();
+      mipped_snapshot->images.assign(2, mipped_storage_descriptor);
+      mipped_prepared.program = std::move(mipped_program);
+      mipped_prepared.snapshot = std::move(mipped_snapshot);
+      mipped_prepared.resources.images.push_back(
+          std::move(plain_mipped_storage_binding));
+      mipped_prepared.resources.images.push_back(
+          std::move(mipped_storage_binding));
+      executor.RebindImages(scheduler.Current(), mipped_prepared);
+      auto &plain_mipped_binding = mipped_prepared.resources.images[0];
+      auto &mipped_binding = mipped_prepared.resources.images[1];
+      plain_mipped_binding.layout = vk::ImageLayout::eGeneral;
+      mipped_binding.layout = vk::ImageLayout::eGeneral;
+      auto plain_view_desc = plain_mipped_binding.desc;
+      plain_view_desc.view_info.level_count = 1;
+      Require(name, "dynamic storage mip views",
+              plain_mipped_binding.desc.view_info.base_level == 1 &&
+                  plain_mipped_binding.desc.view_info.level_count == 3 &&
+                  plain_mipped_binding.mip_views.empty() &&
+                  plain_mipped_binding.image_view == texture_cache.FindTexture(
+                      plain_mipped_binding.image_id, plain_view_desc) &&
+                  mipped_binding.desc.info.resources.levels == 4 &&
+                  mipped_binding.desc.view_info.base_level == 1 &&
+                  mipped_binding.desc.view_info.level_count == 3 &&
+                  mipped_binding.mip_views.size() == 3 &&
+                  mipped_binding.mip_views[0] != nullptr &&
+                  mipped_binding.mip_views[1] != nullptr &&
+                  mipped_binding.mip_views[2] != nullptr &&
+                  mipped_binding.mip_views[0] != mipped_binding.mip_views[1] &&
+                  mipped_binding.mip_views[1] != mipped_binding.mip_views[2] &&
+                  DescriptorCacheTestAccess::MakeImageInfo(mipped_binding, 0)
+                          .imageView == mipped_binding.mip_views[0] &&
+                  DescriptorCacheTestAccess::MakeImageInfo(mipped_binding, 1)
+                          .imageView == mipped_binding.mip_views[1] &&
+                  DescriptorCacheTestAccess::MakeImageInfo(mipped_binding, 2)
+                          .imageView == mipped_binding.mip_views[2],
+              "base-1 through last-3 storage view did not create three "
+              "ordered one-level descriptors");
       auto srgb_storage = storage;
       constexpr auto srgb_format =
           static_cast<uint32_t>(Prospero::BufferFormat::k8_8_8_8Srgb);
@@ -6739,11 +6832,24 @@ public:
       Require(name, "VS-to-PS image acquisition order",
               storage_binding.image_id == storage_id &&
                   sampled_binding.image_id == storage_id &&
+                  texture_cache.GetImage(storage_id).binding.force_general &&
+                  texture_cache.GetImage(storage_id).binding.shader_write &&
                   texture_cache.GetImage(storage_id).IsGpuModified() &&
                   texture_cache.GetImage(storage_id).usage.storage &&
                   texture_cache.GetImage(storage_id).usage.texture,
               "the production graphics binding path did not complete vertex "
               "storage acquisition before pixel sampling");
+      RenderExecutorTestAccess::CommitBindings(executor, scheduler.Current(),
+                                               graphics_bindings.vertex);
+      RenderExecutorTestAccess::CommitBindings(executor, scheduler.Current(),
+                                               *graphics_bindings.pixel);
+      Require(name, "early writable alias retention",
+              texture_cache.GetImage(storage_id).backing.state.access_mask ==
+                  (vk::AccessFlagBits2::eShaderRead |
+                   vk::AccessFlagBits2::eShaderWrite |
+                   vk::AccessFlagBits2::eColorAttachmentRead |
+                   vk::AccessFlagBits2::eColorAttachmentWrite),
+              "a later sampled alias dropped an earlier storage write access");
       RenderExecutorTestAccess::ResetBindings(executor);
 
       auto vertex_sampled_program =
@@ -6782,9 +6888,15 @@ public:
                       writable_alias_bindings.pixel->resources.images[0])
                           .imageLayout == vk::ImageLayout::eGeneral &&
                   texture_cache.GetImage(storage_id).backing.state.layout ==
-                      vk::ImageLayout::eGeneral,
+                      vk::ImageLayout::eGeneral &&
+                  texture_cache.GetImage(storage_id)
+                          .backing.state.access_mask ==
+                      (vk::AccessFlagBits2::eShaderRead |
+                       vk::AccessFlagBits2::eShaderWrite |
+                       vk::AccessFlagBits2::eColorAttachmentRead |
+                       vk::AccessFlagBits2::eColorAttachmentWrite),
               "sampled/storage aliases did not retain the promoted general "
-              "layout in both generated descriptors");
+              "layout and writable access in both generated descriptors");
       RenderExecutorTestAccess::ResetBindings(executor);
 
       const std::array<uint32_t, 3> split_mip_data{0x01020304u, 0x05060708u,
@@ -16382,12 +16494,6 @@ TestCase ImageGetResinfoDmaskMipLevels() {
   return test;
 }
 
-SkippedCase ImageStoreMipWritesExplicitMip2D() {
-  return {"ImageStoreMipWritesExplicitMip2D",
-          "requires per-mip storage-image view descriptors; Vulkan "
-          "OpImageWrite cannot take Lod"};
-}
-
 TestCase ImageSampleAndGather() {
   using O = ShaderOpcode;
 
@@ -16631,6 +16737,37 @@ TestCase ImageStoreVariants() {
   return test;
 }
 
+TestCase ImageStoreMipSelectsPpsa01340Descriptor() {
+  using O = ShaderOpcode;
+
+  std::vector<u32> code;
+  AppendVMovU32(&code, 20, 0);
+  AppendVMovU32(&code, 21, 0);
+  AppendVMovU32(&code, 22, 2);
+  AppendVMovU32(&code, 23, 0);
+  AppendVMovLiteral(&code, 0, 0x3f800000u);
+  AppendVMovLiteral(&code, 1, 0x40000000u);
+  AppendVMovLiteral(&code, 2, 0x40400000u);
+  AppendVMovLiteral(&code, 3, 0x40800000u);
+  code.push_back(EncodeMimg0(0x09, 0xf));
+  code.push_back(EncodeMimg1(0, 20));
+  AppendEnd(&code);
+
+  TestCase test;
+  test.name = "ImageStoreMipSelectsPpsa01340Descriptor";
+  test.code = std::move(code);
+  test.opcodes = {O::VMovB32, O::ImageStoreMip, O::SEndpgm};
+  const std::array<u32, 8> descriptor{
+      0x0294dc00u, 0xc4700000u, 0x00b3c13fu, 0x91b31facu,
+      0x00000000u, 0x00700030u, 0xb07b0000u, 0x0002ac3cu};
+  std::copy(descriptor.begin(), descriptor.end(), test.user_data.begin());
+  test.has_user_data = true;
+  test.compile_only = true;
+  test.expected_storage_mip_descriptors = 3;
+  test.required_spirv = {"OpSwitch", "OpImageWrite"};
+  return test;
+}
+
 TestCase ImageStoreRgbOneUsesInverseSwizzle() {
   using O = ShaderOpcode;
 
@@ -17395,6 +17532,7 @@ std::vector<TestCase> MakeCases() {
   AddCase(ImageSampleA16CompareBiasRdna2AddressOrder);
   AddCase(ImageGatherCompareOpcodes);
   AddCase(ImageStoreVariants);
+  AddCase(ImageStoreMipSelectsPpsa01340Descriptor);
   AddCase(ImageStoreRgbOneUsesInverseSwizzle);
   AddCase(ImageStoreDuplicateSelectorUsesInverseSwizzle);
   AddCase(ImageStoreBgraUsesInverseSwizzle);
@@ -17425,10 +17563,6 @@ std::vector<GraphicsCase> MakeGraphicsCases() {
   };
 }
 
-std::vector<SkippedCase> MakeSkippedCases() {
-  return {ImageStoreMipWritesExplicitMip2D()};
-}
-
 void CheckPs5GameExampleImageClearRuntimeShape() {
   const auto MakeCode = [] {
     std::vector<u32> code;
@@ -17805,8 +17939,6 @@ void CheckRenderTargetFormatContract() {
       resource.atomic = true;
     } else if (std::strcmp(kind, "storage-compare") == 0) {
       resource.depth_compare = true;
-    } else if (std::strcmp(kind, "storage-mip") == 0) {
-      resource.mip_mode = ShaderRecompiler::IR::ImageMipMode::DynamicStorage;
     } else if (std::strcmp(kind, "storage-dimension") == 0) {
       resource.dimension = ShaderRecompiler::Decoder::ImageDimension::Unknown;
     } else {
@@ -18024,6 +18156,13 @@ void CheckSampledColorViews() {
   Require("SampledColorViews", "atomic uint 2D storage resource",
           IsSupportedStorageImageResource(storage_resource),
           "atomic uint storage resource was rejected");
+  storage_resource.atomic = false;
+  storage_resource.mip_mode =
+      ShaderRecompiler::IR::ImageMipMode::DynamicStorage;
+  storage_resource.mip_count = 3;
+  Require("SampledColorViews", "dynamic storage mip resource",
+          IsSupportedStorageImageResource(storage_resource),
+          "dynamic storage mip resource was rejected");
 
   char path[MAX_PATH]{};
   Require("SampledColorViews", "host",
@@ -18033,7 +18172,7 @@ void CheckSampledColorViews() {
        {"sampled-invalid-selector", "sampled-incompatible-format",
         "sampled-invalid-high", "sampled-depth-format", "sampled-depth-swizzle",
         "storage-incompatible-format", "storage-kind", "storage-no-write",
-        "storage-nonuint-atomic", "storage-compare", "storage-mip",
+        "storage-nonuint-atomic", "storage-compare",
         "storage-dimension", "volume-mip-count", "volume-slice-range"}) {
     std::string command =
         std::string("\"") + path + "\" --image-view-death " + kind;
@@ -19389,6 +19528,11 @@ void CheckStorageImageSwizzleSpecializationId() {
   identity_program.info.images.push_back(image);
   auto rgb1_program = identity_program;
   rgb1_program.info.images[0].storage_swizzle = DstSel(4, 5, 6, 1);
+  auto mip1_program = identity_program;
+  mip1_program.info.images[0].mip_mode =
+      ShaderRecompiler::IR::ImageMipMode::DynamicStorage;
+  auto mip3_program = mip1_program;
+  mip3_program.info.images[0].mip_count = 3;
 
   const auto resources =
       std::make_shared<const ShaderRecompiler::IR::ResourceSnapshot>();
@@ -19399,13 +19543,22 @@ void CheckStorageImageSwizzleSpecializationId() {
   auto rgb1_info = identity_info;
   rgb1_info.stage.program =
       std::make_shared<const ShaderRecompiler::IR::Program>(rgb1_program);
+  auto mip1_info = identity_info;
+  mip1_info.stage.program =
+      std::make_shared<const ShaderRecompiler::IR::Program>(mip1_program);
+  auto mip3_info = identity_info;
+  mip3_info.stage.program =
+      std::make_shared<const ShaderRecompiler::IR::Program>(mip3_program);
 
   const auto identity_id = ShaderGetIdCS(regs, identity_info, true);
   const auto rgb1_id = ShaderGetIdCS(regs, rgb1_info, true);
+  const auto mip1_id = ShaderGetIdCS(regs, mip1_info, true);
+  const auto mip3_id = ShaderGetIdCS(regs, mip3_info, true);
   Require("StorageImageSwizzleSpecializationId", "pipeline cache key",
           identity_id != rgb1_id &&
-              identity_id.ids.size() == rgb1_id.ids.size(),
+              identity_id.ids.size() == rgb1_id.ids.size() &&
-          "storage swizzle-specialized SPIR-V variants share a pipeline ID");
+              mip1_id != mip3_id && mip1_id.ids.size() == mip3_id.ids.size(),
+          "storage swizzle or mip-count variants share a pipeline ID");
   std::printf("[host]    %-32s ok\n", "StorageImageSwizzlePipelineId");
 }
 
@@ -20701,6 +20854,15 @@ int main(int argc, char **argv) {
     RunCase(&vulkan, ImageStoreBgraUsesInverseSwizzle());
     return 0;
   }
+  if (argc == 2 && std::strcmp(argv[1], "--storage-mip-host-only") == 0) {
+    VulkanHarness vulkan;
+    vulkan.CheckRenderExecutorStencilBindingDiscovery();
+    return 0;
+  }
+  if (argc == 2 && std::strcmp(argv[1], "--storage-mip-only") == 0) {
+    RunCase(nullptr, ImageStoreMipSelectsPpsa01340Descriptor());
+    return 0;
+  }
   if (argc == 2 && std::strcmp(argv[1], "--storage-yzwx-only") == 0) {
     CheckSampledColorViews();
     CheckBasicStorageTextureDescriptor();
@@ -20783,9 +20945,6 @@ int main(int argc, char **argv) {
   for (const auto &test : tests) {
     RunCase(&vulkan, test);
   }
-  for (const auto &test : MakeSkippedCases()) {
-    std::printf("[skip]    %-32s %s\n", test.name, test.reason);
-  }
   for (const auto &test : graphics_tests) {
     RunGraphicsCase(&vulkan, test);
   }