gemm: Pass implementation function as parameter

Author: colluca

sw/blas/gemm/data/params.json

@@ -1,22 +1,24 @@
-// Copyright 2023 ETH Zurich and University of Bologna.
-// Solderpad Hardware License, Version 0.51, see LICENSE for details.
-// SPDX-License-Identifier: SHL-0.51
-
-// Parameters for a GEMM
+// Copyright 2024 ETH Zurich and University of Bologna.
+// Licensed under the Apache License, Version 2.0, see LICENSE for details.
+// SPDX-License-Identifier: Apache-2.0
 
 {
-    M: 192,
+    prec: "FP32",
+    setup_ssr: 1,
+    parallelize_m: 0,
+    parallelize_k: 0,
+    m_tiles: 2, // number of tiles in M dimension
+    n_tiles: 1, // number of tiles in N dimension
+    k_tiles: 1, // number of tiles in K dimension
+    load_a: 1,
+    load_b: 1,
+    load_c: 1,
+    transa: false,
+    transb: true, // must be true for SIMD
+    M: 16,
     N: 16,
     K: 16,
+    alpha: 1,
     beta: 0,
-    ta: false,
-    tb: true, // must be true for SIMD
-    prec: "FP64",
-    expand: 0,
-    m_tiles: 2, // number of tiles in M dimension
-    k_tiles: 1, // number of tiles in K dimension
-    n_tiles: 1, // number of tiles in N dimension
-    parallelize_k: 0,
-    parallelize_m: 0,
-    implementation: "NAIVE"
+    gemm_fp: "gemm_fp32_opt"
 }

sw/blas/gemm/scripts/datagen.py

@@ -8,21 +8,20 @@
 #          Viviane Potocnik <vivianep@iis.ee.ethz.ch>
 
 import numpy as np
-import argparse
-import pathlib
-import json5
 import sys
 import os
+import re
 import pyflexfloat as ff
 
 sys.path.append(os.path.join(os.path.dirname(__file__), "../../../../util/sim/"))
 import data_utils  # noqa: E402
-from data_utils import emit_license, format_scalar_definition, \
+from data_utils import DataGen, format_array_declaration, format_struct_definition, \
                        format_array_definition, format_ifdef_wrapper  # noqa: E402
 
 
 np.random.seed(42)
 
+
 class GemmDataGen(DataGen):
 
     # AXI splits bursts crossing 4KB address boundaries. To minimize
@@ -41,51 +40,59 @@ def exact_golden_model(self, alpha, a, b, beta, c):
                     result[m][n] += a[m][k] * b[k][n]
         return result
 
-    def validate_config(self, prec, implementation, parallelize_m, parallelize_k, m_tiles, n_tiles, k_tiles, ta,
-                        tb, M, N, K, beta, **kwargs):
+    def infer_implementation(self, gemm_fp):
+        # gemm_fp: "gemm_fp64_opt"
+        # create a regex with fp_<type>_<implementation>
+        prec, impl = re.search(r'gemm_fp(\d+)_(\w+)', gemm_fp).group(1, 2)
+        return (int(prec) / 8), impl
+
+    def validate_config(self, gemm_fp, parallelize_m,
+                        parallelize_k, m_tiles, n_tiles, k_tiles, transa,
+                        transb, M, N, K, beta, **kwargs):
         frac_m = M / m_tiles
         frac_n = N / n_tiles
 
+        dtype, impl = self.infer_implementation(gemm_fp)
+
         assert (M % m_tiles) == 0, 'M is not an integer multiple of tile size'
         assert (N % n_tiles) == 0, 'N is not an integer multiple of tile size'
         assert (K % k_tiles) == 0, 'K is not an integer multiple of tile size'
         assert (frac_m % 8) == 0, 'frac_m is not an integer multiple of the number of cores per' \
-                                ' cluster'
+                                  ' cluster'
         assert not (parallelize_m and parallelize_k), 'Cannot parallelize K and M simultaneously'
-        assert not ta, 'SIMD kernels don\'t support transposed A matrix'
-        assert (prec == "FP64") or (implementation == 'BASELINE') or (implementation == 'NAIVE') \
-               or tb, 'Optimized SIMD kernels only support transposed B matrix'
-        assert not tb or n_tiles == 1, 'Tiling in the N dimension supported only if B is' \
-                                    ' not transposed'
-        assert not tb or k_tiles == 1, 'Tiling in the K dimension supported only if B is' \
-                                    ' not transposed'
-        assert (implementation == 'BASELINE') or (implementation == 'NAIVE') or frac_n >= 8, \
-               'N dimension of tile size must be greater or equal to the unrolling factor (8) ' \
-               'when using optimized kernels'
+        assert not transa, 'SIMD kernels don\'t support transposed A matrix'
+        assert (dtype == 8) or (impl == 'baseline') or (impl == 'naive') \
+            or transb, 'Optimized SIMD kernels only support transposed B matrix'
+        assert not transb or n_tiles == 1, 'Tiling in the N dimension not supported' \
+            ' if B is transposed'
+        assert not transb or k_tiles == 1, 'Tiling in the K dimension not supported' \
+            ' if B is transposed'
+        assert (impl == 'baseline') or (impl == 'naive') or frac_n >= 8, \
+            'N dimension of tile size must be greater or equal to the unrolling factor (8) ' \
+            'when using optimized kernels'
         assert beta == 0 or beta == 1, 'Only values of 0 or 1 supported for beta'
