1#!/usr/bin/env python
2
3"""A shuffle vector fuzz tester.
4
5This is a python program to fuzz test the LLVM shufflevector instruction. It
6generates a function with a random sequnece of shufflevectors, maintaining the
7element mapping accumulated across the function. It then generates a main
8function which calls it with a different value in each element and checks that
9the result matches the expected mapping.
10
11Take the output IR printed to stdout, compile it to an executable using whatever
12set of transforms you want to test, and run the program. If it crashes, it found
13a bug.
14"""
15
16import argparse
17import itertools
18import random
19import sys
20import uuid
21
22def main():
23  element_types=['i8', 'i16', 'i32', 'i64', 'f32', 'f64']
24  parser = argparse.ArgumentParser(description=__doc__)
25  parser.add_argument('-v', '--verbose', action='store_true',
26                      help='Show verbose output')
27  parser.add_argument('--seed', default=str(uuid.uuid4()),
28                      help='A string used to seed the RNG')
29  parser.add_argument('--max-shuffle-height', type=int, default=16,
30                      help='Specify a fixed height of shuffle tree to test')
31  parser.add_argument('--no-blends', dest='blends', action='store_false',
32                      help='Include blends of two input vectors')
33  parser.add_argument('--fixed-bit-width', type=int, choices=[128, 256],
34                      help='Specify a fixed bit width of vector to test')
35  parser.add_argument('--fixed-element-type', choices=element_types,
36                      help='Specify a fixed element type to test')
37  parser.add_argument('--triple',
38                      help='Specify a triple string to include in the IR')
39  args = parser.parse_args()
40
41  random.seed(args.seed)
42
43  if args.fixed_element_type is not None:
44    element_types=[args.fixed_element_type]
45
46  if args.fixed_bit_width is not None:
47    if args.fixed_bit_width == 128:
48      width_map={'i64': 2, 'i32': 4, 'i16': 8, 'i8': 16, 'f64': 2, 'f32': 4}
49      (width, element_type) = random.choice(
50          [(width_map[t], t) for t in element_types])
51    elif args.fixed_bit_width == 256:
52      width_map={'i64': 4, 'i32': 8, 'i16': 16, 'i8': 32, 'f64': 4, 'f32': 8}
53      (width, element_type) = random.choice(
54          [(width_map[t], t) for t in element_types])
55    else:
56      sys.exit(1) # Checked above by argument parsing.
57  else:
58    width = random.choice([2, 4, 8, 16, 32, 64])
59    element_type = random.choice(element_types)
60
61  element_modulus = {
62      'i8': 1 << 8, 'i16': 1 << 16, 'i32': 1 << 32, 'i64': 1 << 64,
63      'f32': 1 << 32, 'f64': 1 << 64}[element_type]
64
65  shuffle_range = (2 * width) if args.blends else width
66
67  # Because undef (-1) saturates and is indistinguishable when testing the
68  # correctness of a shuffle, we want to bias our fuzz toward having a decent
69  # mixture of non-undef lanes in the end. With a deep shuffle tree, the
70  # probabilies aren't good so we need to bias things. The math here is that if
71  # we uniformly select between -1 and the other inputs, each element of the
72  # result will have the following probability of being undef:
73  #
74  #   1 - (shuffle_range/(shuffle_range+1))^max_shuffle_height
75  #
76  # More generally, for any probability P of selecting a defined element in
77  # a single shuffle, the end result is:
78  #
79  #   1 - P^max_shuffle_height
80  #
81  # The power of the shuffle height is the real problem, as we want:
82  #
83  #   1 - shuffle_range/(shuffle_range+1)
84  #
85  # So we bias the selection of undef at any given node based on the tree
86  # height. Below, let 'A' be 'len(shuffle_range)', 'C' be 'max_shuffle_height',
87  # and 'B' be the bias we use to compensate for
88  # C '((A+1)*A^(1/C))/(A*(A+1)^(1/C))':
89  #
90  #   1 - (B * A)/(A + 1)^C = 1 - A/(A + 1)
91  #
92  # So at each node we use:
93  #
94  #   1 - (B * A)/(A + 1)
95  # = 1 - ((A + 1) * A * A^(1/C))/(A * (A + 1) * (A + 1)^(1/C))
96  # = 1 - ((A + 1) * A^((C + 1)/C))/(A * (A + 1)^((C + 1)/C))
97  #
98  # This is the formula we use to select undef lanes in the shuffle.
99  A = float(shuffle_range)
100  C = float(args.max_shuffle_height)
101  undef_prob = 1.0 - (((A + 1.0) * pow(A, (C + 1.0)/C)) /
102                      (A * pow(A + 1.0, (C + 1.0)/C)))
103
104  shuffle_tree = [[[-1 if random.random() <= undef_prob
105                       else random.choice(range(shuffle_range))
106                    for _ in itertools.repeat(None, width)]
107                   for _ in itertools.repeat(None, args.max_shuffle_height - i)]
108                  for i in xrange(args.max_shuffle_height)]
109
110  if args.verbose:
111    # Print out the shuffle sequence in a compact form.
112    print >>sys.stderr, ('Testing shuffle sequence "%s" (v%d%s):' %
113                         (args.seed, width, element_type))
114    for i, shuffles in enumerate(shuffle_tree):
115      print >>sys.stderr, '  tree level %d:' % (i,)
116      for j, s in enumerate(shuffles):
117        print >>sys.stderr, '    shuffle %d: %s' % (j, s)
118    print >>sys.stderr, ''
119
120  # Symbolically evaluate the shuffle tree.
121  inputs = [[int(j % element_modulus)
122             for j in xrange(i * width + 1, (i + 1) * width + 1)]
123            for i in xrange(args.max_shuffle_height + 1)]
124  results = inputs
125  for shuffles in shuffle_tree:
126    results = [[((results[i] if j < width else results[i + 1])[j % width]
127                 if j != -1 else -1)
128                for j in s]
129               for i, s in enumerate(shuffles)]
130  if len(results) != 1:
131    print >>sys.stderr, 'ERROR: Bad results: %s' % (results,)
132    sys.exit(1)
133  result = results[0]
134
135  if args.verbose:
136    print >>sys.stderr, 'Which transforms:'
137    print >>sys.stderr, '  from: %s' % (inputs,)
138    print >>sys.stderr, '  into: %s' % (result,)
139    print >>sys.stderr, ''
140
141  # The IR uses silly names for floating point types. We also need a same-size
142  # integer type.
143  integral_element_type = element_type
144  if element_type == 'f32':
145    integral_element_type = 'i32'
146    element_type = 'float'
147  elif element_type == 'f64':
148    integral_element_type = 'i64'
149    element_type = 'double'
150
151  # Now we need to generate IR for the shuffle function.
152  subst = {'N': width, 'T': element_type, 'IT': integral_element_type}
153  print """
154define internal fastcc <%(N)d x %(T)s> @test(%(arguments)s) noinline nounwind {
155entry:""" % dict(subst,
156                 arguments=', '.join(
157                     ['<%(N)d x %(T)s> %%s.0.%(i)d' % dict(subst, i=i)
158                      for i in xrange(args.max_shuffle_height + 1)]))
159
160  for i, shuffles in enumerate(shuffle_tree):
161   for j, s in enumerate(shuffles):
162    print """
163  %%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>
164""".strip('\n') % dict(subst, i=i, next_i=i + 1, j=j, next_j=j + 1,
165                       S=', '.join(['i32 ' + (str(si) if si != -1 else 'undef')
166                                    for si in s]))
167
168  print """
169  ret <%(N)d x %(T)s> %%s.%(i)d.0
170}
171""" % dict(subst, i=len(shuffle_tree))
172
173  # Generate some string constants that we can use to report errors.
174  for i, r in enumerate(result):
175    if r != -1:
176      s = ('FAIL(%(seed)s): lane %(lane)d, expected %(result)d, found %%d\n\\0A' %
177           {'seed': args.seed, 'lane': i, 'result': r})
178      s += ''.join(['\\00' for _ in itertools.repeat(None, 128 - len(s) + 2)])
179      print """
180@error.%(i)d = private unnamed_addr global [128 x i8] c"%(s)s"
181""".strip() % {'i': i, 's': s}
182
183  # Define a wrapper function which is marked 'optnone' to prevent
184  # interprocedural optimizations from deleting the test.
185  print """
186define internal fastcc <%(N)d x %(T)s> @test_wrapper(%(arguments)s) optnone noinline {
187  %%result = call fastcc <%(N)d x %(T)s> @test(%(arguments)s)
188  ret <%(N)d x %(T)s> %%result
189}
190""" % dict(subst,
191           arguments=', '.join(['<%(N)d x %(T)s> %%s.%(i)d' % dict(subst, i=i)
192                                for i in xrange(args.max_shuffle_height + 1)]))
193
194  # Finally, generate a main function which will trap if any lanes are mapped
195  # incorrectly (in an observable way).
196  print """
197define i32 @main() {
198entry:
199  ; Create a scratch space to print error messages.
200  %%str = alloca [128 x i8]
201  %%str.ptr = getelementptr inbounds [128 x i8], [128 x i8]* %%str, i32 0, i32 0
202
203  ; Build the input vector and call the test function.
204  %%v = call fastcc <%(N)d x %(T)s> @test_wrapper(%(inputs)s)
205  ; We need to cast this back to an integer type vector to easily check the
206  ; result.
207  %%v.cast = bitcast <%(N)d x %(T)s> %%v to <%(N)d x %(IT)s>
208  br label %%test.0
209""" % dict(subst,
210           inputs=', '.join(
211               [('<%(N)d x %(T)s> bitcast '
212                 '(<%(N)d x %(IT)s> <%(input)s> to <%(N)d x %(T)s>)' %
213                 dict(subst, input=', '.join(['%(IT)s %(i)d' % dict(subst, i=i)
214                                              for i in input])))
215                for input in inputs]))
216
217  # Test that each non-undef result lane contains the expected value.
218  for i, r in enumerate(result):
219    if r == -1:
220      print """
221test.%(i)d:
222  ; Skip this lane, its value is undef.
223  br label %%test.%(next_i)d
224""" % dict(subst, i=i, next_i=i + 1)
225    else:
226      print """
227test.%(i)d:
228  %%v.%(i)d = extractelement <%(N)d x %(IT)s> %%v.cast, i32 %(i)d
229  %%cmp.%(i)d = icmp ne %(IT)s %%v.%(i)d, %(r)d
230  br i1 %%cmp.%(i)d, label %%die.%(i)d, label %%test.%(next_i)d
231
232die.%(i)d:
233  ; Capture the actual value and print an error message.
234  %%tmp.%(i)d = zext %(IT)s %%v.%(i)d to i2048
235  %%bad.%(i)d = trunc i2048 %%tmp.%(i)d to i32
236  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)
237  %%length.%(i)d = call i32 @strlen(i8* %%str.ptr)
238  call i32 @write(i32 2, i8* %%str.ptr, i32 %%length.%(i)d)
239  call void @llvm.trap()
240  unreachable
241""" % dict(subst, i=i, next_i=i + 1, r=r)
242
243  print """
244test.%d:
245  ret i32 0
246}
247
248declare i32 @strlen(i8*)
249declare i32 @write(i32, i8*, i32)
250declare i32 @sprintf(i8*, i8*, ...)
251declare void @llvm.trap() noreturn nounwind
252""" % (len(result),)
253
254if __name__ == '__main__':
255  main()
256