llvm-project

1
#!/usr/bin/env python
2

3
"""A shuffle vector fuzz tester.
4

5
This is a python program to fuzz test the LLVM shufflevector instruction. It
6
generates a function with a random sequnece of shufflevectors, maintaining the
7
element mapping accumulated across the function. It then generates a main
8
function which calls it with a different value in each element and checks that
9
the result matches the expected mapping.
10

11
Take the output IR printed to stdout, compile it to an executable using whatever
12
set of transforms you want to test, and run the program. If it crashes, it found
13
a bug.
14
"""
15

16
from __future__ import print_function
17

18
import argparse
19
import itertools
20
import random
21
import sys
22
import uuid
23

24

25
def main():
26
    element_types = ["i8", "i16", "i32", "i64", "f32", "f64"]
27
    parser = argparse.ArgumentParser(description=__doc__)
28
    parser.add_argument(
29
        "-v", "--verbose", action="store_true", help="Show verbose output"
30
    )
31
    parser.add_argument(
32
        "--seed", default=str(uuid.uuid4()), help="A string used to seed the RNG"
33
    )
34
    parser.add_argument(
35
        "--max-shuffle-height",
36
        type=int,
37
        default=16,
38
        help="Specify a fixed height of shuffle tree to test",
39
    )
40
    parser.add_argument(
41
        "--no-blends",
42
        dest="blends",
43
        action="store_false",
44
        help="Include blends of two input vectors",
45
    )
46
    parser.add_argument(
47
        "--fixed-bit-width",
48
        type=int,
49
        choices=[128, 256],
50
        help="Specify a fixed bit width of vector to test",
51
    )
52
    parser.add_argument(
53
        "--fixed-element-type",
54
        choices=element_types,
55
        help="Specify a fixed element type to test",
56
    )
57
    parser.add_argument("--triple", help="Specify a triple string to include in the IR")
58
    args = parser.parse_args()
59

60
    random.seed(args.seed)
61

62
    if args.fixed_element_type is not None:
63
        element_types = [args.fixed_element_type]
64

65
    if args.fixed_bit_width is not None:
66
        if args.fixed_bit_width == 128:
67
            width_map = {"i64": 2, "i32": 4, "i16": 8, "i8": 16, "f64": 2, "f32": 4}
68
            (width, element_type) = random.choice(
69
                [(width_map[t], t) for t in element_types]
70
            )
71
        elif args.fixed_bit_width == 256:
72
            width_map = {"i64": 4, "i32": 8, "i16": 16, "i8": 32, "f64": 4, "f32": 8}
73
            (width, element_type) = random.choice(
74
                [(width_map[t], t) for t in element_types]
75
            )
76
        else:
77
            sys.exit(1)  # Checked above by argument parsing.
78
    else:
79
        width = random.choice([2, 4, 8, 16, 32, 64])
80
        element_type = random.choice(element_types)
81

82
    element_modulus = {
83
        "i8": 1 << 8,
84
        "i16": 1 << 16,
85
        "i32": 1 << 32,
86
        "i64": 1 << 64,
87
        "f32": 1 << 32,
88
        "f64": 1 << 64,
89
    }[element_type]
90

91
    shuffle_range = (2 * width) if args.blends else width
92

93
    # Because undef (-1) saturates and is indistinguishable when testing the
94
    # correctness of a shuffle, we want to bias our fuzz toward having a decent
95
    # mixture of non-undef lanes in the end. With a deep shuffle tree, the
96
    # probabilies aren't good so we need to bias things. The math here is that if
97
    # we uniformly select between -1 and the other inputs, each element of the
98
    # result will have the following probability of being undef:
99
    #
100
    #   1 - (shuffle_range/(shuffle_range+1))^max_shuffle_height
101
    #
102
    # More generally, for any probability P of selecting a defined element in
103
    # a single shuffle, the end result is:
104
    #
105
    #   1 - P^max_shuffle_height
106
    #
107
    # The power of the shuffle height is the real problem, as we want:
108
    #
109
    #   1 - shuffle_range/(shuffle_range+1)
110
    #
111
    # So we bias the selection of undef at any given node based on the tree
112
    # height. Below, let 'A' be 'len(shuffle_range)', 'C' be 'max_shuffle_height',
113
    # and 'B' be the bias we use to compensate for
114
    # C '((A+1)*A^(1/C))/(A*(A+1)^(1/C))':
115
    #
116
    #   1 - (B * A)/(A + 1)^C = 1 - A/(A + 1)
117
    #
118
    # So at each node we use:
119
    #
120
    #   1 - (B * A)/(A + 1)
