PTX Backend

gimmik/__init__.py

@@ -8,6 +8,7 @@
 from gimmik.hip import HIPMatMul
 from gimmik.metal import MetalMatMul
 from gimmik.opencl import OpenCLMatMul
+from gimmik.ptx import PTXMatMul
 
 
 def generate_mm(mat, dtype, platform, alpha=1.0, beta=0.0, funcn='gimmik_mm',
@@ -22,7 +23,8 @@ def generate_mm(mat, dtype, platform, alpha=1.0, beta=0.0, funcn='gimmik_mm',
         'cuda': CUDAMatMul,
         'ispc': ISPCMatMul,
         'hip': HIPMatMul,
-        'opencl': OpenCLMatMul
+        'opencl': OpenCLMatMul,
+        'ptx': PTXMatMul
     }
 
     mm = platmap[platform](alpha*mat, beta, None, n, ldb, ldc)

gimmik/base.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 
 import itertools as it
+import json
+from pathlib import Path
 import pkgutil
 import re
 
@@ -54,6 +56,7 @@ def _chunk(l, chunksz):
 
 class MatMul:
     platform = None
+    _float_suffix = 'f'
 
     def __init__(self, A, beta=0.0, aligne=None, n=None, ldb=None, ldc=None):
         self.A = A
@@ -90,6 +93,9 @@ def __init__(self, A, beta=0.0, aligne=None, n=None, ldb=None, ldc=None):
         self.bix = np.nonzero(np.any(A != 0, axis=0))[0]
         self.bix = {kx: k for k, kx in enumerate(self.bix)}
 
+        # Create config cache
+        self._config_cache = {}
+
     def kernels(self, dtype, kname='gimmik_mm', **kwargs):
         basemeta = self.basemeta
 
@@ -103,14 +109,7 @@ def kernels(self, dtype, kname='gimmik_mm', **kwargs):
             raise ValueError('Invalid floating point data type')
 
         # Common template arguments
-        baseargs = {
-            'dtype': dtype, 'kname': kname,
-            'A': self.A, 'beta': self.beta, 'width': 1,
-            'm': self.m, 'n': self.n, 'k': self.k,
-            'ldb': self.ldb, 'ldc': self.ldc,
-            'afix': self.afix, 'alix': self.alix, 'bix': self.bix,
-            'dot': _dot, 'partition': _partition, 'chunk': _chunk
-        }
+        baseargs = self._base_template_args(dtype, kname)
 
         # Incrementally generate and render the kernels
         gen = self._kernel_generators(dtype, dsize, **kwargs)
@@ -136,17 +135,75 @@ def kernels(self, dtype, kname='gimmik_mm', **kwargs):
         except StopIteration:
             pass
 
+    def _base_template_args(self, dtype, kname):
+        return {
+            'dtype': dtype, 'kname': kname,
+            'A': self.A, 'beta': self.beta, 'width': 1,
+            'm': self.m, 'n': self.n, 'k': self.k,
+            'ldb': self.ldb, 'ldc': self.ldc,
+            'afix': self.afix, 'alix': self.alix, 'bix': self.bix,
+            'dot': _dot, 'partition': _partition, 'chunk': _chunk
+        }
+
     def _process_meta(self, meta):
         pass
 
+    def _get_config(self, key):
+        try:
+            return self._config_cache[key]
+        except KeyError:
+            cfgdir = Path('configs') / self.platform
+            cfgdata = pkgutil.get_data('gimmik', str(cfgdir / f'{key}.json'))
+            self._config_cache[key] = json.loads(cfgdata.decode('utf-8'))
+            return self._config_cache[key]
+
+    def _eval_condition(self, condition, stats):
+        if 'all' in condition:
+            return all(self._eval_condition(c, stats)
+                       for c in condition['all'])
+        if 'any' in condition:
+            return any(self._eval_condition(c, stats)
+                       for c in condition['any'])
+        if 'not' in condition:
+            return not self._eval_condition(condition['not'], stats)
+
+        value = stats[condition['field']]
+        op = next(k for k in condition if k != 'field')
+        expected = condition[op]
+
+        match op:
+            case 'eq':
+                return value == expected
+            case 'ne':
+                return value != expected
+            case 'lt':
+                return value is not None and value < expected
+            case 'lte':
+                return value is not None and value <= expected
+            case 'gt':
+                return value is not None and value > expected
+            case 'gte':
+                return value is not None and value >= expected
+            case 'in':
+                return value in expected
+            case 'is_null':
+                return value is None
+            case 'is_not':
+                return value is not None
+            case 'divisible_by':
+                return value is not None and value % expected == 0
+            case 'is_null_or_divisible_by':
+                return (value is None or value % expected == 0)
+            case _:
+                raise ValueError(f'op `{op}` not supported')
+
     def _render_kernel(self, dtype, tplname, tplargs):
         tpl = _PlatformTemplateLookup(self.platform).get_template(tplname)
         src = tpl.render(**tplargs)
 
-        # At single precision suffix all floating point constants by 'f'
-        if dtype == 'float':
+        if dtype == 'float' and self._float_suffix:
             src = re.sub(r'(?=\d*[.eE])(?=\.?\d)\d*\.?\d*(?:[eE][+-]?\d+)?',
-                         r'\g<0>f', src)
+                         rf'\g<0>{self._float_suffix}', src)
 
         # Cleanup
         src = re.sub(r'^\w+\n$', '', src.strip())

gimmik/configs/ptx/default_double.json

@@ -0,0 +1,41 @@
+{
+  "schema": 1,
+  "cc": [7, 0],
+  "ptx": [7, 0],
+  "kernels": [
+    {
+      "template": "cstream",
+      "family": "sparse",
+      "block": [128, 1, 1],
+      "width": 1,
+      "descriptor": "cstream/x128"
+    },
+    {
+      "template": "bstream",
+      "family": "sparse",
+      "block": [128, 1, 1],
+      "width": 1,
+      "descriptor": "bstream/x128"
+    },
+    {
+      "template": "bstream-msplit",
+      "family": "sparse",
+      "block": [32, 4, 1],
+      "width": 1,
+      "params": {
+        "bsz": 24
+      },
+      "descriptor": "bstream-msplit/m4-b24-x32"
+    },
+    {
+      "template": "cstream-ksplit",
+      "family": "sparse",
+      "block": [32, 2, 1],
+      "width": 1,
+      "params": {
+        "csz": 24
+      },
+      "descriptor": "cstream-ksplit/k2-c24-x32"
+    }
+  ]
+}

gimmik/configs/ptx/default_float.json

@@ -0,0 +1,41 @@
+{
+  "schema": 1,
+  "cc": [7, 0],
+  "ptx": [7, 0],
+  "kernels": [
+    {
+      "template": "cstream",
+      "family": "sparse",
+      "block": [128, 1, 1],
+      "width": 1,
+      "descriptor": "cstream/x128"
+    },
+    {
+      "template": "bstream",
+      "family": "sparse",
+      "block": [128, 1, 1],
+      "width": 1,
+      "descriptor": "bstream/x128"
+    },
+    {
+      "template": "bstream-msplit",
+      "family": "sparse",
+      "block": [32, 4, 1],
+      "width": 1,
+      "params": {
+        "bsz": 24
+      },
+      "descriptor": "bstream-msplit/m4-b24-x32"
+    },
+    {
+      "template": "cstream-ksplit",
+      "family": "sparse",
+      "block": [32, 2, 1],
+      "width": 1,
+      "params": {
+        "csz": 24
+      },
+      "descriptor": "cstream-ksplit/k2-c24-x32"
+    }
+  ]
+}