Skip to content

Add tuned HIP GiMMiK preload-C and width variants with non-temporal loads and stores #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tomjen12 wants to merge 16 commits into
base: master
Choose a base branch
from
Open

Add tuned HIP GiMMiK preload-C and width variants with non-temporal loads and stores #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
8e23d63
Add tuned HIP GiMMiK preload variants
tomjen12 Jun 18, 2026
8f4d03e
Fix HIP GiMMiK block size metadata
tomjen12 Jun 18, 2026
96671a6
Address HIP GiMMiK review comments
tomjen12 Jun 22, 2026
739a82e
Handle ROCm feature suffixes for gfx942 tuning
tomjen12 Jun 22, 2026
0633539
Enable tuned HIP variants on gfx90a
tomjen12 Jun 22, 2026
7b59fb0
Parameterize HIP vector width and refine preload kernels
tomjen12 Jun 23, 2026
e9b921a
Use blockx launch bounds for HIP cstream preload
tomjen12 Jun 23, 2026
2aa2577
Always use non-temporal C accesses for HIP
tomjen12 Jun 24, 2026
be1c1db
feat(hip): add non-temporal B-load (NTB) variants for bstream-msplit
EricKing626 Jun 24, 2026
280e948
Use non-temporal B loads by default for HIP
tomjen12 Jun 25, 2026
c06216d
Make HIP preload-C a template option
tomjen12 Jun 25, 2026
e014e4d
Avoid HIP vector operator+= overloads
tomjen12 Jun 25, 2026
f6bc308
Prune HIP tuned variants to 12
tomjen12 Jun 25, 2026
a3aee45
Remove HIP variant arch gate
tomjen12 Jun 25, 2026
2c7af9b
Restore MI355 HIP baseline variants
tomjen12 Jun 25, 2026
f6075ec
Simplify HIP bstream preload-C beta handling
tomjen12 Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 71 additions & 10 deletions gimmik/hip.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,23 +8,84 @@ class HIPMatMul(MatMul):
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

# B streaming, C accumulation kernel
yield ('bstream', {}, {})
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)

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)

def _process_meta(self, meta):
if self.n is not None:
Expand Down
68 changes: 65 additions & 3 deletions gimmik/kernels/hip/base.mako
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)

@FreddieWitherden FreddieWitherden Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need/benefit from these wrappers? Just calling nt_load/store might be easier.

{
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()}
34 changes: 24 additions & 10 deletions gimmik/kernels/hip/bstream-msplit.mako
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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();
Expand All @@ -51,36 +61,40 @@ ${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
## Accumulate our dot products
% 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:

@FreddieWitherden FreddieWitherden Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This if block is a tad strange since two of the cases reduce to

csub[${j}] += ${jx}*bv;

Suggests we should reorder.

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
## Handle rows of A which are all zero
% 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
Expand Down
37 changes: 26 additions & 11 deletions gimmik/kernels/hip/bstream.mako
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -24,32 +26,45 @@ ${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]);

@FreddieWitherden FreddieWitherden Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does the mul here make sense compared with say doing it inline (so we go from a separate mul and add to an fma)?

% 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:

@FreddieWitherden FreddieWitherden Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above, two cases resulting in the same code.

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]:

@FreddieWitherden FreddieWitherden Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also two branches here resulting in the same code being emitted suggests we can combine the tests with an or.

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

## 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
}
Expand Down
Loading