364 lines · python
1#!/usr/bin/env python2 3"""A shuffle vector fuzz tester.4 5This is a python program to fuzz test the LLVM shufflevector instruction. It6generates a function with a random sequnece of shufflevectors, maintaining the7element mapping accumulated across the function. It then generates a main8function which calls it with a different value in each element and checks that9the result matches the expected mapping.10 11Take the output IR printed to stdout, compile it to an executable using whatever12set of transforms you want to test, and run the program. If it crashes, it found13a bug.14"""15 16from __future__ import print_function17 18import argparse19import itertools20import random21import sys22import uuid23 24 25def 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 width92 93 # Because undef (-1) saturates and is indistinguishable when testing the94 # correctness of a shuffle, we want to bias our fuzz toward having a decent95 # mixture of non-undef lanes in the end. With a deep shuffle tree, the96 # probabilies aren't good so we need to bias things. The math here is that if97 # we uniformly select between -1 and the other inputs, each element of the98 # result will have the following probability of being undef:99 #100 # 1 - (shuffle_range/(shuffle_range+1))^max_shuffle_height101 #102 # More generally, for any probability P of selecting a defined element in103 # a single shuffle, the end result is:104 #105 # 1 - P^max_shuffle_height106 #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 tree112 # height. Below, let 'A' be 'len(shuffle_range)', 'C' be 'max_shuffle_height',113 # and 'B' be the bias we use to compensate for114 # 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 -1135 if random.random() <= undef_prob136 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 = inputs165 for shuffles in shuffle_tree:166 results = [167 [168 (169 (results[i] if j < width else results[i + 1])[j % width]170 if j != -1171 else -1172 )173 for j in s174 ]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-size189 # integer type.190 integral_element_type = element_type191 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 """202define internal fastcc <%(N)d x %(T)s> @test(%(arguments)s) noinline nounwind {203entry:"""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.0238}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 prevent259 # interprocedural optimizations from deleting the test.260 print(261 """262define 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> %%result265}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 mapped279 # incorrectly (in an observable way).280 print(281 """282define i32 @main() {283entry: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 0287 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 the291 ; result.292 %%v.cast = bitcast <%(N)d x %(T)s> %%v to <%(N)d x %(IT)s>293 br label %%test.0294"""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 inputs310 ]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 """320test.%(i)d:321 ; Skip this lane, its value is undef.322 br label %%test.%(next_i)d323"""324 % dict(subst, i=i, next_i=i + 1)325 )326 else:327 print(328 """329test.%(i)d:330 %%v.%(i)d = extractelement <%(N)d x %(IT)s> %%v.cast, i32 %(i)d331 %%cmp.%(i)d = icmp ne %(IT)s %%v.%(i)d, %(r)d332 br i1 %%cmp.%(i)d, label %%die.%(i)d, label %%test.%(next_i)d333 334die.%(i)d:335 ; Capture the actual value and print an error message.336 %%tmp.%(i)d = zext %(IT)s %%v.%(i)d to i2048337 %%bad.%(i)d = trunc i2048 %%tmp.%(i)d to i32338 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 unreachable343"""344 % dict(subst, i=i, next_i=i + 1, r=r)345 )346 347 print(348 """349test.%d:350 ret i32 0351}352 353declare i32 @strlen(i8*)354declare i32 @write(i32, i8*, i32)355declare i32 @sprintf(i8*, i8*, ...)356declare void @llvm.trap() noreturn nounwind357"""358 % (len(result),)359 )360 361 362if __name__ == "__main__":363 main()364