Differential D4969 Diff 21686 source/tools/spirv/compile.py

Changeset View

Standalone View

source/tools/spirv/compile.py

Property	Old Value	New Value
svn:eol-style	null	native \ No newline at end of property

				#!/usr/bin/env python3
				# -- mode: python-mode; python-indent-offset: 4; --
				#
				# Copyright (C) 2023 Wildfire Games.
				#
				# Permission is hereby granted, free of charge, to any person obtaining a copy
				# of this software and associated documentation files (the "Software"), to deal
				# in the Software without restriction, including without limitation the rights
				# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
				# copies of the Software, and to permit persons to whom the Software is
				# furnished to do so, subject to the following conditions:
				#
				# The above copyright notice and this permission notice shall be included in
				# all copies or substantial portions of the Software.
				#
				# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
				# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
				# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
				# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
				# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
				# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
				# THE SOFTWARE.

				import datetime
				import hashlib
				import itertools
				import json
				import os
				import subprocess
				import sys
				import time
				import yaml

				import xml.etree.ElementTree as ET


				def execute(command):
				try:
				process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
				out, err = process.communicate()
				except:
				sys.stderr.write('Failed to run command: {}\n'.format(' '.join(command)))
				raise
				return process.returncode, out, err

				def calculate_hash(path):
				assert os.path.isfile(path)
				with open(path, 'rb') as handle:
				return hashlib.sha1(handle.read()).hexdigest()

				def compare_spirv(path1, path2):
				with open(path1, 'rb') as handle:
				spirv1 = handle.read()
				with open(path2, 'rb') as handle:
				spirv2 = handle.read()
				return spirv1 == spirv2

				def resolve_if(defines, expression):
				for item in expression.strip().split('\|\|'):
				item = item.strip()
				assert len(item) > 1
				name = item
				invert = False
				if name[0] == '!':
				invert = True
				name = item[1:]
				assert item[1].isalpha()
				else:
				assert item[0].isalpha()
				found_define = False
				for define in defines:
				if define['name'] == name:
				assert define['value'] == 'UNDEFINED' or define['value'] == '0' or define['value'] == '1'
				if invert:
				if define['value'] != '1':
				return True
				found_define = True
				else:
				if define['value'] == '1':
				return True
				if invert and not found_define:
				return True
				return False

				def compile_and_reflect(input_root_path, output_root_path, mod_name, stage, path, out_path, defines):
				keep_debug = False
				input_path = os.path.normpath(path)
				fallback_input_path = input_path.replace(mod_name, 'mod')
				if not os.path.isfile(input_path) and os.path.isfile(fallback_input_path):
				input_path = fallback_input_path
				output_path = os.path.normpath(out_path)
				command = [
				'glslc', '-x', 'glsl', '--target-env=vulkan1.0', '-std=450core',
				'-I', os.path.join(input_root_path, 'mod', 'shaders', 'glsl'),
				'-I', os.path.join(input_root_path, mod_name, 'shaders', 'glsl'),
				'-fshader-stage=' + stage, '-O', input_path,
				]
				use_descriptor_indexing = False
				for define in defines:
				if define['value'] == 'UNDEFINED':
				continue
				assert ' ' not in define['value']
				command.append('-D{}={}'.format(define['name'], define['value']))
				if define['name'] == 'USE_DESCRIPTOR_INDEXING':
				use_descriptor_indexing = True
				command.append('-D{}={}'.format('USE_SPIRV', '1'))
				command.append('-DSTAGE_{}={}'.format(stage.upper(), '1'))
				command += ['-o', output_path]
				# Compile the shader with debug information to see names in reflection.
				ret, out, err = execute(command + ['-g'])
				if ret:
				sys.stderr.write('Command returned {}:\nCommand: {}\nInput path: {}\nOutput path: {}\nError: {}\n'.format(
				ret, ' '.join(command), input_path, output_path, err))
				preprocessor_output_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'preprocessed_file.glsl'))
				execute(command[:-2] + ['-g', '-E', '-o', preprocessor_output_path])
				raise ValueError(err)
				ret, out, err = execute(['spirv-reflect', '-y', output_path])
				if ret:
				sys.stderr.write('Command returned {}:\nCommand: {}\nInput path: {}\nOutput path: {}\nError: {}\n'.format(
				ret, ' '.join(command), input_path, output_path, err))
				raise ValueError(err)
				# Reflect the result SPIRV.
				data = yaml.safe_load(out)
				module = data['module']
				interface_variables = []
				if 'all_interface_variables' in data and data['all_interface_variables']:
				interface_variables = data['all_interface_variables']
				push_constants = []
				vertex_attributes = []
				if 'push_constants' in module and module['push_constants']:
				assert len(module['push_constants']) == 1
				def add_push_constants(node, push_constants):
				if ('members' in node) and node['members']:
				for member in node['members']:
				add_push_constants(member, push_constants)
				else:
				assert node['absolute_offset'] + node['size'] <= 128
				push_constants.append({
				'name': node['name'],
				'offset': node['absolute_offset'],
				'size': node['size'],
				})
				assert module['push_constants'][0]['name'] == 'DrawUniforms'
				assert module['push_constants'][0]['size'] <= 128
				add_push_constants(module['push_constants'][0], push_constants)
				descriptor_sets = []
				if 'descriptor_sets' in module and module['descriptor_sets']:
				VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER = 1
				VK_DESCRIPTOR_TYPE_STORAGE_IMAGE = 3
				VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER = 6
				VK_DESCRIPTOR_TYPE_STORAGE_BUFFER = 7
				for descriptor_set in module['descriptor_sets']:
				UNIFORM_SET = 1 if use_descriptor_indexing else 0
				STORAGE_SET = 2
				bindings = []
				if descriptor_set['set'] == UNIFORM_SET:
				assert descriptor_set['binding_count'] > 0
				for binding in descriptor_set['bindings']:
				assert binding['set'] == UNIFORM_SET
				block = binding['block']
				members = []
				for member in block['members']:
				members.append({
				'name': member['name'],
				'offset': member['absolute_offset'],
				'size': member['size'],
				})
				bindings.append({
				'binding': binding['binding'],
				'type': 'uniform',
				'size': block['size'],
				'members': members
				})
				binding = descriptor_set['bindings'][0]
				assert binding['descriptor_type'] == VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER
				elif descriptor_set['set'] == STORAGE_SET:
				assert descriptor_set['binding_count'] > 0
				for binding in descriptor_set['bindings']:
				is_storage_image = binding['descriptor_type'] == VK_DESCRIPTOR_TYPE_STORAGE_IMAGE
				is_storage_buffer = binding['descriptor_type'] == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER
				assert is_storage_image or is_storage_buffer
				assert binding['descriptor_type'] == descriptor_set['bindings'][0]['descriptor_type']
				assert binding['image']['arrayed'] == 0
				assert binding['image']['ms'] == 0
				bindingType = 'storageImage'
				if is_storage_buffer:
				bindingType = 'storageBuffer'
				bindings.append({
				'binding': binding['binding'],
				'type': bindingType,
				'name': binding['name'],
				})
				else:
				if use_descriptor_indexing:
				if descriptor_set['set'] == 0:
				assert descriptor_set['binding_count'] >= 1
				for binding in descriptor_set['bindings']:
				assert binding['descriptor_type'] == VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
				assert binding['array']['dims'][0] == 16384
				if binding['binding'] == 0:
				assert binding['name'] == 'textures2D'
				elif binding['binding'] == 1:
				assert binding['name'] == 'texturesCube'
				elif binding['binding'] == 2:
				assert binding['name'] == 'texturesShadow'
				else:
				assert False
				else:
				assert descriptor_set['binding_count'] > 0
				for binding in descriptor_set['bindings']:
				assert binding['descriptor_type'] == VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
				assert binding['image']['sampled'] == 1
				assert binding['image']['arrayed'] == 0
				assert binding['image']['ms'] == 0
				sampler_type = 'sampler{}D'.format(binding['image']['dim'] + 1)
				if binding['image']['dim'] == 3:
				sampler_type = 'samplerCube'
				bindings.append({
				'binding': binding['binding'],
				'type': sampler_type,
				'name': binding['name'],
				})
				descriptor_sets.append({
				'set': descriptor_set['set'],
				'bindings': bindings,
				})
				if stage == 'vertex':
				for variable in interface_variables:
				if variable['storage_class'] == 1:
				# Input.
				vertex_attributes.append({
				'name': variable['name'],
				'location': variable['location'],
				})
				# Compile the final version without debug information.
				if not keep_debug:
				ret, out, err = execute(command)
				if ret:
				sys.stderr.write('Command returned {}:\nCommand: {}\nInput path: {}\nOutput path: {}\nError: {}\n'.format(
				ret, ' '.join(command), input_path, output_path, err))
				raise ValueError(err)
				return {
				'push_constants': push_constants,
				'vertex_attributes': vertex_attributes,
				'descriptor_sets': descriptor_sets,
				}


				def output_xml_tree(tree, path):
				''' We use a simple custom printer to have the same output for all platforms.'''
				with open(path, 'wt') as handle:
				handle.write('<?xml version="1.0" encoding="utf-8"?>\n')
				handle.write('<!-- DO NOT EDIT: GENERATED BY SCRIPT {} -->\n'.format(os.path.basename(__file__)))
				def output_xml_node(node, handle, depth):
				indent = '\t' * depth
				attributes = ''
				for attribute_name in sorted(node.attrib.keys()):
				attributes += ' {}="{}"'.format(attribute_name, node.attrib[attribute_name])
				if node:
				handle.write('{}<{}{}>\n'.format(indent, node.tag, attributes))
				for child in node:
				output_xml_node(child, handle, depth + 1)
				handle.write('{}</{}>\n'.format(indent, node.tag))
				else:
				handle.write('{}<{}{}/>\n'.format(indent, node.tag, attributes))
				output_xml_node(tree.getroot(), handle, 0)


				def build(rules, input_root_path, output_root_path, mod_name, program_name):
				sys.stdout.write('Program "{}" ({})\n'.format(program_name, mod_name))
				if rules and program_name not in rules:
				sys.stdout.write(' Skip.\n')
				return
				sys.stdout.write(' Building.\n')

				rebuild = False

				defines = []
				program_defines = []
				shaders = []

				tree = ET.parse(os.path.join(input_root_path, mod_name, 'shaders', 'glsl', program_name + '.xml'))
				root = tree.getroot()
				for element in root:
				element_tag = element.tag
				if element_tag == 'defines':
				for child in element:
				values = []
				for value in child:
				values.append({
				'name': child.attrib['name'],
				'value': value.text,
				})
				defines.append(values)
				elif element_tag == 'define':
				program_defines.append({'name': element.attrib['name'], 'value': element.attrib['value']})
				elif element_tag == 'vertex':
				streams = []
				for shader_child in element:
				assert shader_child.tag == 'stream'
				streams.append({
				'name': shader_child.attrib['name'],
				'attribute': shader_child.attrib['attribute'],
				})
				if 'if' in shader_child.attrib:
				streams[-1]['if'] = shader_child.attrib['if']
				shaders.append({
				'type': 'vertex',
				'file': element.attrib['file'],
				'streams': streams,
				})
				elif element_tag == 'fragment':
				shaders.append({
				'type': 'fragment',
				'file': element.attrib['file'],
				})
				elif element_tag == 'compute':
				shaders.append({
				'type': 'compute',
				'file': element.attrib['file'],
				})
				else:
				raise ValueError('Unsupported element tag: "{}"'.format(element_tag))

				stage_extension = {
				'vertex': '.vs',
				'fragment': '.fs',
				'geometry': '.gs',
				'compute': '.cs',
				}

				output_spirv_root_path = os.path.join(output_root_path, 'shaders', 'spirv')
				if not os.path.isdir(output_spirv_root_path):
				os.mkdir(output_spirv_root_path)

				root = ET.Element('programs')

				if 'combinations' in rules[program_name]:
				combinations = rules[program_name]['combinations']
				else:
				combinations = list(itertools.product(*defines))

				hashed_cache = {}

				for index, combination in enumerate(combinations):
				assert index < 10000
				program_path = 'spirv/' + program_name + ('_%04d' % index) + '.xml'

				programs_element = ET.SubElement(root, 'program')
				programs_element.set('type', 'spirv')
				programs_element.set('file', program_path)

				defines_element = ET.SubElement(programs_element, 'defines')
				for define in combination:
				if define['value'] == 'UNDEFINED':
				continue
				define_element = ET.SubElement(defines_element, 'define')
				define_element.set('name', define['name'])
				define_element.set('value', define['value'])

				if not rebuild and os.path.isfile(os.path.join(output_root_path, 'shaders', program_path)):
				continue

				program_root = ET.Element('program')
				program_root.set('type', 'spirv')
				for shader in shaders:
				extension = stage_extension[shader['type']]
				file_name = program_name + ('_%04d' % index) + extension + '.spv'
				output_spirv_path = os.path.join(output_spirv_root_path, file_name)
				reflection = compile_and_reflect(
				input_root_path,
				output_root_path,
				mod_name,
				shader['type'],
				os.path.join(input_root_path, mod_name, 'shaders', shader['file']),
				output_spirv_path,
				combination + program_defines)

				spirv_hash = calculate_hash(output_spirv_path)
				if spirv_hash not in hashed_cache:
				hashed_cache[spirv_hash] = [file_name]
				else:
				found_candidate = False
				for candidate_name in hashed_cache[spirv_hash]:
				candidate_path = os.path.join(output_spirv_root_path, candidate_name)
				if compare_spirv(output_spirv_path, candidate_path):
				found_candidate = True
				file_name = candidate_name
				break
				if found_candidate:
				os.remove(output_spirv_path)
				else:
				hashed_cache[spirv_hash].append(file_name)

				shader_element = ET.SubElement(program_root, shader['type'])
				shader_element.set('file', 'spirv/' + file_name)
				if shader['type'] == 'vertex':
				for stream in shader['streams']:
				if 'if' in stream and not resolve_if(combination, stream['if']):
				continue

				found_vertex_attribute = False
				for vertex_attribute in reflection['vertex_attributes']:
				if vertex_attribute['name'] == stream['attribute']:
				found_vertex_attribute = True
				break
				if not found_vertex_attribute and stream['attribute'] == 'a_tangent':
				continue
				if not found_vertex_attribute:
				sys.stderr.write('Vertex attribute not found: {}\n'.format(stream['attribute']))
				assert found_vertex_attribute

				stream_element = ET.SubElement(shader_element, 'stream')
				stream_element.set('name', stream['name'])
				stream_element.set('attribute', stream['attribute'])
				for vertex_attribute in reflection['vertex_attributes']:
				if vertex_attribute['name'] == stream['attribute']:
				stream_element.set('location', vertex_attribute['location'])
				break

				for push_constant in reflection['push_constants']:
				push_constant_element = ET.SubElement(shader_element, 'push_constant')
				push_constant_element.set('name', push_constant['name'])
				push_constant_element.set('size', push_constant['size'])
				push_constant_element.set('offset', push_constant['offset'])
				descriptor_sets_element = ET.SubElement(shader_element, 'descriptor_sets')
				for descriptor_set in reflection['descriptor_sets']:
				descriptor_set_element = ET.SubElement(descriptor_sets_element, 'descriptor_set')
				descriptor_set_element.set('set', descriptor_set['set'])
				for binding in descriptor_set['bindings']:
				binding_element = ET.SubElement(descriptor_set_element, 'binding')
				binding_element.set('type', binding['type'])
				binding_element.set('binding', binding['binding'])
				if binding['type'] == 'uniform':
				binding_element.set('size', binding['size'])
				for member in binding['members']:
				member_element = ET.SubElement(binding_element, 'member')
				member_element.set('name', member['name'])
				member_element.set('size', member['size'])
				member_element.set('offset', member['offset'])
				elif binding['type'].startswith('sampler'):
				binding_element.set('name', binding['name'])
				elif binding['type'].startswith('storage'):
				binding_element.set('name', binding['name'])
				program_tree = ET.ElementTree(program_root)
				output_xml_tree(program_tree, os.path.join(output_root_path, 'shaders', program_path))


				tree = ET.ElementTree(root)
				output_xml_tree(tree, os.path.join(output_root_path, 'shaders', 'spirv', program_name + '.xml'))


				if __name__ == '__main__':
				assert len(sys.argv) >= 4
				# A path to a folder with mods like binaries/data/mods.
				input_root_path = sys.argv[1]
				# A path to a mod folder like binaries/data/mods/spirv.
				output_root_path = sys.argv[2]
				# A path to JSON with rules
				rules_path = sys.argv[3]
				rules = None
				if os.path.isfile(rules_path):
				with open(rules_path, 'rt') as handle:
				rules = json.load(handle)

				if len(sys.argv) == 4:
				for mod_name in ['mod', 'public']:
				sys.stdout.write('Building SPIRV for "{}"\n'.format(mod_name))
				mod_shaders_path = os.path.join(input_root_path, mod_name, 'shaders', 'glsl')
				if not os.path.isdir(mod_shaders_path):
				continue
				for file_name in os.listdir(mod_shaders_path):
				name, ext = os.path.splitext(file_name)
				if ext.lower() == '.xml':
				build(rules, input_root_path, output_root_path, mod_name, name)
				elif len(sys.argv) == 5:
				build(rules, input_root_path, output_root_path, sys.argv[1], sys.argv[2], sys.argv[4])