src/dawn/native/vulkan/ShaderModuleVk.cpp - dawn - Git at Google

 // Copyright 2018 The Dawn Authors
 //
 // 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.

 #include "dawn/native/vulkan/ShaderModuleVk.h"

 #include "dawn/native/SpirvValidation.h"
 #include "dawn/native/TintUtils.h"
 #include "dawn/native/vulkan/BindGroupLayoutVk.h"
 #include "dawn/native/vulkan/DeviceVk.h"
 #include "dawn/native/vulkan/FencedDeleter.h"
 #include "dawn/native/vulkan/PipelineLayoutVk.h"
 #include "dawn/native/vulkan/UtilsVulkan.h"
 #include "dawn/native/vulkan/VulkanError.h"
 #include "dawn/platform/DawnPlatform.h"
 #include "dawn/platform/tracing/TraceEvent.h"

 #include <tint/tint.h>
 #include <spirv-tools/libspirv.hpp>

 namespace dawn::native::vulkan {

     ShaderModule::ConcurrentTransformedShaderModuleCache::ConcurrentTransformedShaderModuleCache(
         Device* device)
         : mDevice(device) {
     }

     ShaderModule::ConcurrentTransformedShaderModuleCache::
         ~ConcurrentTransformedShaderModuleCache() {
         std::lock_guard<std::mutex> lock(mMutex);
         for (const auto& [_, moduleAndSpirv] : mTransformedShaderModuleCache) {
             mDevice->GetFencedDeleter()->DeleteWhenUnused(moduleAndSpirv.first);
         }
     }

     std::optional<ShaderModule::ModuleAndSpirv>
     ShaderModule::ConcurrentTransformedShaderModuleCache::Find(
         const PipelineLayoutEntryPointPair& key) {
         std::lock_guard<std::mutex> lock(mMutex);
         auto iter = mTransformedShaderModuleCache.find(key);
         if (iter != mTransformedShaderModuleCache.end()) {
             return std::make_pair(iter->second.first, iter->second.second.get());
         }
         return {};
     }

     ShaderModule::ModuleAndSpirv ShaderModule::ConcurrentTransformedShaderModuleCache::AddOrGet(
         const PipelineLayoutEntryPointPair& key,
         VkShaderModule module,
         std::vector<uint32_t>&& spirv) {
         ASSERT(module != VK_NULL_HANDLE);
         std::lock_guard<std::mutex> lock(mMutex);
         auto iter = mTransformedShaderModuleCache.find(key);
         if (iter == mTransformedShaderModuleCache.end()) {
             mTransformedShaderModuleCache.emplace(
                 key, std::make_pair(module, std::unique_ptr<Spirv>(new Spirv(spirv))));
         } else {
             mDevice->GetFencedDeleter()->DeleteWhenUnused(module);
         }
         // Now the key should exist in the map, so find it again and return it.
         iter = mTransformedShaderModuleCache.find(key);
         return std::make_pair(iter->second.first, iter->second.second.get());
     }

     // static
     ResultOrError<Ref<ShaderModule>> ShaderModule::Create(Device* device,
                                                           const ShaderModuleDescriptor* descriptor,
                                                           ShaderModuleParseResult* parseResult) {
         Ref<ShaderModule> module = AcquireRef(new ShaderModule(device, descriptor));
         DAWN_TRY(module->Initialize(parseResult));
         return module;
     }

     ShaderModule::ShaderModule(Device* device, const ShaderModuleDescriptor* descriptor)
         : ShaderModuleBase(device, descriptor),
           mTransformedShaderModuleCache(
               std::make_unique<ConcurrentTransformedShaderModuleCache>(device)) {
     }

     MaybeError ShaderModule::Initialize(ShaderModuleParseResult* parseResult) {
         if (GetDevice()->IsRobustnessEnabled()) {
             ScopedTintICEHandler scopedICEHandler(GetDevice());

             tint::transform::Robustness robustness;
             tint::transform::DataMap transformInputs;

             tint::Program program;
             DAWN_TRY_ASSIGN(program, RunTransforms(&robustness, parseResult->tintProgram.get(),
                                                    transformInputs, nullptr, nullptr));
             // Rather than use a new ParseResult object, we just reuse the original parseResult
             parseResult->tintProgram = std::make_unique<tint::Program>(std::move(program));
         }

         return InitializeBase(parseResult);
     }

     void ShaderModule::DestroyImpl() {
         ShaderModuleBase::DestroyImpl();
         // Remove reference to internal cache to trigger cleanup.
         mTransformedShaderModuleCache = nullptr;
     }

     ShaderModule::~ShaderModule() = default;

     ResultOrError<ShaderModule::ModuleAndSpirv> ShaderModule::GetHandleAndSpirv(
         const char* entryPointName,
         PipelineLayout* layout) {
         TRACE_EVENT0(GetDevice()->GetPlatform(), General, "ShaderModuleVk::GetHandleAndSpirv");

         // If the shader was destroyed, we should never call this function.
         ASSERT(IsAlive());

         ScopedTintICEHandler scopedICEHandler(GetDevice());

         // Check to see if we have the handle and spirv cached already.
         auto cacheKey = std::make_pair(layout, entryPointName);
         auto handleAndSpirv = mTransformedShaderModuleCache->Find(cacheKey);
         if (handleAndSpirv.has_value()) {
             return std::move(*handleAndSpirv);
         }

         // Creation of module and spirv is deferred to this point when using tint generator

         // Remap BindingNumber to BindingIndex in WGSL shader
         using BindingRemapper = tint::transform::BindingRemapper;
         using BindingPoint = tint::transform::BindingPoint;
         BindingRemapper::BindingPoints bindingPoints;
         BindingRemapper::AccessControls accessControls;

         const BindingInfoArray& moduleBindingInfo = GetEntryPoint(entryPointName).bindings;

         for (BindGroupIndex group : IterateBitSet(layout->GetBindGroupLayoutsMask())) {
             const BindGroupLayout* bgl = ToBackend(layout->GetBindGroupLayout(group));
             const auto& groupBindingInfo = moduleBindingInfo[group];
             for (const auto& it : groupBindingInfo) {
                 BindingNumber binding = it.first;
                 BindingIndex bindingIndex = bgl->GetBindingIndex(binding);
                 BindingPoint srcBindingPoint{static_cast<uint32_t>(group),
                                              static_cast<uint32_t>(binding)};

                 BindingPoint dstBindingPoint{static_cast<uint32_t>(group),
                                              static_cast<uint32_t>(bindingIndex)};
                 if (srcBindingPoint != dstBindingPoint) {
                     bindingPoints.emplace(srcBindingPoint, dstBindingPoint);
                 }
             }
         }

         tint::transform::Manager transformManager;
         transformManager.append(std::make_unique<tint::transform::BindingRemapper>());
         // Many Vulkan drivers can't handle multi-entrypoint shader modules.
         transformManager.append(std::make_unique<tint::transform::SingleEntryPoint>());

         tint::transform::DataMap transformInputs;
         transformInputs.Add<BindingRemapper::Remappings>(std::move(bindingPoints),
                                                          std::move(accessControls),
                                                          /* mayCollide */ false);
         transformInputs.Add<tint::transform::SingleEntryPoint::Config>(entryPointName);

         // Transform external textures into the binding locations specified in the bgl
         // TODO(dawn:1082): Replace this block with ShaderModuleBase::AddExternalTextureTransform.
         tint::transform::MultiplanarExternalTexture::BindingsMap newBindingsMap;
         for (BindGroupIndex i : IterateBitSet(layout->GetBindGroupLayoutsMask())) {
             BindGroupLayoutBase* bgl = layout->GetBindGroupLayout(i);

             ExternalTextureBindingExpansionMap expansions =
                 bgl->GetExternalTextureBindingExpansionMap();

             std::map<BindingNumber, dawn_native::ExternalTextureBindingExpansion>::iterator it =
                 expansions.begin();

             while (it != expansions.end()) {
                 newBindingsMap[{static_cast<uint32_t>(i),
                                 static_cast<uint32_t>(bgl->GetBindingIndex(it->second.plane0))}] = {
                     {static_cast<uint32_t>(i),
                      static_cast<uint32_t>(bgl->GetBindingIndex(it->second.plane1))},
                     {static_cast<uint32_t>(i),
                      static_cast<uint32_t>(bgl->GetBindingIndex(it->second.params))}};
                 it++;
             }
         }

         if (!newBindingsMap.empty()) {
             transformManager.Add<tint::transform::MultiplanarExternalTexture>();
             transformInputs.Add<tint::transform::MultiplanarExternalTexture::NewBindingPoints>(
                 newBindingsMap);
         }

         tint::Program program;
         {
             TRACE_EVENT0(GetDevice()->GetPlatform(), General, "RunTransforms");
             DAWN_TRY_ASSIGN(program, RunTransforms(&transformManager, GetTintProgram(),
                                                    transformInputs, nullptr, nullptr));
         }

 #if TINT_BUILD_SPV_WRITER
         tint::writer::spirv::Options options;
         options.emit_vertex_point_size = true;
         options.disable_workgroup_init = GetDevice()->IsToggleEnabled(Toggle::DisableWorkgroupInit);
         options.use_zero_initialize_workgroup_memory_extension =
             GetDevice()->IsToggleEnabled(Toggle::VulkanUseZeroInitializeWorkgroupMemoryExtension);

         Spirv spirv;
         {
             TRACE_EVENT0(GetDevice()->GetPlatform(), General, "tint::writer::spirv::Generate()");
             auto result = tint::writer::spirv::Generate(&program, options);
             DAWN_INVALID_IF(!result.success, "An error occured while generating SPIR-V: %s.",
                             result.error);

             spirv = std::move(result.spirv);
         }

         DAWN_TRY(
             ValidateSpirv(GetDevice(), spirv, GetDevice()->IsToggleEnabled(Toggle::DumpShaders)));

         VkShaderModuleCreateInfo createInfo;
         createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
         createInfo.pNext = nullptr;
         createInfo.flags = 0;
         createInfo.codeSize = spirv.size() * sizeof(uint32_t);
         createInfo.pCode = spirv.data();

         Device* device = ToBackend(GetDevice());

         VkShaderModule newHandle = VK_NULL_HANDLE;
         {
             TRACE_EVENT0(GetDevice()->GetPlatform(), General, "vkCreateShaderModule");
             DAWN_TRY(CheckVkSuccess(device->fn.CreateShaderModule(
                                         device->GetVkDevice(), &createInfo, nullptr, &*newHandle),
                                     "CreateShaderModule"));
         }
         ModuleAndSpirv moduleAndSpirv;
         if (newHandle != VK_NULL_HANDLE) {
             moduleAndSpirv =
                 mTransformedShaderModuleCache->AddOrGet(cacheKey, newHandle, std::move(spirv));
         }

         SetDebugName(ToBackend(GetDevice()), VK_OBJECT_TYPE_SHADER_MODULE,
                      reinterpret_cast<uint64_t&>(moduleAndSpirv.first), "Dawn_ShaderModule",
                      GetLabel());

         return std::move(moduleAndSpirv);
 #else
         return DAWN_INTERNAL_ERROR("TINT_BUILD_SPV_WRITER is not defined.");
 #endif
     }

 }  // namespace dawn::native::vulkan
	// Copyright 2018 The Dawn Authors
	//
	// 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.

	#include "dawn/native/vulkan/ShaderModuleVk.h"

	#include "dawn/native/SpirvValidation.h"
	#include "dawn/native/TintUtils.h"
	#include "dawn/native/vulkan/BindGroupLayoutVk.h"
	#include "dawn/native/vulkan/DeviceVk.h"
	#include "dawn/native/vulkan/FencedDeleter.h"
	#include "dawn/native/vulkan/PipelineLayoutVk.h"
	#include "dawn/native/vulkan/UtilsVulkan.h"
	#include "dawn/native/vulkan/VulkanError.h"
	#include "dawn/platform/DawnPlatform.h"
	#include "dawn/platform/tracing/TraceEvent.h"

	#include <tint/tint.h>
	#include <spirv-tools/libspirv.hpp>

	namespace dawn::native::vulkan {

	ShaderModule::ConcurrentTransformedShaderModuleCache::ConcurrentTransformedShaderModuleCache(
	Device* device)
	: mDevice(device) {
	}

	ShaderModule::ConcurrentTransformedShaderModuleCache::
	~ConcurrentTransformedShaderModuleCache() {
	std::lock_guard<std::mutex> lock(mMutex);
	for (const auto& [_, moduleAndSpirv] : mTransformedShaderModuleCache) {
	mDevice->GetFencedDeleter()->DeleteWhenUnused(moduleAndSpirv.first);
	}
	}

	std::optional<ShaderModule::ModuleAndSpirv>
	ShaderModule::ConcurrentTransformedShaderModuleCache::Find(
	const PipelineLayoutEntryPointPair& key) {
	std::lock_guard<std::mutex> lock(mMutex);
	auto iter = mTransformedShaderModuleCache.find(key);
	if (iter != mTransformedShaderModuleCache.end()) {
	return std::make_pair(iter->second.first, iter->second.second.get());
	}
	return {};
	}

	ShaderModule::ModuleAndSpirv ShaderModule::ConcurrentTransformedShaderModuleCache::AddOrGet(
	const PipelineLayoutEntryPointPair& key,
	VkShaderModule module,
	std::vector<uint32_t>&& spirv) {
	ASSERT(module != VK_NULL_HANDLE);
	std::lock_guard<std::mutex> lock(mMutex);
	auto iter = mTransformedShaderModuleCache.find(key);
	if (iter == mTransformedShaderModuleCache.end()) {
	mTransformedShaderModuleCache.emplace(
	key, std::make_pair(module, std::unique_ptr<Spirv>(new Spirv(spirv))));
	} else {
	mDevice->GetFencedDeleter()->DeleteWhenUnused(module);
	}
	// Now the key should exist in the map, so find it again and return it.
	iter = mTransformedShaderModuleCache.find(key);
	return std::make_pair(iter->second.first, iter->second.second.get());
	}

	// static
	ResultOrError<Ref<ShaderModule>> ShaderModule::Create(Device* device,
	const ShaderModuleDescriptor* descriptor,
	ShaderModuleParseResult* parseResult) {
	Ref<ShaderModule> module = AcquireRef(new ShaderModule(device, descriptor));
	DAWN_TRY(module->Initialize(parseResult));
	return module;
	}

	ShaderModule::ShaderModule(Device* device, const ShaderModuleDescriptor* descriptor)
	: ShaderModuleBase(device, descriptor),
	mTransformedShaderModuleCache(
	std::make_unique<ConcurrentTransformedShaderModuleCache>(device)) {
	}

	MaybeError ShaderModule::Initialize(ShaderModuleParseResult* parseResult) {
	if (GetDevice()->IsRobustnessEnabled()) {
	ScopedTintICEHandler scopedICEHandler(GetDevice());

	tint::transform::Robustness robustness;
	tint::transform::DataMap transformInputs;

	tint::Program program;
	DAWN_TRY_ASSIGN(program, RunTransforms(&robustness, parseResult->tintProgram.get(),
	transformInputs, nullptr, nullptr));
	// Rather than use a new ParseResult object, we just reuse the original parseResult
	parseResult->tintProgram = std::make_unique<tint::Program>(std::move(program));
	}

	return InitializeBase(parseResult);
	}

	void ShaderModule::DestroyImpl() {
	ShaderModuleBase::DestroyImpl();
	// Remove reference to internal cache to trigger cleanup.
	mTransformedShaderModuleCache = nullptr;
	}

	ShaderModule::~ShaderModule() = default;

	ResultOrError<ShaderModule::ModuleAndSpirv> ShaderModule::GetHandleAndSpirv(
	const char* entryPointName,
	PipelineLayout* layout) {
	TRACE_EVENT0(GetDevice()->GetPlatform(), General, "ShaderModuleVk::GetHandleAndSpirv");

	// If the shader was destroyed, we should never call this function.
	ASSERT(IsAlive());

	ScopedTintICEHandler scopedICEHandler(GetDevice());

	// Check to see if we have the handle and spirv cached already.
	auto cacheKey = std::make_pair(layout, entryPointName);
	auto handleAndSpirv = mTransformedShaderModuleCache->Find(cacheKey);
	if (handleAndSpirv.has_value()) {
	return std::move(*handleAndSpirv);
	}

	// Creation of module and spirv is deferred to this point when using tint generator

	// Remap BindingNumber to BindingIndex in WGSL shader
	using BindingRemapper = tint::transform::BindingRemapper;
	using BindingPoint = tint::transform::BindingPoint;
	BindingRemapper::BindingPoints bindingPoints;
	BindingRemapper::AccessControls accessControls;

	const BindingInfoArray& moduleBindingInfo = GetEntryPoint(entryPointName).bindings;

	for (BindGroupIndex group : IterateBitSet(layout->GetBindGroupLayoutsMask())) {
	const BindGroupLayout* bgl = ToBackend(layout->GetBindGroupLayout(group));
	const auto& groupBindingInfo = moduleBindingInfo[group];
	for (const auto& it : groupBindingInfo) {
	BindingNumber binding = it.first;
	BindingIndex bindingIndex = bgl->GetBindingIndex(binding);
	BindingPoint srcBindingPoint{static_cast<uint32_t>(group),
	static_cast<uint32_t>(binding)};

	BindingPoint dstBindingPoint{static_cast<uint32_t>(group),
	static_cast<uint32_t>(bindingIndex)};
	if (srcBindingPoint != dstBindingPoint) {
	bindingPoints.emplace(srcBindingPoint, dstBindingPoint);
	}
	}
	}

	tint::transform::Manager transformManager;
	transformManager.append(std::make_unique<tint::transform::BindingRemapper>());
	// Many Vulkan drivers can't handle multi-entrypoint shader modules.
	transformManager.append(std::make_unique<tint::transform::SingleEntryPoint>());

	tint::transform::DataMap transformInputs;
	transformInputs.Add<BindingRemapper::Remappings>(std::move(bindingPoints),
	std::move(accessControls),
	/* mayCollide */ false);
	transformInputs.Add<tint::transform::SingleEntryPoint::Config>(entryPointName);

	// Transform external textures into the binding locations specified in the bgl
	// TODO(dawn:1082): Replace this block with ShaderModuleBase::AddExternalTextureTransform.
	tint::transform::MultiplanarExternalTexture::BindingsMap newBindingsMap;
	for (BindGroupIndex i : IterateBitSet(layout->GetBindGroupLayoutsMask())) {
	BindGroupLayoutBase* bgl = layout->GetBindGroupLayout(i);

	ExternalTextureBindingExpansionMap expansions =
	bgl->GetExternalTextureBindingExpansionMap();

	std::map<BindingNumber, dawn_native::ExternalTextureBindingExpansion>::iterator it =
	expansions.begin();

	while (it != expansions.end()) {
	newBindingsMap[{static_cast<uint32_t>(i),
	static_cast<uint32_t>(bgl->GetBindingIndex(it->second.plane0))}] = {
	{static_cast<uint32_t>(i),
	static_cast<uint32_t>(bgl->GetBindingIndex(it->second.plane1))},
	{static_cast<uint32_t>(i),
	static_cast<uint32_t>(bgl->GetBindingIndex(it->second.params))}};
	it++;
	}
	}

	if (!newBindingsMap.empty()) {
	transformManager.Add<tint::transform::MultiplanarExternalTexture>();
	transformInputs.Add<tint::transform::MultiplanarExternalTexture::NewBindingPoints>(
	newBindingsMap);
	}

	tint::Program program;
	{
	TRACE_EVENT0(GetDevice()->GetPlatform(), General, "RunTransforms");
	DAWN_TRY_ASSIGN(program, RunTransforms(&transformManager, GetTintProgram(),
	transformInputs, nullptr, nullptr));
	}

	#if TINT_BUILD_SPV_WRITER
	tint::writer::spirv::Options options;
	options.emit_vertex_point_size = true;
	options.disable_workgroup_init = GetDevice()->IsToggleEnabled(Toggle::DisableWorkgroupInit);
	options.use_zero_initialize_workgroup_memory_extension =
	GetDevice()->IsToggleEnabled(Toggle::VulkanUseZeroInitializeWorkgroupMemoryExtension);

	Spirv spirv;
	{
	TRACE_EVENT0(GetDevice()->GetPlatform(), General, "tint::writer::spirv::Generate()");
	auto result = tint::writer::spirv::Generate(&program, options);
	DAWN_INVALID_IF(!result.success, "An error occured while generating SPIR-V: %s.",
	result.error);

	spirv = std::move(result.spirv);
	}

	DAWN_TRY(
	ValidateSpirv(GetDevice(), spirv, GetDevice()->IsToggleEnabled(Toggle::DumpShaders)));

	VkShaderModuleCreateInfo createInfo;
	createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
	createInfo.pNext = nullptr;
	createInfo.flags = 0;
	createInfo.codeSize = spirv.size() * sizeof(uint32_t);
	createInfo.pCode = spirv.data();

	Device* device = ToBackend(GetDevice());

	VkShaderModule newHandle = VK_NULL_HANDLE;
	{
	TRACE_EVENT0(GetDevice()->GetPlatform(), General, "vkCreateShaderModule");
	DAWN_TRY(CheckVkSuccess(device->fn.CreateShaderModule(
	device->GetVkDevice(), &createInfo, nullptr, &*newHandle),
	"CreateShaderModule"));
	}
	ModuleAndSpirv moduleAndSpirv;
	if (newHandle != VK_NULL_HANDLE) {
	moduleAndSpirv =
	mTransformedShaderModuleCache->AddOrGet(cacheKey, newHandle, std::move(spirv));
	}

	SetDebugName(ToBackend(GetDevice()), VK_OBJECT_TYPE_SHADER_MODULE,
	reinterpret_cast<uint64_t&>(moduleAndSpirv.first), "Dawn_ShaderModule",
	GetLabel());

	return std::move(moduleAndSpirv);
	#else
	return DAWN_INTERNAL_ERROR("TINT_BUILD_SPV_WRITER is not defined.");
	#endif
	}

	} // namespace dawn::native::vulkan