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