Define datatypes for attention dynamically based on chip type (#1894)

Set DATA_TYPES_ATTENTION dynamically based on chip type in perfRunner
---------

Signed-off-by: Djordje Antic <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: dorde-antic

mlir/utils/performance/perfRunner.py

@@ -35,7 +35,8 @@
 LAYOUTS = ['NHWC', 'NCHW']
 
 DATA_TYPES_GEMM = ['f32', 'f16', 'bf16', 'i8', 'fp8']
-DATA_TYPES_ATTENTION = ['i8', 'f32', 'f16', 'bf16']
+DATA_TYPES_ATTENTION_WMMA = ['i8', 'f16', 'bf16']
+DATA_TYPES_ATTENTION_MFMA = ['i8', 'f32', 'f16', 'bf16']
 DATA_TYPES_GEMM_GEMM = ['f32', 'f16', 'bf16']
 DATA_TYPES_CONV_GEMM = ['f32', 'f16', 'bf16']
 OUTPUT_DATA_TYPES_MAP = {'f32': 'f32', 'f16': 'f16', 'bf16': 'bf16', 'i8': 'i32', 'fp8':'f32',
@@ -117,6 +118,43 @@ def find_mlir_build_dir() -> str:
     build_dir = Path(rocmlir_gen_path).parent.parent
     return str(build_dir)
 
+def hip_check(call_result):
+    err = call_result[0]
+    result = call_result[1:]
+    if len(result) == 1:
+        result = result[0]
+    if isinstance(err, hip.hipError_t) and err != hip.hipError_t.hipSuccess:
+        raise RuntimeError(str(err))
+    return result
+
+def getArch() -> str:
+    agents = set()
+    device_count = hip_check(hip.hipGetDeviceCount())
+    for device in range(device_count):
+        props = hip.hipDeviceProp_t()
+        hip_check(hip.hipGetDeviceProperties(props,device))
+        agent = props.gcnArchName.decode('utf-8')
+        agents.add(agent)
+    if(len(agents) > 1):
+        print(f"WARNING: Found {len(agents)} different kinds of agents on the same machine :  {', '.join(agents)}")
+        print("WARNING: Using the first agent by default. If you want to use a different agent, please set the HIP_VISIBLE_DEVICES environment variable.")
+    # select first agent by default
+    return list(agents)[0]
+
+def getChip():
+    arch = getArch()
+    chip = GFX_CHIP_RE.search(arch).group(0)
+    return chip
+
+DATA_TYPES_ATTENTION = None
+
+def initializeDataTypesAttention():
+    global DATA_TYPES_ATTENTION
+    if getChip().startswith('gfx9'):
+        DATA_TYPES_ATTENTION = DATA_TYPES_ATTENTION_MFMA
+    else:
+        DATA_TYPES_ATTENTION = DATA_TYPES_ATTENTION_WMMA
+        
 def create_paths(config_file_path, mlir_build_dir_path) -> Paths:
     """Creates the composite Paths structure using build dir paths"""
 
@@ -686,6 +724,8 @@ def getGemmGemmConfigurations(fileName):
     return configs
 
 def getAttentionConfigurations(fileName):
+    if DATA_TYPES_ATTENTION is None:
+        initializeDataTypesAttention()
     bool_space = ['false', 'true']
     default_test_space = {
         "-t": DATA_TYPES_ATTENTION,
@@ -1740,29 +1780,6 @@ def tuneMLIRKernels(configs, arch, numCU):
                 print("MIOpen tuning timed out")
                 _, errs = p1.communicate()
 
-def hip_check(call_result):
-    err = call_result[0]
-    result = call_result[1:]
-    if len(result) == 1:
-        result = result[0]
-    if isinstance(err, hip.hipError_t) and err != hip.hipError_t.hipSuccess:
-        raise RuntimeError(str(err))
-    return result
-
-def getArch() -> str:
-    agents = set()
-    device_count = hip_check(hip.hipGetDeviceCount())
-    for device in range(device_count):
-        props = hip.hipDeviceProp_t()
-        hip_check(hip.hipGetDeviceProperties(props,device))
-        agent = props.gcnArchName.decode('utf-8')
-        agents.add(agent)
-    if(len(agents) > 1):
-        print(f"WARNING: Found {len(agents)} different kinds of agents on the same machine :  {', '.join(agents)}")
-        print("WARNING: Using the first agent by default. If you want to use a different agent, please set the HIP_VISIBLE_DEVICES environment variable.")
-    # select first agent by default
-    return list(agents)[0]
-
 def parseDataTypes(data_types):
     if not data_types:
         return DATA_TYPES_GEMM, OUTPUT_DATA_TYPES_MAP
@@ -1780,11 +1797,6 @@ def parseDataTypes(data_types):
             outMap[dt[0]] = 'f32'
     return datatypes, outMap
 
-def getChip():
-    arch = getArch()
-    chip = GFX_CHIP_RE.search(arch).group(0)
-    return chip
-
 def getNumCU(chip):
     try:
         rocminfo = subprocess.check_output("/opt/rocm/bin/rocminfo",
@@ -1842,6 +1854,7 @@ def main(args=None):
     arch = getArch()
     chip = getChip() 
     numCU = getNumCU(chip)
+    initializeDataTypesAttention()
 
     root_dir = str(subprocess.check_output(['git', 'rev-parse', '--show-toplevel']).decode().strip())
     default_conv_configs = root_dir + '/mlir/utils/jenkins/performance/configs/tier1-conv-configs'