MoltenVK/MoltenVK/GPUObjects/MVKShaderModule.h - external/github.com/KhronosGroup/MoltenVK - Git at Google

 /*
  * MVKShaderModule.h
  *
  * Copyright (c) 2015-2022 The Brenwill Workshop Ltd. (http://www.brenwill.com)
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 #pragma once

 #include "MVKDevice.h"
 #include "MVKSync.h"
 #include "MVKSmallVector.h"
 #include <MoltenVKShaderConverter/SPIRVToMSLConverter.h>
 #include <MoltenVKShaderConverter/GLSLToSPIRVConverter.h>
 #include <mutex>

 #import <Metal/Metal.h>

 class MVKPipelineCache;
 class MVKShaderCacheIterator;
 class MVKShaderLibraryCache;
 class MVKShaderModule;

 using namespace mvk;


 #pragma mark -
 #pragma mark MVKShaderLibrary

 /** A MTLFunction and corresponding result information resulting from a shader conversion. */
 typedef struct MVKMTLFunction {
 	SPIRVToMSLConversionResults shaderConversionResults;
 	MTLSize threadGroupSize;
 	inline id<MTLFunction> getMTLFunction() { return _mtlFunction; }

 	MVKMTLFunction(id<MTLFunction> mtlFunc, const SPIRVToMSLConversionResults scRslts, MTLSize tgSize);
 	MVKMTLFunction(const MVKMTLFunction& other);
 	MVKMTLFunction& operator=(const MVKMTLFunction& other);
 	MVKMTLFunction() {}
 	~MVKMTLFunction();

 private:
 	id<MTLFunction> _mtlFunction = nil;

 } MVKMTLFunction;

 /** A MVKMTLFunction indicating an invalid MTLFunction. The mtlFunction member is nil. */
 const MVKMTLFunction MVKMTLFunctionNull(nil, SPIRVToMSLConversionResults(), MTLSizeMake(1, 1, 1));

 /** Wraps a single MTLLibrary. */
 class MVKShaderLibrary : public MVKBaseObject {

 public:

 	/** Returns the Vulkan API opaque object controlling this object. */
 	MVKVulkanAPIObject* getVulkanAPIObject() override { return _owner->getVulkanAPIObject(); };

 	/**
 	 * Sets the entry point function name.
 	 *
 	 * This is usually set automatically during shader conversion from SPIR-V to MSL.
 	 * For a library that was created directly from MSL, this function can be used to
 	 * set the name of the function if it has a different name than the default main0().
 	 */
 	void setEntryPointName(std::string& funcName);

     /**
 	 * Sets the number of threads in a single compute kernel workgroup, per dimension.
 	 *
 	 * This is usually set automatically during shader conversion from SPIR-V to MSL.
 	 * For a library that was created directly from MSL, this function can be used to
 	 * set the workgroup size..
 	 */
     void setWorkgroupSize(uint32_t x, uint32_t y, uint32_t z);

 	/** Constructs an instance from the specified MSL source code. */
 	MVKShaderLibrary(MVKVulkanAPIDeviceObject* owner,
 					 const std::string& mslSourceCode,
 					 const SPIRVToMSLConversionResults& shaderConversionResults);

 	/** Constructs an instance from the specified compiled MSL code data. */
 	MVKShaderLibrary(MVKVulkanAPIDeviceObject* owner,
 					 const void* mslCompiledCodeData,
 					 size_t mslCompiledCodeLength);

 	MVKShaderLibrary(const MVKShaderLibrary& other);

 	MVKShaderLibrary& operator=(const MVKShaderLibrary& other);

 	~MVKShaderLibrary() override;

 protected:
 	friend MVKShaderCacheIterator;
 	friend MVKShaderLibraryCache;
 	friend MVKShaderModule;

 	MVKMTLFunction getMTLFunction(const VkSpecializationInfo* pSpecializationInfo, MVKShaderModule* shaderModule);
 	void handleCompilationError(NSError* err, const char* opDesc);
     MTLFunctionConstant* getFunctionConstant(NSArray<MTLFunctionConstant*>* mtlFCs, NSUInteger mtlFCID);

 	MVKVulkanAPIDeviceObject* _owner;
 	id<MTLLibrary> _mtlLibrary;
 	SPIRVToMSLConversionResults _shaderConversionResults;
 	std::string _msl;
 };


 #pragma mark -
 #pragma mark MVKShaderLibraryCache

 /** Represents a cache of shader libraries for one shader module. */
 class MVKShaderLibraryCache : public MVKBaseObject {

 public:

 	/** Returns the Vulkan API opaque object controlling this object. */
 	MVKVulkanAPIObject* getVulkanAPIObject() override { return _owner->getVulkanAPIObject(); };

 	/**
 	 * Returns a shader library from the shader conversion configuration sourced from the shader module,
 	 * lazily creating the shader library from source code in the shader module, if needed.
 	 *
 	 * If pWasAdded is not nil, this function will set it to true if a new shader library was created,
 	 * and to false if an existing shader library was found and returned.
 	 */
 	MVKShaderLibrary* getShaderLibrary(SPIRVToMSLConversionConfiguration* pShaderConfig,
 									   MVKShaderModule* shaderModule,
 									   bool* pWasAdded = nullptr);

 	MVKShaderLibraryCache(MVKVulkanAPIDeviceObject* owner) : _owner(owner) {};

 	~MVKShaderLibraryCache() override;

 protected:
 	friend MVKShaderCacheIterator;
 	friend MVKPipelineCache;
 	friend MVKShaderModule;

 	MVKShaderLibrary* findShaderLibrary(SPIRVToMSLConversionConfiguration* pShaderConfig);
 	MVKShaderLibrary* addShaderLibrary(SPIRVToMSLConversionConfiguration* pShaderConfig,
 									   const std::string& mslSourceCode,
 									   const SPIRVToMSLConversionResults& shaderConversionResults);
 	void merge(MVKShaderLibraryCache* other);

 	MVKVulkanAPIDeviceObject* _owner;
 	MVKSmallVector<std::pair<SPIRVToMSLConversionConfiguration, MVKShaderLibrary*>> _shaderLibraries;
 };


 #pragma mark -
 #pragma mark MVKShaderModule

 typedef struct MVKShaderModuleKey {
 	std::size_t codeSize;
 	std::size_t codeHash;

 	bool operator==(const MVKShaderModuleKey& rhs) const {
 		return ((codeSize == rhs.codeSize) && (codeHash == rhs.codeHash));
 	}
 	MVKShaderModuleKey(std::size_t codeSize, std::size_t codeHash) : codeSize(codeSize), codeHash(codeHash) {}
 	MVKShaderModuleKey() :  MVKShaderModuleKey(0, 0) {}
 } MVKShaderModuleKey;

 /**
  * Hash structure implementation for MVKShaderModuleKey in std namespace,
  * so MVKShaderModuleKey can be used as a key in a std::map and std::unordered_map.
  */
 namespace std {
 	template <>
 	struct hash<MVKShaderModuleKey> {
 		std::size_t operator()(const MVKShaderModuleKey& k) const { return k.codeHash; }
 	};
 }

 /** Represents a Vulkan shader module. */
 class MVKShaderModule : public MVKVulkanAPIDeviceObject {

 public:

 	/** Returns the Vulkan type of this object. */
 	VkObjectType getVkObjectType() override { return VK_OBJECT_TYPE_SHADER_MODULE; }

 	/** Returns the debug report object type of this object. */
 	VkDebugReportObjectTypeEXT getVkDebugReportObjectType() override { return VK_DEBUG_REPORT_OBJECT_TYPE_SHADER_MODULE_EXT; }

 	/** Returns the Metal shader function, possibly specialized. */
 	MVKMTLFunction getMTLFunction(SPIRVToMSLConversionConfiguration* pShaderConfig,
 								  const VkSpecializationInfo* pSpecializationInfo,
 								  MVKPipelineCache* pipelineCache);

 	/** Convert the SPIR-V to MSL, using the specified shader conversion configuration. */
 	bool convert(SPIRVToMSLConversionConfiguration* pShaderConfig);

 	/** Returns the original SPIR-V code that was specified when this object was created. */
 	const std::vector<uint32_t>& getSPIRV() { return _spvConverter.getSPIRV(); }

 	/**
 	 * Returns the Metal Shading Language source code as converted by the most recent
 	 * call to convert() function, or set directly using the setMSL() function.
 	 */
 	const std::string& getMSL() { return _spvConverter.getMSL(); }

 	/** Returns information about the shader conversion results. */
 	const SPIRVToMSLConversionResults& getConversionResults() { return _spvConverter.getConversionResults(); }

     /** Sets the number of threads in a single compute kernel workgroup, per dimension. */
     void setWorkgroupSize(uint32_t x, uint32_t y, uint32_t z);

 	/** Returns a key as a means of identifying this shader module in a pipeline cache. */
 	MVKShaderModuleKey getKey() { return _key; }

 	MVKShaderModule(MVKDevice* device, const VkShaderModuleCreateInfo* pCreateInfo);

 	~MVKShaderModule() override;

 protected:
 	friend MVKShaderCacheIterator;

 	void propagateDebugName() override {}
 	MVKGLSLConversionShaderStage getMVKGLSLConversionShaderStage(SPIRVToMSLConversionConfiguration* pShaderConfig);

 	MVKShaderLibraryCache _shaderLibraryCache;
 	SPIRVToMSLConverter _spvConverter;
 	GLSLToSPIRVConverter _glslConverter;
 	MVKShaderLibrary* _directMSLLibrary;
 	MVKShaderModuleKey _key;
     std::mutex _accessLock;
 };


 #pragma mark -
 #pragma mark MVKShaderLibraryCompiler

 /**
  * Creates a MTLLibrary from source code.
  *
  * Instances of this class are one-shot, and can only be used for a single library compilation.
  */
 class MVKShaderLibraryCompiler : public MVKMetalCompiler {

 public:

 	/**
 	 * Returns a new (retained) MTLLibrary object compiled from the MSL source code.
 	 *
 	 * If the Metal library compiler does not return within MVKConfiguration::metalCompileTimeout
 	 * nanoseconds, an error will be generated and logged, and nil will be returned.
 	 */
 	id<MTLLibrary> newMTLLibrary(NSString* mslSourceCode,
 								 const SPIRVToMSLConversionResults& shaderConversionResults);


 #pragma mark Construction

 	MVKShaderLibraryCompiler(MVKVulkanAPIDeviceObject* owner) : MVKMetalCompiler(owner) {
 		_compilerType = "Shader library";
 		_pPerformanceTracker = &_owner->getDevice()->_performanceStatistics.shaderCompilation.mslCompile;
 	}

 	~MVKShaderLibraryCompiler() override;

 protected:
 	bool compileComplete(id<MTLLibrary> mtlLibrary, NSError *error);
 	void handleError() override;

 	id<MTLLibrary> _mtlLibrary = nil;
 };


 #pragma mark -
 #pragma mark MVKFunctionSpecializer

 /**
  * Compiles a specialized MTLFunction.
  *
  * Instances of this class are one-shot, and can only be used for a single function compilation.
  */
 class MVKFunctionSpecializer : public MVKMetalCompiler {

 public:

 	/**
 	 * Returns a new (retained) MTLFunction object compiled from the MTLLibrary and specialization constants.
 	 *
 	 * If the Metal function compiler does not return within MVKConfiguration::metalCompileTimeout
 	 * nanoseconds, an error will be generated and logged, and nil will be returned.
 	 */
 	id<MTLFunction> newMTLFunction(id<MTLLibrary> mtlLibrary, NSString* funcName, MTLFunctionConstantValues* constantValues);


 #pragma mark Construction

 	MVKFunctionSpecializer(MVKVulkanAPIDeviceObject* owner) : MVKMetalCompiler(owner) {
 		_compilerType = "Function specialization";
 		_pPerformanceTracker = &_owner->getDevice()->_performanceStatistics.shaderCompilation.functionSpecialization;
 	}

 	~MVKFunctionSpecializer() override;

 protected:
 	bool compileComplete(id<MTLFunction> mtlFunction, NSError *error);

 	id<MTLFunction> _mtlFunction = nil;
 };
	/*
	* MVKShaderModule.h
	*
	* Copyright (c) 2015-2022 The Brenwill Workshop Ltd. (http://www.brenwill.com)
	*
	* Licensed under the Apache License, Version 2.0 (the "License");
	* you may not use this file except in compliance with the License.
	* You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	#pragma once

	#include "MVKDevice.h"
	#include "MVKSync.h"
	#include "MVKSmallVector.h"
	#include <MoltenVKShaderConverter/SPIRVToMSLConverter.h>
	#include <MoltenVKShaderConverter/GLSLToSPIRVConverter.h>
	#include <mutex>

	#import <Metal/Metal.h>

	class MVKPipelineCache;
	class MVKShaderCacheIterator;
	class MVKShaderLibraryCache;
	class MVKShaderModule;

	using namespace mvk;


	#pragma mark -
	#pragma mark MVKShaderLibrary

	/** A MTLFunction and corresponding result information resulting from a shader conversion. */
	typedef struct MVKMTLFunction {
	SPIRVToMSLConversionResults shaderConversionResults;
	MTLSize threadGroupSize;
	inline id<MTLFunction> getMTLFunction() { return _mtlFunction; }

	MVKMTLFunction(id<MTLFunction> mtlFunc, const SPIRVToMSLConversionResults scRslts, MTLSize tgSize);
	MVKMTLFunction(const MVKMTLFunction& other);
	MVKMTLFunction& operator=(const MVKMTLFunction& other);
	MVKMTLFunction() {}
	~MVKMTLFunction();

	private:
	id<MTLFunction> _mtlFunction = nil;

	} MVKMTLFunction;

	/** A MVKMTLFunction indicating an invalid MTLFunction. The mtlFunction member is nil. */
	const MVKMTLFunction MVKMTLFunctionNull(nil, SPIRVToMSLConversionResults(), MTLSizeMake(1, 1, 1));

	/** Wraps a single MTLLibrary. */
	class MVKShaderLibrary : public MVKBaseObject {

	public:

	/** Returns the Vulkan API opaque object controlling this object. */
	MVKVulkanAPIObject* getVulkanAPIObject() override { return _owner->getVulkanAPIObject(); };

	/**
	* Sets the entry point function name.
	*
	* This is usually set automatically during shader conversion from SPIR-V to MSL.
	* For a library that was created directly from MSL, this function can be used to
	* set the name of the function if it has a different name than the default main0().
	*/
	void setEntryPointName(std::string& funcName);

	/**
	* Sets the number of threads in a single compute kernel workgroup, per dimension.
	*
	* This is usually set automatically during shader conversion from SPIR-V to MSL.
	* For a library that was created directly from MSL, this function can be used to
	* set the workgroup size..
	*/
	void setWorkgroupSize(uint32_t x, uint32_t y, uint32_t z);

	/** Constructs an instance from the specified MSL source code. */
	MVKShaderLibrary(MVKVulkanAPIDeviceObject* owner,
	const std::string& mslSourceCode,
	const SPIRVToMSLConversionResults& shaderConversionResults);

	/** Constructs an instance from the specified compiled MSL code data. */
	MVKShaderLibrary(MVKVulkanAPIDeviceObject* owner,
	const void* mslCompiledCodeData,
	size_t mslCompiledCodeLength);

	MVKShaderLibrary(const MVKShaderLibrary& other);

	MVKShaderLibrary& operator=(const MVKShaderLibrary& other);

	~MVKShaderLibrary() override;

	protected:
	friend MVKShaderCacheIterator;
	friend MVKShaderLibraryCache;
	friend MVKShaderModule;

	MVKMTLFunction getMTLFunction(const VkSpecializationInfo* pSpecializationInfo, MVKShaderModule* shaderModule);
	void handleCompilationError(NSError* err, const char* opDesc);
	MTLFunctionConstant* getFunctionConstant(NSArray<MTLFunctionConstant> mtlFCs, NSUInteger mtlFCID);

	MVKVulkanAPIDeviceObject* _owner;
	id<MTLLibrary> _mtlLibrary;
	SPIRVToMSLConversionResults _shaderConversionResults;
	std::string _msl;
	};


	#pragma mark -
	#pragma mark MVKShaderLibraryCache

	/** Represents a cache of shader libraries for one shader module. */
	class MVKShaderLibraryCache : public MVKBaseObject {

	public:

	/** Returns the Vulkan API opaque object controlling this object. */
	MVKVulkanAPIObject* getVulkanAPIObject() override { return _owner->getVulkanAPIObject(); };

	/**
	* Returns a shader library from the shader conversion configuration sourced from the shader module,
	* lazily creating the shader library from source code in the shader module, if needed.
	*
	* If pWasAdded is not nil, this function will set it to true if a new shader library was created,
	* and to false if an existing shader library was found and returned.
	*/
	MVKShaderLibrary* getShaderLibrary(SPIRVToMSLConversionConfiguration* pShaderConfig,
	MVKShaderModule* shaderModule,
	bool* pWasAdded = nullptr);

	MVKShaderLibraryCache(MVKVulkanAPIDeviceObject* owner) : _owner(owner) {};

	~MVKShaderLibraryCache() override;

	protected:
	friend MVKShaderCacheIterator;
	friend MVKPipelineCache;
	friend MVKShaderModule;

	MVKShaderLibrary* findShaderLibrary(SPIRVToMSLConversionConfiguration* pShaderConfig);
	MVKShaderLibrary* addShaderLibrary(SPIRVToMSLConversionConfiguration* pShaderConfig,
	const std::string& mslSourceCode,
	const SPIRVToMSLConversionResults& shaderConversionResults);
	void merge(MVKShaderLibraryCache* other);

	MVKVulkanAPIDeviceObject* _owner;
	MVKSmallVector<std::pair<SPIRVToMSLConversionConfiguration, MVKShaderLibrary*>> _shaderLibraries;
	};


	#pragma mark -
	#pragma mark MVKShaderModule

	typedef struct MVKShaderModuleKey {
	std::size_t codeSize;
	std::size_t codeHash;

	bool operator==(const MVKShaderModuleKey& rhs) const {
	return ((codeSize == rhs.codeSize) && (codeHash == rhs.codeHash));
	}
	MVKShaderModuleKey(std::size_t codeSize, std::size_t codeHash) : codeSize(codeSize), codeHash(codeHash) {}
	MVKShaderModuleKey() : MVKShaderModuleKey(0, 0) {}
	} MVKShaderModuleKey;

	/**
	* Hash structure implementation for MVKShaderModuleKey in std namespace,
	* so MVKShaderModuleKey can be used as a key in a std::map and std::unordered_map.
	*/
	namespace std {
	template <>
	struct hash<MVKShaderModuleKey> {
	std::size_t operator()(const MVKShaderModuleKey& k) const { return k.codeHash; }
	};
	}

	/** Represents a Vulkan shader module. */
	class MVKShaderModule : public MVKVulkanAPIDeviceObject {

	public:

	/** Returns the Vulkan type of this object. */
	VkObjectType getVkObjectType() override { return VK_OBJECT_TYPE_SHADER_MODULE; }

	/** Returns the debug report object type of this object. */
	VkDebugReportObjectTypeEXT getVkDebugReportObjectType() override { return VK_DEBUG_REPORT_OBJECT_TYPE_SHADER_MODULE_EXT; }

	/** Returns the Metal shader function, possibly specialized. */
	MVKMTLFunction getMTLFunction(SPIRVToMSLConversionConfiguration* pShaderConfig,
	const VkSpecializationInfo* pSpecializationInfo,
	MVKPipelineCache* pipelineCache);

	/** Convert the SPIR-V to MSL, using the specified shader conversion configuration. */
	bool convert(SPIRVToMSLConversionConfiguration* pShaderConfig);

	/** Returns the original SPIR-V code that was specified when this object was created. */
	const std::vector<uint32_t>& getSPIRV() { return _spvConverter.getSPIRV(); }

	/**
	* Returns the Metal Shading Language source code as converted by the most recent
	* call to convert() function, or set directly using the setMSL() function.
	*/
	const std::string& getMSL() { return _spvConverter.getMSL(); }

	/** Returns information about the shader conversion results. */
	const SPIRVToMSLConversionResults& getConversionResults() { return _spvConverter.getConversionResults(); }

	/** Sets the number of threads in a single compute kernel workgroup, per dimension. */
	void setWorkgroupSize(uint32_t x, uint32_t y, uint32_t z);

	/** Returns a key as a means of identifying this shader module in a pipeline cache. */
	MVKShaderModuleKey getKey() { return _key; }

	MVKShaderModule(MVKDevice* device, const VkShaderModuleCreateInfo* pCreateInfo);

	~MVKShaderModule() override;

	protected:
	friend MVKShaderCacheIterator;

	void propagateDebugName() override {}
	MVKGLSLConversionShaderStage getMVKGLSLConversionShaderStage(SPIRVToMSLConversionConfiguration* pShaderConfig);

	MVKShaderLibraryCache _shaderLibraryCache;
	SPIRVToMSLConverter _spvConverter;
	GLSLToSPIRVConverter _glslConverter;
	MVKShaderLibrary* _directMSLLibrary;
	MVKShaderModuleKey _key;
	std::mutex _accessLock;
	};


	#pragma mark -
	#pragma mark MVKShaderLibraryCompiler

	/**
	* Creates a MTLLibrary from source code.
	*
	* Instances of this class are one-shot, and can only be used for a single library compilation.
	*/
	class MVKShaderLibraryCompiler : public MVKMetalCompiler {

	public:

	/**
	* Returns a new (retained) MTLLibrary object compiled from the MSL source code.
	*
	* If the Metal library compiler does not return within MVKConfiguration::metalCompileTimeout
	* nanoseconds, an error will be generated and logged, and nil will be returned.
	*/
	id<MTLLibrary> newMTLLibrary(NSString* mslSourceCode,
	const SPIRVToMSLConversionResults& shaderConversionResults);


	#pragma mark Construction

	MVKShaderLibraryCompiler(MVKVulkanAPIDeviceObject* owner) : MVKMetalCompiler(owner) {
	_compilerType = "Shader library";
	_pPerformanceTracker = &_owner->getDevice()->_performanceStatistics.shaderCompilation.mslCompile;
	}

	~MVKShaderLibraryCompiler() override;

	protected:
	bool compileComplete(id<MTLLibrary> mtlLibrary, NSError *error);
	void handleError() override;

	id<MTLLibrary> _mtlLibrary = nil;
	};


	#pragma mark -
	#pragma mark MVKFunctionSpecializer

	/**
	* Compiles a specialized MTLFunction.
	*
	* Instances of this class are one-shot, and can only be used for a single function compilation.
	*/
	class MVKFunctionSpecializer : public MVKMetalCompiler {

	public:

	/**
	* Returns a new (retained) MTLFunction object compiled from the MTLLibrary and specialization constants.
	*
	* If the Metal function compiler does not return within MVKConfiguration::metalCompileTimeout
	* nanoseconds, an error will be generated and logged, and nil will be returned.
	*/
	id<MTLFunction> newMTLFunction(id<MTLLibrary> mtlLibrary, NSString* funcName, MTLFunctionConstantValues* constantValues);


	#pragma mark Construction

	MVKFunctionSpecializer(MVKVulkanAPIDeviceObject* owner) : MVKMetalCompiler(owner) {
	_compilerType = "Function specialization";
	_pPerformanceTracker = &_owner->getDevice()->_performanceStatistics.shaderCompilation.functionSpecialization;
	}

	~MVKFunctionSpecializer() override;

	protected:
	bool compileComplete(id<MTLFunction> mtlFunction, NSError *error);

	id<MTLFunction> _mtlFunction = nil;
	};