diff --git a/gimmik/hip.py b/gimmik/hip.py index a58c8fa..9e19715 100644 --- a/gimmik/hip.py +++ b/gimmik/hip.py @@ -2,31 +2,178 @@ 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', {}, {}) + 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: + wargs = ({'dtype': f'{dtype}{width}', 'width': width, + 'sdtype': dtype} if width > 1 else {}) + wpfx = f'w{width}-' if width > 1 else '' + + for MT, NT, KT, blockx, blocky in packed_mfma_tiles: + if NT % width: + raise ValueError('mfma-tile-gemm width expects NT divisible by width') + + 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 + 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 + width_meta = {'width': width} if width > 1 else {} + yield from emit( + 'mfma-tile-gemm', + args, + { + '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 + ) + + 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) diff --git a/gimmik/kernels/hip/base.mako b/gimmik/kernels/hip/base.mako index 874fbbd..d67ee25 100644 --- a/gimmik/kernels/hip/base.mako +++ b/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()} diff --git a/gimmik/kernels/hip/bstream-msplit.mako b/gimmik/kernels/hip/bstream-msplit.mako index 6359ca1..69db5ed 100644 --- a/gimmik/kernels/hip/bstream-msplit.mako +++ b/gimmik/kernels/hip/bstream-msplit.mako @@ -3,6 +3,7 @@ <% mx = partition(A, into=msplit, by='rows') bchunks = chunk(bix, bsz) +preload = context.get('preload', False) %> __global__ __launch_bounds__(${blockx*msplit}) void @@ -12,7 +13,7 @@ ${kname}(int n, ${dtype}* __restrict__ c, int ldc) { % if width > 1: - n = ((n + ${width} - 1) / ${width}) * ${width}; + n = (n + ${width} - 1) / ${width}; ldb /= ${width}; ldc /= ${width}; % endif @@ -34,9 +35,18 @@ ${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) { % for kx in bchunks[0]: % if loop.index % msplit == cid: - bsub[0][${loop.index}][threadIdx.x] = b[i + ${kx}*ldb]; + bsub[0][${loop.index}][threadIdx.x] = load_b(&b[i + ${kx}*ldb]); % endif % endfor + + % if preload and beta != 0: + ## Preload C values for active rows owned by this m-split lane + % for j, jx in enumerate(mx[cid]): + % if afix[jx] != -1: + csub[${j}] = ${'' if beta == 1 else f'{beta}*'}load_c(&c[i + ${jx}*ldc]); + % endif + % endfor + % endif } % endfor __syncthreads(); @@ -51,7 +61,7 @@ ${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) % if not loop.parent.last: % for kx in bchunks[bb + 1]: % if loop.index % msplit == cid: - bsub[${(bb + 1) % 2}][${loop.index}][threadIdx.x] = b[i + ${kx}*ldb]; + bsub[${(bb + 1) % 2}][${loop.index}][threadIdx.x] = load_b(&b[i + ${kx}*ldb]); % endif % endfor % endif @@ -59,18 +69,22 @@ ${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) % for kx in bchunks[bb]: bv = bsub[${bb % 2}][${loop.index}][threadIdx.x]; % for j, jx in enumerate(A[mcx, kx]): - % if jx != 0 and kx == afix[mcx[j]]: + % if preload and beta != 0 and jx != 0: + csub[${j}] += ${jx}*bv; + % elif jx != 0 and kx == afix[mcx[j]]: csub[${j}] = ${jx}*bv; % elif jx != 0: csub[${j}] += ${jx}*bv; % endif ## If we're done with this dot product then store to global - % if kx == alix[mcx[j]] and beta == 0: - c[i + ${mcx[j]}*ldc] = csub[${j}]; + % if preload and kx == alix[mcx[j]]: + store_c(&c[i + ${mcx[j]}*ldc], csub[${j}]); + % elif kx == alix[mcx[j]] and beta == 0: + store_c(&c[i + ${mcx[j]}*ldc], csub[${j}]); % elif kx == alix[mcx[j]] and beta == 1: - c[i + ${mcx[j]}*ldc] += csub[${j}]; + store_c(&c[i + ${mcx[j]}*ldc], load_c(&c[i + ${mcx[j]}*ldc]) + csub[${j}]); % elif kx == alix[mcx[j]]: - c[i + ${mcx[j]}*ldc] = csub[${j}] + ${beta}*c[i + ${mcx[j]}*ldc]; + store_c(&c[i + ${mcx[j]}*ldc], csub[${j}] + ${beta}*load_c(&c[i + ${mcx[j]}*ldc])); % endif % endfor % endfor @@ -78,9 +92,9 @@ ${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) % if loop.parent.last: % for j, jx in enumerate(afix): % if jx == -1 and j % msplit == cid and beta == 0: - c[i + ${j}*ldc] = make_zero(); + store_c(&c[i + ${j}*ldc], make_zero()); % elif jx == -1 and j % msplit == cid and beta != 1: - c[i + ${j}*ldc] *= ${beta}; + store_c(&c[i + ${j}*ldc], ${beta}*load_c(&c[i + ${j}*ldc])); % endif % endfor % endif diff --git a/gimmik/kernels/hip/bstream.mako b/gimmik/kernels/hip/bstream.mako index 2f6dc62..689b94e 100644 --- a/gimmik/kernels/hip/bstream.mako +++ b/gimmik/kernels/hip/bstream.mako @@ -1,13 +1,15 @@ <%inherit file='base'/> -__global__ __launch_bounds__(128) void +<% preload = context.get('preload', False) %> + +__global__ __launch_bounds__(${blockx}) void % if n is None: ${kname}(int n, const ${dtype}* __restrict__ b, int ldb, ${dtype}* __restrict__ c, int ldc) { % if width > 1: - n = ((n + ${width} - 1) / ${width}) * ${width}; + n = (n + ${width} - 1) / ${width}; ldb /= ${width}; ldc /= ${width}; % endif @@ -24,22 +26,35 @@ ${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) { ${dtype} bv, csub[${m}]; -## Iterare through the used rows of B +% if preload and beta != 0: +## Preload C values for rows which will receive a non-zero dot product +% for j, jx in enumerate(afix): + % if jx != -1: + csub[${j}] = ${'' if beta == 1 else f'{beta}*'}load_c(&c[i + ${j}*ldc]); + % endif +% endfor +% endif + +## Iterate through the used rows of B % for kx in bix: - bv = b[i + ${kx}*ldb]; + bv = load_b(&b[i + ${kx}*ldb]); % for j, jx in enumerate(A[:, kx]): - % if jx != 0 and kx == afix[j]: + % if preload and beta != 0 and jx != 0: + csub[${j}] += ${jx}*bv; + % elif jx != 0 and kx == afix[j]: csub[${j}] = ${jx}*bv; % elif jx != 0: csub[${j}] += ${jx}*bv; % endif ## - % if kx == alix[j] and beta == 0: - c[i + ${j}*ldc] = csub[${j}]; + % if preload and kx == alix[j]: + store_c(&c[i + ${j}*ldc], csub[${j}]); + % elif kx == alix[j] and beta == 0: + store_c(&c[i + ${j}*ldc], csub[${j}]); % elif kx == alix[j] and beta == 1: - c[i + ${j}*ldc] += csub[${j}]; + store_c(&c[i + ${j}*ldc], load_c(&c[i + ${j}*ldc]) + csub[${j}]); % elif kx == alix[j]: - c[i + ${j}*ldc] = csub[${j}] + ${beta}*c[i + ${j}*ldc]; + store_c(&c[i + ${j}*ldc], csub[${j}] + ${beta}*load_c(&c[i + ${j}*ldc])); % endif % endfor % endfor @@ -47,9 +62,9 @@ ${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) ## Handle rows of A which are all zero % for j, jx in enumerate(afix): % if jx == -1 and beta == 0: - c[i + ${j}*ldc] = make_zero(); + store_c(&c[i + ${j}*ldc], make_zero()); % elif jx == -1 and beta != 1: - c[i + ${j}*ldc] *= ${beta}; + store_c(&c[i + ${j}*ldc], ${beta}*load_c(&c[i + ${j}*ldc])); % endif % endfor } diff --git a/gimmik/kernels/hip/cstream-ksplit.mako b/gimmik/kernels/hip/cstream-ksplit.mako index bae2d2a..12c59ba 100644 --- a/gimmik/kernels/hip/cstream-ksplit.mako +++ b/gimmik/kernels/hip/cstream-ksplit.mako @@ -4,6 +4,7 @@ kparts = partition(A, ksplit, by='cols') cchunks = chunk(range(m), csz) loaded = set() +preload = context.get('preload', False) %> __global__ __launch_bounds__(${blockx*ksplit}) void @@ -13,7 +14,7 @@ ${kname}(int n, ${dtype}* __restrict__ c, int ldc) { % if width > 1: - n = ((n + ${width} - 1) / ${width}) * ${width}; + n = (n + ${width} - 1) / ${width}; ldb /= ${width}; ldc /= ${width}; % endif @@ -43,14 +44,31 @@ ${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) bv[${loop.index}] = b[i + ${kx}*ldb]; <% loaded.add(kx) %> % endif % endfor - % if (dotex := dot(lambda kx: f'bv[{kx}]', A[j, kbx])) != '0.0': + <% + nzixs = [(l_idx, kbx[l_idx]) for l_idx in A[j, kbx].nonzero()[0]] + has_dotp = A[j].any() + if nzixs: + first_l_idx, first_kx = nzixs[0] + dotex = f"{A[j, first_kx]}*bv[{first_l_idx}]" + for l_idx, kx in nzixs[1:]: + dotex = f"{dotex} + {A[j, kx]}*bv[{l_idx}]" + else: + dotex = 'make_zero()' + %> dotp = ${dotex}; - % else: - dotp = make_zero(); - % endif ## Save to a register % if loop.index % ksplit == bid: + % if preload and beta == 0: cv[${loop.index // ksplit}] = dotp; + % elif preload and beta == 1 and has_dotp: + cv[${loop.index // ksplit}] = load_c(&c[i + ${j}*ldc]); + cv[${loop.index // ksplit}] += dotp; + % elif preload and has_dotp: + cv[${loop.index // ksplit}] = ${beta}*load_c(&c[i + ${j}*ldc]); + cv[${loop.index // ksplit}] += dotp; + % elif not preload: + cv[${loop.index // ksplit}] = dotp; + % endif ## Save to shared memory % else: csub[${bid - (bid > loop.index % ksplit)}][${loop.index}][threadIdx.x] = dotp; @@ -66,14 +84,32 @@ ${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) ## Sum and output the final set of dot products % for j in cchunk: % if loop.index % ksplit == bid: - dotp = cv[${loop.index // ksplit}] + ${' + '.join(f'csub[{i}][{loop.index}][threadIdx.x]' - for i in range(ksplit - 1))}; - % if beta == 0: - c[i + ${j}*ldc] = dotp; + <% has_dotp = A[j].any() %> + <% + sum_expr = f"cv[{loop.index // ksplit}]" + for s_idx in range(ksplit - 1): + sum_expr = f"{sum_expr} + csub[{s_idx}][{loop.index}][threadIdx.x]" + %> + % if preload and beta == 0: + dotp = ${sum_expr}; + store_c(&c[i + ${j}*ldc], dotp); + % elif preload and beta == 1 and has_dotp: + dotp = ${sum_expr}; + store_c(&c[i + ${j}*ldc], dotp); + % elif preload and beta != 1 and has_dotp: + dotp = ${sum_expr}; + store_c(&c[i + ${j}*ldc], dotp); + % elif preload and beta != 1: + store_c(&c[i + ${j}*ldc], ${beta}*load_c(&c[i + ${j}*ldc])); + % elif beta == 0: + dotp = ${sum_expr}; + store_c(&c[i + ${j}*ldc], dotp); % elif beta == 1: - c[i + ${j}*ldc] += dotp; + dotp = ${sum_expr}; + store_c(&c[i + ${j}*ldc], load_c(&c[i + ${j}*ldc]) + dotp); % else: - c[i + ${j}*ldc] = dotp + ${beta}*c[i + ${j}*ldc]; + dotp = ${sum_expr}; + store_c(&c[i + ${j}*ldc], dotp + ${beta}*load_c(&c[i + ${j}*ldc])); % endif % endif % endfor diff --git a/gimmik/kernels/hip/cstream.mako b/gimmik/kernels/hip/cstream.mako index f75301d..2ee9574 100644 --- a/gimmik/kernels/hip/cstream.mako +++ b/gimmik/kernels/hip/cstream.mako @@ -1,15 +1,17 @@ <%inherit file='base'/> -<% ksplit = 2 if m < 36 else 1 %> +<% +preload = context.get('preload', False) +%> -__global__ __launch_bounds__(128) void +__global__ __launch_bounds__(${blockx}) void % if n is None: ${kname}(int n, const ${dtype}* __restrict__ b, int ldb, ${dtype}* __restrict__ c, int ldc) { % if width > 1: - n = ((n + ${width} - 1) / ${width}) * ${width}; + n = (n + ${width} - 1) / ${width}; ldb /= ${width}; ldc /= ${width}; % endif @@ -26,17 +28,39 @@ ${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) if (i < n) { % for j, jx in enumerate(A): - % if (dotex := dot(lambda kx: f'b[i + {kx}*ldb]', jx, maxsplit=ksplit)) != '0.0': + <% + nzixs = [kx for kx, val in enumerate(jx) if val != 0] + if nzixs: + first_kx = nzixs[0] + dotex = f"{jx[first_kx]}*b[i + {first_kx}*ldb]" + for kx in nzixs[1:]: + dotex = f"{dotex} + {jx[kx]}*b[i + {kx}*ldb]" + else: + dotex = 'make_zero()' + %> dotp = ${dotex}; + % if preload and nzixs: + % if beta == 0: + store_c(&c[i + ${j}*ldc], dotp); + % elif beta == 1: + dotp = load_c(&c[i + ${j}*ldc]) + dotp; + store_c(&c[i + ${j}*ldc], dotp); + % else: + dotp = ${beta}*load_c(&c[i + ${j}*ldc]) + dotp; + store_c(&c[i + ${j}*ldc], dotp); + % endif + % elif preload: + % if beta == 0: + store_c(&c[i + ${j}*ldc], make_zero()); + % elif beta != 1: + store_c(&c[i + ${j}*ldc], ${beta}*load_c(&c[i + ${j}*ldc])); + % endif + % elif beta == 0: + store_c(&c[i + ${j}*ldc], dotp); + % elif beta == 1 and nzixs: + store_c(&c[i + ${j}*ldc], load_c(&c[i + ${j}*ldc]) + dotp); % else: - dotp = make_zero(); - % endif - % if beta == 0: - c[i + ${j}*ldc] = dotp; - % elif beta == 1 and dotex != '0.0': - c[i + ${j}*ldc] += dotp; - % else: - c[i + ${j}*ldc] = dotp + ${beta}*c[i + ${j}*ldc]; + store_c(&c[i + ${j}*ldc], dotp + ${beta}*load_c(&c[i + ${j}*ldc])); % endif % endfor } diff --git a/gimmik/kernels/hip/mfma-tile-gemm.mako b/gimmik/kernels/hip/mfma-tile-gemm.mako new file mode 100644 index 0000000..b37e6d2 --- /dev/null +++ b/gimmik/kernels/hip/mfma-tile-gemm.mako @@ -0,0 +1,419 @@ +<%inherit file='base'/> +## +## Dense FP64 tiled GEMM optimized for AMD wave-level MFMA: +## 1. Uses v_mfma_f64_16x16x4_f64 for 16x16 FP64 wave tiles. +## 2. Packs baked A in MFMA lane-major order so each lane can vector-load +## the two FP64 A operands used by a pair of consecutive MFMA K groups. +## 3. Loads A directly from global/constant memory into VGPRs; A does not +## use LDS. +## 4. Loads B with normal cached global loads, stages B through LDS, and +## reads LDS in the MFMA operand layout. +## 5. Uses PGR2/double buffering for both operands: A is double-buffered in +## VGPR pairs, while B is double-buffered as global->VGPR->LDS. +## 6. Remaps block ids to a column-major/B-reuse tile order so neighboring +## workgroups compute different M tiles for the same N tile, improving B +## cache reuse. +## +<% + sdtype = context.get('sdtype', dtype) + mfma_k = 4 + a_pair_groups = 2 + + nthreads = blockx * blocky + if nthreads % 64: + raise ValueError('mfma-tile-gemm expects a whole number of wave64 waves') + if MT % 16: + raise ValueError('mfma-tile-gemm expects MT to be a multiple of 16') + if NT % 16: + raise ValueError('mfma-tile-gemm expects NT to be a multiple of 16') + if KT % (mfma_k * a_pair_groups): + raise ValueError('mfma-tile-gemm expects KT to cover whole A pairs') + if blocky <= 0: + raise ValueError('mfma-tile-gemm expects blocky > 0') + if width not in (1, 2): + raise ValueError('mfma-tile-gemm expects width to be 1 or 2') + + valid_m_tiles = min(-(-m // 16), MT // 16) + n_tiles = NT // 16 + k_groups = KT // mfma_k + k_group_pairs = k_groups // 2 + nwaves = nthreads // 64 + mtpg = -(-valid_m_tiles // nwaves) + b_tile_iters = -(-(KT * NT) // nthreads) + c_epilogue_indices = [ + (j, t, reg) + for j in range(mtpg) + for t in range(n_tiles) + for reg in range(4) + ] +%> +typedef ${sdtype} gimmik_f64x4 __attribute__((ext_vector_type(4))); +typedef ${sdtype} gimmik_a_f64x2 __attribute__((ext_vector_type(2))); + +// A is packed as row16-tile, K-group-pair, lane, pair-element. +__device__ static const __align__(16) ${sdtype} ${kname}_Ag[${m_pad * k_pad}] = { + ${', '.join(a_hex)} +}; + +## --------------------------------------------------------------------------- +## Tile-level global prefetch and LDS staging helpers +## --------------------------------------------------------------------------- + +<%def name="a_operand_prefetch_tile(row_offset_expr, slot)"> +% for j in range(mtpg): +<% + mt = f'wmt + {j}' + full_valid_mt = (nwaves - 1) * mtpg + j < valid_m_tiles +%> +% for kgp in range(k_group_pairs): + { +% if full_valid_mt: + a_pair_${slot}_${j}_${kgp} = *(const gimmik_a_f64x2*)&${kname}_Ag[ + (((row_base / 16 + (${mt}))*${k_pad // 8} + + ((${row_offset_expr}) / ${KT})*${k_group_pairs} + ${kgp})*64 + + lane)*2 + ]; +% else: + if (${mt} < ${valid_m_tiles}) + { + a_pair_${slot}_${j}_${kgp} = *(const gimmik_a_f64x2*)&${kname}_Ag[ + (((row_base / 16 + (${mt}))*${k_pad // 8} + + ((${row_offset_expr}) / ${KT})*${k_group_pairs} + ${kgp})*64 + + lane)*2 + ]; + } + else + { + a_pair_${slot}_${j}_${kgp} = {(${sdtype})0.0, (${sdtype})0.0}; + } +% endif + } +% endfor +% endfor + + +<%def name="b_prefetch_tile_frag(row_offset_expr, slot, pp)"> + { + const int idx = tid + ${pp * nthreads}; +% if (pp + 1) * nthreads <= KT * NT: + const int kk = idx / ${NT}; + const int cc = idx % ${NT}; + const int krow = ${row_offset_expr} + kk; + const int col = col_base + cc; + b_next_${slot}_${pp} = fast_b_tile ? b[krow*ldb + col] + : ((krow < ${k} && col < n) ? b[krow*ldb + col] : make_zero()); +% else: + if (idx < ${KT * NT}) + { + const int kk = idx / ${NT}; + const int cc = idx % ${NT}; + const int krow = ${row_offset_expr} + kk; + const int col = col_base + cc; + b_next_${slot}_${pp} = fast_b_tile ? b[krow*ldb + col] + : ((krow < ${k} && col < n) ? b[krow*ldb + col] : make_zero()); + } +% endif + } + + +<%def name="b_prefetch_tile(row_offset_expr, slot)"> + { + const bool fast_b_tile = (${row_offset_expr} + ${KT} <= ${k}) && fast_col; +% for pp in range(b_tile_iters): +${b_prefetch_tile_frag(row_offset_expr, slot, pp)} +% endfor + } + + +<%def name="b_write_tile_frag(buf_expr, slot, pp)"> + { + const int idx = tid + ${pp * nthreads}; +% if (pp + 1) * nthreads <= KT * NT: + ${kname}_Bs[${buf_expr} + idx] = b_next_${slot}_${pp}; +% else: + if (idx < ${KT * NT}) + ${kname}_Bs[${buf_expr} + idx] = b_next_${slot}_${pp}; +% endif + } + + +<%def name="b_write_tile(buf_expr, slot)"> +% for pp in range(b_tile_iters): +${b_write_tile_frag(buf_expr, slot, pp)} +% endfor + + +<%def name="b_prefetch_write_tile(row_offset_expr, buf_expr, slot)"> +${b_prefetch_tile(row_offset_expr, slot)} +${b_write_tile(buf_expr, slot)} + + +## --------------------------------------------------------------------------- +## C epilogue helpers +## --------------------------------------------------------------------------- + +<%def name="c_epilogue_coords(j, t, reg)"> + const int mt = wmt + ${j}; + const int row = row_base + mt*16 + ${4 * reg} + g; + const int col = col_base + ${t * 16} + p; + + +<%def name="c_epilogue_acc_expr(j, t, reg)"> +% if width == 1: +acc_${j}_${t}[${reg}]\ +% else: +make_${dtype}(acc_0_${j}_${t}[${reg}], acc_1_${j}_${t}[${reg}])\ +% endif + + +<%def name="store_c_epilogue_beta1(guarded)"> +% for j, t, reg in c_epilogue_indices: + ${dtype} c_old_${j}_${t}_${reg}; +% if guarded: + bool c_valid_${j}_${t}_${reg}; +% endif +% endfor + +% for j, t, reg in c_epilogue_indices: + { +${c_epilogue_coords(j, t, reg)} +% if guarded: + c_valid_${j}_${t}_${reg} = + mt < ${valid_m_tiles} && (fast_row || row < ${m}) && (fast_col || col < n); + if (c_valid_${j}_${t}_${reg}) +% endif + c_old_${j}_${t}_${reg} = load_c(&c[row*ldc + col]); + } +% endfor + +% for j, t, reg in c_epilogue_indices: + { +${c_epilogue_coords(j, t, reg)} +% if guarded: + if (c_valid_${j}_${t}_${reg}) +% endif + store_c(&c[row*ldc + col], c_old_${j}_${t}_${reg} + ${c_epilogue_acc_expr(j, t, reg)}); + } +% endfor + + +<%def name="store_c_epilogue_scalar(guarded)"> +% for j in range(mtpg): +<% + full_valid_mt = (nwaves - 1) * mtpg + j < valid_m_tiles +%> +% for t in range(n_tiles): +% for reg in range(4): + { +${c_epilogue_coords(j, t, reg)} +% if guarded: + if (mt < ${valid_m_tiles} && (fast_row || row < ${m}) && (fast_col || col < n)) +% elif not full_valid_mt: + if (mt < ${valid_m_tiles}) +% endif +% if width == 1: +% if beta == 0: + store_c(&c[row*ldc + col], acc_${j}_${t}[${reg}]); +% else: + store_c(&c[row*ldc + col], ${beta}*load_c(&c[row*ldc + col]) + acc_${j}_${t}[${reg}]); +% endif +% else: +% if beta == 0: + store_c(&c[row*ldc + col], ${c_epilogue_acc_expr(j, t, reg)}); +% else: + store_c(&c[row*ldc + col], ${beta}*load_c(&c[row*ldc + col]) + ${c_epilogue_acc_expr(j, t, reg)}); +% endif +% endif + } +% endfor +% endfor +% endfor + + +<%def name="store_c_epilogue(guarded)"> +% if beta == 1: +${store_c_epilogue_beta1(guarded)} +% else: +${store_c_epilogue_scalar(guarded)} +% endif + + +__global__ __launch_bounds__(${blockx * blocky}) void +% if n is None: +${kname}(int n, + const ${dtype}* __restrict__ b, int ldb, + ${dtype}* __restrict__ c, int ldc) +{ +% if width > 1: + n = (n + ${width} - 1) / ${width}; + ldb /= ${width}; + ldc /= ${width}; +% endif +% else: +${kname}(const ${dtype}* __restrict__ b, ${dtype}* __restrict__ c) +{ + const int n = ${-(-n // width)}; + const ${'long long' if k * ldb >= width*2**31 else 'int'} ldb = ${ldb // width}; + const ${'long long' if m * ldc >= width*2**31 else 'int'} ldc = ${ldc // width}; +% endif + const int tid = threadIdx.y*${blockx} + threadIdx.x; + const int lane = tid & 63; + const int wave = tid >> 6; + const int wmt = wave*${mtpg}; + + const int g = lane / 16; // MFMA K group / C row group + const int p = lane % 16; // MFMA row/column position + + const int logical_bid = blockIdx.y*gridDim.x + blockIdx.x; + const int m_tile = logical_bid % gridDim.y; + const int n_tile = logical_bid / gridDim.y; + const int row_base = m_tile * ${MT}; + const int col_base = n_tile * ${NT}; + const bool fast_col = col_base + ${NT} <= n; + const bool fast_row = row_base + ${MT} <= ${m}; + const bool fast_tile = fast_col & fast_row & (${nwaves * mtpg} <= ${valid_m_tiles}); + + __shared__ __align__(16) ${dtype} ${kname}_Bs[${2 * KT * NT}]; + +% for j in range(mtpg): +% for t in range(n_tiles): +% for w in range(width): +% if width == 1: + gimmik_f64x4 acc_${j}_${t} = {0.0, 0.0, 0.0, 0.0}; +% else: + gimmik_f64x4 acc_${w}_${j}_${t} = {0.0, 0.0, 0.0, 0.0}; +% endif +% endfor +% endfor +% endfor + +% for j in range(mtpg): +% for kgp in range(k_group_pairs): + gimmik_a_f64x2 a_pair_0_${j}_${kgp}; + gimmik_a_f64x2 a_pair_1_${j}_${kgp}; +% endfor +% endfor +% for pp in range(b_tile_iters): + ${dtype} b_next_0_${pp}; + ${dtype} b_next_1_${pp}; +% endfor + +${a_operand_prefetch_tile('0', 0)} +% if KT < k_pad: +${a_operand_prefetch_tile(str(KT), 1)} +% endif +${b_prefetch_write_tile('0', '0', 0)} +% if KT < k_pad: +${b_prefetch_tile(str(KT), 1)} +% endif + __syncthreads(); + + for (int k0 = 0; k0 < ${k_pad}; k0 += ${KT}) + { + const int curbuf = (k0 / ${KT}) & 1; + const int nextbuf = curbuf ^ 1; + const int k_next = k0 + ${KT}; + const int k_next2 = k0 + ${2 * KT}; + + if (k_next < ${k_pad}) + { + if ((k0 / ${KT}) & 1) + { +${b_write_tile('nextbuf*' + str(KT * NT), 0)} + } + else + { +${b_write_tile('nextbuf*' + str(KT * NT), 1)} + } + } + + if (k_next2 < ${k_pad}) + { + if ((k0 / ${KT}) & 1) + { +${b_prefetch_tile('k_next2', 1)} + } + else + { +${b_prefetch_tile('k_next2', 0)} + } + } +% for kgp in range(k_group_pairs): +% for which in range(2): +<% + kg = 2*kgp + which + acomp = 'xy'[which] +%> +% for t in range(n_tiles): + const ${dtype} bv_${kg}_${t} = + ${kname}_Bs[curbuf*${KT * NT} + (${kg * 4} + g)*${NT} + ${t * 16} + p]; +% endfor +% for j in range(mtpg): +<% + mt = f'wmt + {j}' + full_valid_mt = (nwaves - 1) * mtpg + j < valid_m_tiles +%> +% if full_valid_mt: + { + const ${sdtype} av = + curbuf ? a_pair_1_${j}_${kgp}.${acomp} : a_pair_0_${j}_${kgp}.${acomp}; +% for t in range(n_tiles): +% if width == 1: + acc_${j}_${t} = + __builtin_amdgcn_mfma_f64_16x16x4f64(av, bv_${kg}_${t}, acc_${j}_${t}, 0, 0, 0); +% else: +% for w in range(width): + acc_${w}_${j}_${t} = + __builtin_amdgcn_mfma_f64_16x16x4f64(av, bv_${kg}_${t}.${'xy'[w]}, acc_${w}_${j}_${t}, 0, 0, 0); +% endfor +% endif +% endfor + } +% else: + if (${mt} < ${valid_m_tiles}) + { + const ${sdtype} av = + curbuf ? a_pair_1_${j}_${kgp}.${acomp} : a_pair_0_${j}_${kgp}.${acomp}; +% for t in range(n_tiles): +% if width == 1: + acc_${j}_${t} = + __builtin_amdgcn_mfma_f64_16x16x4f64(av, bv_${kg}_${t}, acc_${j}_${t}, 0, 0, 0); +% else: +% for w in range(width): + acc_${w}_${j}_${t} = + __builtin_amdgcn_mfma_f64_16x16x4f64(av, bv_${kg}_${t}.${'xy'[w]}, acc_${w}_${j}_${t}, 0, 0, 0); +% endfor +% endif +% endfor + } +% endif +% endfor +% endfor +% endfor + + if (k_next < ${k_pad}) + { + if (k_next2 < ${k_pad}) + { + if ((k0 / ${KT}) & 1) + { +${a_operand_prefetch_tile('k_next2', 1)} + } + else + { +${a_operand_prefetch_tile('k_next2', 0)} + } + } + __syncthreads(); + } + } + + if (fast_tile) + { +${store_c_epilogue(False)} + } + else + { +${store_c_epilogue(True)} + } +}