121
    # = 1 - ((A + 1) * A * A^(1/C))/(A * (A + 1) * (A + 1)^(1/C))
122
    # = 1 - ((A + 1) * A^((C + 1)/C))/(A * (A + 1)^((C + 1)/C))
123
    #
124
    # This is the formula we use to select undef lanes in the shuffle.
125
    A = float(shuffle_range)
126
    C = float(args.max_shuffle_height)
127
    undef_prob = 1.0 - (
128
        ((A + 1.0) * pow(A, (C + 1.0) / C)) / (A * pow(A + 1.0, (C + 1.0) / C))
129
    )
130

131
    shuffle_tree = [
132
        [
133
            [
134
                -1
135
                if random.random() <= undef_prob
136
                else random.choice(range(shuffle_range))
137
                for _ in itertools.repeat(None, width)
138
            ]
139
            for _ in itertools.repeat(None, args.max_shuffle_height - i)
140
        ]
141
        for i in range(args.max_shuffle_height)
142
    ]
143

144
    if args.verbose:
145
        # Print out the shuffle sequence in a compact form.
146
        print(
147
            (
148
                'Testing shuffle sequence "%s" (v%d%s):'
149
                % (args.seed, width, element_type)
150
            ),
151
            file=sys.stderr,
152
        )
153
        for i, shuffles in enumerate(shuffle_tree):
154
            print("  tree level %d:" % (i,), file=sys.stderr)
155
            for j, s in enumerate(shuffles):
156
                print("    shuffle %d: %s" % (j, s), file=sys.stderr)
157
        print("", file=sys.stderr)
158

159
    # Symbolically evaluate the shuffle tree.
160
    inputs = [
161
        [int(j % element_modulus) for j in range(i * width + 1, (i + 1) * width + 1)]
162
        for i in range(args.max_shuffle_height + 1)
163
    ]
164
    results = inputs
165
    for shuffles in shuffle_tree:
166
        results = [
167
            [
168
                (
169
                    (results[i] if j < width else results[i + 1])[j % width]
170
                    if j != -1
171
                    else -1
172
                )
173
                for j in s
174
            ]
175
            for i, s in enumerate(shuffles)
176
        ]
177
    if len(results) != 1:
178
        print("ERROR: Bad results: %s" % (results,), file=sys.stderr)
179
        sys.exit(1)
180
    result = results[0]
181

182
    if args.verbose:
183
        print("Which transforms:", file=sys.stderr)
184
        print("  from: %s" % (inputs,), file=sys.stderr)
185
        print("  into: %s" % (result,), file=sys.stderr)
186
        print("", file=sys.stderr)
187

188
    # The IR uses silly names for floating point types. We also need a same-size
189
    # integer type.
190
    integral_element_type = element_type
191
    if element_type == "f32":
192
        integral_element_type = "i32"
193
        element_type = "float"
194
    elif element_type == "f64":
195
        integral_element_type = "i64"
196
        element_type = "double"
197

198
    # Now we need to generate IR for the shuffle function.
199
    subst = {"N": width, "T": element_type, "IT": integral_element_type}
200
    print(
201
        """
202
define internal fastcc <%(N)d x %(T)s> @test(%(arguments)s) noinline nounwind {
203
entry:"""
204
        % dict(
205
            subst,
206
            arguments=", ".join(
207
                [
208
                    "<%(N)d x %(T)s> %%s.0.%(i)d" % dict(subst, i=i)
209
                    for i in range(args.max_shuffle_height + 1)
210
                ]
211
            ),
212
        )
213
    )
214

215
    for i, shuffles in enumerate(shuffle_tree):
216
        for j, s in enumerate(shuffles):
217
            print(
218
                """
219
  %%s.%(next_i)d.%(j)d = shufflevector <%(N)d x %(T)s> %%s.%(i)d.%(j)d, <%(N)d x %(T)s> %%s.%(i)d.%(next_j)d, <%(N)d x i32> <%(S)s>
220
""".strip(
221
                    "\n"
222
                )
223
                % dict(
224
                    subst,
225
                    i=i,
226
                    next_i=i + 1,
227
                    j=j,
228
                    next_j=j + 1,
229
                    S=", ".join(
230
                        ["i32 " + (str(si) if si != -1 else "undef") for si in s]
231
                    ),
232
                )
233
            )
234

235
    print(
236
        """
237
  ret <%(N)d x %(T)s> %%s.%(i)d.0
238
}
239
"""
240
        % dict(subst, i=len(shuffle_tree))
241
    )
242

243
    # Generate some string constants that we can use to report errors.
244
    for i, r in enumerate(result):
245
        if r != -1:
