#!/usr/bin/env python3


# Usage  : nvdisasm [options] <image file>

# Options
# =======

# --base-address <value>                         (-base)
#         Specify the logical base address of the image to disassemble. This option
#         is only valid when disassembling a raw instruction binary (see option '--binary'),
#         and is ignored when disassembling an Elf file.
#         Default value:  0x0.

# --binary <SMxy>                                (-b)
#         When this option is specified, the input file is assumed to contain a raw
#         instruction binary, that is, a sequence of binary instruction encodings as
#         they occur in instruction memory. The value of this option must be the asserted
#         architecture of the raw binary.
#         Allowed values for this option:  'SM35','SM37','SM50','SM52','SM53','SM60',
#         'SM61','SM62','SM70','SM72','SM75','SM80','SM86','SM87','SM89','SM90'.

# --cuda-function-index <symbol index>,...       (-fun)
#         Restrict the output to the CUDA functions represented by symbols with the
#         given indices. The CUDA function for a given symbol is the enclosing section.
#         This only restricts executable sections; all other sections will still be
#         printed.

# --help                                         (-h)
#         Print this help information on this tool.

# --life-range-mode                              (-lrm)
#         This option implies option --print-life-ranges, and determines how register
#         live range info should be printed:
#         - count         : Not at all, leaving only the # column (number of live registers)
#         - wide          : Columns spaced out for readability (default)
#         - narrow        : A one-character column for each register, economizing on
#         table width
#         Allowed values for this option:  'count','narrow','wide'.

# --no-dataflow                                  (-ndf)
#         Disable dataflow analyzer after disassembly. Dataflow analysis is normally
#         enabled to perform branch stack analysis and annotate all instructions that
#         jump via the GPU branch stack with inferred branch target labels. However,
#         it may occasionally fail when certain restrictions on the input nvelf/cubin
#         are not met.

# --no-vliw                                      (-novliw)
#         Conventional mode; disassemble paired instructions in normal syntax, instead
#         of VLIW syntax.

# --options-file <file>,...                      (-optf)
#         Include command line options from specified file.

# --output-control-flow-graph                    (-cfg)
#         When specified output the control flow graph, where each node is a hyperblock,
#         in a format consumable by graphviz tools (such as dot).

# --output-control-flow-graph-with-basic-blocks  (-bbcfg)
#         When specified output the control flow graph, where each node is a basicblock,
#         in a format consumable by graphviz tools (such as dot).

# --print-code                                   (-c)
#         Only print code sections.

# --print-instr-offsets-cfg                      (-poff)
#         When specified, print instruction offsets in the control flow graph. This
#         should be used along with the option --output-control-flow-graph or --output-control-flow-graph-with-basic-blocks.

# --print-instruction-encoding                   (-hex)
#         When specified, print the encoding bytes after each disassembled operation.

# --print-life-ranges                            (-plr)
#         Print register life range information in a trailing column in the produced
#         disassembly.

# --print-line-info                              (-g)
#         Annotate disassembly with source line information obtained from .debug_line
#         section, if present.

# --print-line-info-inline                       (-gi)
#         Annotate disassembly with source line information obtained from .debug_line
#         section along with function inlining info, if present.

# --print-line-info-ptx                          (-gp)
#         Annotate disassembly with source line information obtained from .nv_debug_line_sass
#         section, if present.

# --print-raw                                    (-raw)
#         Print the disassembly without any attempt to beautify it.

# --separate-functions                           (-sf)
#         Separate the code corresponding with function symbols by some new lines to
#         let them stand out in the printed disassembly.

# --version                                      (-V)
#         Print version information on this tool.

import argparse
import os
import subprocess
import datetime

parser = argparse.ArgumentParser(description='nvdisasm extracts information from standalone cubin files and presents them in human readable format.')

