rocm_agent_enumerator

#!/usr/bin/env python3

import os
import re
import subprocess
import sys

# get current working directory
CWD = os.path.dirname(os.path.realpath(__file__))

ISA_TO_ID = {
  # Kaveri - Temporary
  "gfx700" : [0x130f],
  # Hawaii
  "gfx701" : [0x67a0, 0x67a1, 0x67a2, 0x67a8, 0x67a9, 0x67aa, 0x67b0, 0x67b1,
              0x67b8, 0x67b9, 0x67ba, 0x67be],
  # Carrizo
  "gfx801" : [0x9870, 0x9874, 0x9875, 0x9876, 0x9877, 0x98e4],
  # Tonga
  "gfx802" : [0x6920, 0x6921, 0x6928, 0x6929, 0x692b, 0x692f, 0x6930, 0x6938,
              0x6939],
  # Fiji
  "gfx803" : [0x7300, 0x730f,
  # Polaris10
              0x67c0, 0x67c1, 0x67c2, 0x67c4, 0x67c7, 0x67c8, 0x67c9, 0x67ca,
              0x67cc, 0x67cf,
  # Polaris11
              0x67d0, 0x67df, 0x67e0, 0x67e1, 0x67e3, 0x67e7, 0x67e8, 0x67e9,
              0x67eb, 0x67ef, 0x67ff,
  # Polaris12
              0x6980, 0x6981, 0x6985, 0x6986, 0x6987, 0x6995, 0x6997, 0x699f],
  # Vega10
  "gfx900" : [0x6860, 0x6861, 0x6862, 0x6863, 0x6864, 0x6867, 0x6868, 0x686c,
              0x687f]
}

def staticVars(**kwargs):
  def deco(func):
    for k in kwargs:
      setattr(func, k, kwargs[k])
    return func
  return deco

@staticVars(search_term=re.compile("gfx[0-9a-fA-F]+"))
def getGCNISA(line, match_from_beginning = False):
 if match_from_beginning is True:
   result = getGCNISA.search_term.match(line)
 else:
   result = getGCNISA.search_term.search(line)

 if result is not None:
   return result.group(0)
 return None

@staticVars(search_name=re.compile("gfx[0-9a-fA-F]+:[a-zA-Z]+-"))
def getGCNArchName(line):
 result = getGCNArchName.search_name.search(line)

 if result is not None:
   return result.group(0)
 return None

def readFromTargetLstFile():
  target_list = []

  # locate target.lst using environment variable or
  # it should be placed at the same directory with this script
  target_lst_path = os.environ.get("ROCM_TARGET_LST");
  if target_lst_path == None:
    target_lst_path = os.path.join(CWD, "target.lst")
  if os.path.isfile(target_lst_path):
    target_lst_file = open(target_lst_path, 'r')
    for line in target_lst_file:
      # for target.lst match from beginning so targets can be disabled by
      # commenting it out
      target = getGCNISA(line, match_from_beginning = True)
      if target is not None:
        target_list.append(target)

  return target_list

def readFromROCMINFO(search_arch_name = False):
  target_list = []
  # locate rocminfo binary which should be placed at the same directory with
  # this script
  rocminfo_executable = os.path.join(CWD, "rocminfo")

  try:
    # run rocminfo
    rocminfo_output = subprocess.Popen(rocminfo_executable, stdout=subprocess.PIPE).communicate()[0].decode("utf-8").split('\n')
  except:
    rocminfo_output = []

  # search AMDGCN gfx ISA
  if search_arch_name is True:
    line_search_term = re.compile("\A\s+Name:\s+(amdgcn-amd-amdhsa--gfx\d+)")
  else:
    line_search_term = re.compile("\A\s+Name:\s+(gfx\d+)")
  for line in rocminfo_output:
    if line_search_term.match(line) is not None:
      if search_arch_name is True:
        target = getGCNArchName(line)
      else:
        target = getGCNISA(line)
      if target is not None:
        target_list.append(target)

  return target_list

def readFromLSPCI():
  target_list = []

  try:
    # run lspci
    lspci_output = subprocess.Popen(["/usr/bin/lspci", "-n", "-d", "1002:"], stdout=subprocess.PIPE).communicate()[0].split('\n')
  except:
    lspci_output = []

  target_search_term = re.compile("1002:\w+")
  for line in lspci_output:
    search_result = target_search_term.search(line)
    if search_result is not None:
      device_id = int(search_result.group(0).split(':')[1], 16)
      # try lookup from ISA_TO_ID dict
      for target in ISA_TO_ID.keys():
        for target_device_id in ISA_TO_ID[target]:
          if device_id == target_device_id:
            target_list.append(target)
            break

  return target_list

def main():
  if len(sys.argv) == 2 and sys.argv[1] == '-name' :
    """ Prints the list of available AMD GCN target names extracted from rocminfo, a tool
        shipped with this script to enumerate GPU agents available on a working ROCm stack."""
    target_list = readFromROCMINFO(True) 
  else:
    """Prints the list of available AMD GCN ISA

    The program collects the list in 3 different ways, in the order of
    precendence:

    1. ROCM_TARGET_LST : a user defined environment variable, set to the path and
                       filename where to find the "target.lst" file. This can be
                       used in an install environment with sandbox, where
                       execution of "rocminfo" is not possible.
    2. target.lst : user-supplied text file. This is used in a container setting
                  where ROCm stack may usually not available.
    3. rocminfo : a tool shipped with this script to enumerate GPU agents
                available on a working ROCm stack.
    4. lspci : enumerate PCI bus and locate supported devices from a hard-coded
             lookup table.
    """
    target_list = readFromTargetLstFile()

    if len(target_list) == 0:
      target_list = readFromROCMINFO()

    if len(target_list) == 0:
      target_list = readFromLSPCI()

  # workaround to cope with existing rocm_agent_enumerator behavior where gfx000
  # would always be returned
    print("gfx000")

  for gfx in target_list:
    print(gfx)

if __name__ == "__main__":
  main()