1903 lines · python
1# This test generates all variants of wmma intrinsics and verifies that LLVM2# generates correct instructions for them. This is the test generator only. The3# test scripts themselves are in wmma-ptx*-sm*.py files.4 5# RUN: true6 7from __future__ import print_function8 9import argparse10from itertools import product11from string import Template12 13 14class MMAType:15 def __init__(self, ptx_type):16 self.ptx_type = ptx_type17 self.llvm_type = {18 "e4m3": "i32",19 "e5m2": "i32",20 "e3m2": "i32",21 "e2m3": "i32",22 "e2m1": "i32",23 "f16": "<2 x half>",24 "f32": "float",25 "f64": "double",26 "s32": "i32",27 "b16": "i32",28 "b8": "i32",29 "b8x16.b6x16_p32": "i32",30 "b8x16.b4x16_p64": "i32",31 "s8": "i32",32 "u8": "i32",33 "s4": "i32",34 "u4": "i32",35 "b1": "i32",36 "bf16": "i32",37 "tf32": "i32",38 }[ptx_type]39 40 self.ptx_reg_pattern = {41 "f16": "%r[0-9]+",42 "f32": "%r[0-9]+",43 "f64": "%rd[0-9]+",44 }.get(ptx_type, "%r[0-9]+")45 46 def __repr__(self):47 return "%s/%s" % (self.ptx_type, self.llvm_type)48 49 50class MMAFrag:51 def __init__(self, geom, frag, ptx_elt_type, is_mma_sparse=False):52 self.geom = geom53 self.frag = frag54 self.mma_type = MMAType(ptx_elt_type)55 self.nregs = {56 # u8/s8 -> s32 @ m16n16k16/m8n32k16/m32n8k1657 "m16n16k16:a:u8": 2,58 "m16n16k16:a:s8": 2,59 "m16n16k16:b:u8": 2,60 "m16n16k16:b:s8": 2,61 "m16n16k16:c:s32": 8,62 "m16n16k16:d:s32": 8,63 "m8n32k16:a:u8": 1,64 "m8n32k16:a:s8": 1,65 "m8n32k16:b:u8": 4,66 "m8n32k16:b:s8": 4,67 "m8n32k16:c:s32": 8,68 "m8n32k16:d:s32": 8,69 "m32n8k16:a:u8": 4,70 "m32n8k16:a:s8": 4,71 "m32n8k16:b:u8": 1,72 "m32n8k16:b:s8": 1,73 "m32n8k16:c:s32": 8,74 "m32n8k16:d:s32": 8,75 "m8n8k16:a:u8": 1,76 "m8n8k16:a:s8": 1,77 "m8n8k16:b:u8": 1,78 "m8n8k16:b:s8": 1,79 "m8n8k16:c:s32": 2,80 "m8n8k16:d:s32": 2,81 "m16n8k16:a:u8": 2,82 "m16n8k16:a:s8": 2,83 "m16n8k16:b:u8": 1,84 "m16n8k16:b:s8": 1,85 "m16n8k16:c:s32": 4,86 "m16n8k16:d:s32": 4,87 "m16n8k32:a:u8": 2 if is_mma_sparse else 4,88 "m16n8k32:a:s8": 2 if is_mma_sparse else 4,89 "m16n8k32:b:u8": 2,90 "m16n8k32:b:s8": 2,91 "m16n8k32:c:s32": 4,92 "m16n8k32:d:s32": 4,93 # e4m3/e5m2/e3m2/e2m3/e2m1 -> f16/f32 @ m16n8k16/m16n8k3294 "m16n8k16:a:e4m3": 2,95 "m16n8k16:a:e5m2": 2,96 "m16n8k32:a:e4m3": 4,97 "m16n8k32:a:e5m2": 4,98 "m16n8k32:a:e3m2": 4,99 "m16n8k32:a:e2m3": 4,100 "m16n8k32:a:e2m1": 4,101 "m16n8k16:b:e4m3": 1,102 "m16n8k16:b:e5m2": 1,103 "m16n8k32:b:e4m3": 2,104 "m16n8k32:b:e5m2": 2,105 "m16n8k32:b:e3m2": 2,106 "m16n8k32:b:e2m3": 2,107 "m16n8k32:b:e2m1": 2,108 # mma sp109 "m16n8k32:a:bf16": 4,110 "m16n8k32:a:f16": 4,111 "m16n8k32:b:bf16": 4,112 "m16n8k32:b:f16": 4,113 "m16n8k32:c:f16": 2,114 "m16n8k32:c:f32": 4,115 "m16n8k32:d:f16": 2,116 "m16n8k32:d:f32": 4,117 "m16n8k16:a:tf32": 4,118 "m16n8k16:b:tf32": 4,119 "m16n8k16:c:tf32": 4,120 "m16n8k16:d:tf32": 4,121 "m16n8k64:a:u8": 4,122 "m16n8k64:a:s8": 4,123 "m16n8k64:a:e4m3": 4,124 "m16n8k64:a:e5m2": 4,125 "m16n8k64:a:e3m2": 4,126 "m16n8k64:a:e2m3": 4,127 "m16n8k64:a:e2m1": 4,128 "m16n8k64:b:u8": 4,129 "m16n8k64:b:s8": 4,130 "m16n8k64:b:e4m3": 4,131 "m16n8k64:b:e5m2": 4,132 "m16n8k64:b:e3m2": 4,133 "m16n8k64:b:e2m3": 4,134 "m16n8k64:b:e2m1": 4 if is_mma_sparse else 2,135 "m16n8k64:c:f16": 2,136 "m16n8k64:c:f32": 4,137 "m16n8k64:d:f16": 2,138 "m16n8k64:d:f32": 4,139 "m16n8k128:a:u4": 4,140 "m16n8k128:a:s4": 4,141 "m16n8k128:a:e2m1": 4,142 "m16n8k128:b:u4": 4,143 "m16n8k128:b:s4": 4,144 "m16n8k128:b:e2m1": 4,145 "m16n8k128:c:s32": 4,146 "m16n8k128:c:f32": 4,147 "m16n8k128:d:s32": 4,148 "m16n8k128:d:f32": 4,149 # u4/s4 -> s32 @ m8n8k32 (u4/s4)150 "m8n8k32:a:u4": 1,151 "m8n8k32:a:s4": 1,152 "m8n8k32:b:u4": 1,153 "m8n8k32:b:s4": 1,154 "m8n8k32:c:s32": 2,155 "m8n8k32:d:s32": 2,156 "m16n8k32:a:u4": 2,157 "m16n8k32:a:s4": 2,158 "m16n8k32:b:u4": 1,159 "m16n8k32:b:s4": 1,160 "m16n8k32:c:s32": 4,161 "m16n8k32:d:s32": 4,162 "m16n8k64:a:u4": 2 if is_mma_sparse else 4,163 "m16n8k64:a:s4": 2 if is_mma_sparse else 4,164 "m16n8k64:b:u4": 2,165 "m16n8k64:b:s4": 2,166 "m16n8k64:c:s32": 4,167 "m16n8k64:d:s32": 4,168 # b1 -> s32 @ m8n8k128(b1)169 "m8n8k128:a:b1": 1,170 "m8n8k128:b:b1": 1,171 "m8n8k128:c:s32": 2,172 "m8n8k128:d:s32": 2,173 "m16n8k128:a:b1": 2,174 "m16n8k128:b:b1": 1,175 "m16n8k128:c:s32": 4,176 "m16n8k128:d:s32": 4,177 "m16n8k256:a:b1": 4,178 "m16n8k256:b:b1": 2,179 "m16n8k256:c:s32": 4,180 "m16n8k256:d:s32": 4,181 # bf16 -> s32 @ m16n16k16/m8n32k16/m32n8k16182 "m16n16k16:a:bf16": 4,183 "m16n16k16:b:bf16": 4,184 "m8n32k16:a:bf16": 2,185 "m8n32k16:b:bf16": 8,186 "m32n8k16:a:bf16": 8,187 "m32n8k16:b:bf16": 2,188 "m16n8k16:a:bf16": 2 if is_mma_sparse else 4,189 "m16n8k16:b:bf16": 2,190 "m16n8k16:c:f32": 4,191 "m16n8k16:d:f32": 4,192 "m16n8k8:a:bf16": 2,193 "m16n8k8:b:bf16": 1,194 "m16n8k8:c:f32": 4,195 "m16n8k8:d:f32": 4,196 "m8n8k4:a:f64": 1,197 "m8n8k4:b:f64": 1,198 "m8n8k4:c:f64": 2,199 "m8n8k4:d:f64": 2,200 "m16n8k4:a:f64": 2,201 "m16n8k4:b:f64": 1,202 "m16n8k4:c:f64": 4,203 "m16n8k4:d:f64": 4,204 "m16n8k8:a:f64": 4,205 "m16n8k8:b:f64": 2,206 "m16n8k8:c:f64": 4,207 "m16n8k8:d:f64": 4,208 "m16n8k16:a:f64": 8,209 "m16n8k16:b:f64": 4,210 "m16n8k16:c:f64": 4,211 "m16n8k16:d:f64": 4,212 # tf32 -> s32 @ m16n16k8213 "m16n16k8:a:tf32": 4,214 "m16n16k8:b:tf32": 4,215 "m16n8k4:a:tf32": 2,216 "m16n8k4:b:tf32": 1,217 "m16n8k4:c:f32": 4,218 "m16n8k4:d:f32": 4,219 "m16n8k8:a:tf32": 2 if is_mma_sparse else 4,220 "m16n8k8:b:tf32": 2,221 "m16n8k8:c:f32": 4,222 "m16n8k8:d:f32": 4,223 "m8n8k4:a:f16": 2,224 "m8n8k4:b:f16": 2,225 "m16n8k8:a:f16": 2,226 "m16n8k8:b:f16": 1,227 "m16n8k8:c:f16": 2,228 "m16n8k8:d:f16": 2,229 "m16n8k8:c:f32": 4,230 "m16n8k8:d:f32": 4,231 "m16n8k16:a:f16": 2 if is_mma_sparse else 4,232 "m16n8k16:b:f16": 2,233 "m16n8k16:c:f16": 2,234 "m16n8k16:d:f16": 2,235 "m16n8k16:c:f32": 4,236 "m16n8k16:d:f32": 4,237 # ldmatrix238 "m8n8:x1:b16": 1,239 "m8n8:x2:b16": 2,240 "m8n8:x4:b16": 4,241 "m16n16:x1:b8": 2,242 "m16n16:x2:b8": 4,243 "m16n16:x1:b8x16.b6x16_p32": 2,244 "m16n16:x2:b8x16.b6x16_p32": 4,245 "m16n16:x1:b8x16.b4x16_p64": 2,246 "m16n16:x2:b8x16.b4x16_p64": 4,247 "m8n16:x1:b8x16.b6x16_p32": 1,248 "m8n16:x2:b8x16.b6x16_p32": 2,249 "m8n16:x4:b8x16.b6x16_p32": 4,250 "m8n16:x1:b8x16.b4x16_p64": 1,251 "m8n16:x2:b8x16.b4x16_p64": 2,252 "m8n16:x4:b8x16.b4x16_p64": 4,253 # stmatrix254 "m8n8:x1:b16": 1,255 "m8n8:x2:b16": 2,256 "m8n8:x4:b16": 4,257 "m16n8:x1:b8": 1,258 "m16n8:x2:b8": 2,259 "m16n8:x4:b8": 4,260 }.get(261 "%s:%s:%s" % (geom, frag, ptx_elt_type),262 {263 # All other FP shape/fragment/type combinations have the same size264 "a:f16": 8,265 "b:f16": 8,266 "c:f16": 4,267 "d:f16": 4,268 "c:f32": 8,269 "d:f32": 8,270 }.get("%s:%s" % (frag, ptx_elt_type), None),271 )272 assert self.nregs273 274 def __repr__(self):275 return "%s:%s:%s%s" % (276 self.geom,277 self.frag,278 self.mma_type,279 "" if self.nregs == 1 else ("*%d" % self.nregs),280 )281 282 283class MMAOp:284 def __init__(self, a, b, c, d):285 self.a = a286 self.b = b287 self.c = c288 self.d = d289 290 def __repr__(self):291 return "{A:%s, B:%s, C:%s, D:%s}" % (self.a, self.b, self.c, self.d)292 293 294def make_mma_ops(geoms, types_a, types_b, types_c, types_d, is_mma_sparse=False):295 ops = []296 for geom, type_a, type_c in product(geoms, types_a, types_c):297 for type_b, type_d in product(298 types_b if types_b else [type_a], types_d if types_d else [type_c]299 ):300 ops.append(301 MMAOp(302 MMAFrag(geom, "a", type_a, is_mma_sparse),303 MMAFrag(geom, "b", type_b, is_mma_sparse),304 MMAFrag(geom, "c", type_c, is_mma_sparse),305 MMAFrag(geom, "d", type_d, is_mma_sparse),306 )307 )308 return ops309 310 311def make_ldst_ops(geoms, frags, types):312 return [313 MMAFrag(geom, frag, ptx_type)314 for (geom, frag, ptx_type) in product(geoms, frags, types)315 ]316 317 318def make_ldmatrix_ops(geoms, frags, types):319 return [320 MMAFrag(geom, frag, ptx_type)321 for (geom, frag, ptx_type) in product(geoms, frags, types)322 ]323 324 325def make_stmatrix_ops(geoms, frags, types):326 return [327 MMAFrag(geom, frag, ptx_type)328 for (geom, frag, ptx_type) in product(geoms, frags, types)329 ]330 331 332def get_wmma_ops():333 return (334 make_mma_ops(["m16n16k8"], ["tf32"], [], ["f32"], [])335 + make_mma_ops(["m16n16k16", "m32n8k16", "m8n32k16"], ["bf16"], [], ["f32"], [])336 + make_mma_ops(["m8n8k4"], ["f64"], [], ["f64"], [])337 + make_mma_ops(338 ["m16n16k16", "m32n8k16", "m8n32k16"],339 ["f16"],340 [],341 ["f16", "f32"],342 ["f16", "f32"],343 )344 + make_mma_ops(345 ["m16n16k16", "m32n8k16", "m8n32k16"], ["s8", "u8"], [], ["s32"], []346 )347 + make_mma_ops(["m8n8k32"], ["s4", "u4"], [], ["s32"], [])348 + make_mma_ops(["m8n8k128"], ["b1"], [], ["s32"], [])349 )350 351 352def get_mma_ops():353 return (354 make_mma_ops(355 ["m8n8k4", "m16n8k4", "m16n8k8", "m16n8k16"], ["f64"], [], ["f64"], []356 )357 + make_mma_ops(["m16n8k4", "m16n8k8"], ["tf32"], [], ["f32"], [])358 + make_mma_ops(["m16n8k16", "m16n8k8"], ["bf16"], [], ["f32"], [])359 + make_mma_ops(360 ["m8n8k4", "m16n8k8", "m16n8k16"],361 ["f16"],362 [],363 ["f16", "f32"],364 ["f16", "f32"],365 )366 + make_mma_ops(367 ["m8n8k16", "m16n8k16", "m16n8k32"], ["s8", "u8"], ["s8", "u8"], ["s32"], []368 )369 + make_mma_ops(370 ["m8n8k32", "m16n8k32", "m16n8k64"], ["s4", "u4"], ["s4", "u4"], ["s32"], []371 )372 + make_mma_ops(["m8n8k128", "m16n8k128", "m16n8k256"], ["b1"], [], ["s32"], [])373 + make_mma_ops(374 ["m16n8k16"],375 ["e4m3", "e5m2"],376 ["e4m3", "e5m2"],377 ["f16", "f32"],378 ["f16", "f32"],379 )380 + make_mma_ops(381 ["m16n8k32"],382 ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],383 ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],384 ["f16", "f32"],385 ["f16", "f32"],386 )387 )388 389 390def get_ldst_ops(kind):391 ldst_ops = (392 make_ldst_ops(393 ["m16n16k16", "m32n8k16", "m8n32k16"],394 ["a", "b"],395 ["f16", "u8", "s8", "bf16"],396 )397 + make_ldst_ops(398 ["m16n16k16", "m32n8k16", "m8n32k16"], ["c", "d"], ["f16", "f32", "s32"]399 )400 + make_ldst_ops(["m8n8k32"], ["a", "b"], ["s4", "u4"])401 + make_ldst_ops(["m8n8k128"], ["a", "b"], ["b1"])402 + make_ldst_ops(["m8n8k32", "m8n8k128"], ["c", "d"], ["s32"])403 + make_ldst_ops(["m8n8k4"], ["a", "b", "c", "d"], ["f64"])404 + make_ldst_ops(["m16n16k8"], ["a", "b"], ["tf32"])405 + make_ldst_ops(["m16n16k8"], ["c", "d"], ["f32"])406 )407 return [x for x in ldst_ops if (x.frag == "d") == (kind == "store")]408 409 410def get_ldmatrix_ops():411 return (412 make_ldmatrix_ops(["m8n8"], ["x1", "x2", "x4"], ["b16"])413 + make_ldmatrix_ops(414 ["m16n16"], ["x1", "x2"], ["b8", "b8x16.b6x16_p32", "b8x16.b4x16_p64"]415 )416 + make_ldmatrix_ops(417 ["m8n16"], ["x1", "x2", "x4"], ["b8x16.b6x16_p32", "b8x16.b4x16_p64"]418 )419 )420 421 422def get_stmatrix_ops():423 return make_stmatrix_ops(["m8n8"], ["x1", "x2", "x4"], ["b16"]) + make_stmatrix_ops(424 ["m16n8"], ["x1", "x2", "x4"], ["b8"]425 )426 427 428def is_wmma_geom_supported(geom):429 # geometries for FP and ints.430 if geom in ["m8n32k16", "m32n8k16"]:431 return ptx_version >= 61432 # geometries for sub-ints.433 if geom in ["m8n8k32", "m8n8k128"]:434 return ptx_version >= 63 and gpu_arch >= 75435 if geom == "m16n16k16":436 return ptx_version >= 60437 if geom == "m16n8k8":438 return ptx_version >= 65439 if geom in ["m16n16k8", "m8n8k4"]:440 return ptx_version >= 70441 assert False # Unexpected geometry.442 443 444def is_mma_geom_supported(geom):445 # geometries for FP and ints.446 if geom == "m8n8k4":447 return ptx_version >= 64448 if geom in ["m16n8k8", "m8n8k16", "m8n8k32"]:449 return ptx_version >= 65450 if geom in [451 "m16n8k16",452 "m16n8k4",453 "m16n8k32",454 "m16n8k64",455 "m8n8k128",456 "m16n8k128",457 "m16n8k256",458 ]:459 return ptx_version >= 70460 assert False # Unexpected geometry.461 462 463def is_ldmatrix_geom_supported(geom):464 if geom in ["m8n8"]:465 return ptx_version >= 65 and gpu_arch >= 75466 elif geom in ["m16n16"]:467 return ptx_version >= 86 and gpu_arch >= 100 and aa468 elif geom in ["m8n16"]:469 return ptx_version >= 86 and gpu_arch >= 100 and aa470 assert False # Unexpected geometry.471 472 473def is_stmatrix_geom_supported(geom):474 if geom in ["m8n8"]:475 return ptx_version >= 78 and gpu_arch >= 90476 elif geom in ["m16n8"]:477 return ptx_version >= 86 and gpu_arch >= 100 and aa478 assert False # Unexpected geometry.479 480 481def is_ldmatrix_trans_supported(geom, trans):482 if geom in ["m8n8"]:483 return True484 elif geom in ["m16n16"]:485 return trans == ".trans"486 elif geom in ["m8n16"]:487 return trans == ""488 assert False # Unexpected geometry.489 490 491def is_stmatrix_trans_supported(geom, trans):492 if geom in ["m8n8"]:493 return True494 elif geom in ["m16n8"]:495 return trans == ".trans"496 assert False # Unexpected geometry.497 498 499def is_type_supported(ptx_type):500 if ptx_type in ["s8", "u8", "s32"]:501 return ptx_version >= 63 and gpu_arch >= 72502 if ptx_type in ["s4", "u4", "b1"]:503 return ptx_version >= 63 and gpu_arch >= 75504 if ptx_type == "b16":505 return ptx_version >= 65 and gpu_arch >= 75506 if ptx_type in ["bf16", "tf32", "f64"]:507 return ptx_version >= 70508 if ptx_type in ["e4m3", "e5m2"]:509 return ptx_version >= 84 and gpu_arch >= 89510 if ptx_type in ["e3m2", "e2m3", "e2m1"]:511 return ptx_version >= 87 and gpu_arch >= 120 and aa512 return ptx_version >= 60 and gpu_arch >= 70513 514 515def is_wmma_variant_supported(op, layout_a, layout_b, rnd, satf):516 if not (517 is_type_supported(op.a.mma_type.ptx_type) and is_wmma_geom_supported(op.a.geom)518 ):519 return False520 521 # rnd is only supported for FP64 WMMA522 if rnd and op.a.mma_type.ptx_type != "f64":523 return False524 525 if satf:526 # satfinite for floating points was removed in PTX 6.5527 if op.a.mma_type.ptx_type == "f16" and ptx_version >= 65:528 return False529 if not op.a.mma_type.ptx_type in ["f16", "s8", "u8", "s4", "u4"]:530 return False531 532 # sub-integer require row/col layout.533 if op.a.mma_type.ptx_type in ["s4", "u4", "b1"]:534 return layout_a == "row" and layout_b == "col"535 return True536 537 538def is_mma_variant_supported(op, layout_a, layout_b, kind, satf):539 if not (540 is_type_supported(op.a.mma_type.ptx_type) and is_mma_geom_supported(op.a.geom)541 ):542 return False543 544 if satf and op.a.mma_type.ptx_type not in ["s8", "u8", "s4", "u4"]:545 return False546 547 # If the type of C is f32 then so must the type of D548 if (549 op.a.geom == "m8n8k4"550 and op.c.mma_type.ptx_type == "f32"551 and op.d.mma_type.ptx_type != "f32"552 ):553 return False554 555 # A and B type must be the same. C and D type must be the same556 if op.a.geom == "m16n8k8" and (557 op.a.mma_type.ptx_type != op.b.mma_type.ptx_type558 or op.c.mma_type.ptx_type != op.d.mma_type.ptx_type559 ):560 return False561 562 if (563 op.a.geom != "m8n8k4"564 and op.a.mma_type.ptx_type == "f64"565 and (ptx_version < 78 or gpu_arch < 90)566 ):567 return False568 569 # C and D type must be the same570 if (571 op.a.geom in ["m16n8k16", "m16n8k32"]572 and op.c.mma_type.ptx_type != op.d.mma_type.ptx_type573 ):574 return False575 576 if (577 op.a.geom in ["m16n8k16", "m16n8k32"]578 and any(579 x in ["e4m3", "e5m2"]580 for x in (op.a.mma_type.ptx_type, op.b.mma_type.ptx_type)581 )582 and ptx_version < 87583 ):584 return False585 586 if kind != "" and not (ptx_version >= 87 and gpu_arch >= 120 and aa):587 return False588 589 if kind != "" and (590 op.a.geom != "m16n8k32"591 or op.a.mma_type.ptx_type not in ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"]592 ):593 return False594 595 if (596 kind == ""597 and op.a.geom in ["m16n8k16", "m16n8k32"]598 and any(599 x in ["e3m2", "e2m3", "e2m1"]600 for x in (op.a.mma_type.ptx_type, op.b.mma_type.ptx_type)601 )602 ):603 return False604 605 # Require row/col layout for all MMA except m8n8k4 on FP16606 if not (op.a.geom == "m8n8k4" and op.a.mma_type.ptx_type == "f16"):607 return layout_a == "row" and layout_b == "col"608 609 return True610 611 612def is_ldst_variant_supported(frag, layout):613 if not (614 is_type_supported(frag.mma_type.ptx_type) and is_wmma_geom_supported(frag.geom)615 ):616 return False617 if frag.mma_type.ptx_type in ["s4", "u4", "b1"]:618 # sub-integer require sm_75 and ptx63, row/col layout for a/b.619 return (620 (frag.frag == "a" and layout == "row")621 or (frag.frag == "b" and layout == "col")622 or frag.frag in ["c", "d"]623 )624 return True625 626 627def is_ldmatrix_variant_supported(frag, trans):628 if not (629 is_type_supported(frag.mma_type.ptx_type)630 and is_ldmatrix_geom_supported(frag.geom)631 and is_ldmatrix_trans_supported(frag.geom, trans)632 ):633 return False634 return frag.frag in ["x1", "x2", "x4"]635 636 637def is_stmatrix_variant_supported(frag, trans):638 if not (639 is_type_supported(frag.mma_type.ptx_type)640 and is_stmatrix_geom_supported(frag.geom)641 and is_stmatrix_trans_supported(frag.geom, trans)642 ):643 return False644 return frag.frag in ["x1", "x2", "x4"]645 646 647def make_wmma_slice_ty(frag):648 return [frag.mma_type.llvm_type] * frag.nregs649 650 651def make_wmma_ld_ret_ty(frag):652 results = make_wmma_slice_ty(frag)653 if len(results) == 1:654 return "%s" % results[0]655 return "{%s}" % ", ".join(results)656 657 658# returns address space659def get_aspace(space):660 space_map = {661 ".global": 1,662 ".shared": 3,663 ".const": 4,664 ".local": 5,665 ".param": 101,666 "": 0,667 ".generic": 0,668 }669 return space_map[space]670 671 672def get_pspace(space):673 return "p%di8" % get_aspace(space)674 675 676def check_pattern(frag):677 return "{{%s}}" % ", *".join([frag.mma_type.ptx_reg_pattern] * frag.nregs)678 679 680def gen_wmma_load_tests():681 load_template = """682declare ${ret_ty} @${intrinsic}(i8 ${as}* %src ${extra_args});683 684; CHECK-LABEL: .func {{.*}}test_${function}(685define ${ret_ty} @test_${function}(i8 ${as}* %src ${extra_args}) {686; CHECK: ${instruction}687; CHECK: {${check_result}}688; CHECK: [%rd{{[0-9]+}}]${stride_pattern}689 %v0 = call ${ret_ty} @${intrinsic}(i8 ${as}* %src ${extra_args});690 ret ${ret_ty} %v0;691}692 693; CHECK-LABEL: .func{{.*}}test_${function}_o(694define ${ret_ty} @test_${function}_o(i8 ${as}* %src ${extra_args}) {695; CHECK: ${instruction}696; CHECK: {${check_result}}697; CHECK: [%rd{{[0-9]+}}+128]${stride_pattern}698 %src1 = getelementptr i8, i8 ${as}* %src, i32 128;699 %v0 = call ${ret_ty} @${intrinsic}(i8 ${as}* %src1 ${extra_args});700 ret ${ret_ty} %v0;701}702"""703 intrinsic_template = (704 "llvm.nvvm.wmma.${geom}.load.${abc}.${layout}${stride}.${itype}.${pspace}"705 )706 instruction_template = (707 "wmma.load.${abc}.sync${aligned}.${layout}.${geom}${space}.${itype}"708 )709 710 generated_items = []711 712 for frag, layout, space, stride in product(713 get_ldst_ops("load"),714 ["row", "col"],715 ["", ".shared", ".global"],716 ["", ".stride"],717 ):718 if not is_ldst_variant_supported(frag, layout):719 continue720 721 params = {722 "abc": frag.frag,723 "aligned": ".aligned" if ptx_version >= 63 else "",724 "layout": layout,725 "space": space,726 "stride": stride,727 "itype": frag.mma_type.ptx_type,728 "pspace": get_pspace(space),729 "as": "addrspace(%d)" % get_aspace(space),730 "geom": frag.geom,731 }732 733 test_params = params734 test_params["intrinsic"] = Template(intrinsic_template).substitute(params)735 test_params["function"] = test_params["intrinsic"].replace(".", "_")736 test_params["instruction"] = Template(instruction_template).substitute(params)737 test_params["ret_ty"] = make_wmma_ld_ret_ty(frag)738 test_params["check_result"] = check_pattern(frag)739 740 if stride:741 test_params["extra_args"] = ", i32 %stride"742 test_params["stride_pattern"] = ", %r{{[0-9]+}}"743 else:744 test_params["extra_args"] = ""745 test_params["stride_pattern"] = ""746 747 print(Template(load_template).substitute(test_params))748 749 generated_items.append((test_params["intrinsic"], test_params["instruction"]))750 751 return generated_items752 753 754def make_wmma_slice_args(frag):755 return ", ".join(756 [757 "%s %%%s%d" % (t, frag.frag, i)758 for i, t in enumerate(make_wmma_slice_ty(frag))759 ]760 )761 762 763def gen_wmma_store_tests():764 store_template = """765declare void @${intrinsic}(i8 ${as}* %src, ${args}${extra_args});766 767; CHECK-LABEL: .func {{.*}}test_${function}(768define void @test_${function}(i8 ${as}* %src, ${args}${extra_args}) {769; CHECK: ${instruction} {{.*}}[%rd{{[0-9+]}}770; CHECK: {${check_args}}771; CHECK: ${stride_pattern}772 call void @${intrinsic}(i8 ${as}* %src, ${args} ${extra_args});773 ret void774}775 776; CHECK-LABEL: .func{{.*}}test_${function}_o(777define void @test_${function}_o(i8 ${as}* %src, ${args}${extra_args}) {778; CHECK: ${instruction} {{.*}}[%rd{{[0-9+]}}+128]779; CHECK: ${check_args}780; CHECK: ${stride_pattern}781 %src1 = getelementptr i8, i8 ${as}* %src, i32 128;782 call void @${intrinsic}(i8 ${as}* %src1, ${args}${extra_args});783 ret void784}785"""786 intrinsic_template = (787 "llvm.nvvm.wmma.${geom}.store.${abc}.${layout}${stride}.${itype}.${pspace}"788 )789 instruction_template = (790 "wmma.store.${abc}.sync${aligned}.${layout}.${geom}${space}.${itype}"791 )792 793 generated_items = []794 795 for frag, layout, space, stride in product(796 get_ldst_ops("store"),797 ["row", "col"],798 ["", ".shared", ".global"],799 ["", ".stride"],800 ):801 802 if not is_ldst_variant_supported(frag, layout):803 continue804 805 params = {806 "abc": frag.frag,807 "aligned": ".aligned" if ptx_version >= 63 else "",808 "layout": layout,809 "space": space,810 "stride": stride,811 "itype": frag.mma_type.ptx_type,812 "pspace": get_pspace(space),813 "as": "addrspace(%d)" % get_aspace(space),814 "geom": frag.geom,815 }816 817 test_params = params818 test_params["intrinsic"] = Template(intrinsic_template).substitute(params)819 test_params["function"] = test_params["intrinsic"].replace(".", "_")820 test_params["instruction"] = Template(instruction_template).substitute(params)821 test_params["ret_ty"] = make_wmma_ld_ret_ty(frag)822 test_params["check_args"] = check_pattern(frag)823 if stride:824 test_params["extra_args"] = ", i32 %stride"825 test_params["stride_pattern"] = ", %r{{[0-9]+}};"826 else:827 test_params["extra_args"] = ""828 test_params["stride_pattern"] = ";"829 test_params["args"] = make_wmma_slice_args(frag)830 831 print(Template(store_template).substitute(test_params))832 generated_items.append((test_params["intrinsic"], test_params["instruction"]))833 834 return generated_items835 836 837def gen_ldmatrix_tests():838 ldmatrix_template = """839declare ${ret_ty} @${intrinsic}(i8 ${as}* %src);840 841; CHECK-LABEL: .func {{.*}}test_${function}(842define ${ret_ty} @test_${function}(i8 ${as}* %src) {843; CHECK: ${instruction}844; CHECK: {${check_result}}845; CHECK: [%rd{{[0-9]+}}]846 %v0 = call ${ret_ty} @${intrinsic}(i8 ${as}* %src);847 ret ${ret_ty} %v0;848}849 850; CHECK-LABEL: .func{{.*}}test_${function}_o(851define ${ret_ty} @test_${function}_o(i8 ${as}* %src) {852; CHECK: ${instruction}853; CHECK: {${check_result}}854; CHECK: [%rd{{[0-9]+}}+128]855 %src1 = getelementptr i8, i8 ${as}* %src, i32 128;856 %v0 = call ${ret_ty} @${intrinsic}(i8 ${as}* %src1);857 ret ${ret_ty} %v0;858}859"""860 intrinsic_template = (861 "llvm.nvvm.ldmatrix.sync.aligned.${geom}.${frag}${trans}.${itype}.${pspace}"862 )863 instruction_template = (864 "ldmatrix.sync.aligned.${geom}.${frag}${trans}${space}.${itype}"865 )866 867 generated_items = []868 869 for frag, space, trans in product(870 get_ldmatrix_ops(),871 ["", ".shared"],872 ["", ".trans"],873 ):874 if not is_ldmatrix_variant_supported(frag, trans):875 continue876 877 params = {878 "frag": frag.frag,879 "space": space,880 "trans": trans,881 "itype": frag.mma_type.ptx_type,882 "pspace": get_pspace(space),883 "as": "addrspace(%d)" % get_aspace(space),884 "geom": frag.geom,885 }886 887 test_params = params888 test_params["intrinsic"] = Template(intrinsic_template).substitute(params)889 test_params["function"] = test_params["intrinsic"].replace(".", "_")890 test_params["instruction"] = Template(instruction_template).substitute(params)891 test_params["ret_ty"] = make_wmma_ld_ret_ty(frag)892 test_params["check_result"] = check_pattern(frag)893 894 print(Template(ldmatrix_template).substitute(test_params))895 896 generated_items.append((test_params["intrinsic"], test_params["instruction"]))897 898 return generated_items899 900 901def gen_stmatrix_tests():902 stmatrix_template = """903declare void @${intrinsic}(i8 ${as}* %dst, ${args});904 905; CHECK-LABEL: .func {{.*}}test_${function}(906define void @test_${function}(i8 ${as}* %dst, ${args}) {907; CHECK: ${instruction} {{.*}}[%rd{{[0-9+]}}]908; CHECK: {${check_args}}909 call void @${intrinsic}(i8${as}* %dst, ${args});910 ret void911}912 913; CHECK-LABEL: .func{{.*}}test_${function}_o(914define void @test_${function}_o(i8 ${as}* %dst, ${args}) {915; CHECK: ${instruction} {{.*}}[%rd{{[0-9+]}}+128],916; CHECK: {${check_args}}917 %dst1 = getelementptr i8, i8 ${as}* %dst, i32 128;918 call void @${intrinsic}(i8 ${as}* %dst1, ${args});919 ret void920}921"""922 intrinsic_template = (923 "llvm.nvvm.stmatrix.sync.aligned.${geom}.${frag}${trans}.${itype}.${pspace}"924 )925 instruction_template = (926 "stmatrix.sync.aligned.${geom}.${frag}${trans}${space}.${itype}"927 )928 generated_items = []929 930 for frag, space, trans in product(931 get_stmatrix_ops(),932 ["", ".shared"],933 ["", ".trans"],934 ):935 if not is_stmatrix_variant_supported(frag, trans):936 continue937 938 params = {939 "frag": frag.frag,940 "space": space,941 "trans": trans,942 "itype": frag.mma_type.ptx_type,943 "pspace": get_pspace(space),944 "as": "addrspace(%d)" % get_aspace(space),945 "geom": frag.geom,946 }947 948 test_params = params949 test_params["intrinsic"] = Template(intrinsic_template).substitute(params)950 test_params["function"] = test_params["intrinsic"].replace(".", "_")951 test_params["instruction"] = Template(instruction_template).substitute(params)952 test_params["args"] = make_wmma_slice_args(frag)953 test_params["check_args"] = check_pattern(frag)954 955 print(Template(stmatrix_template).substitute(test_params))956 generated_items.append((test_params["intrinsic"], test_params["instruction"]))957 958 return generated_items959 960def mma_signature(op):961 if op.a.mma_type.ptx_type in ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"]:962 # FP8/F8F6F4 ops identified by inputs, accumulator & result types.963 return "%s.%s.%s.%s" % (964 op.d.mma_type.ptx_type,965 op.a.mma_type.ptx_type,966 op.b.mma_type.ptx_type,967 op.c.mma_type.ptx_type,968 )969 elif op.a.mma_type.ptx_type == "f16":970 # FP16 ops identified by accumulator & result type.971 return "%s.%s" % (op.d.mma_type.ptx_type, op.c.mma_type.ptx_type)972 elif op.a.mma_type.ptx_type != op.b.mma_type.ptx_type:973 # other ops are identified by input types.974 return "%s.%s" % (op.a.mma_type.ptx_type, op.b.mma_type.ptx_type)975 else:976 # if input types are the same, it only appears once.977 return op.a.mma_type.ptx_type978 979 980def mma_ptx_signature(op):981 # Encode all four types as D.A.B.C982 return ".".join(x.mma_type.ptx_type for x in (op.d, op.a, op.b, op.c))983 984 985def wmma_signature(op):986 if op.a.mma_type.ptx_type == "f16":987 # FP16 ops identified by accumulator & result type.988 return "%s.%s" % (op.d.mma_type.ptx_type, op.c.mma_type.ptx_type)989 else:990 # other ops are identified by input type.991 return op.a.mma_type.ptx_type992 993 994def wmma_ptx_signature(op):995 if op.a.mma_type.ptx_type == "f16":996 # FP16 instructions use D.C997 return "%s.%s" % (op.d.mma_type.ptx_type, op.c.mma_type.ptx_type)998 else:999 # other instructions encode all four types as D.A.B.C1000 return ".".join(x.mma_type.ptx_type for x in (op.d, op.a, op.b, op.c))1001 1002 1003def common_mma_test_gen(params, op, intrinsic_template, instruction_template):1004 mma_template = """1005declare ${ret_ty} @${intrinsic}(1006 ${args});1007 1008; CHECK-LABEL: .func {{.*}}test_${function}(1009define ${ret_ty} @test_${function}(1010 ${args}) {1011; CHECK: ${instruction}1012; CHECK-NEXT: ${check_d}1013; CHECK-NEXT: ${check_a}1014; CHECK-NEXT: ${check_b}1015; CHECK-NEXT: ${check_c}1016 %r = call ${ret_ty} @${intrinsic}(1017 ${args});1018 ret ${ret_ty} %r;1019}1020"""1021 1022 test_params = params1023 test_params["intrinsic"] = (1024 Template(intrinsic_template)1025 .substitute(params)1026 .replace("::", ".")1027 .replace("_", ".")1028 )1029 test_params["function"] = test_params["intrinsic"].replace(".", "_")1030 test_params["instruction"] = Template(instruction_template).substitute(params)1031 test_params["ret_ty"] = make_wmma_ld_ret_ty(op.d)1032 test_params["check_a"] = check_pattern(op.a)1033 test_params["check_b"] = check_pattern(op.b)1034 test_params["check_c"] = check_pattern(op.c)1035 test_params["check_d"] = check_pattern(op.d)1036 args = ",\n ".join(make_wmma_slice_args(frag) for frag in (op.a, op.b, op.c))1037 test_params["args"] = args1038 print(Template(mma_template).substitute(test_params))1039 return (test_params["intrinsic"], test_params["instruction"])1040 1041 1042def get_b1_ops(ptx_type):1043 if ptx_type != "b1":1044 return [""]1045 if ptx_version >= 71:1046 return [".xor.popc", ".and.popc"]1047 return [".xor.popc"]1048 1049 1050def gen_wmma_mma_tests():1051 wmma_intrinsic_template = "llvm.nvvm.wmma.${geom}.mma${b1op}.${alayout}.${blayout}${rnd}.${intrinsic_signature}${satf}"1052 wmma_instruction_template = "wmma.mma${b1op}.sync${aligned}.${alayout}.${blayout}.${geom}${rnd}.${ptx_signature}${satf}"1053 1054 generated_items = []1055 1056 for op, alayout, blayout, rnd, satf in product(1057 get_wmma_ops(),1058 ["row", "col"],1059 ["row", "col"],1060 [".rn", ".rz", ".rm", ".rp", ""],1061 [".satfinite", ""],1062 ):1063 1064 if not is_wmma_variant_supported(op, alayout, blayout, rnd, satf):1065 continue1066 1067 for b1op in get_b1_ops(op.a.mma_type.ptx_type):1068 params = {1069 "aligned": ".aligned" if ptx_version >= 63 else "",1070 "alayout": alayout,1071 "blayout": blayout,1072 "intrinsic_signature": wmma_signature(op),1073 "ptx_signature": wmma_ptx_signature(op),1074 "satf": satf,1075 "rnd": rnd,1076 "geom": op.a.geom,1077 "b1op": b1op,1078 }1079 1080 intrinsic_template = wmma_intrinsic_template1081 instruction_template = wmma_instruction_template1082 1083 generated_items.append(1084 common_mma_test_gen(1085 params, op, intrinsic_template, instruction_template1086 )1087 )1088 1089 return generated_items1090 1091 1092def gen_mma_tests():1093 mma_intrinsic_template = "llvm.nvvm.mma${b1op}.${geom}.${alayout}.${blayout}${kind}${satf}.${intrinsic_signature}"1094 mma_instruction_template = "mma.sync${aligned}.${geom}.${alayout}.${blayout}${kind}${satf}.${ptx_signature}${b1op}"1095 1096 generated_items = []1097 1098 for op, alayout, blayout, kind, satf in product(1099 get_mma_ops(),1100 ["row", "col"],1101 ["row", "col"],1102 ["", ".kind::f8f6f4"],1103 [".satfinite", ""],1104 ):1105 1106 if not is_mma_variant_supported(op, alayout, blayout, kind, satf):1107 continue1108 1109 for b1op in get_b1_ops(op.a.mma_type.ptx_type):1110 params = {1111 "aligned": ".aligned" if ptx_version >= 63 else "",1112 "alayout": alayout,1113 "blayout": blayout,1114 "intrinsic_signature": mma_signature(op),1115 "ptx_signature": mma_ptx_signature(op),1116 "satf": satf,1117 "geom": op.a.geom,1118 "b1op": b1op,1119 "kind": kind,1120 }1121 1122 intrinsic_template = mma_intrinsic_template1123 instruction_template = mma_instruction_template1124 1125 generated_items.append(1126 common_mma_test_gen(1127 params, op, intrinsic_template, instruction_template1128 )1129 )1130 1131 return generated_items1132 1133 1134def get_mma_block_scale_ops():1135 return make_mma_ops(["m16n8k64"], ["e2m1"], [], ["f32"], []) + make_mma_ops(1136 ["m16n8k32"],1137 ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],1138 ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],1139 ["f32"],1140 [],1141 )1142 1143 1144def is_mma_block_scale_geom_supported(geom):1145 # geometries for FP.1146 if geom in [1147 "m16n8k32",1148 "m16n8k64",1149 ]:1150 return True1151 raise ValueError(f"Unexpected MMA block scale geometry: {geom}")1152 1153 1154def is_mma_block_scale_variant_supported(op, kind, scale_vec_size, stype):1155 if not (1156 is_type_supported(op.a.mma_type.ptx_type)1157 and is_mma_block_scale_geom_supported(op.a.geom)1158 ):1159 return False1160 1161 if (1162 op.a.geom == "m16n8k64"1163 and kind == "mxf4"1164 and stype == "ue8m0"1165 and scale_vec_size in ["", ".scale_vec::2X"]1166 ):1167 return True1168 1169 if (1170 op.a.geom == "m16n8k64"1171 and kind == "mxf4nvf4"1172 and stype == "ue8m0"1173 and scale_vec_size == ".scale_vec::2X"1174 ):1175 return True1176 1177 if (1178 op.a.geom == "m16n8k64"1179 and kind == "mxf4nvf4"1180 and stype == "ue4m3"1181 and scale_vec_size == ".scale_vec::4X"1182 ):1183 return True1184 1185 if (1186 op.a.geom == "m16n8k32"1187 and kind == "mxf8f6f4"1188 and stype == "ue8m0"1189 and scale_vec_size in ["", ".scale_vec::1X"]1190 ):1191 return True1192 1193 return False1194 1195 1196def common_mma_block_scale_test_gen(1197 params, op, intrinsic_template, instruction_template1198):1199 mma_block_scale_template = """1200declare ${ret_ty} @${intrinsic}(1201 ${args});1202 1203; CHECK-LABEL: .func {{.*}}test_${function}(1204define ${ret_ty} @test_${function}(1205 ${args}) {1206; CHECK: ${instruction}1207; CHECK-NEXT: ${check_d}1208; CHECK-NEXT: ${check_a}1209; CHECK-NEXT: ${check_b}1210; CHECK-NEXT: ${check_c}1211; CHECK-NEXT: ${check_scale_a_data}1212; CHECK-NEXT: ${check_byte_id_a}1213; CHECK-NEXT: ${check_thread_id_a}1214; CHECK-NEXT: ${check_scale_b_data}1215; CHECK-NEXT: ${check_byte_id_b}1216; CHECK-NEXT: ${check_thread_id_b}1217 %r = call ${ret_ty} @${intrinsic}(1218 ${args});1219 ret ${ret_ty} %r;1220}1221"""1222 1223 test_params = params1224 test_params["intrinsic"] = Template(intrinsic_template).substitute(params)1225 test_params["function"] = test_params["intrinsic"].replace(".", "_")1226 test_params["instruction"] = Template(instruction_template).substitute(params)1227 test_params["ret_ty"] = make_wmma_ld_ret_ty(op.d)1228 test_params["check_a"] = check_pattern(op.a)1229 test_params["check_b"] = check_pattern(op.b)1230 test_params["check_c"] = check_pattern(op.c)1231 test_params["check_d"] = check_pattern(op.d)1232 test_params["check_scale_a_data"] = "{{%r[0-9]+}}"1233 test_params["check_byte_id_a"] = "{{%r[0-9]+}}"1234 test_params["check_thread_id_a"] = "{{%r[0-9]+}}"1235 test_params["check_scale_b_data"] = "{{%r[0-9]+}}"1236 test_params["check_byte_id_b"] = "{{%r[0-9]+}}"1237 test_params["check_thread_id_b"] = "{{%r[0-9]+}}"1238 args = ",\n ".join(1239 list(make_wmma_slice_args(frag) for frag in (op.a, op.b, op.c))1240 + ["i32 %scale_a_data", "i16 %byte_id_a, i16 %thread_id_a"]1241 + ["i32 %scale_b_data", "i16 %byte_id_b, i16 %thread_id_b"]1242 )1243 test_params["args"] = args1244 print(Template(mma_block_scale_template).substitute(test_params))1245 return (test_params["intrinsic"], test_params["instruction"])1246 1247 1248def gen_mma_block_scale_tests():1249 if not (ptx_version >= 88 and gpu_arch >= 120 and aa):1250 return []1251 1252 mma_block_scale_intrinsic_template = "llvm.nvvm.mma.block.scale.${geom}.row.col.${kind}${scale}.${intrinsic_signature}.${stype}"1253 mma_block_scale_instruction_template = "mma.sync.aligned.${geom}.row.col.kind::${kind}.block_scale${scale_vec_size}.${ptx_signature}.${stype}"1254 1255 generated_items = []1256 1257 for op, kind, scale_vec_size, stype in product(1258 get_mma_block_scale_ops(),1259 ["mxf4", "mxf4nvf4", "mxf8f6f4"],1260 ["", ".scale_vec::1X", ".scale_vec::2X", ".scale_vec::4X"],1261 ["ue8m0", "ue4m3"],1262 ):1263 if not is_mma_block_scale_variant_supported(op, kind, scale_vec_size, stype):1264 continue1265 1266 params = {1267 "intrinsic_signature": mma_signature(op),1268 "ptx_signature": mma_ptx_signature(op),1269 "geom": op.a.geom,1270 "kind": kind,1271 "scale_vec_size": scale_vec_size,1272 "scale": scale_vec_size.replace("_vec::", ".").lower(),1273 "stype": stype,1274 }1275 1276 intrinsic_template = mma_block_scale_intrinsic_template1277 instruction_template = mma_block_scale_instruction_template1278 1279 generated_items.append(1280 common_mma_block_scale_test_gen(1281 params, op, intrinsic_template, instruction_template1282 )1283 )1284 1285 return generated_items1286 1287 1288def get_mma_sp_ops():1289 return (1290 make_mma_ops(["m16n8k16", "m16n8k32"], ["bf16"], [], ["f32"], [], True)1291 + make_mma_ops(["m16n8k8", "m16n8k16"], ["tf32"], [], ["f32"], [], True)1292 + make_mma_ops(1293 ["m16n8k16", "m16n8k32"],1294 ["f16"],1295 [],1296 ["f16", "f32"],1297 ["f16", "f32"],1298 True,1299 )1300 + make_mma_ops(1301 ["m16n8k64", "m16n8k128"], ["s4", "u4"], ["s4", "u4"], ["s32"], [], True1302 )1303 + make_mma_ops(1304 ["m16n8k32", "m16n8k64"], ["s8", "u8"], ["s8", "u8"], ["s32"], [], True1305 )1306 + make_mma_ops(1307 ["m16n8k64"],1308 ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],1309 ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],1310 ["f16", "f32"],1311 ["f16", "f32"],1312 True,1313 )1314 )1315 1316 1317def is_mma_sp_geom_supported(geom):1318 # geometries for FP and ints.1319 if geom in [1320 "m16n8k16",1321 "m16n8k32",1322 "m16n8k8",1323 "m16n8k64",1324 "m16n8k128",1325 ]:1326 return ptx_version >= 711327 raise ValueError(f"Unexpected sparse MMA geometry: {geom}")1328 1329 1330def is_mma_sp_variant_supported(op, metadata, kind, satf):1331 if metadata != "sp" and (ptx_version < 85 or gpu_arch < 80):1332 return False1333 1334 if kind != "" and (ptx_version < 87 or gpu_arch < 120 or not aa):1335 return False1336 1337 if not (1338 is_type_supported(op.a.mma_type.ptx_type)1339 and is_mma_sp_geom_supported(op.a.geom)1340 ):1341 return False1342 1343 is_int = op.a.mma_type.ptx_type in ["s8", "u8", "s4", "u4"]1344 1345 if satf and not is_int:1346 return False1347 1348 # A and B type must be the same1349 if (1350 op.a.mma_type.ptx_type in ["f16", "bf16", "tf32"]1351 and op.a.mma_type.ptx_type != op.b.mma_type.ptx_type1352 ):1353 return False1354 1355 # C and D type must be the same for m16n8k16/m16n8k32/m16n8k641356 if (1357 op.a.geom in ["m16n8k16", "m16n8k32", "m16n8k64"]1358 and op.c.mma_type.ptx_type != op.d.mma_type.ptx_type1359 ):1360 return False1361 1362 if kind == "" and (1363 op.a.mma_type.ptx_type in ["e3m2", "e2m3", "e2m1"]1364 or op.b.mma_type.ptx_type in ["e3m2", "e2m3", "e2m1"]1365 ):1366 return False1367 1368 if (1369 kind == ""1370 and op.a.geom == "m16n8k64"1371 and (op.c.mma_type.ptx_type == "f16" or op.d.mma_type.ptx_type == "f16")1372 ):1373 return False1374 1375 if kind != "" and (metadata == "sp" or op.a.geom != "m16n8k64" or is_int):1376 return False1377 1378 return True1379 1380 1381def sp_selector_gen(op, block_scale=False):1382 if block_scale:1383 # PTX ISA 9.0 has the sparsity selector equal to 0 only1384 return range(1)1385 1386 # (geom, type) -> allowed selector range1387 range_01 = {1388 ("m16n8k32", "bf16"),1389 ("m16n8k32", "f16"),1390 ("m16n8k16", "tf32"),1391 ("m16n8k32", "u8"),1392 ("m16n8k32", "s8"),1393 ("m16n8k64", "u4"),1394 ("m16n8k64", "s4"),1395 }1396 1397 if (op.a.geom, op.a.mma_type.ptx_type) in range_01:1398 return range(2)1399 if op.a.geom == "m16n8k64" and op.a.mma_type.ptx_type in [1400 "u8",1401 "s8",1402 "e4m3",1403 "e5m2",1404 "e3m2",1405 "e2m3",1406 "e2m1",1407 ]:1408 return range(1)1409 if op.a.geom == "m16n8k128" and op.a.mma_type.ptx_type in [1410 "u4",1411 "s4",1412 ]:1413 return range(1)1414 return range(4)1415 1416 1417def common_mma_sp_test_gen(params, op, intrinsic_template, instruction_template):1418 mma_sp_decl_template = """1419declare ${ret_ty} @${intrinsic}(1420 ${args});1421"""1422 1423 mma_sp_test_template = """1424; CHECK-LABEL: .func {{.*}}test_${function}_${selector}(1425define ${ret_ty} @test_${function}_${selector}(1426 ${args}) {1427; CHECK: ${instruction}1428; CHECK-NEXT: ${check_d}1429; CHECK-NEXT: ${check_a}1430; CHECK-NEXT: ${check_b}1431; CHECK-NEXT: ${check_c}1432; CHECK-NEXT: ${check_metadata}1433; CHECK-NEXT: ${check_selector}1434 %r = call ${ret_ty} @${intrinsic}(1435 ${call_args});1436 ret ${ret_ty} %r;1437}1438"""1439 1440 test_params = params1441 test_params["intrinsic"] = (1442 Template(intrinsic_template)1443 .substitute(params)1444 .replace("::", ".")1445 .replace("_", ".")1446 )1447 test_params["function"] = test_params["intrinsic"].replace(".", "_")1448 test_params["instruction"] = Template(instruction_template).substitute(params)1449 test_params["ret_ty"] = make_wmma_ld_ret_ty(op.d)1450 test_params["check_a"] = check_pattern(op.a)1451 test_params["check_b"] = check_pattern(op.b)1452 test_params["check_c"] = check_pattern(op.c)1453 test_params["check_d"] = check_pattern(op.d)1454 test_params["check_metadata"] = "{{%r[0-9]+}}"1455 args = ",\n ".join(1456 list(make_wmma_slice_args(frag) for frag in (op.a, op.b, op.c))1457 + ["i32 %metadata", "i32 %selector"]1458 )1459 test_params["args"] = args1460 1461 print(Template(mma_sp_decl_template).substitute(test_params))1462 1463 for selector in [str(r) for r in sp_selector_gen(op)]:1464 test_params["selector"] = selector1465 test_params["check_selector"] = "{{" + test_params["selector"] + "}}"1466 test_params["call_args"] = test_params["args"].replace(1467 "%selector", test_params["selector"]1468 )1469 1470 print(Template(mma_sp_test_template).substitute(test_params))1471 1472 return (test_params["intrinsic"], test_params["instruction"])1473 1474 1475def gen_mma_sp_tests():1476 if ptx_version < 71 or gpu_arch < 80:1477 return []1478 1479 mma_sp_intrinsic_template = (1480 "llvm.nvvm.mma.${metadata}.${geom}.row.col${kind}${satf}.${intrinsic_signature}"1481 )1482 mma_sp_instruction_template = (1483 "mma.${metadata}.sync.aligned.${geom}.row.col${kind}${satf}.${ptx_signature}"1484 )1485 1486 generated_items = []1487 1488 for op, metadata, kind, satf in product(1489 get_mma_sp_ops(),1490 ["sp::ordered_metadata", "sp"],1491 ["", ".kind::f8f6f4"],1492 [".satfinite", ""],1493 ):1494 if not is_mma_sp_variant_supported(op, metadata, kind, satf):1495 continue1496 1497 params = {1498 "intrinsic_signature": mma_signature(op),1499 "ptx_signature": mma_ptx_signature(op),1500 "satf": satf,1501 "geom": op.a.geom,1502 "metadata": metadata,1503 "kind": kind,1504 }1505 1506 intrinsic_template = mma_sp_intrinsic_template1507 instruction_template = mma_sp_instruction_template1508 1509 generated_items.append(1510 common_mma_sp_test_gen(params, op, intrinsic_template, instruction_template)1511 )1512 1513 return generated_items1514 1515 1516def get_mma_sp_block_scale_ops():1517 return make_mma_ops(["m16n8k128"], ["e2m1"], [], ["f32"], [], True) + make_mma_ops(1518 ["m16n8k64"],1519 ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],1520 ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],1521 ["f32"],1522 [],1523 True,1524 )1525 1526 1527def is_mma_sp_block_scale_geom_supported(geom):1528 # geometries for FP.1529 if geom in [1530 "m16n8k64",1531 "m16n8k128",1532 ]:1533 return True1534 raise ValueError(f"Unexpected sparse MMA block scale geometry: {geom}")1535 1536 1537def is_mma_sp_block_scale_variant_supported(op, kind, scale_vec_size, stype):1538 if not (1539 is_type_supported(op.a.mma_type.ptx_type)1540 and is_mma_sp_block_scale_geom_supported(op.a.geom)1541 ):1542 return False1543 1544 if (1545 op.a.geom == "m16n8k128"1546 and kind == "mxf4"1547 and stype == "ue8m0"1548 and scale_vec_size in ["", ".scale_vec::2X"]1549 ):1550 return True1551 1552 if (1553 op.a.geom == "m16n8k128"1554 and kind == "mxf4nvf4"1555 and stype == "ue8m0"1556 and scale_vec_size == ".scale_vec::2X"1557 ):1558 return True1559 1560 if (1561 op.a.geom == "m16n8k128"1562 and kind == "mxf4nvf4"1563 and stype == "ue4m3"1564 and scale_vec_size == ".scale_vec::4X"1565 ):1566 return True1567 1568 if (1569 op.a.geom == "m16n8k64"1570 and kind == "mxf8f6f4"1571 and stype == "ue8m0"1572 and scale_vec_size in ["", ".scale_vec::1X"]1573 ):1574 return True1575 1576 return False1577 1578 1579def common_mma_sp_block_scale_test_gen(1580 params, op, intrinsic_template, instruction_template1581):1582 mma_sp_block_scale_decl_template = """1583declare ${ret_ty} @${intrinsic}(1584 ${args});1585"""1586 1587 mma_sp_block_scale_test_template = """1588; CHECK-LABEL: .func {{.*}}test_${function}_${selector}(1589define ${ret_ty} @test_${function}_${selector}(1590 ${args}) {1591; CHECK: ${instruction}1592; CHECK-NEXT: ${check_d}1593; CHECK-NEXT: ${check_a}1594; CHECK-NEXT: ${check_b}1595; CHECK-NEXT: ${check_c}1596; CHECK-NEXT: ${check_metadata}1597; CHECK-NEXT: ${check_selector}1598; CHECK-NEXT: ${check_scale_a_data}1599; CHECK-NEXT: ${check_byte_id_a}1600; CHECK-NEXT: ${check_thread_id_a}1601; CHECK-NEXT: ${check_scale_b_data}1602; CHECK-NEXT: ${check_byte_id_b}1603; CHECK-NEXT: ${check_thread_id_b}1604 %r = call ${ret_ty} @${intrinsic}(1605 ${call_args});1606 ret ${ret_ty} %r;1607}1608"""1609 1610 test_params = params1611 test_params["intrinsic"] = Template(intrinsic_template).substitute(params)1612 test_params["function"] = test_params["intrinsic"].replace(".", "_")1613 test_params["instruction"] = Template(instruction_template).substitute(params)1614 test_params["ret_ty"] = make_wmma_ld_ret_ty(op.d)1615 test_params["check_a"] = check_pattern(op.a)1616 test_params["check_b"] = check_pattern(op.b)1617 test_params["check_c"] = check_pattern(op.c)1618 test_params["check_d"] = check_pattern(op.d)1619 test_params["check_metadata"] = "{{%r[0-9]+}}"1620 test_params["check_scale_a_data"] = "{{%r[0-9]+}}"1621 test_params["check_byte_id_a"] = "{{%r[0-9]+}}"1622 test_params["check_thread_id_a"] = "{{%r[0-9]+}}"1623 test_params["check_scale_b_data"] = "{{%r[0-9]+}}"1624 test_params["check_byte_id_b"] = "{{%r[0-9]+}}"1625 test_params["check_thread_id_b"] = "{{%r[0-9]+}}"1626 args = ",\n ".join(1627 list(make_wmma_slice_args(frag) for frag in (op.a, op.b, op.c))1628 + ["i32 %metadata", "i32 %selector"]1629 + ["i32 %scale_a_data", "i16 %byte_id_a, i16 %thread_id_a"]1630 + ["i32 %scale_b_data", "i16 %byte_id_b, i16 %thread_id_b"]1631 )1632 test_params["args"] = args1633 1634 print(Template(mma_sp_block_scale_decl_template).substitute(test_params))1635 1636 for selector in [str(r) for r in sp_selector_gen(op, True)]:1637 test_params["selector"] = selector1638 test_params["check_selector"] = "{{" + test_params["selector"] + "}}"1639 test_params["call_args"] = test_params["args"].replace(1640 "%selector", test_params["selector"]1641 )1642 1643 print(Template(mma_sp_block_scale_test_template).substitute(test_params))1644 1645 return (test_params["intrinsic"], test_params["instruction"])1646 1647 1648def gen_mma_sp_block_scale_tests():1649 if not (ptx_version >= 88 and gpu_arch >= 120 and aa):1650 return []1651 1652 mma_sp_block_scale_intrinsic_template = "llvm.nvvm.mma.sp.ordered.metadata.block.scale.${geom}.row.col.${kind}${scale}.${intrinsic_signature}.${stype}"1653 mma_sp_block_scale_instruction_template = "mma.sp::ordered_metadata.sync.aligned.${geom}.row.col.kind::${kind}.block_scale${scale_vec_size}.${ptx_signature}.${stype}"1654 1655 generated_items = []1656 1657 for op, kind, scale_vec_size, stype in product(1658 get_mma_sp_block_scale_ops(),1659 ["mxf4", "mxf4nvf4", "mxf8f6f4"],1660 ["", ".scale_vec::1X", ".scale_vec::2X", ".scale_vec::4X"],1661 ["ue8m0", "ue4m3"],1662 ):1663 if not is_mma_sp_block_scale_variant_supported(op, kind, scale_vec_size, stype):1664 continue1665 1666 params = {1667 "intrinsic_signature": mma_signature(op),1668 "ptx_signature": mma_ptx_signature(op),1669 "geom": op.a.geom,1670 "kind": kind,1671 "scale_vec_size": scale_vec_size,1672 "scale": scale_vec_size.replace("_vec::", ".").lower(),1673 "stype": stype,1674 }1675 1676 intrinsic_template = mma_sp_block_scale_intrinsic_template1677 instruction_template = mma_sp_block_scale_instruction_template1678 1679 generated_items.append(1680 common_mma_sp_block_scale_test_gen(1681 params, op, intrinsic_template, instruction_template1682 )1683 )1684 1685 return generated_items1686 1687 1688# Append complete list of intrinsics and instructions we've generated tests for.1689# Generate set of checks to verify that that we did generate sensible set of1690# tests for the given combination of PTX and SM variants.1691#1692def gen_check_unsupported_ops(items):1693 print(1694 "; Complete list of intrinsics supported by PTX%d on sm_%d"1695 % (ptx_version, gpu_arch)1696 )1697 print("; INTRINSICS: {{^; INTRINSICS_LIST_BEGIN}}")1698 print(1699 """1700 1701; NOEXTGEOM-NOT: {{m8n32|m32n8}}1702; NOINT-NOT: .{{s32|s8}}1703; NOSUBINT-NOT: {{s4|u4|b1}}1704; NOMMA-NOT: .m8n8k4.1705; NOALTFLOAT-NOT: .{{bf16|tf32}}1706; NODOUBLE-NOT: .f641707; NOLDMATRIX-NOT: ldmatrix.sync.aligned1708; NOSTMATRIX-NOT: stmatrix.sync.aligned1709 1710; M16N16-DAG: m16n16k16.load.{{[ab].*}}.f16.p1711; M16N16-DAG: m16n16k16.{{load|store}}.{{[cd].*\.(f16|f32)}}.p1712; M16N16-DAG: m16n16k16.mma.{{.*}}.f16.f321713; M16N16-DAG: m16n16k16.mma.{{.*}}.f32.f161714; M16N16-DAG: m16n16k16.mma.{{.*}}.f16.f161715; M16N16-DAG: m16n16k16.mma.{{.*}}.f32.f321716 1717; PTX60 adds support for m32n8k16/m8n32k16 geometries.1718; EXTGEOM-DAG: m32n8k16.load.{{[ab].*}}.f16.p1719; EXTGEOM-DAG: m32n8k16.{{load|store}}.{{[cd].*\.(f16|f32)}}.p1720; EXTGEOM-DAG: m32n8k16.mma.{{.*}}.f16.f321721; EXTGEOM-DAG: m32n8k16.mma.{{.*}}.f32.f161722; EXTGEOM-DAG: m32n8k16.mma.{{.*}}.f16.f161723; EXTGEOM-DAG: m32n8k16.mma.{{.*}}.f32.f321724 1725; EXTGEOM-DAG: m8n32k16.load.{{[ab].*}}.f16.p1726; EXTGEOM-DAG: m8n32k16.{{load|store}}.{{[cd].*\.(f16|f32)}}.p1727; EXTGEOM-DAG: m8n32k16.mma.{{.*}}.f16.f321728; EXTGEOM-DAG: m8n32k16.mma.{{.*}}.f32.f161729; EXTGEOM-DAG: m8n32k16.mma.{{.*}}.f16.f161730; EXTGEOM-DAG: m8n32k16.mma.{{.*}}.f32.f321731 1732; INT-DAG: m16n16k16.load.{{[ab].*}}.s8.p1733; INT-DAG: m8n32k16.load.{{[ab].*}}.s8.p1734; INT-DAG: m32n8k16.load.{{[ab].*}}.s8.p1735; INT-DAG: m16n16k16.load.{{[ab].*}}.u8.p1736; INT-DAG: m8n32k16.load.{{[ab].*}}.u8.p1737; INT-DAG: m32n8k16.load.{{[ab].*}}.u8.p1738; INT-DAG: m32n8k16.{{load|store}}.{{[cd].*\.s32}}.p1739; INT-DAG: m16n16k16.mma.{{.*}}.u81740; INT-DAG: m16n16k16.mma.{{.*}}.s81741; INT-DAG: m8n32k16.mma.{{.*}}.u81742; INT-DAG: m8n32k16.mma.{{.*}}.s81743; INT-DAG: m32n8k16.mma.{{.*}}.u81744; INT-DAG: m32n8k16.mma.{{.*}}.s81745 1746; SUBINT-DAG: m8n8k128.load.{{[ab].*}}.b1.p1747; SUBINT-DAG: m8n8k32.load.{{[ab].*}}.s4.p1748; SUBINT-DAG: m8n8k32.load.{{[ab].*}}.u4.p1749; SUBINT-DAG: m8n8k128.{{load|store}}.{{[cd].*\.s32}}.p1750; SUBINT-DAG: m8n8k32.{{load|store}}.{{[cd].*\.s32}}.p1751; SUBINT-DAG: m8n8k32.mma.{{.*}}.u41752; SUBINT-DAG: m8n8k32.mma.{{.*}}.s41753; SUBINT-DAG: m8n8k128.mma.{{.*}}.b11754 1755; ALTFLOAT-DAG: m16n16k16.load.{{[ab].*}}.bf16.p1756; ALTFLOAT-DAG: m8n32k16.load.{{[ab].*}}.bf16.p1757; ALTFLOAT-DAG: m32n8k16.load.{{[ab].*}}.bf16.p1758; ALTFLOAT-DAG: m16n16k8.load.{{[ab].*}}.tf32.p1759; ALTFLOAT-DAG: m16n16k16.mma.{{.*}}.bf161760; ALTFLOAT-DAG: m8n32k16.mma.{{.*}}.bf161761; ALTFLOAT-DAG: m32n8k16.mma.{{.*}}.bf161762; ALTFLOAT-DAG: m16n16k8.mma.{{.*}}.tf321763 1764; DOUBLE-DAG: m8n8k4.load.{{[abc].*}}.f64.p1765; DOUBLE-DAG: m8n8k4.store.d.{{.*}}.f64.p1766; DOUBLE-DAG: m8n8k4.mma.{{.*}}.f641767 1768; MMA-DAG: mma.m8n8k4.{{.*}}.f16.f321769; MMA-DAG: mma.m8n8k4.{{.*}}.f32.f161770; MMA-DAG: mma.m8n8k4.{{.*}}.f16.f161771; MMA-DAG: mma.m8n8k4.{{.*}}.f32.f321772 1773; PTX65MMA-DAG: mma.m16n8k8.row.col.f16.f161774; PTX65MMA-DAG: mma.m16n8k8.row.col.f32.f321775; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.u8.u81776; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.s8.s81777; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.s8.u81778; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.u8.s81779; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.u4.u41780; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.s4.s41781; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.s4.u41782; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.u4.s41783 1784; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x1.b161785; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x2.b161786; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x4.b161787; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x1.trans.b161788; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x2.trans.b161789; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x4.trans.b161790; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x1.shared.b161791; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x2.shared.b161792; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x4.shared.b161793; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x1.trans.shared.b161794; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x2.trans.shared.b161795; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x4.trans.shared.b161796 1797; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m16n16.x1.trans.shared.b81798; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m16n16.x2.trans.shared.b81799; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m16n16.x1.trans.b8x16.b6x16_p321800; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m16n16.x1.trans.b8x16.b4x16_p641801; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m16n16.x2.trans.b8x16.b6x16_p321802; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m16n16.x2.trans.b8x16.b4x16_p641803; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m8n16.x1.b8x16.b6x16_p321804; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m8n16.x1.b8x16.b4x16_p641805; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m8n16.x2.b8x16.b6x16_p321806; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m8n16.x2.b8x16.b4x16_p641807; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m8n16.x4.b8x16.b6x16_p321808; PTX86LDMATRIX-DAG: ldmatrix.sync.aligned.m8n16.x4.b8x16.b4x16_p641809 1810; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x1.b161811; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x2.b161812; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x4.b161813; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x1.trans.b161814; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x2.trans.b161815; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x4.trans.b161816; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x1.shared.b161817; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x2.shared.b161818; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x4.shared.b161819; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x1.trans.shared.b161820; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x2.trans.shared.b161821; PTX78STMATRIX-DAG: stmatrix.sync.aligned.m8n8.x4.trans.shared.b161822 1823; PTX86STMATRIX-DAG: stmatrix.sync.aligned.m16n8.x1.trans.b81824; PTX86STMATRIX-DAG: stmatrix.sync.aligned.m16n8.x2.trans.b81825; PTX86STMATRIX-DAG: stmatrix.sync.aligned.m16n8.x4.trans.b81826; PTX86STMATRIX-DAG: stmatrix.sync.aligned.m16n8.x1.trans.shared.b81827; PTX86STMATRIX-DAG: stmatrix.sync.aligned.m16n8.x2.trans.shared.b81828; PTX86STMATRIX-DAG: stmatrix.sync.aligned.m16n8.x4.trans.shared.b81829 1830; PTX71MMA-DAG: mma.m8n8k4.row.col.f641831; PTX71MMA-DAG: mma.m16n8k4.row.col.tf321832; PTX71MMA-DAG: mma.m16n8k8.row.col.tf321833; PTX71MMA-DAG: mma.m16n8k16.row.col.bf161834; PTX71MMA-DAG: mma.m16n8k8.row.col.bf161835; PTX71MMA-DAG: mma.m16n8k16.row.col.f16.f161836; PTX71MMA-DAG: mma.m16n8k16.row.col.f32.f321837; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.u8.u81838; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.s8.s81839; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.s8.u81840; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.u8.s81841; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u8.u81842; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s8.s81843; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s8.u81844; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u8.s81845; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u4.u41846; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s4.s41847; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s4.u41848; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u4.s41849; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.u4.u41850; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.s4.s41851; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.s4.u41852; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.u4.s41853; PTX71MMA-DAG: mma.and.popc.m8n8k128.row.col.b11854; PTX71MMA-DAG: mma.xor.popc.m8n8k128.row.col.b11855; PTX71MMA-DAG: mma.and.popc.m16n8k128.row.col.b11856; PTX71MMA-DAG: mma.xor.popc.m16n8k128.row.col.b11857; PTX71MMA-DAG: mma.and.popc.m16n8k256.row.col.b11858; PTX71MMA-DAG: mma.xor.popc.m16n8k256.row.col.b11859;1860 1861"""1862 )1863 1864 print("; INTRINSICS_LIST_BEGIN")1865 for intrinsic, instruction in sorted(items):1866 print("; ", intrinsic, " -> ", instruction, "")1867 print("; INTRINSICS_LIST_END")1868 print("; INTRINSICS: ; INTRINSICS_LIST_END")1869 1870 1871def gen_tests():1872 items = gen_wmma_load_tests()1873 items += gen_wmma_store_tests()1874 items += gen_ldmatrix_tests()1875 items += gen_stmatrix_tests()1876 items += gen_wmma_mma_tests()1877 items += gen_mma_tests()1878 items += gen_mma_block_scale_tests()1879 items += gen_mma_sp_tests()1880 items += gen_mma_sp_block_scale_tests()1881 gen_check_unsupported_ops(items)1882 1883 1884def main():1885 global ptx_version1886 global gpu_arch1887 global aa1888 parser = argparse.ArgumentParser()1889 parser.add_argument("--ptx", type=int, default=60)1890 parser.add_argument("--gpu-arch", type=int, default=70)1891 parser.add_argument("--aa", action="store_true")1892 args = parser.parse_args()1893 1894 ptx_version = args.ptx1895 gpu_arch = args.gpu_arch1896 aa = args.aa1897 1898 gen_tests()1899 1900 1901if __name__ == "__main__":1902 main()1903