#!/usr/bin/perl
# Copyright (c) 2020 - 2021 Advanced Micro Devices, Inc. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

use strict;
use File::Copy;
use File::Spec;
use File::Basename;
use File::Which;
use Cwd 'realpath';
use Getopt::Std;
use List::Util qw(max);
use URI::Encode;

sub usage {
  print("Usage: $0 [-v|h] executable...\n");
  print("List the URIs of the code objects embedded in the specfied host executables.\n");
  print("-v \tVerbose output (includes Entry ID)\n");
  print("-h \tShow this help message\n");
  exit;
}

# sub to read a qword. 1st arg is a FP, 2nd arg is ref to destination var.
sub readq {
 my ($input_fp, $qword) = @_;
 read($input_fp, my $bytes, 8) == 8 or die("Error: Failed to read 8 bytes\n");
 ${$qword} = unpack("Q<", $bytes);
}

# sub to move address to next alignment boundary
#  first arg is address to move
#  second arg is alignment requirement/boundary
sub align_up {
	my ($address, $alignment) = @_;
	$address = int(($address + ($alignment - 1)) / $alignment) * $alignment;
}

# Process options
my %options=();
getopts('vhd', \%options);

if (defined $options{h}) {
  usage();
}

my $verbose = $options{v};
my $debug = $options{d};

my $num_bundles = 1;
my $bundle_alignment = 4096;

# look for objdump
my $objdump = which("objdump");
(-f $objdump) || die("Error: Can't find objdump command\n");

# for each argument (which should be an executable):
foreach my $executable_file(@ARGV) {

  # debug message
  print("Reading input file \"$executable_file\" ...\n") if ($debug);

  # verify/open file specified.
  open (INPUT_FP, "<", $executable_file) || die("Error: failed to open file: $executable_file\n");
  binmode INPUT_FP;

  # kernel section information
  my $escaped_name=quotemeta($executable_file);
  my $bundle_section_name = ".hip_fatbin";
  my $bundle_section_size = hex(`$objdump -h $escaped_name | grep $bundle_section_name | awk '{print \$3}'`);
  my $bundle_section_offset =  hex(`$objdump -h $escaped_name | grep $bundle_section_name | awk '{print \$6}'`);

  $bundle_section_size or die("Error: No kernel section found\n");

  my $bundle_section_end = $bundle_section_offset + $bundle_section_size;

  if ($debug) {
    printf("Code Objects Bundle section size: %x\n",$bundle_section_size);
    printf("Code Objects Bundle section offset: %x\n",$bundle_section_offset);
    printf("Code Objects Bundle section end: %x\n\n",$bundle_section_end);
  }

  my $current_bundle_offset = $bundle_section_offset;
  printf("Current Bundle offset: 0x%X\n",$current_bundle_offset) if ($debug);

  # move fp to current_bundle_offset.
  seek(INPUT_FP, $current_bundle_offset, 0);

  while ($current_bundle_offset < $bundle_section_end) {

    # skip OFFLOAD_BUNDLER_MAGIC_STR
    my $magic_str;
    my $read_bytes = read(INPUT_FP, $magic_str, 24);
    if (($read_bytes != 24) || ($magic_str ne "__CLANG_OFFLOAD_BUNDLE__")) {
      print(STDERR "Error: Offload bundle magic string not detected\n") if ($debug);
    }

    # read number of bundle entries, which are code objects.
    my $num_codeobjects;
    readq(\*INPUT_FP,\$num_codeobjects);

    # header with current bundle number and number of embedded code objcts in that bundle.
    #    print("Bundle Number: $num_bundles with $num_codeobjects Code Objects:\n") if ($very_verbose);

    my $end_of_current_bundle = $current_bundle_offset;

    # Column Header.
    printf("%-8s%-40s%35s\n","Bundle#","Entry ID:","URI:") if ($verbose);

    # for each Bundle entry (code object)  ....
    for (my $iter = 0; $iter < $num_codeobjects; $iter++) {

      print("\nEntry #$iter\n") if $debug;

      # read bundle entry (code object) offset
      my $entry_offset;
      my $abs_offset;
      readq(*INPUT_FP,\$entry_offset);
      printf("entry_offset: 0x%X\n",$entry_offset) if $debug;

      # read bundle entry (code object) size
      my $entry_size;
      readq(*INPUT_FP,\$entry_size);
      printf("entry_size: 0x%X\n",$entry_size) if $debug;

      # read triple size
      my $triple_size;
      readq(*INPUT_FP,\$triple_size);
      printf("triple_size: 0x%X\n",$triple_size) if $debug;

      # read triple string
      my $triple;
      my $read_bytes = read(INPUT_FP, $triple, $triple_size);
      $read_bytes == $triple_size or die("Error: Fail to parse triple\n");
      print("triple: $triple\n") if $debug;

      # because the bundle entry's offset is relative to the beginning of the bundled code object section.
      $abs_offset = int($current_bundle_offset + $entry_offset); 

      # and we need to keep track of where we are in the current bundle.
      $end_of_current_bundle = int($abs_offset + $entry_size);

      printf("abs_offset: 0x%X\n",$abs_offset) if $debug;

      my $obj_uri_encode = URI::Encode->new();
      my $encoded_executable_file = $obj_uri_encode->encode($executable_file);

      printf("%-8s%-40s%35s%s%s%s%s%s%s\n",$num_bundles,$triple,"file:\/\/",$encoded_executable_file,"\#offset=",$abs_offset, "\&size=",$entry_size);

      printf("end_of_current_bundle: 0x%X\n",$end_of_current_bundle) if $debug;
      printf("Hex values: file:\/\/$encoded_executable_file#offset=0x%X$abs_offset\&size=0x%X\n", $abs_offset, $entry_size) if $debug;

    } # End of for each Bundle entry (code object) ...

    printf("\n") if ($verbose);

    # we've finished listing this current bundle ...
    printf("current_bundle_offset: %x \n",$current_bundle_offset) if ($debug);
    printf("bundle_section_end: %x \n", $bundle_section_end) if ($debug);

    # move current_bundle_offset to next alignment boundary.
    $current_bundle_offset = align_up($end_of_current_bundle,$bundle_alignment);
    printf("Adjusting for alignment of next bundle: current_bundle_offset: %x \n\n\n", $current_bundle_offset) if ($debug);

    # seek to the end of the current bundle:
    seek(INPUT_FP, $current_bundle_offset, 0);

    # increment the number of bundles listed.
    $num_bundles = $num_bundles+1;

  }  # End of while loop

} # End of for each command line argument

exit(0);
