Enable kernels to be generated with constant dimensions.

This can reduce the size of the generated code and enable
further optimisations on the part of the compiler.

gimmik/__init__.py

@@ -9,7 +9,8 @@
 from gimmik._version import __version__
 
 
-def generate_mm(mat, dtype, platform, alpha=1.0, beta=0.0, funcn='gimmik_mm'):
+def generate_mm(mat, dtype, platform, alpha=1.0, beta=0.0, funcn='gimmik_mm',
+                n=None, ldb=None, ldc=None):
     # Data type
     dtype = np.dtype(dtype).type
     if dtype == np.float32:
@@ -19,11 +20,22 @@ def generate_mm(mat, dtype, platform, alpha=1.0, beta=0.0, funcn='gimmik_mm'):
     else:
         raise ValueError('Invalid floating point data type')
 
+    if 0 < (n is None) + (ldb is None) + (ldc is None) < 3:
+        raise ValueError('Must provide all of (n, ldb, ldc) or none')
+
     # Multiply the matrix through by alpha
     mat = alpha*mat
 
     # Template arguments
-    tplargs = {'dtype': dtype, 'mat': mat, 'beta': beta, 'funcn': funcn}
+    tplargs = {
+        'dtype': dtype,
+        'mat': mat,
+        'beta': beta,
+        'funcn': funcn,
+        'n': n,
+        'ldb': ldb,
+        'ldc': ldc
+    }
 
     # Load and render the template
     tpl = pkgutil.get_data(__name__, f'kernels/{platform}.mako')

gimmik/kernels/c-omp.mako

@@ -1,14 +1,22 @@
 # -*- coding: utf-8 -*-
 
 void
-${funcn}(int ncol,
+% if n is None:
+${funcn}(int n,
          const ${dtype}* restrict b, int ldb,
          ${dtype}* restrict c, int ldc)
 {
+% else:
+${funcn}(const ${dtype}* restrict b, ${dtype}* restrict c)
+{
+    const int n = ${n};
+    const int ldb = ${ldb};
+    const int ldc = ${ldc};
+% endif
     ${dtype} dotp;
 
     #pragma omp parallel for simd private(dotp)
-    for (int i = 0; i < ncol; i++)
+    for (int i = 0; i < n; i++)
     {
     % for j, jx in enumerate(mat):
         dotp = ${' + '.join(f'{kx}*b[i + {k}*ldb]'

gimmik/kernels/c.mako

@@ -1,10 +1,18 @@
 # -*- coding: utf-8 -*-
 
 void
+% if n is None:
 ${funcn}(int n,
          const ${dtype}* restrict b, int ldb,
          ${dtype}* restrict c, int ldc)
 {
+% else:
+${funcn}(const ${dtype}* restrict b, ${dtype}* restrict c)
+{
+    const int n = ${n};
+    const int ldb = ${ldb};
+    const int ldc = ${ldc};
+% endif
     ${dtype} dotp;
 
     #pragma omp simd

gimmik/kernels/cuda.mako

@@ -1,10 +1,18 @@
 # -*- coding: utf-8 -*-
 
 __global__ void
+% if n is None:
 ${funcn}(int n,
          const ${dtype}* __restrict__ b, int ldb,
          ${dtype}* __restrict__ c, int ldc)
 {
+% else:
+${funcn}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c)
+{
+    const int n = ${n};
+    const int ldb = ${ldb};
+    const int ldc = ${ldc};
+% endif
     int i = blockDim.x*blockIdx.x + threadIdx.x;
     ${dtype} dotp;
 

gimmik/kernels/hip.mako

@@ -1,10 +1,18 @@
 # -*- coding: utf-8 -*-
 
 __global__ __launch_bounds__(128) void
+% if n is None:
 ${funcn}(int n,
          const ${dtype}* __restrict__ b, int ldb,
          ${dtype}* __restrict__ c, int ldc)
 {
+% else:
+${funcn}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c)
+{
+    const int n = ${n};
+    const int ldb = ${ldb};
+    const int ldc = ${ldc};
+% endif
     int i = blockDim.x*blockIdx.x + threadIdx.x;
     ${dtype} dotp;
 

gimmik/kernels/ispc.mako

@@ -1,10 +1,18 @@
 # -*- coding: utf-8 -*-
 
 export void
+% if n is None:
 ${funcn}(uniform int n,
          const uniform ${dtype} b[], uniform int ldb,
          ${dtype} uniform c[], uniform int ldc)
 {
+% else:
+${funcn}(const uniform ${dtype} b[], ${dtype} uniform c[])
+{
+    const uniform int n = ${n};
+    const uniform int ldb = ${ldb};
+    const uniform int ldc = ${ldc};
+% endif
     ${dtype} dotp;
 
     foreach (i = 0 ... n)

gimmik/kernels/opencl.mako

@@ -7,10 +7,17 @@
 % endif
 
 __kernel void
+% if n is None:
 ${funcn}(int n,
          __global const ${dtype}* restrict b, int ldb,
          __global ${dtype}* restrict c, int ldc)
 {
+% else:
+${funcn}(__global const ${dtype}* restrict b, __global ${dtype}* restrict c)
+    const int n = ${n};
+    const int ldb = ${ldb};
+    const int ldc = ${ldc};
+% endif
     int i = get_global_id(0);
     ${dtype} dotp;