-        assert not (prec == "FP64" and implementation == "BASELINE"), 'No baseline implemented' \
-                                                                  ' for FP64 (switch to NAIVE)'
-        assert not (((prec == "FP64") or (prec == "FP32")) and implementation == "OPT_EX"), \
+        assert not (dtype == 8 and impl == "baseline"), 'No baseline implemented' \
+            ' for FP64 (switch to NAIVE)'
+        assert not (((dtype == 8) or (dtype == 4)) and impl == "OPT_EX"), \
             'Expanding GEMM kernels' \
             ' not supported for FP64 and FP32'
-        assert not (((prec == "FP16") or (prec == "FP8")) and implementation == "NAIVE"), \
+        assert not (((dtype == 2) or (dtype == 1)) and impl == "NAIVE"), \
             'FP16 and FP8 not supported' \
             ' in naive implementation'
-        assert not (prec == "FP8" and implementation == "OPT"), 'FP8 not supported in' \
-                                                                ' optimized implementation' \
-                                                                ' (switch to OPT_EX)'
-
+        assert not (dtype == 1 and impl == "OPT"), 'FP8 not supported in' \
+            ' optimized implementation' \
+            ' (switch to OPT_EX)'
 
     def emit_header(self, **kwargs):
         header = [super().emit_header()]
 
         # Validate parameters
         self.validate_config(**kwargs)
 
-        # Generate random input matrices
-        prec = kwargs['prec']
         M, N, K = kwargs['M'], kwargs['N'], kwargs['K']
 
+        prec, _ = self.infer_implementation(kwargs['gemm_fp'])
+
         ff_desc = data_utils.ff_desc_from_precision_t(prec)
         ctype = data_utils.ctype_from_precision_t(prec)
 
@@ -95,28 +102,34 @@ def emit_header(self, **kwargs):
         result = self.exact_golden_model(1, a, b, kwargs['beta'], c)
 
         # Store matrices in transposed form if requested
