gimmik 4.0

What does GiMMiK do?

Consider matrix multiplication of the form

C = α∙A×B + β∙C

GiMMiK generates kernels specialised to a given operator matrix. The generated code is fully unrolled, with each kernel computing a single column of the output matrix. GiMMiK is designed for block-by-panel matrix multiplication where the operator matrix is small. It removes sparsity from the operator matrix and attempts to reduce common subexpressions.

How do I install GiMMiK?

Install the latest release from PyPI with

python -m pip install gimmik

To install from a source checkout, run

python -m pip install .

GiMMiK requires Python 3.10 or newer, Mako 1.0 or newer, and NumPy 2.2 or newer.

How do I use GiMMiK?

Create a platform-specific matrix multiplication generator and request its kernels. Multiplying the operator matrix applies the corresponding alpha factor.

import numpy as np

from gimmik import CUDAMatMul

mat = np.array([[1.0, 0.0], [2.0, 3.0]])

# Generate a CUDA kernel for C = 2*mat*B
mm = CUDAMatMul(2.0*mat, beta=0.0)
src, metadata = next(mm.kernels(np.float32))

The kernels method is a generator which yields a sequence of candidate kernels for the operator matrix, each paired with a dictionary of metadata describing how it should be launched. Rather than taking the first candidate, callers are expected to benchmark the sequence and keep whichever kernel performs best on their hardware.

Only kernels which suit the operator, the target and the arguments given are yielded, so the sequence may be empty. This is not an error: it means GiMMiK has nothing to offer for the problem, and a vendor GEMM should be preferred. Callers which benchmark the sequence therefore need to handle having found no kernel at all.

The number of columns of B may either be baked into the kernel, by passing n together with the leading dimensions ldb and ldc, or left as a runtime argument by passing none of them. Passing them is a licence to specialise rather than an instruction to, and an individual kernel may still take its sizes at runtime; consult meta['args'] to see which it did.

Call signatures

Kernels do not all take the same arguments. Callers declare which call signatures they are able to invoke, and GiMMiK only yields kernels which match.

from gimmik import SIG_ABC, SIG_BC

mm.kernels(np.float32, sigs={SIG_BC, SIG_ABC})

Arguments named _desc are TMA tensor map descriptors rather than pointers, and arise only on PTX. Every one of them is described by meta['operands'], covered below.

The default is sigs={'bc'}, so a caller which asks for nothing keeps the classic two-pointer kernels and can never be handed one it cannot call. Do not infer the signature from what was passed to the constructor: consult meta['sig'] for the label and meta['args'] for the ordered tuple of argument names, and build the call from those.

The signature of a kernel is not a property of the target alone. The same template can take C as a descriptor or as a plain pointer depending on beta, since with beta=0 there is nothing to read back and the result is written straight out. Two generators built from the same operator and target, differing only in beta, may therefore offer different signatures.

MatMul.sigs gives the signatures a platform is capable of emitting, and available_sigs reports which it would actually offer for a given operator and target. This distinguishes “GiMMiK has nothing for this problem” from “you declined the only kernels which suited it”.

mm.available_sigs(np.float32, gpu_family=9)

Passing A in a buffer

Kernels with the abc signature read the operator matrix from a buffer rather than from constants baked into the source. GiMMiK owns the layout, which is padded and swizzled to suit the target, and hands the caller the bytes to upload.

src, meta = next(kgen)

if meta['sig'] == 'abc':
    apack = mm.pack_a(meta)

pack_a returns a one dimensional contiguous array which the caller copies into a buffer of at least spec['nbytes'] bytes, aligned to spec['align'], where spec is mm.operands(meta)['a']. Its data type is spec['dtype'], which need not be the type the kernel was generated for. The buffer must stay valid and unmodified for as long as the kernel is used; kernels only read it.

Packing is a pure function, and an operator may be named explicitly. Because A is no longer part of the source, one compiled kernel can serve many operators of the same shape and sparsity pattern.

other = mm.pack_a(meta, a=2.0*mat2)

Note that with abc the generated source no longer identifies the operator, so a cache of bound kernels must not be keyed on the source text alone. A supplied A must have the same shape as the one the generator was built with, and must be zero wherever the kernel elides an all zero tile; both are checked. Metadata may only be used with the generator which produced it.

When benchmarking a sequence which mixes signatures, bind a real A buffer to the abc candidates. They pay to read A on every launch where a bc kernel does not, so timings are otherwise not comparable.

Passing operands as tensor maps

An argument named b_desc or c_desc is not a pointer but the address of a TMA tensor map, which the caller encodes on the host before launch. GiMMiK states everything it requires of that map in meta['operands'], keyed by the argument name, so nothing has to be inferred from the kernel. Ask for the descriptions through operands, whose sizes let a kernel which takes them at runtime be described as precisely as one which baked them in.

for name, spec in mm.operands(meta, n, ldb, ldc).items():
    if spec['kind'] == 'tensormap':
        ...     # encode a map for spec['operand'] and pass its address

Dimension zero is the contiguous one, matching the order cuTensorMapEncodeTiled expects: for b that is (n, k) and for c it is (n, m). The modes are reported as names rather than as driver constants so the description stays independent of any one API; treat them as values to be translated, not as defaults to assume. A caller may give a larger global_dim than requested provided the strides still match the ldb and ldc the kernel was generated for.

Metadata reference

Every kernel carries these four keys, and nothing else is guaranteed.

The remaining keys depend on the platform and on the kernel. Launch geometry is named in the terms of the platform’s own API.

operands maps an argument name to a description of what that argument needs to be. An argument absent from it is an ordinary pointer or scalar, and for many kernels the mapping is empty. Buffer operands carry kind, dtype, align, nbytes and a private token; tensor map operands carry the fields tabulated above. variant is intended for logging and cache keys, and its form is not guaranteed.

Metadata is only meaningful to the generator which produced it, and is consumed by pack_a; do not construct it by hand or move it between generators.

The launch key describes the whole geometry as a function of n, and every kernel carries it; it is empty only on platforms which have no launch geometry at all. Every parameter the platform needs to launch the kernel appears in it, including those which do not vary with n, so a caller has one place to read rather than a description for some and metadata for the rest. Each entry maps a geometry parameter onto a tuple of axes, where an integer axis is a constant and a mapping axis is ceil(n / div) * mul, with mul defaulting to one. A kernel with n baked in has a fixed geometry, which the metadata additionally carries resolved.

mm = CUDAMatMul(2.0*mat, beta=0.0)
src, metadata = next(mm.kernels(np.float32))

# A grid of ceil(n / 128) by 1 by 1, of blocks of 128 threads
metadata['launch'] == {'grid': ({'div': 128}, 1, 1),
                       'block': (128, 1, 1)}

Being plain data, the description can be serialised alongside the kernel source, allowing an application compiled ahead of time to work out its own geometry with no access to the generator. Callers which have a generator to hand may instead evaluate the description with launch_config, which works for every kernel and is the one way of asking which does not need to know whether n was baked in.

# Geometry needed to multiply a B with 40000 columns
cfg = mm.launch_config(metadata, 40000)
grid, block = cfg['grid'], cfg['block']

The available generators are CMatMul, COpenMPMatMul, CUDAMatMul, HIPMatMul, ISPCMatMul, MetalMatMul, OpenCLMatMul, and PTXMatMul.

Who uses GiMMiK?

GiMMiK was developed to improve the performance of the PyFR framework. It was originally developed as part of Bartosz Wozniak’s master’s thesis in the Department of Computing at Imperial College London and is currently maintained by Freddie Witherden.