VK_KHR_shader_expect_assume
This document proposes adding support for expect/assume SPIR-V instructions to guide shader program optimizations.
1. Problem Statement
Shader writers or generators as well as other SPIR-V producers (e.g. Machine Learning compilers) often have access to information that could enable the SPIR-V consumers in Vulkan implementations to make better optimization decisions, such as knowledge of the likely value of objects or whether a given condition holds, but which they cannot communicate to a Vulkan SPIR-V consumer using existing features.
2. Solution Space
SPIR-V already provides some mechanisms for producers to give hints to consumers in a limited number of scenarios:
-
OpBranchConditionalcan accept branch weights that enable producers to indicate the likelihood of each path. This does not however generalize toOpSwitchconstructs. -
Various so called Loop Controls make it possible for producers to provide metadata about the iteration count of loops or desired unrolling behavior.
There is however no exposed generic mechanism for SPIR-V producers to communicate optimization information to consumers. SPIR-V does support dedicated instructions, introduced by the SPV_KHR_expect_assume extension, that make it possible for producers to communicate to consumers the likely value of an object or whether a given condition holds, but this extension is currently not exposed in Vulkan.
3. Proposal
Expose the SPV_KHR_expect_assume extension in Vulkan.
The SPV_KHR_expect_assume extension introduces two new instructions:
-
OpExpectKHRmakes it possible to state the most probable value of its input. -
OpAssumeTrueKHRenables the optimizer to assume that the provided condition is always true.
4. Examples
As an illustration, consider the following pseudocode example:
c = 20
d = 2
b = c / d
if (a - b > 0) {
...
} else {
...
}The writer or producer may know that a > 10. This knowledge makes it possible
to completely remove the else branch. In this case, the producer could perform
that optimization alone. However, if the producer only knows that a is greater
than some value provided, say with a specialization constant, it can no longer
perform the optimization. Adding that information to the SPIR-V module would
enable the SPIR-V consumer to do it.
Another possible use could be to provide guarantees that a particular value is not NaN or infinite:
value = load(...)
assume(!isnan(value))