Add indexed GiMMiK candidate rendering

gimmik/base.py

@@ -90,6 +90,58 @@ 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)}
 
+    def _process_dtype(self, dtype):
+        dtype = np.dtype(dtype).type
+        if dtype == np.float32:
+            return 'float', 4
+        elif dtype == np.float64:
+            return 'double', 8
+        else:
+            raise ValueError('Invalid floating point data type')
+
+    def _base_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 _candidate_specs(self, dtype, dsize, **kwargs):
+        yield from self._kernel_generators(dtype, dsize, **kwargs)
+
+    def _render_candidate_spec(self, dtype, kname, spec):
+        name, exargs, exmeta = spec
+
+        # Merge in the base arguments and metadata
+        args = self._base_args(dtype, kname) | exargs
+        meta = self.basemeta | exmeta
+
+        # Render the kernel template
+        src = self._render_kernel(dtype, name, args)
+
+        # Post-process the metadata
+        meta['tplname'] = name
+        self._process_meta(meta)
+
+        return src, meta
+
+    def candidate_count(self, dtype, **kwargs):
+        dtype, dsize = self._process_dtype(dtype)
+
+        return sum(1 for _ in self._candidate_specs(dtype, dsize, **kwargs))
+
+    def render_candidate(self, idx, dtype, kname='gimmik_mm', **kwargs):
+        dtype, dsize = self._process_dtype(dtype)
+
+        for i, spec in enumerate(self._candidate_specs(dtype, dsize, **kwargs)):
+            if i == idx:
+                return self._render_candidate_spec(dtype, kname, spec)
+
+        raise IndexError(idx)
+
     def kernels(self, dtype, kname='gimmik_mm', **kwargs):
         basemeta = self.basemeta
 

gimmik/hip.py

@@ -2,31 +2,202 @@
 
 from gimmik.base import MatMul
 
+import numpy as np
+
 
 class HIPMatMul(MatMul):
     platform = 'hip'
     basemeta = {'block': (128, 1, 1), 'width': 1, 'shared': 0}
 
-    def _kernel_generators(self, dtype, dsize, *, gcn_arch=None, warp_size=64):
-        # B loading, C streaming kernel
-        yield ('cstream', {}, {})
+    def _candidate_specs(self, dtype, dsize, *, gcn_arch=None, warp_size=64):
+        max_block_threads = 1024
+        max_shared = 64*1024
+
+        def emit(name, args, meta):
+            block = meta.get('block', self.basemeta['block'])
+            shared = meta.get('shared', self.basemeta['shared'])
+            threads = block[0]*block[1]*block[2]
+
+            if threads <= max_block_threads and shared <= max_shared:
+                yield (name, args, meta)
 
-        # B streaming, C accumulation kernel
-        yield ('bstream', {}, {})
+        def emit_preload(name, args, meta):
+            yield from emit(name, args | {'preload': True}, meta)
 
-        # Four-way m-split B streaming, C accumulation kernel
         ms, bsz, blkx = 4, 24, 64
         args = {'msplit': ms, 'bsz': bsz, 'blockx': blkx}
-        meta = {'block': (blkx, ms, 1), 'shared': 2*bsz*blkx*dsize}
-        yield ('bstream-msplit', args, meta)
+        meta = {
+            'block': (blkx, ms, 1), 'shared': 2*bsz*blkx*dsize,
+            'desc': f'bstream-msplit/m{ms}-b{bsz}-x{blkx}'
+        }
+        yield from emit('bstream-msplit', args, meta)
 
-        # Two-way k-split B loading, C streaming kernel
         ks, csz, blkx = 2, 24, 64
         args = {'ksplit': ks, 'csz': csz, 'blockx': blkx}
