teeny_kernels/math/matmul.rs
1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17// import torch
18
19// from triton import Config, autotune, cdiv, heuristics, jit
20// from triton import language as tl
21// from .matmul_perf_model import early_config_prune, estimate_matmul_time
22
23// _ordered_datatypes = [torch.int8, torch.float16, torch.bfloat16, torch.float32]
24
25// def upcast_if_fp8(a):
26// if "fp8" in str(a):
27// return torch.float16
28// return a
29
30// def get_higher_dtype(a, b):
31// a = upcast_if_fp8(a)
32// b = upcast_if_fp8(b)
33// if a is b:
34// return a
35
36// assert a in _ordered_datatypes
37// assert b in _ordered_datatypes
38
39// for d in _ordered_datatypes:
40// if a is d:
41// return b
42// if b is d:
43// return a
44
45// def init_to_zero(name):
46// return lambda nargs: nargs[name].zero_()
47
48// def get_configs_io_bound():
49// configs = []
50// for num_stages in [2, 3, 4, 5, 6]:
51// for block_m in [16, 32]:
52// for block_k in [32, 64]:
53// for block_n in [32, 64, 128, 256]:
54// num_warps = 2 if block_n <= 64 else 4
55// configs.append(
56// Config(
57// {
58// "BLOCK_M": block_m,
59// "BLOCK_N": block_n,
60// "BLOCK_K": block_k,
61// "SPLIT_K": 1,
62// },
63// num_stages=num_stages,
64// num_warps=num_warps,
65// )
66// )
67// # split_k
68// for split_k in [2, 4, 8, 16]:
69// configs.append(
70// Config(
71// {
72// "BLOCK_M": block_m,
73// "BLOCK_N": block_n,
74// "BLOCK_K": block_k,
75// "SPLIT_K": split_k,
76// },
77// num_stages=num_stages,
78// num_warps=num_warps,
79// pre_hook=init_to_zero("C"),
80// )
81// )
82// return configs
83
84// @autotune(
85// configs=[
86// # basic configs for compute-bound matmuls
87// Config(
88// {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1},
89// num_stages=3,
90// num_warps=8,
91// ),
92// Config(
93// {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1},
94// num_stages=3,
95// num_warps=8,
96// ),
97// Config(
98// {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1},
99// num_stages=4,
100// num_warps=4,
101// ),
102// Config(
103// {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1},
104// num_stages=4,
105// num_warps=4,
106// ),
107// Config(
108// {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1},
109// num_stages=4,
110// num_warps=4,
111// ),
112// Config(
113// {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1},
114// num_stages=4,
115// num_warps=4,
116// ),
117// Config(
118// {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1},
119// num_stages=4,
120// num_warps=4,
121// ),
122// Config(
123// {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1},
124// num_stages=4,
125// num_warps=4,
126// ),
127// Config(
128// {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1},
129// num_stages=5,
130// num_warps=2,
131// ),
132// # good for int8
133// Config(
134// {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1},
135// num_stages=3,
136// num_warps=8,
137// ),
138// Config(
139// {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1},
140// num_stages=3,
141// num_warps=8,
142// ),
143// Config(
144// {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1},
145// num_stages=4,
146// num_warps=4,
147// ),
148// Config(
149// {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1},
150// num_stages=4,
151// num_warps=4,
152// ),
153// Config(
154// {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1},
155// num_stages=4,
156// num_warps=4,
157// ),
158// Config(
159// {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1},
160// num_stages=4,
161// num_warps=4,
162// ),
163// Config(
164// {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1},
165// num_stages=4,
166// num_warps=4,
167// ),
168// Config(
169// {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1},
170// num_stages=4,
171// num_warps=4,
172// ),
173// Config(
174// {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1},
175// num_stages=5,
176// num_warps=2,
177// ),
178// ]
179// + get_configs_io_bound(),
180// key=["M", "N", "K"],
181// prune_configs_by={
182// "early_config_prune": early_config_prune,
183// "perf_model": estimate_matmul_time,
184// "top_k": 10,
185// },
186// )
187// @heuristics(
188// {
189// "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0,
190// }
191// )
192// @jit
193
194// fn kernel(
195// A: &DenseTensor<DynamicShape, f32>,
196// B: &DenseTensor<DynamicShape, f32>,
197// C: &DenseTensor<DynamicShape, f32>,
198// M: usize,
199// N: usize,
200// K: usize,
201// stride_am: usize,
202// stride_ak: usize,
203// stride_bk: usize,
204// stride_bn: usize,
205// stride_cm: usize,
206// stride_cn: usize,
207// acc_dtype: usize,
208// input_precision: usize,
209// fp8_fast_accum: usize,
210// BLOCK_M: usize,
211// BLOCK_N: usize,
212// BLOCK_K: usize,
213// GROUP_M: usize,
214// SPLIT_K: usize,
215// EVEN_K: usize,
216// AB_DTYPE: usize,
217// ) {
218// // matrix multiplication
219// let pid = triton::program_id(0);
220// let pid_z = triton::program_id(1);
221// let grid_m = triton::cdiv(M, BLOCK_M);
222// let grid_n = triton::cdiv(N, BLOCK_N);
223// // re-order program ID for better L2 performance
224// let width = GROUP_M * grid_n;
225// let group_id = pid / width;
226// let group_size = min(grid_m - group_id * GROUP_M, GROUP_M);
227// let pid_m = group_id * GROUP_M + (pid % group_size);
228// let pid_n = (pid % width) / group_size;
229// // do matrix multiplication
230// let rm = pid_m * BLOCK_M + triton::arange(0, BLOCK_M, 1);
231// let rn = pid_n * BLOCK_N + triton::arange(0, BLOCK_N, 1);
232// let ram = triton::max_contiguous(triton::multiple_of(rm % M, BLOCK_M), BLOCK_M);
233// let rbn = triton::max_contiguous(triton::multiple_of(rn % N, BLOCK_N), BLOCK_N);
234// let rk = pid_z * BLOCK_K + triton::arange(0, BLOCK_K, 1);
235// // pointers
236// let A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak);
237// let B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn);
238// let acc = triton::zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype);
239// for k in range(0, triton::cdiv(K, BLOCK_K * SPLIT_K)) {
240// if EVEN_K {
241// let a = triton::load(A);
242// let b = triton::load(B);
243// } else {
244// let k_remaining = K - k * (BLOCK_K * SPLIT_K);
245// let _0 = triton::zeros((1, 1), dtype=C.dtype.element_ty);
246// let a = triton::load(A, mask=rk[None, :] < k_remaining, other=_0);
247// let b = triton::load(B, mask=rk[:, None] < k_remaining, other=_0);
248// }
249// if AB_DTYPE is not None {
250// a = a.to(AB_DTYPE);
251// b = b.to(AB_DTYPE);
252// }
253// if fp8_fast_accum {
254// acc = triton::dot(
255// a, b, acc, out_dtype=acc_dtype, input_precision=input_precision
256// )
257// } else {
258// acc += triton::dot(a, b, out_dtype=acc_dtype, input_precision=input_precision)
259// }
260// let A += BLOCK_K * SPLIT_K * stride_ak;
261// let B += BLOCK_K * SPLIT_K * stride_bk;
262// }
263
264// let acc = acc.to(C.dtype.element_ty)
265
266// // rematerialize rm and rn to save registers
267// let rm = pid_m * BLOCK_M + triton::arange(0, BLOCK_M, 1);
268// let rn = pid_n * BLOCK_N + triton::arange(0, BLOCK_N, 1);
269// let C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn);
270// let mask = (rm < M)[:, None] & (rn < N)[None, :];
271
272// // handles write-back with reduction-splitting
273// if SPLIT_K == 1 {
274// triton::store(C, acc, mask=mask)
275// } else {
276// triton::atomic_add(C, acc, mask=mask)
277// }
278// }
279
280// fn matmul(a: &DenseTensor<DynamicShape, f32>, b: &DenseTensor<DynamicShape, f32>, acc_dtype: usize,
281// input_precision: usize, fp8_fast_accum: usize, output_dtype: usize) {
282// let device = a.device
283// // handle non-contiguous inputs if necessary
284// if a.stride(0) > 1 and a.stride(1) > 1 {
285// a = a.contiguous()
286// }
287// if b.stride(0) > 1 and b.stride(1) > 1 {
288// b = b.contiguous()
289// }
290// // checks constraints
291// assert a.shape[1] == b.shape[0], "incompatible dimensions {a.shape} and {b.shape}";
292// let M = a.shape[0];
293// let K = a.shape[1];
294// let N = b.shape[1];
295
296// // common type between a and b
297// let ab_dtype = get_higher_dtype(a.dtype, b.dtype)
298
299// // allocates output
300// if output_dtype is None {
301// output_dtype = ab_dtype
302// }
303
304// let c = torch.empty((M, N), device=device, dtype=output_dtype)
305
306// // Allowed types for acc_type given the types of a and b.
307// let supported_acc_dtypes = {
308// torch.float16: (torch.float32, torch.float16),
309// torch.bfloat16: (torch.float32, torch.bfloat16),
310// torch.float32: (torch.float32,),
311// torch.int8: (torch.int32,),
312// }
313
314// if acc_dtype is None {
315// acc_dtype = supported_acc_dtypes[ab_dtype][0]
316// } else {
317// assert isinstance(acc_dtype, torch.dtype), "acc_dtype must be a torch.dtype"
318// assert (
319// acc_dtype in supported_acc_dtypes[a.dtype]
320// ), "acc_dtype not compatible with the type of a"
321// assert (
322// acc_dtype in supported_acc_dtypes[b.dtype]
323// ), "acc_dtype not compatible with the type of b"
324
325// def to_tl_type(ty):
326// return getattr(tl, str(ty).split(".")[-1])
327
328// acc_dtype = to_tl_type(acc_dtype)
329// ab_dtype = to_tl_type(ab_dtype)
330// output_dtype = to_tl_type(output_dtype)
331
332// // Tensor cores support input with mixed float8 types.
333// if a.dtype in [tl.float8e4nv, tl.float8e5] and b.dtype in [
334// tl.float8e4nv,
335// tl.float8e5,
336// ] {
337// ab_dtype = None
338// }
339// // launch kernel
340// grid = lambda META: (
341// triton::cdiv(M, META["BLOCK_M"]) * triton::cdiv(N, META["BLOCK_N"]),
342// META["SPLIT_K"],
343// )
344
345// kernel(
346// a,
347// b,
348// c,
349// M,
350// N,
351// K,
352// a.stride(0),
353// a.stride(1),
354// b.stride(0),
355// b.stride(1),
356// c.stride(0),
357// c.stride(1),
358// acc_dtype,
359// input_precision,
360// fp8_fast_accum,
361// GROUP_M=8,
362// AB_DTYPE=ab_dtype,
363// )
364// return c
365// }