From: Mikko Rasa <tdb@tdb.fi>
Date: Fri, 19 Nov 2021 19:00:52 +0000 (+0200)
Subject: Create specialized copies of SPIR-V modules
X-Git-Url: http://git.tdb.fi/?p=libs%2Fgl.git;a=commitdiff_plain;h=d8d4c2aeaf109ab7ba55888bef5364fe11b5b02d

Create specialized copies of SPIR-V modules

It turns out Vulkan does want all declared resources in the descriptor
sets, even if they're unused.  Manual specialization avoids the need to
deal with them when binding stuff.
---

diff --git a/source/core/module.cpp b/source/core/module.cpp
index 0c2df8e1..636fd13c 100644
--- a/source/core/module.cpp
+++ b/source/core/module.cpp
@@ -37,9 +37,15 @@ enum SpirVConstants
 	OP_ACCESS_CHAIN = 65,
 	OP_DECORATE = 71,
 	OP_MEMBER_DECORATE = 72,
+	OP_SELECTION_MERGE = 247,
 	OP_LABEL = 248,
 	OP_BRANCH = 249,
 	OP_BRANCH_CONDITIONAL = 250,
+	OP_SWITCH = 251,
+	OP_KILL = 252,
+	OP_RETURN = 253,
+	OP_RETURN_VALUE = 254,
+	OP_UNREACHABLE = 255,
 
 	DECO_SPEC_ID = 1,
 	DECO_ARRAY_STRIDE = 6,
@@ -246,6 +252,191 @@ void SpirVModule::reflect()
 	}
 }
 
