Fix use of deleted shader during compute execution

ComputeProgram::run() uses the `shader` member, which was previously a
weak pointer to a SpirvShader object. If the corresponding shader module
is destroyed, and there is no (more) pipeline cache to hold a shared_ptr
reference to it, we'd access a deleted object.

The `pipelineLayout` member is safe-ish because we explicitly reference
count it in the vk::Pipeline class, since graphics pipelines also use it
to (re)compile routines at draw execution time.

The ComputeProgramKey was modified to use unique integer identifiers
instead of pointers to the shader and pipeline layout. This prevents a
false cache hit when these objects are destroyed and new ones get
allocated at the same address. Note that this means we lose the nice
property of constructing new cache entries only from cache key data,
but this is a better compromise than the previous buggy code.

Bug: b/197982536
Change-Id: Ie75bb0e8df21ad25c2a7fed29adbf086fab072e3
Reviewed-on: https://swiftshader-review.googlesource.com/c/SwiftShader/+/57168
Tested-by: Nicolas Capens <nicolascapens@google.com>
Kokoro-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Alexis Hétu <sugoi@google.com>

src/Pipeline/ComputeProgram.cpp

diff --git a/src/Pipeline/ComputeProgram.cpp b/src/Pipeline/ComputeProgram.cpp
index 43926f4..ad946d0 100644
--- a/src/Pipeline/ComputeProgram.cpp
+++ b/src/Pipeline/ComputeProgram.cpp
@@ -37,7 +37,7 @@
 
 namespace sw {
 
-ComputeProgram::ComputeProgram(vk::Device *device, SpirvShader const *shader, vk::PipelineLayout const *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets)
+ComputeProgram::ComputeProgram(vk::Device *device, std::shared_ptr<SpirvShader> shader, vk::PipelineLayout const *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets)
     : device(device)
     , shader(shader)
     , pipelineLayout(pipelineLayout)
@@ -70,7 +70,7 @@
 	routine->subgroupsPerWorkgroup = *Pointer<Int>(data + OFFSET(Data, subgroupsPerWorkgroup));
 	routine->invocationsPerSubgroup = *Pointer<Int>(data + OFFSET(Data, invocationsPerSubgroup));
 
-	routine->setInputBuiltin(shader, spv::BuiltInNumWorkgroups, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
+	routine->setInputBuiltin(shader.get(), spv::BuiltInNumWorkgroups, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
 		for(uint32_t component = 0; component < builtin.SizeInComponents; component++)
 		{
 			value[builtin.FirstComponent + component] =
@@ -78,7 +78,7 @@
 		}
 	});
 
-	routine->setInputBuiltin(shader, spv::BuiltInWorkgroupId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
+	routine->setInputBuiltin(shader.get(), spv::BuiltInWorkgroupId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
 		for(uint32_t component = 0; component < builtin.SizeInComponents; component++)
 		{
 			value[builtin.FirstComponent + component] =
@@ -86,7 +86,7 @@
 		}
 	});
 
-	routine->setInputBuiltin(shader, spv::BuiltInWorkgroupSize, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
+	routine->setInputBuiltin(shader.get(), spv::BuiltInWorkgroupSize, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
 		for(uint32_t component = 0; component < builtin.SizeInComponents; component++)
 		{
 			value[builtin.FirstComponent + component] =
@@ -94,17 +94,17 @@
 		}
 	});
 
-	routine->setInputBuiltin(shader, spv::BuiltInNumSubgroups, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
+	routine->setInputBuiltin(shader.get(), spv::BuiltInNumSubgroups, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
 		ASSERT(builtin.SizeInComponents == 1);
 		value[builtin.FirstComponent] = As<SIMD::Float>(SIMD::Int(routine->subgroupsPerWorkgroup));
 	});
 
-	routine->setInputBuiltin(shader, spv::BuiltInSubgroupSize, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
+	routine->setInputBuiltin(shader.get(), spv::BuiltInSubgroupSize, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
 		ASSERT(builtin.SizeInComponents == 1);
 		value[builtin.FirstComponent] = As<SIMD::Float>(SIMD::Int(routine->invocationsPerSubgroup));
 	});
 
-	routine->setImmutableInputBuiltins(shader);
+	routine->setImmutableInputBuiltins(shader.get());
 }
 
 void ComputeProgram::setSubgroupBuiltins(Pointer<Byte> data, SpirvRoutine *routine, Int workgroupID[3], SIMD::Int localInvocationIndex, Int subgroupIndex)
