1#!/usr/bin/env python 2 3import subprocess 4import os.path 5import yaml 6import io 7import re 8 9def parseKernelUsages(usageStr, usageDict): 10 demangler = 'c++filt -p' 11 12 def getKernelMem(usages): 13 match = re.search(r"([0-9]+) bytes cmem\[0\]", usages) 14 return match.group(1) if match else None 15 def getSharedMem(usages): 16 match = re.search(r"([0-9]+) bytes smem", usages) 17 return match.group(1) if match else None 18 def getRegisters(usages): 19 match = re.search(r"[Uu]sed ([0-9]+) registers", usages) 20 return match.group(1) if match else None 21 def demangle(fn): 22 expr = re.compile("__omp_offloading_[a-zA-Z0-9]*_[a-zA-Z0-9]*_(_Z.*_)_l[0-9]*$") 23 match = expr.search(fn) 24 function = match.group(1) if match else fn 25 output = subprocess.run(demangler.split(' ') + [function], check=True, stdout=subprocess.PIPE) 26 return output.stdout.decode('utf-8').strip() 27 def getLine(fn): 28 expr = re.compile("__omp_offloading_[a-zA-Z0-9]*_[a-zA-Z0-9]*_.*_l([0-9]*)$") 29 match = expr.search(fn) 30 return match.group(1) if match else 0 31 32 expr = re.compile("Function properties for \'?([a-zA-Z0-9_]*)\'?\n(.*,.*)\n") 33 for (fn, usages) in expr.findall(usageStr): 34 info = usageDict[fn] if fn in usageDict else dict() 35 info["Name"] = demangle(fn) 36 info["DebugLoc"] = {"File" : "unknown", "Line": getLine(fn), "Column" : 0} 37 info["Usage"] = {"Registers" : getRegisters(usages), "Shared" : getSharedMem(usages), "Kernel" : getKernelMem(usages)} 38 usageDict[fn] = info 39 40def getKernelUsage(stderr, fname='usage.yaml'): 41 remarks = [line for line in stderr.split('\n') if re.search(r"^remark:", line)] 42 ptxas = '\n'.join([line.split(':')[1].strip() for line in stderr.split('\n') if re.search(r"^ptxas info *:", line)]) 43 nvlink = '\n'.join([line.split(':')[1].strip() for line in stderr.split('\n') if re.search(r"^nvlink info *:", line)]) 44 45 if os.path.exists(fname): 46 with io.open(fname, 'r', encoding = 'utf-8') as f: 47 usage = yaml.load(f, Loader=yaml.Loader) 48 else: 49 usage = dict() 50 51 parseKernelUsages(ptxas, usage) 52 parseKernelUsages(nvlink, usage) 53 54 return usage 55