-        meta = {'block': (blkx, ks, 1), 'shared': (ks - 1)*csz*blkx*dsize}
-        yield ('cstream-ksplit', args, meta)
+        meta = {
+            'block': (blkx, ks, 1), 'shared': (ks - 1)*csz*blkx*dsize,
+            'desc': f'cstream-ksplit/k{ks}-c{csz}-x{blkx}'
+        }
+        yield from emit('cstream-ksplit', args, meta)
+
+        # Tuned HIP variants
+        msplits, ksplits = [8, 4], [4, 2]
+        bsz, csz, blkx = 8, 8, 64
+        widths = [1]
+        if self.aligne is not None and self.aligne % 2 == 0:
+            widths.insert(0, 2)
+
+        for width in widths:
+            wargs = ({'dtype': f'{dtype}{width}', 'width': width}
+                     if width > 1 else {})
+            wmeta = {'width': width} if width > 1 else {}
+            wpfx = f'w{width}-' if width > 1 else ''
+
+            for ms in msplits:
+                # m-split B streaming, C accumulation kernel
+                args = {'msplit': ms, 'bsz': bsz, 'blockx': blkx} | wargs
+                shared = 2*bsz*blkx*dsize*width
+                meta = {
+                    'block': (blkx, ms, 1), 'shared': shared,
+                    'desc': f'bstream-msplit/{wpfx}m{ms}-b{bsz}-x{blkx}'
+                } | wmeta
+                yield from emit('bstream-msplit', args, meta)
+
+            for ms in msplits:
+                # m-split B streaming, C preloading, C accumulation kernel
+                args = {'msplit': ms, 'bsz': bsz, 'blockx': blkx} | wargs
+                shared = 2*bsz*blkx*dsize*width
+                meta = {
+                    'block': (blkx, ms, 1), 'shared': shared,
+                    'desc': (
+                        f'bstream-msplit-preload-c/'
+                        f'{wpfx}m{ms}-b{bsz}-x{blkx}'
+                    )
+                } | wmeta
+                yield from emit_preload('bstream-msplit', args, meta)
+
+            for ks in ksplits:
+                # k-split B loading, C preloading, C streaming kernel
+                args = {'ksplit': ks, 'csz': csz, 'blockx': blkx} | wargs
+                shared = (ks - 1)*csz*blkx*dsize*width
+                meta = {
+                    'block': (blkx, ks, 1), 'shared': shared,
+                    'desc': (
+                        f'cstream-ksplit-preload-c/'
+                        f'{wpfx}k{ks}-c{csz}-x{blkx}'
+                    )
+                } | wmeta
+                yield from emit_preload('cstream-ksplit', args, meta)
+
+        if dsize == 8:
+            # ── mfma-tile-gemm ────────────────────────────────────────────
+            # Packed Direct-A MFMA path with B-reuse workgroup mapping and
+            # cached B loads.  NT is the scalar output-column tile; width
+            # converts it to vector columns before rendering the kernel.
+            packed_mfma_tiles = [
+                (64, 64, 8, 64, 4),
+                (128, 64, 8, 64, 4),
+                (64, 128, 8, 64, 4),
+                (128, 128, 8, 64, 4),
+            ]
+
+            widths = [2] if self.aligne is not None and self.aligne % 2 == 0 else [1]
+
+            for width in widths:
+                for MT, NT, KT, blockx, blocky in packed_mfma_tiles:
+                    if NT % width:
+                        raise ValueError('mfma-tile-gemm width expects NT divisible by width')
+
+                    block = (blockx, blocky, 1)
+                    shared = 2*(KT*NT)*dsize
+                    threads = block[0]*block[1]*block[2]
+
+                    if threads <= max_block_threads and shared <= max_shared:
+                        yield ('mfma-tile-gemm',
+                               (width, MT, NT, KT, blockx, blocky, dsize))
+
+    def _render_candidate_spec(self, dtype, kname, spec):
+        if len(spec) == 2 and spec[0] == 'mfma-tile-gemm':
+            spec = self._expand_mfma_candidate_spec(dtype, spec)
+
+        return super()._render_candidate_spec(dtype, kname, spec)
+
+    def _expand_mfma_candidate_spec(self, dtype, spec):
+        name, mspec = spec
+        width, MT, NT, KT, blockx, blocky, dsize = mspec
+        vNT = NT // width
+        block = (blockx, blocky, 1)
+        a_packed_hex, m_pad, k_pad = self._dense_mfma_lane_bake(MT, KT)
+        direct_a_shared = 2*(KT*NT)*dsize
+
+        wargs = ({'dtype': f'{dtype}{width}', 'width': width,
+                  'sdtype': dtype} if width > 1 else {})
+        width_meta = {'width': width} if width > 1 else {}
+        wpfx = f'w{width}-' if width > 1 else ''
+        bpfx = f'b{blocky}-' if blocky != 4 else ''
+
+        args = {
+            'MT': MT, 'NT': vNT, 'KT': KT,
+            'blockx': block[0], 'blocky': block[1],
+            'a_hex': a_packed_hex, 'm_pad': m_pad,
+            'k_pad': k_pad,
+        } | wargs
+        meta = {
+            'block': block, 'shared': direct_a_shared,
+            'bm': MT, 'ncols': vNT,
+            'desc': (
+                f'mfma-tile-gemm/'
+                f'{wpfx}{bpfx}mt{MT}-nt{NT}-kt{KT}'
+            ),
+        } | width_meta
+
+        return name, args, meta
+
+    def _kernel_generators(self, dtype, dsize, *, gcn_arch=None, warp_size=64):
+        for spec in self._candidate_specs(dtype, dsize, gcn_arch=gcn_arch,
+                                          warp_size=warp_size):
+            if len(spec) == 2 and spec[0] == 'mfma-tile-gemm':
+                spec = self._expand_mfma_candidate_spec(dtype, spec)
+
+            yield spec
+
+    def _dense_mfma_lane_bake(self, BM, BK):
+        # Pack A so each lane can vector-load the two FP64 operands it consumes
+        # across a pair of consecutive 16x16x4 MFMA K groups:
+        #   Apg[row16_tile][kg_pair][lane][which]
+        # where lane = g*16 + p and which selects kg_pair*2 + {0, 1}.
+        if BK % 8:
+            raise ValueError('mfma lane-packed A expects BK to be a multiple of 8')
+
+        m, k = self.A.shape
+        m_pad = -(-m // BM) * BM
+        k_pad = -(-k // BK) * BK
+        a_pad = np.zeros((m_pad, k_pad), dtype=np.float64)
+        a_pad[:m, :k] = self.A
+
+        packed = []
+        kg_pairs = BK // 8
+        for row16 in range(m_pad // 16):
+            row_base = row16*16
+            for ktile in range(k_pad // BK):
+                k_base = ktile*BK
+                for kgp in range(kg_pairs):
+                    for lane in range(64):
+                        g = lane // 16
+                        p = lane % 16
+                        row = row_base + p
+                        for which in range(2):
+                            kg = 2*kgp + which
+                            packed.append(a_pad[row, k_base + kg*4 + g])
+
+        return [float(x).hex() for x in packed], m_pad, k_pad
 
     def _process_meta(self, meta):
+        bm = meta.get('bm')
+        if bm is not None:
+            meta['grid_y'] = -(-self.A.shape[0] // bm)
+
         if self.n is not None:
-            div = meta['block'][0]*meta['width']
-            meta['grid'] = (-(-self.n // div), 1, 1)
+            div = meta.get('ncols', meta['block'][0])*meta['width']
+            gy = meta.get('grid_y', 1)
+            meta['grid'] = (-(-self.n // div), gy, 1)

gimmik/kernels/hip/base.mako

@@ -1,12 +1,74 @@
 % if dtype.endswith('4'):
-static inline __device__ ${dtype} make_zero()
+inline __device__ ${dtype} operator+(${dtype} a, ${dtype} b)
+{ return make_${dtype}(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); }
+
+inline __device__ ${dtype} operator*(${dtype[:-1]} a, ${dtype} b)
+{ return make_${dtype}(a*b.x, a*b.y, a*b.z, a*b.w); }
+
+inline __device__ ${dtype} make_zero()
 { return make_${dtype}(0, 0, 0, 0); }
 % elif dtype.endswith('2'):
-static inline __device__ ${dtype} make_zero()
+inline __device__ ${dtype} operator+(${dtype} a, ${dtype} b)
+{ return make_${dtype}(a.x + b.x, a.y + b.y); }
+
+inline __device__ ${dtype} operator*(${dtype[:-1]} a, ${dtype} b)
+{ return make_${dtype}(a*b.x, a*b.y); }
+
+inline __device__ ${dtype} make_zero()
 { return make_${dtype}(0, 0); }
 % else:
-static inline __device__ ${dtype} make_zero()
+inline __device__ ${dtype} make_zero()
 { return 0; }
 % endif
 
+static inline __device__ void
+nt_store(${dtype}* p, ${dtype} v)
+{
+% if dtype.endswith('4'):
+    __builtin_nontemporal_store(v.x, &p->x);
+    __builtin_nontemporal_store(v.y, &p->y);
+    __builtin_nontemporal_store(v.z, &p->z);
+    __builtin_nontemporal_store(v.w, &p->w);
+% elif dtype.endswith('2'):
+    __builtin_nontemporal_store(v.x, &p->x);
+    __builtin_nontemporal_store(v.y, &p->y);
+% else:
+    __builtin_nontemporal_store(v, p);
+% endif
+}
+
+static inline __device__ ${dtype}
+nt_load(const ${dtype}* p)
+{
+% if dtype.endswith('4'):
+    return make_${dtype}(__builtin_nontemporal_load(&p->x),
+                         __builtin_nontemporal_load(&p->y),
+                         __builtin_nontemporal_load(&p->z),
+                         __builtin_nontemporal_load(&p->w));
+% elif dtype.endswith('2'):
+    return make_${dtype}(__builtin_nontemporal_load(&p->x),
+                         __builtin_nontemporal_load(&p->y));
+% else:
+    return __builtin_nontemporal_load(p);
+% endif
+}
+
+static inline __device__ void
+store_c(${dtype}* p, ${dtype} v)
+{
+    nt_store(p, v);
+}
+
+static inline __device__ ${dtype}
+load_c(const ${dtype}* p)
+{
+    return nt_load(p);
+}
+
+static inline __device__ ${dtype}
+load_b(const ${dtype}* p)
+{
+    return nt_load(p);
+}
+
 ${next.body()}