group = parser.add_mutually_exclusive_group()
group.add_argument('--version', '-V', action='store_true', help='Display version information')
group.add_argument('-hex', '--print-instruction-encoding', action='store_true', help='When specified, print the encoding bytes after each disassembled operation.')
group.add_argument('-g', '--print-line-info', action='store_true', help='Annotate disassembly with source line information obtained from .debug_line section, if present.')
group.add_argument('-gi', '--print-line-info-inline', action='store_true', help='Annotate disassembly with source line information obtained from .debug_line section along with function inlining info, if present.')
group.add_argument('-gp', '--print-line-info-ptx', action='store_true', help='Annotate disassembly with source line information obtained from .nv_debug_line_sass section, if present.')
group.add_argument('-raw', '--print-raw', action='store_true', help='Print the disassembly without any attempt to beautify it.')
# group.add_argument('-plr', '--print-life-ranges', action='store_true', help='Print register life range information in a trailing column in the produced disassembly.')
# group.add_argument('-poff', '--print-instr-offsets-cfg', action='store_true', help='When specified, print instruction offsets in the control flow graph. This should be used along with the option --output-control-flow-graph or --output-control-flow-graph-with-basic-blocks.')
# group.add_argument('-c', '--print-code', action='store_true', help='Only print code sections.')
# group.add_argument('-b', '--binary', metavar='<SMxy>', help='When this option is specified, the input file is assumed to contain a raw instruction binary, that is, a sequence of binary instruction encodings as they occur in instruction memory. The value of this option must be the asserted architecture of the raw binary. Allowed values for this option:  SM35,SM37,SM50,SM52,SM53,SM60,SM61,SM62,SM70,SM72,SM75,SM80,SM86,SM87,SM89,SM90.')
# group.add_argument('-novliw', '--no-vliw', action='store_true', help='Conventional mode; disassemble paired instructions in normal syntax, instead of VLIW syntax.')
# group.add_argument('-ndf', '--no-dataflow', action='store_true', help='Disable dataflow analyzer after disassembly. Dataflow analysis is normally enabled to perform branch stack analysis and annotate all instructions that jump via the GPU branch stack with inferred branch target labels. However, it may occasionally fail when certain restrictions on the input nvelf/cubin are not met.')
# group.add_argument('-sf', '--separate-functions', action='store_true', help='Separate the code corresponding with function symbols by some new lines to let them stand out in the printed disassembly.')
# group.add_argument('-cfg', '--output-control-flow-graph', action='store_true', help='When specified output the control flow graph, where each node is a hyperblock, in a format consumable by graphviz tools (such as dot).')
# group.add_argument('-bbcfg', '--output-control-flow-graph-with-basic-blocks', action='store_true', help='When specified output the control flow graph, where each node is a basicblock, in a format consumable by graphviz tools (such as dot).')
# group.add_argument('-lrm', '--life-range-mode', metavar='<mode>', help='This option implies option --print-life-ranges, and determines how register live range info should be printed: count : Not at all, leaving only the # column (number of live registers); wide : Columns spaced out for readability (default); narrow : A one-character column for each register, economizing on table width. Allowed values for this option:  count,narrow,wide.')
# group.add_argument('-optf', '--options-file', metavar='<file>', help='Include command line options from specified file.')
# group.add_argument('-base', '--base-address', metavar='<value>', help='Specify the logical base address of the image to disassemble. This option is only valid when disassembling a raw instruction binary (see option --binary), and is ignored when disassembling an Elf file. Default value:  0x0.')


parser.add_argument('infile', metavar='<input cubin file>', type=str, help='Specify the input cubin file')

run_on_shell=True
DEBUG = False

def output_cmd(cmd):
    if DEBUG:
        print("cmd: {}".format(cmd))

def check_ROCM_PATH():
    if 'ROCM_PATH' in os.environ:
        return os.environ['ROCM_PATH']
    else:
        exit('Error: ROCM_PATH is not set')

def exec_cmd(cmd):
    output_cmd(cmd)
    status = subprocess.run(cmd, shell=run_on_shell, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding='utf-8')
    if status.returncode != 0:
        print("Error: {}\n{}".format(cmd, status.stderr))
        exit(1)
    return status.stdout

def check_device_file(input):
    cmd = "readelf -h {}".format(input)
    output = exec_cmd(cmd)
    lines = output.split('\n')
    for line in lines:
        if line.strip().startswith('Machine:'):
            if 'GPU' in line:
                return True
    return False

def print_realelf_all_info(input):
    cmd = "readelf -a {}".format(input)
    output = exec_cmd(cmd)
    print(output)

def main():
    args = parser.parse_args()

    if args.version:
        print('nvdisasm: version 1.0')
        return

    input = args.infile
    if not check_device_file(input):
        print("nvdisasm fatal : {} is not a supported Elf file".format(input))
        exit(0)

    if args.print_raw or args.print_instruction_encoding or args.print_line_info or args.print_line_info_inline or args.print_line_info_ptx:
        print_realelf_all_info(input)
        return
    
#  ./nvdisasm --help
#  ./nvdisasm -hex hipv4-amdgcn-amd-amdhsa--gfx906

main()