+SpirVModule *SpirVModule::specialize(const map<string, int> &spec_values) const
+{
+	vector<uint8_t> flags(code[3], 1);
+
+	std::map<unsigned, int> spec_values_by_id;
+	for(const Constant &c: spec_constants)
+	{
+		auto i = spec_values.find(c.name);
+		if(i!=spec_values.end())
+		{
+			flags[c.id] = (i->second ? 5 : 3);
+			spec_values_by_id[c.constant_id] = i->second;
+		}
+	}
+
+	for(const Variable &v: variables)
+		flags[v.id] = 0;
+	for(const InstructionBlock &b: blocks)
+		flags[b.id] = 0;
+	for(const InstructionBlock *b: collect_visited_blocks(spec_values_by_id))
+	{
+		flags[b->id] = 1;
+		for(const Variable *v: b->accessed_variables)
+			flags[v->id] = 1;
+	}
+
+	std::vector<uint32_t> new_code;
+	new_code.reserve(code.size());
+
+	auto op = code.begin()+5;
+	new_code.insert(new_code.begin(), code.begin(), op);
+
+	bool skip_block = false;
+	while(op!=code.end())
+	{
+		unsigned word_count = *op>>16;
+		unsigned opcode = *op&0xFFFF;
+
+		bool copy = !skip_block;
+		if(skip_block)
+		{
+			skip_block = (opcode!=OP_BRANCH && opcode!=OP_BRANCH_CONDITIONAL && opcode!=OP_SWITCH &&
+				opcode!=OP_KILL && opcode!=OP_RETURN && opcode!=OP_RETURN_VALUE && opcode!=OP_UNREACHABLE);
+		}
+		else
+		{
+			if(opcode==OP_NAME)
+				copy = flags[*(op+1)];
+			else if(opcode==OP_ENTRY_POINT)
+			{
+				unsigned start = new_code.size();
+				new_code.push_back(opcode);
+				new_code.push_back(*(op+1));
+				new_code.push_back(*(op+2));
+
+				unsigned i=3;
+				while(i<word_count)
+				{
+					unsigned word = *(op+i++);
+					new_code.push_back(word);
+					if(!(word&(word>>8)&(word>>16)&(word>>24)&0xFF))
+						break;
+				}
+
+				for(; i<word_count; ++i)
+				{
+					unsigned id = *(op+i);
+					if(flags[id])
+						new_code.push_back(id);
+				}
+
+				new_code[start] |= (new_code.size()-start)<<16;
+
+				copy = false;
+			}
+			else if(opcode==OP_SPEC_CONSTANT_TRUE || opcode==OP_SPEC_CONSTANT_FALSE)
+			{
+				unsigned id = *(op+2);
+				if(flags[id]&2)
+				{
+					new_code.push_back(0x30000 | (flags[id]&4 ? OP_CONSTANT_TRUE : OP_CONSTANT_FALSE));
+					new_code.push_back(*(op+1));
+					new_code.push_back(id);
+
+					copy = false;
+				}
+			}
+			else if(opcode==OP_VARIABLE)
+				copy = flags[*(op+2)];
+			else if(opcode==OP_DECORATE)
+			{
+				unsigned id = *(op+1);
+				copy = flags[id];
+				if(copy && *(op+2)==DECO_SPEC_ID)
+					copy = !(flags[id]&2);
+			}
+			else if(opcode==OP_LABEL)
+			{
+				copy = flags[*(op+1)];
+				skip_block = !copy;
+			}
+			else if(opcode==OP_SELECTION_MERGE)
+			{
+				unsigned next_opcode = *(op+word_count)&0xFFFF;
+				if(next_opcode==OP_BRANCH_CONDITIONAL)
+				{
+					unsigned true_id = *(op+word_count+2);
+					unsigned false_id = *(op+word_count+3);
+					if(!flags[true_id] || !flags[false_id])
+					{
+						new_code.push_back(0x20000 | OP_BRANCH);
+						new_code.push_back(flags[true_id] ? true_id : false_id);
+						copy = false;
+
+						/* Skip the branch instruction when it's encountered on the
+						next iteration */
+						skip_block = true;
+					}
+				}
+			}
+		}
+
+		if(copy)
+		{
+			for(unsigned i=0; i<word_count; ++i)
+				new_code.push_back(*(op+i));
+		}
+
+		op += word_count;
+	}
+
+	SpirVModule *spec_mod = new SpirVModule;
+	spec_mod->code = move(new_code);
+	spec_mod->reflect();
+	spec_mod->create();
+
+	return spec_mod;
+}
+
+vector<const SpirVModule::InstructionBlock *> SpirVModule::collect_visited_blocks(const map<unsigned, int> &spec_values) const
+{
+	vector<uint8_t> visited(blocks.size(), 4);
+	for(unsigned i=0; i<blocks.size(); ++i)
+	{
+		const InstructionBlock &b = blocks[i];
+
+		bool cond = true;
+		if(b.condition)
+		{
+			cond = b.condition->i_value;
+			auto j = spec_values.find(b.condition->constant_id);
+			if(j!=spec_values.end())
+				cond = j->second;
+			if(b.negate_condition)
+				cond = !cond;
+		}
+
+		visited[i] |= cond*2;
+		for(const InstructionBlock *s: b.successors)
+			visited[s-blocks.data()] &= 3;
+	}
+
+	for(unsigned i=0; i<blocks.size(); ++i)
+		if(visited[i]&4)
+			collect_visited_blocks(i, visited);
+
+	vector<const SpirVModule::InstructionBlock *> result;
+	for(unsigned i=0; i<blocks.size(); ++i)
+		if(visited[i]&1)
+			result.push_back(&blocks[i]);
+
+	return result;
+}
+
+void SpirVModule::collect_visited_blocks(unsigned i, vector<uint8_t> &visited) const
+{
+	visited[i] |= 1;
+	for(const InstructionBlock *s: blocks[i].successors)
+	{
+		unsigned j = s-blocks.data();
+		if((visited[j]&3)==2)
+			collect_visited_blocks(j, visited);
+	}
+}
+
 
 bool SpirVModule::Variable::operator==(const Variable &other) const
 {
@@ -358,7 +549,9 @@ void SpirVModule::Reflection::reflect_member_name(CodeIterator op)
 void SpirVModule::Reflection::reflect_entry_point(CodeIterator op)
 {
 	CodeIterator op_end = get_op_end(op);
-	EntryPoint &entry = entry_points[*(op+2)];
+	unsigned id = *(op+2);
+	EntryPoint &entry = entry_points[id];
+	entry.id = id;
 	entry.stage = static_cast<Stage>(*(op+1));  // Execution model in SPIR-V spec
 	op += 3;
 	entry.name = read_string(op, op_end);
@@ -444,6 +637,7 @@ void SpirVModule::Reflection::reflect_struct_type(CodeIterator op)
 	unsigned id = *(op+1);
 	Structure &strct = structs[id];
 	strct.name = names[id];
+	strct.id = id;
 	types[id].struct_type = &strct;
 	op += 2;
 
@@ -472,6 +666,7 @@ void SpirVModule::Reflection::reflect_constant(CodeIterator op)
 	unsigned id = *(op+2);
 	Constant &cnst = constants[id];
 	cnst.name = names[id];
+	cnst.id = id;
 	cnst.type = types[*(op+1)].type;
 	if(opcode==OP_CONSTANT_TRUE || opcode==OP_SPEC_CONSTANT_TRUE)
 		cnst.i_value = true;
@@ -488,6 +683,7 @@ void SpirVModule::Reflection::reflect_variable(CodeIterator op)
 	unsigned id = *(op+2);
 	Variable &var = variables[id];
 	var.name = names[id];
+	var.id = id;
 	const TypeInfo &type = types[*(op+1)];
 	var.storage = type.storage;
 	var.type = type.type;
@@ -570,7 +766,9 @@ void SpirVModule::Reflection::reflect_member_decorate(CodeIterator op)
 
 void SpirVModule::Reflection::reflect_label(CodeIterator op)
 {
-	current_block = &blocks[*(op+1)];
+	unsigned id = *(op+1);
+	current_block = &blocks[id];
+	current_block->id = id;
 }
 
 void SpirVModule::Reflection::reflect_branch(CodeIterator op)
diff --git a/source/core/module.h b/source/core/module.h
index 855d5195..2a4cac6b 100644
--- a/source/core/module.h
+++ b/source/core/module.h
@@ -129,6 +129,7 @@ public:
 	struct EntryPoint
 	{
 		std::string name;
+		unsigned id = 0;
 		Stage stage = VERTEX;
 		std::vector<const Variable *> globals;
 	};
@@ -148,16 +149,18 @@ public:
 	struct Structure
 	{
 		std::string name;
-		std::vector<StructMember> members;
+		unsigned id = 0;
 		unsigned size = 0;
+		std::vector<StructMember> members;
 	};
 
 	struct Variable
 	{
 		std::string name;
-		StorageClass storage = static_cast<StorageClass>(-1);
+		unsigned id = 0;
 		DataType type = VOID;
 		const Structure *struct_type = 0;
+		StorageClass storage = static_cast<StorageClass>(-1);
 		unsigned array_size = 0;
 		int location = -1;
 		unsigned descriptor_set = 0;
@@ -170,6 +173,7 @@ public:
 	struct Constant
 	{
 		std::string name;
+		unsigned id = 0;
 		int constant_id = -1;
 		DataType type = VOID;
 		union
@@ -181,8 +185,9 @@ public:
 
 	struct InstructionBlock
 	{
-		const Constant *condition = 0;
+		unsigned id = 0;
 		bool negate_condition = false;
+		const Constant *condition = 0;
 		std::vector<const Variable *> accessed_variables;
 		std::vector<const InstructionBlock *> successors;
 	};
@@ -264,6 +269,13 @@ public:
 	const std::vector<Variable> &get_variables() const { return variables; }
 	const std::vector<Constant> &get_spec_constants() const { return spec_constants; }
 	const std::vector<InstructionBlock> &get_blocks() const { return blocks; }
+
+	/** Creates a new module which is a specialized version of this one. */
+	SpirVModule *specialize(const std::map<std::string, int> &) const;
+
+private:
+	std::vector<const InstructionBlock *> collect_visited_blocks(const std::map<unsigned, int> &) const;
+	void collect_visited_blocks(unsigned, std::vector<std::uint8_t> &) const;
 };
 
 } // namespace GL
diff --git a/source/core/program.cpp b/source/core/program.cpp
index 16b2fa79..471803c5 100644
--- a/source/core/program.cpp
+++ b/source/core/program.cpp
@@ -12,12 +12,26 @@ Program::Program(const Module &mod, const map<string, int> &spec_values)
 	add_stages(mod, spec_values);
 }
 
+Program::Program(Program &&other):
+	ProgramBackend(move(other)),
+	reflect_data(move(other.reflect_data)),
+	specialized_spirv(other.specialized_spirv)
+{
+	other.specialized_spirv = 0;
+}
+
+Program::~Program()
+{
+	delete specialized_spirv;
+}
+
 void Program::add_stages(const Module &mod, const map<string, int> &spec_values)
 {
 	if(has_stages())
 		throw invalid_operation("Program::add_stages");
 
 	reflect_data = ReflectData();
+	const Module *final_module = &mod;
 
 	TransientData transient;
 	switch(mod.get_format())
@@ -26,20 +40,24 @@ void Program::add_stages(const Module &mod, const map<string, int> &spec_values)
 		add_glsl_stages(static_cast<const GlslModule &>(mod), spec_values, transient);
 		break;
 	case Module::SPIR_V:
-		add_spirv_stages(static_cast<const SpirVModule &>(mod), spec_values);
+		if(!spec_values.empty())
+		{
+			specialized_spirv = static_cast<const SpirVModule &>(mod).specialize(spec_values);
+			final_module = specialized_spirv;
+		}
+		add_spirv_stages(*static_cast<const SpirVModule *>(final_module), spec_values);
 		break;
 	default:
 		throw invalid_argument("Program::add_stages");
 	}
 
-	finalize(mod, transient);
+	finalize(*final_module, transient);
 
-	if(mod.get_format()==Module::SPIR_V)
+	if(final_module->get_format()==Module::SPIR_V)
 	{
-		const SpirVModule &spirv_mod = static_cast<const SpirVModule &>(mod);
-		vector<uint8_t> used_variables = collect_used_variables(spirv_mod, spec_values);
-		collect_uniforms(spirv_mod, used_variables);
-		collect_attributes(spirv_mod, used_variables);
+		const SpirVModule &spirv_mod = *static_cast<const SpirVModule *>(final_module);
+		collect_uniforms(spirv_mod);
+		collect_attributes(spirv_mod);
 		collect_builtins(spirv_mod);
 	}
 
@@ -51,79 +69,15 @@ void Program::add_stages(const Module &mod, const map<string, int> &spec_values)
 		require_type(a.type);
 }
 