-        a = a.T if kwargs['ta'] else a
-        b = b.T if kwargs['tb'] else b
-
-        header += [format_scalar_definition('uint32_t', 'M', M)]
-        header += [format_scalar_definition('uint32_t', 'N', N)]
-        header += [format_scalar_definition('uint32_t', 'K', K)]
-        header += [format_scalar_definition('uint32_t', 'TA', int(kwargs['ta']))]
-        header += [format_scalar_definition('uint32_t', 'TB', int(kwargs['tb']))]
-        header += [format_scalar_definition('uint32_t', 'BETA', kwargs['beta'])]
-        header += [format_scalar_definition('uint32_t', 'dtype_size', prec)]
-        header += [format_scalar_definition('uint32_t', 'expand', int(kwargs['expand']))]
-        header += [format_scalar_definition('uint32_t', 'm_tiles', kwargs['m_tiles'])]
-        header += [format_scalar_definition('uint32_t', 'n_tiles', kwargs['n_tiles'])]
-        header += [format_scalar_definition('uint32_t', 'k_tiles', kwargs['k_tiles'])]
-        header += [format_scalar_definition('uint32_t', 'parallelize_m', kwargs['parallelize_m'])]
-        header += [format_scalar_definition('uint32_t', 'parallelize_k', kwargs['parallelize_k'])]
-        header += [format_scalar_definition('implementation_t', 'implementation', kwargs['implementation'])]
-        header += [format_array_definition(ctype, 'a', a.flatten(), alignment=self.BURST_ALIGNMENT,
+        a = a.T if kwargs['transa'] else a
+        b = b.T if kwargs['transb'] else b
+
+        a_uid = 'a'
+        b_uid = 'b'
+        c_uid = 'c'
+
+        cfg = {
+            'prec': prec,
+            **kwargs,
+            'a': a_uid,
+            'b': b_uid,
+            'c': c_uid,
+        }
+
+        a = a.flatten()
+        b = b.flatten()
+        c = c.flatten()
+
+        header += [format_array_declaration(ctype, a_uid, a.shape)]
+        header += [format_array_declaration(ctype, b_uid, b.shape)]
+        header += [format_array_declaration(ctype, c_uid, c.shape)]
+        header += [format_struct_definition('gemm_args_t', 'args', cfg)]
+        header += [format_array_definition(ctype, a_uid, a,
                                            section=kwargs['section'])]
-        header += [format_array_definition(ctype, 'b', b.flatten(), alignment=self.BURST_ALIGNMENT,
+        header += [format_array_definition(ctype, b_uid, b,
                                            section=kwargs['section'])]
-        header += [format_array_definition(ctype, 'c', c.flatten(), alignment=self.BURST_ALIGNMENT,
+        header += [format_array_definition(ctype, c_uid, c,
                                            section=kwargs['section'])]
         result_def = format_array_definition(ctype, 'result', result.flatten())
         header += [format_ifdef_wrapper('BIST', result_def)]
@@ -125,5 +138,5 @@ def emit_header(self, **kwargs):
         return header
 
 
-if __name__ == '__main__':
-    main()
+if __name__ == "__main__":
+    sys.exit(GemmDataGen().main())

sw/blas/gemm/scripts/verify.py

@@ -8,7 +8,7 @@
 import numpy as np
 import sys
 from pathlib import Path
-from datagen import golden_model
+from datagen import GemmDataGen
 
 sys.path.append(str(Path(__file__).parent / '../../../../util/sim/'))
 from verif_utils import Verifier  # noqa: E402
@@ -19,27 +19,54 @@ class GemmVerifier(Verifier):
 
     OUTPUT_UIDS = ['c']
     ERR_THRESHOLD = {
-        0: {8: 1e-6, 4: 1e-6, 2: 1e-2, 1: 1e-4},
-        1: {8: 0, 4: 0, 2: 0, 1: 0}
+        1: 1e-4,
+        2: 1e-2,
+        4: 1e-6,
+        8: 1e-6
     }
 
     def __init__(self):
         super().__init__()
-        self.prec = self.get_input_from_symbol('dtype_size', 'uint32_t')[0]
-        self.baseline = self.get_input_from_symbol('baseline', 'uint32_t')[0]
+        self.func_args = {
+            'alpha': 'd',
+            'prec': 'I',
+            'setup_ssr': 'I',
+            'parallelize_m': 'I',
+            'parallelize_k': 'I',
+            'm_tiles': 'I',
+            'n_tiles': 'I',
+            'k_tiles': 'I',
+            'load_a': 'I',
+            'load_b': 'I',
+            'load_c': 'I',
+            'transa': 'I',
+            'transb': 'I',
+            'M': 'I',
+            'N': 'I',
+            'K': 'I',
+            'a': 'I',
+            'b': 'I',
+            'beta': 'I',
+            'c': 'I',
+            'gemm_fp': 'I'
+        }
+        self.func_args = self.get_input_from_symbol('args', self.func_args)
 
     def get_actual_results(self):
-        return self.get_output_from_symbol(self.OUTPUT_UIDS[0], ctype_from_precision_t(self.prec))
+        prec = self.func_args['prec']
+        return self.get_output_from_symbol(self.OUTPUT_UIDS[0], ctype_from_precision_t(prec))
 
     def get_expected_results(self):
-        a = self.get_input_from_symbol('a', ctype_from_precision_t(self.prec))
-        b = self.get_input_from_symbol('b', ctype_from_precision_t(self.prec))
-        c = self.get_input_from_symbol('c', ctype_from_precision_t(self.prec))
-        beta = self.get_input_from_symbol('BETA', 'uint32_t')[0]
-        m = self.get_input_from_symbol('M', 'uint32_t')[0]
-        n = self.get_input_from_symbol('N', 'uint32_t')[0]
-        k = self.get_input_from_symbol('K', 'uint32_t')[0]
-        tb = self.get_input_from_symbol('TB', 'uint32_t')[0]
+        prec = self.func_args['prec']
+        a = self.get_input_from_symbol('a', ctype_from_precision_t(prec))
+        b = self.get_input_from_symbol('b', ctype_from_precision_t(prec))
+        c = self.get_input_from_symbol('c', ctype_from_precision_t(prec))
+        beta = self.func_args['beta']
+        m = self.func_args['M']
+        n = self.func_args['N']
+        k = self.func_args['K']
+        tb = self.func_args['transb']
+
         a = np.reshape(a, (m, k))
         if tb:
             b = np.reshape(b, (n, k))
@@ -50,7 +77,8 @@ def get_expected_results(self):
         return GemmDataGen().exact_golden_model(1, a, b, beta, c).flatten()
 
     def check_results(self, *args):
-        return super().check_results(*args, rtol=self.ERR_THRESHOLD[self.baseline][self.prec])
+        prec = self.func_args['prec']
+        return super().check_results(*args, rtol=self.ERR_THRESHOLD[prec])
 
 
 if __name__ == "__main__":