@@ -142,17 +142,17 @@
 	routine->globalInvocationID[Y] = globalInvocationID[Y];
 	routine->globalInvocationID[Z] = globalInvocationID[Z];
 
-	routine->setInputBuiltin(shader, spv::BuiltInLocalInvocationIndex, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
+	routine->setInputBuiltin(shader.get(), spv::BuiltInLocalInvocationIndex, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
 		ASSERT(builtin.SizeInComponents == 1);
 		value[builtin.FirstComponent] = As<SIMD::Float>(localInvocationIndex);
 	});
 
-	routine->setInputBuiltin(shader, spv::BuiltInSubgroupId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
+	routine->setInputBuiltin(shader.get(), spv::BuiltInSubgroupId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
 		ASSERT(builtin.SizeInComponents == 1);
 		value[builtin.FirstComponent] = As<SIMD::Float>(SIMD::Int(subgroupIndex));
 	});
 
-	routine->setInputBuiltin(shader, spv::BuiltInLocalInvocationId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
+	routine->setInputBuiltin(shader.get(), spv::BuiltInLocalInvocationId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
 		for(uint32_t component = 0; component < builtin.SizeInComponents; component++)
 		{
 			value[builtin.FirstComponent + component] =
@@ -160,7 +160,7 @@
 		}
 	});
 
-	routine->setInputBuiltin(shader, spv::BuiltInGlobalInvocationId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
+	routine->setInputBuiltin(shader.get(), spv::BuiltInGlobalInvocationId, [&](const SpirvShader::BuiltinMapping &builtin, Array<SIMD::Float> &value) {
 		for(uint32_t component = 0; component < builtin.SizeInComponents; component++)
 		{
 			value[builtin.FirstComponent + component] =

src/Pipeline/ComputeProgram.hpp

diff --git a/src/Pipeline/ComputeProgram.hpp b/src/Pipeline/ComputeProgram.hpp
index 839f494..79a0bcb 100644
--- a/src/Pipeline/ComputeProgram.hpp
+++ b/src/Pipeline/ComputeProgram.hpp
@@ -46,7 +46,7 @@
                            int32_t subgroupCount)>
 {
 public:
-	ComputeProgram(vk::Device *device, SpirvShader const *spirvShader, vk::PipelineLayout const *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets);
+	ComputeProgram(vk::Device *device, std::shared_ptr<SpirvShader> spirvShader, vk::PipelineLayout const *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets);
 
 	virtual ~ComputeProgram();
 
@@ -81,8 +81,8 @@
 	};
 
 	vk::Device *const device;
-	SpirvShader const *const shader;
-	vk::PipelineLayout const *const pipelineLayout;
+	const std::shared_ptr<SpirvShader> shader;
+	const vk::PipelineLayout *const pipelineLayout;  // Reference held by vk::Pipeline
 	const vk::DescriptorSet::Bindings &descriptorSets;
 };
 

src/Vulkan/VkPipeline.cpp

diff --git a/src/Vulkan/VkPipeline.cpp b/src/Vulkan/VkPipeline.cpp
index 8ab06d9..aae2104 100644
--- a/src/Vulkan/VkPipeline.cpp
+++ b/src/Vulkan/VkPipeline.cpp
@@ -128,15 +128,16 @@
 	                                         code, key.getRenderPass(), key.getSubpassIndex(), robustBufferAccess, dbgctx);
 }
 
