1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
|
import tempfile
import unittest
from tools.gen_vulkan_spv import DEFAULT_ENV, SPVGenerator
####################
# Data for testing #
####################
test_shader = """
#version 450 core
#define FORMAT ${FORMAT}
#define PRECISION ${PRECISION}
#define OP(X) ${OPERATOR}
$def is_int(dtype):
$ return dtype in {"int", "int32", "int8"}
$def is_uint(dtype):
$ return dtype in {"uint", "uint32", "uint8"}
$if is_int(DTYPE):
#define VEC4_T ivec4
$elif is_uint(DTYPE):
#define VEC4_T uvec4
$else:
#define VEC4_T vec4
$if not INPLACE:
$if is_int(DTYPE):
layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly iimage3D uOutput;
layout(set = 0, binding = 1) uniform PRECISION isampler3D uInput;
$elif is_uint(DTYPE):
layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly uimage3D uOutput;
layout(set = 0, binding = 1) uniform PRECISION usampler3D uInput;
$else:
layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
$else:
$if is_int(DTYPE):
layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict iimage3D uOutput;
$elif is_uint(DTYPE):
layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict uimage3D uOutput;
$else:
layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict image3D uOutput;
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
void main() {
const ivec3 pos = ivec3(gl_GlobalInvocationID);
$if not INPLACE:
VEC4_T v = texelFetch(uInput, pos, 0);
$else:
VEC4_T v = imageLoad(uOutput, pos);
$for i in range(ITER[0]):
for (int i = 0; i < ${ITER[1]}; ++i) {
v = OP(v + i);
}
imageStore(uOutput, pos, OP(v));
}
"""
test_params_yaml = """
test_shader:
parameter_names_with_default_values:
DTYPE: float
INPLACE: false
OPERATOR: X + 3
ITER: !!python/tuple [3, 5]
generate_variant_forall:
INPLACE:
- VALUE: false
SUFFIX: ""
- VALUE: true
SUFFIX: inplace
DTYPE:
- VALUE: int8
- VALUE: float
shader_variants:
- NAME: test_shader_1
- NAME: test_shader_3
OPERATOR: X - 1
ITER: !!python/tuple [3, 2]
generate_variant_forall:
DTYPE:
- VALUE: float
- VALUE: int
"""
##############
# Unit Tests #
##############
class TestVulkanSPVCodegen(unittest.TestCase):
def setUp(self) -> None:
self.tmpdir = tempfile.TemporaryDirectory()
with open(f"{self.tmpdir.name}/test_shader.glsl,", "w") as f:
f.write(test_shader)
with open(f"{self.tmpdir.name}/test_params.yaml", "w") as f:
f.write(test_params_yaml)
self.tmpoutdir = tempfile.TemporaryDirectory()
self.generator = SPVGenerator(
src_dir_paths=self.tmpdir.name, env=DEFAULT_ENV, glslc_path=None
)
def cleanUp(self) -> None:
self.tmpdir.cleanup()
self.tmpoutdir.cleanup()
def testOutputMap(self) -> None:
# Each shader variant will produce variants generated based on all possible combinations
# of the DTYPE and INPLACE parameters. test_shader_3 has fewer generated variants due to
# a custom specified generate_variant_forall field.
expected_output_shaders = {
"test_shader_1_float",
"test_shader_1_inplace_float",
"test_shader_1_inplace_int8",
"test_shader_1_int8",
"test_shader_3_float",
"test_shader_3_int",
}
actual_output_shaders = set(self.generator.output_shader_map.keys())
self.assertEqual(expected_output_shaders, actual_output_shaders)
|
