utils/shuffle_fuzz.py

   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 import argparse
  17 import itertools
  18 import random
  19 import sys
  20 import uuid
  21
  22 def 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 """
 154 define internal fastcc <%(N)d x %(T)s> @test(%(arguments)s) noinline nounwind {
 155 entry:""" % 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\\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 """
 186 define 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 """
 197 define i32 @main() {
 198 entry:
 199   ; Create a scratch space to print error messages.
 200   %%str = alloca [128 x i8]
 201   %%str.ptr = getelementptr inbounds [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 """
 221 test.%(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 """
 227 test.%(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
 232 die.%(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]* @error.%(i)d, i32 0, i32 0), i32 %%bad.%(i)d)
 237   %%length.%(i)d = call i32 @strlen(i8* %%str.ptr)
 238   %%size.%(i)d = add i32 %%length.%(i)d, 1
 239   call i32 @write(i32 2, i8* %%str.ptr, i32 %%size.%(i)d)
 240   call void @llvm.trap()
 241   unreachable
 242 """ % dict(subst, i=i, next_i=i + 1, r=r)
 243
 244   print """
 245 test.%d:
 246   ret i32 0
 247 }
 248
 249 declare i32 @strlen(i8*)
 250 declare i32 @write(i32, i8*, i32)
 251 declare i32 @sprintf(i8*, i8*, ...)
 252 declare void @llvm.trap() noreturn nounwind
 253 """ % (len(result),)
 254
 255 if __name__ == '__main__':
 256   main()