246
            s = (
247
                "FAIL(%(seed)s): lane %(lane)d, expected %(result)d, found %%d\n\\0A"
248
                % {"seed": args.seed, "lane": i, "result": r}
249
            )
250
            s += "".join(["\\00" for _ in itertools.repeat(None, 128 - len(s) + 2)])
251
            print(
252
                """
253
@error.%(i)d = private unnamed_addr global [128 x i8] c"%(s)s"
254
""".strip()
255
                % {"i": i, "s": s}
256
            )
257

258
    # Define a wrapper function which is marked 'optnone' to prevent
259
    # interprocedural optimizations from deleting the test.
260
    print(
261
        """
262
define internal fastcc <%(N)d x %(T)s> @test_wrapper(%(arguments)s) optnone noinline {
263
  %%result = call fastcc <%(N)d x %(T)s> @test(%(arguments)s)
264
  ret <%(N)d x %(T)s> %%result
265
}
266
"""
267
        % dict(
268
            subst,
269
            arguments=", ".join(
270
                [
271
                    "<%(N)d x %(T)s> %%s.%(i)d" % dict(subst, i=i)
272
                    for i in range(args.max_shuffle_height + 1)
273
                ]
274
            ),
275
        )
276
    )
277

278
    # Finally, generate a main function which will trap if any lanes are mapped
279
    # incorrectly (in an observable way).
280
    print(
281
        """
282
define i32 @main() {
283
entry:
284
  ; Create a scratch space to print error messages.
285
  %%str = alloca [128 x i8]
286
  %%str.ptr = getelementptr inbounds [128 x i8], [128 x i8]* %%str, i32 0, i32 0
287

288
  ; Build the input vector and call the test function.
289
  %%v = call fastcc <%(N)d x %(T)s> @test_wrapper(%(inputs)s)
290
  ; We need to cast this back to an integer type vector to easily check the
291
  ; result.
292
  %%v.cast = bitcast <%(N)d x %(T)s> %%v to <%(N)d x %(IT)s>
293
  br label %%test.0
294
"""
295
        % dict(
296
            subst,
297
            inputs=", ".join(
298
                [
299
                    (
300
                        "<%(N)d x %(T)s> bitcast "
301
                        "(<%(N)d x %(IT)s> <%(input)s> to <%(N)d x %(T)s>)"
302
                        % dict(
303
                            subst,
304
                            input=", ".join(
305
                                ["%(IT)s %(i)d" % dict(subst, i=i) for i in input]
306
                            ),
307
                        )
308
                    )
309
                    for input in inputs
310
                ]
311
            ),
312
        )
313
    )
314

315
    # Test that each non-undef result lane contains the expected value.
316
    for i, r in enumerate(result):
317
        if r == -1:
318
            print(
319
                """
320
test.%(i)d:
321
  ; Skip this lane, its value is undef.
322
  br label %%test.%(next_i)d
323
"""
324
                % dict(subst, i=i, next_i=i + 1)
325
            )
326
        else:
327
            print(
328
                """
329
test.%(i)d:
330
  %%v.%(i)d = extractelement <%(N)d x %(IT)s> %%v.cast, i32 %(i)d
331
  %%cmp.%(i)d = icmp ne %(IT)s %%v.%(i)d, %(r)d
332
  br i1 %%cmp.%(i)d, label %%die.%(i)d, label %%test.%(next_i)d
333

334
die.%(i)d:
335
  ; Capture the actual value and print an error message.
336
  %%tmp.%(i)d = zext %(IT)s %%v.%(i)d to i2048
337
  %%bad.%(i)d = trunc i2048 %%tmp.%(i)d to i32
338
  call i32 (i8*, i8*, ...) @sprintf(i8* %%str.ptr, i8* getelementptr inbounds ([128 x i8], [128 x i8]* @error.%(i)d, i32 0, i32 0), i32 %%bad.%(i)d)
339
  %%length.%(i)d = call i32 @strlen(i8* %%str.ptr)
340
  call i32 @write(i32 2, i8* %%str.ptr, i32 %%length.%(i)d)
341
  call void @llvm.trap()
342
  unreachable
343
"""
344
                % dict(subst, i=i, next_i=i + 1, r=r)
345
            )
346

347
    print(
348
        """
349
test.%d:
350
  ret i32 0
351
}
352

353
declare i32 @strlen(i8*)
354
declare i32 @write(i32, i8*, i32)
355
declare i32 @sprintf(i8*, i8*, ...)
356
declare void @llvm.trap() noreturn nounwind
357
"""
358
        % (len(result),)
359
    )
360

361

362
if __name__ == "__main__":
363
    main()
364