-std::shared_ptr<sw::ComputeProgram> createProgram(vk::Device *device, const vk::PipelineCache::ComputeProgramKey &key)
+std::shared_ptr<sw::ComputeProgram> createProgram(vk::Device *device, std::shared_ptr<sw::SpirvShader> shader, const vk::PipelineLayout *layout)
 {
 	MARL_SCOPED_EVENT("createProgram");
 
-	vk::DescriptorSet::Bindings descriptorSets;  // FIXME(b/129523279): Delay code generation until invoke time.
+	vk::DescriptorSet::Bindings descriptorSets;  // TODO(b/129523279): Delay code generation until dispatch time.
 	// TODO(b/119409619): use allocator.
-	auto program = std::make_shared<sw::ComputeProgram>(device, key.getShader(), key.getLayout(), descriptorSets);
+	auto program = std::make_shared<sw::ComputeProgram>(device, shader, layout, descriptorSets);
 	program->generate();
 	program->finalize("ComputeProgram");
+
 	return program;
 }
 
@@ -285,16 +286,16 @@
 			return createShader(shaderKey, module, robustBufferAccess, device->getDebuggerContext());
 		});
 
-		const PipelineCache::ComputeProgramKey programKey(shader.get(), layout);
+		const PipelineCache::ComputeProgramKey programKey(shader->getSerialID(), layout->identifier);
 		program = pPipelineCache->getOrCreateComputeProgram(programKey, [&] {
-			return createProgram(device, programKey);
+			return createProgram(device, shader, layout);
 		});
 	}
 	else
 	{
 		shader = createShader(shaderKey, module, robustBufferAccess, device->getDebuggerContext());
-		const PipelineCache::ComputeProgramKey programKey(shader.get(), layout);
-		program = createProgram(device, programKey);
+		const PipelineCache::ComputeProgramKey programKey(shader->getSerialID(), layout->identifier);
+		program = createProgram(device, shader, layout);
 	}
 }
 

src/Vulkan/VkPipelineCache.cpp

diff --git a/src/Vulkan/VkPipelineCache.cpp b/src/Vulkan/VkPipelineCache.cpp
index 0aa72d6..dc4afe3 100644
--- a/src/Vulkan/VkPipelineCache.cpp
+++ b/src/Vulkan/VkPipelineCache.cpp
@@ -75,6 +75,16 @@
 	return (specializationInfo < other.specializationInfo);
 }
 
+PipelineCache::ComputeProgramKey::ComputeProgramKey(uint64_t shaderIdentifier, uint32_t pipelineLayoutIdentifier)
+    : shaderIdentifier(shaderIdentifier)
+    , pipelineLayoutIdentifier(pipelineLayoutIdentifier)
+{}
+
+bool PipelineCache::ComputeProgramKey::operator<(const ComputeProgramKey &other) const
+{
+	return std::tie(shaderIdentifier, pipelineLayoutIdentifier) < std::tie(other.shaderIdentifier, other.pipelineLayoutIdentifier);
+}
+
 PipelineCache::PipelineCache(const VkPipelineCacheCreateInfo *pCreateInfo, void *mem)
     : dataSize(ComputeRequiredAllocationSize(pCreateInfo))
     , data(reinterpret_cast<uint8_t *>(mem))

src/Vulkan/VkPipelineCache.hpp

diff --git a/src/Vulkan/VkPipelineCache.hpp b/src/Vulkan/VkPipelineCache.hpp
index 933bfbf..9c479e4 100644
--- a/src/Vulkan/VkPipelineCache.hpp
+++ b/src/Vulkan/VkPipelineCache.hpp
@@ -90,22 +90,13 @@
 
 	struct ComputeProgramKey
 	{
-		ComputeProgramKey(const sw::SpirvShader *shader, const vk::PipelineLayout *layout)
-		    : shader(shader)
-		    , layout(layout)
-		{}
+		ComputeProgramKey(uint64_t shaderIdentifier, uint32_t pipelineLayoutIdentifier);
 
-		bool operator<(const ComputeProgramKey &other) const
-		{
-			return std::tie(shader, layout) < std::tie(other.shader, other.layout);
-		}
-
-		const sw::SpirvShader *getShader() const { return shader; }
-		const vk::PipelineLayout *getLayout() const { return layout; }
+		bool operator<(const ComputeProgramKey &other) const;
 
 	private:
-		const sw::SpirvShader *shader;
-		const vk::PipelineLayout *layout;
+		const uint64_t shaderIdentifier;
+		const uint32_t pipelineLayoutIdentifier;
 	};
 
 	// getOrCreateComputeProgram() queries the cache for a compute program with