From c89c8a67313882b9529275aba5ee628d9a380667 Mon Sep 17 00:00:00 2001
From: Freddie Witherden <freddie@witherden.org>
Date: Fri, 17 Dec 2021 16:00:49 -0600
Subject: [PATCH] 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         | 16 ++++++++++++++--
 gimmik/kernels/c-omp.mako  | 12 ++++++++++--
 gimmik/kernels/c.mako      |  8 ++++++++
 gimmik/kernels/cuda.mako   |  8 ++++++++
 gimmik/kernels/hip.mako    |  8 ++++++++
 gimmik/kernels/ispc.mako   |  8 ++++++++
 gimmik/kernels/opencl.mako |  7 +++++++
 7 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/gimmik/__init__.py b/gimmik/__init__.py
index c22e57d..7ee8a28 100644
--- a/gimmik/__init__.py
+++ b/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')
diff --git a/gimmik/kernels/c-omp.mako b/gimmik/kernels/c-omp.mako
index 74d8cd0..1f77d94 100644
--- a/gimmik/kernels/c-omp.mako
+++ b/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]'
diff --git a/gimmik/kernels/c.mako b/gimmik/kernels/c.mako
index c844985..aa42194 100644
--- a/gimmik/kernels/c.mako
+++ b/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
diff --git a/gimmik/kernels/cuda.mako b/gimmik/kernels/cuda.mako
index 17bf159..c1b7564 100644
--- a/gimmik/kernels/cuda.mako
+++ b/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;
 
diff --git a/gimmik/kernels/hip.mako b/gimmik/kernels/hip.mako
index f9b4b0c..92c2897 100644
--- a/gimmik/kernels/hip.mako
+++ b/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;
 
diff --git a/gimmik/kernels/ispc.mako b/gimmik/kernels/ispc.mako
index 5cebafd..ac186a8 100644
--- a/gimmik/kernels/ispc.mako
+++ b/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)
diff --git a/gimmik/kernels/opencl.mako b/gimmik/kernels/opencl.mako
index 16edf9d..110bc2d 100644
--- a/gimmik/kernels/opencl.mako
+++ b/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;