-vector<uint8_t> Program::collect_used_variables(const SpirVModule &mod, const map<string, int> &spec_values)
-{
-	std::map<unsigned, int> spec_values_by_id;
-	for(const SpirVModule::Constant &c: mod.get_spec_constants())
-	{
-		auto i = spec_values.find(c.name);
-		if(i!=spec_values.end())
-			spec_values_by_id[c.constant_id] = i->second;
-	}
-
-	const vector<SpirVModule::InstructionBlock> &blocks = mod.get_blocks();
-	vector<uint8_t> visited(blocks.size(), 4);
-	for(unsigned i=0; i<blocks.size(); ++i)
-	{
-		const SpirVModule::InstructionBlock &b = blocks[i];
-
-		bool cond = true;
-		if(b.condition)
-		{
-			cond = b.condition->i_value;
-			auto j = spec_values_by_id.find(b.condition->constant_id);
-			if(j!=spec_values_by_id.end())
-				cond = j->second;
-			if(b.negate_condition)
-				cond = !cond;
-		}
-
-		visited[i] |= cond*2;
-		for(const SpirVModule::InstructionBlock *s: b.successors)
-			visited[s-blocks.data()] &= 3;
-	}
-
-	for(unsigned i=0; i<blocks.size(); ++i)
-		if(visited[i]&4)
-			collect_visited_blocks(blocks, i, visited);
-
-	const vector<SpirVModule::Variable> &variables = mod.get_variables();
-	vector<uint8_t> used(variables.size());
-	for(unsigned i=0; i<blocks.size(); ++i)
-		if(visited[i]&1)
-		{
-			for(const SpirVModule::Variable *v: blocks[i].accessed_variables)
-				used[v-variables.data()] = 1;
-		}
-
-	return used;
-}
-
-void Program::collect_visited_blocks(const vector<SpirVModule::InstructionBlock> &blocks, unsigned i, vector<uint8_t> &visited)
-{
-	visited[i] |= 1;
-	for(const SpirVModule::InstructionBlock *s: blocks[i].successors)
-	{
-		unsigned j = s-blocks.data();
-		if((visited[j]&3)==2)
-			collect_visited_blocks(blocks, j, visited);
-	}
-}
-
-void Program::collect_uniforms(const SpirVModule &mod, const vector<uint8_t> &used_variables)
+void Program::collect_uniforms(const SpirVModule &mod)
 {
 	// Prepare the default block
 	reflect_data.uniform_blocks.emplace_back();
 	vector<vector<string> > block_uniform_names(1);
 
-	const vector<SpirVModule::Variable> &variables = mod.get_variables();
 	unsigned n_descriptor_sets = 0;
-	for(unsigned i=0; i<variables.size(); ++i)
+	for(const SpirVModule::Variable &v: mod.get_variables())
 	{
-		if(!used_variables[i])
-			continue;
-
-		const SpirVModule::Variable &v = variables[i];
 		if((v.storage==SpirVModule::UNIFORM || v.storage==SpirVModule::PUSH_CONSTANT) && v.struct_type)
 		{
 			reflect_data.uniform_blocks.emplace_back();
@@ -221,14 +175,13 @@ void Program::collect_block_uniforms(const SpirVModule::Structure &strct, const
 	}
 }
 
-void Program::collect_attributes(const SpirVModule &mod, const vector<uint8_t> &used_variables)
+void Program::collect_attributes(const SpirVModule &mod)
 {
-	const vector<SpirVModule::Variable> &variables = mod.get_variables();
 	for(const SpirVModule::EntryPoint &e: mod.get_entry_points())
 		if(e.stage==SpirVModule::VERTEX && e.name=="main")
 		{
 			for(const SpirVModule::Variable *v: e.globals)
-				if(v->storage==SpirVModule::INPUT && used_variables[v-variables.data()])
+				if(v->storage==SpirVModule::INPUT)
 				{
 					reflect_data.attributes.emplace_back();
 					ReflectData::AttributeInfo &info = reflect_data.attributes.back();
diff --git a/source/core/program.h b/source/core/program.h
index 7f6ad54e..b6b639bf 100644
--- a/source/core/program.h
+++ b/source/core/program.h
@@ -51,6 +51,7 @@ private:
 	};
 
 	ReflectData reflect_data;
+	SpirVModule *specialized_spirv = 0;
 
 public:
 	/// Constructs an empty Program with no shader stages attached.
@@ -59,13 +60,14 @@ public:
 	/// Constructs a Program from a Module, with specialization constants.
 	Program(const Module &, const std::map<std::string, int> & = std::map<std::string, int>());
 
+	Program(Program &&);
+	~Program();
+
 	void add_stages(const Module &, const std::map<std::string, int> & = std::map<std::string, int>());
 private:
-	static std::vector<std::uint8_t> collect_used_variables(const SpirVModule &, const std::map<std::string, int> &);
-	static void collect_visited_blocks(const std::vector<SpirVModule::InstructionBlock> &, unsigned, std::vector<std::uint8_t> &);
-	void collect_uniforms(const SpirVModule &, const std::vector<std::uint8_t> &);
+	void collect_uniforms(const SpirVModule &);
 	void collect_block_uniforms(const SpirVModule::Structure &, const std::string &, unsigned, std::vector<std::string> &);
-	void collect_attributes(const SpirVModule &, const std::vector<std::uint8_t> &);
+	void collect_attributes(const SpirVModule &);
 	void collect_builtins(const SpirVModule &);
 	void collect_builtins(const SpirVModule::Structure &);