python3 find_def_use.py -f a.s -v v7 -l 634 -d

find_def_use.py

import argparse
import re

pattern = re.compile(r"^v((\d+)|\[(\d+):(\d+)\])$")


def parse_instr(instr: str):
    if not instr.startswith(('v_', 's_', 'buffer_', 'scratch_', 'tensor_', 'ds_', 'global_')):
        # skip non-vector instr
        return None, None
    instr_name = [x.strip() for x in instr.split()][0]
    line_split = [x.strip() for x in instr[len(instr_name):].split(',')]
    if len(line_split) <= 1:
        return None, None
    regs = []
    for r in line_split:
        if not r.startswith('v'):
            continue
        if '/*' in r and '*/' in r:
            r = r.split('/*')[-1].split('*/')[0]
        regs.append(r)
    return instr_name, regs

def split_instr(instr: str):
    if '::' in instr:
        return [x.strip() for x in instr.split('::')]
    else:
        return [instr]

def has_reg(reg_no: int, target: str):
    m = pattern.match(target)
    if not m:
        return False
    if m.group(2):
        # v123
        if reg_no == int(m.group(2)):
            return True
    else:
        # v[1:2]
        lb = int(m.group(3))
        ub = int(m.group(4))
        if lb <= reg_no <= ub:
            return True

def reg0_is_not_used(instr: str):
    return instr.startswith(('scratch_store_', 'global_load_async_to_lds_', 'tensor_load_to_lds'))

def main(args):
    with open(args.filename, 'r') as f:
        asm = [x.strip() for x in f.readlines()]
    v: str = args.value
    if not v.startswith('v'):
        raise ValueError(f"value should start with v, but got {v}")
    
    start_line = int(args.start_line)
    if start_line > len(asm):
        raise ValueError(f"Start line number {start_line} out of bound")
    vn = int(v[1:])

    start_instr = asm[start_line - 1]
    instr, regs = parse_instr(start_instr)
    assert instr, f"Can't parse instr {start_instr}"
    found = False
    for r in regs:
        if has_reg(vn, r):
            found = True
            break
    if not found:
        raise ValueError(f"Reg {v} is not in starting line {start_instr}")

    print(f'Start at line {start_line}: {start_instr}')

    if args.find_def:
        # reverse to find 
        for i in range(start_line - 2, 0, -1):
            line = asm[i]

            instrs = split_instr(line)
            for instr in instrs:
                instr_name, regs = parse_instr(instr)
                if not instr_name:
                    continue

                if len(regs) == 0:
                    continue
                if has_reg(vn, regs[0]) and not reg0_is_not_used(instr_name):
                    print(f'Found def at {i + 1}: {line}')
                    return
    else:
        for i in range(start_line, len(asm)):
            line = asm[i]

            instrs = split_instr(line)
            for instr in instrs:
                instr_name, regs = parse_instr(instr)
                if not instr_name:
                    continue
                for ele_idx, ele in enumerate(regs):
                    if has_reg(vn, ele):
                        if ele_idx == 0 and not reg0_is_not_used(instr_name):
                            print(f'Found redefinition at line {i + 1}: {line}, exit')
                            return
                        else:
                            print(f'{i + 1}: {line}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--filename', type=str)
    parser.add_argument('-v', '--value', type=str)
    parser.add_argument('-l', '--start_line', type=str)
    parser.add_argument('-d', '--find_def', action='store_true')
    args = parser.parse_args()
    main(args)

How to use

Find immediate def

python3 find_def_use.py -f <filename> -v <reg> -l <lino> -d
# example
python3 find_def_use.py -f a.s -v v7 -l 2306 -d

Find immediate use until modified

python3 find_def_use.py -f <filename> -v <reg> -l <lino>
# example
python3 find_def_use.py -f a.s -v v7 -l 2306