From d758cdb8197c797b481f87aec3cab2bac6a08a17 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Mon, 23 Feb 2026 15:21:11 +0100 Subject: [PATCH 01/40] fewer tests --- tests/conftest.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 345e0b7..60a2840 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,21 +11,21 @@ def channels_out(request): return request.param -@pytest.fixture(params=[1, 2, 3, 4]) +@pytest.fixture(params=[1, 2, 3]) def bsize(request): return request.param -@pytest.fixture(params=[2, 3, 7, 13, 16, 32, 37, 53, 111, 128]) +@pytest.fixture(params=[2, 3, 13, 16, 32, 37, 53, 64]) def other_1(request): return request.param -@pytest.fixture(params=[2, 3, 7, 13, 16, 32, 37, 53, 111, 128]) +@pytest.fixture(params=[2, 3, 13, 16, 32, 37, 53, 64]) def other_2(request): return request.param -@pytest.fixture(params=[2, 3, 7, 13, 16, 32, 37, 53, 111, 128]) +@pytest.fixture(params=[2, 3, 13, 16, 32, 37, 53, 64]) def other_3(request): return request.param From 07c186527d72dfe12e766df1c3fcadb730287df2 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Tue, 24 Feb 2026 11:21:25 +0100 Subject: [PATCH 02/40] conv3d, utils --- kerops/kernels/conv.py | 137 ++++++++++++++++++ kerops/ops/addition.py | 7 +- kerops/ops/avgpool.py | 4 +- kerops/ops/bnrelu.py | 7 +- kerops/ops/conv/__init__.py | 1 + kerops/ops/conv/conv.py | 44 ++++++ kerops/ops/conv/dwconv.py | 9 +- kerops/ops/conv/dwconv_wgrad.py | 9 +- .../ops/linear/linear_bias_relu_linear_add.py | 5 +- .../linear_bias_relu_linear_backward.py | 5 +- kerops/ops/linear/relu_linear_add.py | 5 +- kerops/ops/linear/relu_linear_backward.py | 5 +- kerops/ops/quantization.py | 7 +- kerops/ops/stats.py | 7 +- kerops/utils.py | 7 + 15 files changed, 223 insertions(+), 36 deletions(-) create mode 100644 kerops/kernels/conv.py create mode 100644 kerops/ops/conv/conv.py diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py new file mode 100644 index 0000000..50790e2 --- /dev/null +++ b/kerops/kernels/conv.py @@ -0,0 +1,137 @@ +import triton +from triton import language as tl + +""" +- Conv3dV2: D_cell -> W_cell -> H_cell changed to W_cell -> H_cell -> D_cell +- Conv3dV3: `tl.max_contiguous(tl.multiple_of(channels_offset, CHANNELS), CHANNELS)` +- Conv3dV4: Input channels BLOCKCING +- Conv3dV5: output blocking by HW: 2x2 tile + +No impact: + - `x @ w` and `w @ x` orientation via ORDER parameter + - computations of multiple h, w in one block (WHEN THEY ARE FAR AWAY FROM EACH OTHER) + - Boundary checks outside the loops + - INTERIOR and BOUNDARY masking inside one kernel + - `x: [D_BLOCK, CHANNELS] -> x: [NEAR, D_BLOCK, CHANNELS]` where NEAR=4 represents neighbours; `w: [4, CHANNELS_IN, CHANNELS_OUT]` + - Output channels BLOCKCING (why?) + - Output channels BLOCKCING via expanding grid +""" + +@triton.jit +def _Conv_cl3d_impl_V5( + input_ptr, + weight_ptr, + output_ptr, + H, + W, + D, + D_BLOCK: tl.constexpr, + ACCTYPE: tl.constexpr, + IN_CHANNELS: tl.constexpr, + OUT_CHANNELS: tl.constexpr, + CIN_BLOCK: tl.constexpr, +): + W_cell = tl.program_id(0) + H_cell = tl.program_id(1) + D_cell = tl.program_id(2) + + CIN_STEPS: tl.constexpr = IN_CHANNELS // CIN_BLOCK + + in_channels_offset = tl.arange(0, CIN_BLOCK) + out_channels_offset = tl.arange(0, OUT_CHANNELS) + d_offset = tl.arange(0, D_BLOCK) + d_offset_shifted = d_offset[:, None] + D_cell * D_BLOCK + + input_offset = d_offset[:, None] * IN_CHANNELS + tl.max_contiguous(tl.multiple_of(in_channels_offset, CIN_BLOCK), CIN_BLOCK)[None, :] + output_offset = d_offset[:, None] * OUT_CHANNELS + tl.max_contiguous(tl.multiple_of(out_channels_offset, OUT_CHANNELS), OUT_CHANNELS)[None, :] + weight_offset = in_channels_offset[:, None] * OUT_CHANNELS + tl.max_contiguous(tl.multiple_of(out_channels_offset, OUT_CHANNELS), OUT_CHANNELS)[None, :] + + input_ptr += D_cell * D_BLOCK * IN_CHANNELS + input_ptr += W_cell * 2 * IN_CHANNELS * D + input_ptr += H_cell * 2 * IN_CHANNELS * D * W + + output_ptr += D_cell * D_BLOCK * OUT_CHANNELS + output_ptr += W_cell * 2 * OUT_CHANNELS * D + output_ptr += H_cell * 2 * OUT_CHANNELS * D * W + + acc00 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + acc01 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + acc10 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + acc11 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + + for h_block in tl.static_range(0, 2): + for w_block in tl.static_range(0, 2): + for cin in tl.static_range(0, CIN_STEPS): + for dd in tl.static_range(-1, 2): # MB other order? + w_ptr = ( + weight_ptr + + (dd + 1) * IN_CHANNELS * OUT_CHANNELS + + w_block * IN_CHANNELS * OUT_CHANNELS * 3 + + h_block * IN_CHANNELS * OUT_CHANNELS * 9 + + cin * CIN_BLOCK * OUT_CHANNELS + ) + + weights = [ + [ + tl.load(w_ptr + weight_offset), + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 3), + ], + [ + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 9), + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 12), + ] + ] + + i_ptr = ( + input_ptr + + (h_block * 2 - 1) * IN_CHANNELS * D * W + + (w_block * 2 - 1) * IN_CHANNELS * D + + dd * IN_CHANNELS + + cin * CIN_BLOCK + ) + mask = ((d_offset_shifted + dd) < D) & ((d_offset_shifted + dd) >= 0) + + m00 = ((H_cell * 2 + h_block * 2 - 1) < H) & ((H_cell * 2 + h_block * 2 - 1) >= 0) & ((W_cell * 2 + w_block * 2 - 1) < W) & ((W_cell * 2 + w_block * 2 - 1) >= 0) + m01 = ((H_cell * 2 + h_block * 2 - 1) < H) & ((H_cell * 2 + h_block * 2 - 1) >= 0) & ((W_cell * 2 + w_block * 2) < W) & ((W_cell * 2 + w_block * 2) >= 0) + m10 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2 - 1) < W) & ((W_cell * 2 + w_block * 2 - 1) >= 0) + m11 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2) < W) & ((W_cell * 2 + w_block * 2) >= 0) + + xs = [ + [ + tl.load(i_ptr + input_offset, mask=mask & m00, other=0.0), + tl.load(i_ptr + input_offset + IN_CHANNELS * D, mask=mask & m01, other=0.0) + ], + [ + tl.load(i_ptr + input_offset + IN_CHANNELS * D * W, mask=mask & m10, other=0.0), + tl.load(i_ptr + input_offset + IN_CHANNELS * D + IN_CHANNELS * D * W, mask=mask & m11, other=0.0) + ] + ] + + for h in tl.static_range(0, 2): + for w in tl.static_range(0, 2): + # h_weight_idx = 2 * h_block + h - acc_abs_h + 1 - h_block <-- weights window shift + # <---x_h-------> ^-- +1 since weight indexed from 0 + + # acc00 + if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) < 3): + acc00 += tl.dot(xs[h][w], weights[h_block + h][w_block + w]) + + # acc01 + if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) > 0): + acc01 += tl.dot(xs[h][w], weights[h_block + h][w_block + w - 1]) + + # acc10 + if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) < 3): + acc10 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w]) + + # acc11 + if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) > 0): + acc11 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w - 1]) + + + # ADD BOUNDARY CHECK + omask = d_offset_shifted < D + tl.store(output_ptr + output_offset, acc00, mask=omask & ((W_cell * 2) < W) & ((H_cell * 2) < H)) + tl.store(output_ptr + output_offset + OUT_CHANNELS * D, acc01, mask=omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2) < H)) + tl.store(output_ptr + output_offset + OUT_CHANNELS * D * W, acc10, mask=omask & ((W_cell * 2) < W) & ((H_cell * 2 + 1) < H)) + tl.store(output_ptr + output_offset + OUT_CHANNELS * D * W + OUT_CHANNELS * D, acc11, mask=omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2 + 1) < H)) diff --git a/kerops/ops/addition.py b/kerops/ops/addition.py index 532d1b4..22ac7f1 100644 --- a/kerops/ops/addition.py +++ b/kerops/ops/addition.py @@ -1,11 +1,12 @@ from functools import reduce -from math import ceil, floor, log2 +from math import floor, log2 import torch from triton import next_power_of_2 from ..kernels.addition import _AddStats_cl3d_backward_impl, _AddStats_cl3d_impl from ..settings import ConfigurableArg, configure, get_l1_cache +from ...utils import cdiv @configure(l1_cache_bytes=get_l1_cache, num_warps=8) @@ -24,7 +25,7 @@ def AddStats(x, y, inplace=False, *, l1_cache_bytes: ConfigurableArg, num_warps: other = min(MAX_SIZE // num_channels, numel_no_channels) other = int(2 ** (floor(log2(other)))) BLOCK_SIZE = num_channels * other - grid_size = ceil(numel / BLOCK_SIZE) + grid_size = cdiv(numel, BLOCK_SIZE) if inplace: output = x @@ -68,7 +69,7 @@ def AddStatsBackward( other = min(MAX_SIZE // num_channels, numel_no_channels) other = int(2 ** (floor(log2(other)))) BLOCK_SIZE = num_channels * other - grid_size = ceil(numel / BLOCK_SIZE) + grid_size = cdiv(numel, BLOCK_SIZE) output_grad = torch.empty_like(add_grad) diff --git a/kerops/ops/avgpool.py b/kerops/ops/avgpool.py index 19df6e9..1ac44b7 100644 --- a/kerops/ops/avgpool.py +++ b/kerops/ops/avgpool.py @@ -1,11 +1,11 @@ from functools import reduce -from math import ceil import torch from triton import next_power_of_2 from ..kernels.avgpool import _AvgPoolCeilStats_cl3d_backward_impl, _AvgPoolCeilStats_cl3d_impl from ..settings import ConfigurableArg, configure, get_l1_cache +from ...utils import cdiv @configure( @@ -26,7 +26,7 @@ def AvgPoolCeilStats(x, *, l1_cache_bytes: ConfigurableArg, num_warps: Configura BLOCK_SIZE = next_power_of_2(input_d * num_channels) almost_half_d = BLOCK_SIZE // (2 * num_channels) - out_shape = [x.shape[0]] + [ceil(sh / 2) for sh in x.shape[2:]] + [x.shape[1]] + out_shape = [x.shape[0]] + [cdiv(sh, 2) for sh in x.shape[2:]] + [x.shape[1]] output = torch.empty(out_shape, dtype=torch.float16, device=x.device).permute(0, 4, 1, 2, 3) grid_batch, _, grid_H, grid_W, _ = output.shape diff --git a/kerops/ops/bnrelu.py b/kerops/ops/bnrelu.py index f51d14a..de45391 100644 --- a/kerops/ops/bnrelu.py +++ b/kerops/ops/bnrelu.py @@ -1,11 +1,12 @@ from functools import reduce -from math import ceil, floor, log2 +from math import floor, log2 import torch from triton import next_power_of_2 from ..kernels.bnrelu import _ApplyBNReLU_cl3d_backward_impl, _ApplyBNReLU_cl3d_impl from ..settings import ConfigurableArg, configure, get_l1_cache +from ...utils import cdiv @configure(l1_cache_bytes=get_l1_cache, num_warps=8) @@ -23,7 +24,7 @@ def ApplyBNReLU(x, weight, bias, *, l1_cache_bytes: ConfigurableArg, num_warps: other = min(MAX_SIZE // num_channels, numel_no_channels) other = int(2 ** (floor(log2(other)))) BLOCK_SIZE = num_channels * other - grid_size = ceil(numel / BLOCK_SIZE) + grid_size = cdiv(numel, BLOCK_SIZE) output = torch.empty_like(x) @@ -59,7 +60,7 @@ def ApplyBNReLUBackward(x, weight, bias, grad, *, l1_cache_bytes: ConfigurableAr other = min(MAX_SIZE // num_channels, numel_no_channels) other = int(2 ** (floor(log2(other)))) BLOCK_SIZE = num_channels * other - grid_size = ceil(numel / BLOCK_SIZE) + grid_size = cdiv(numel, BLOCK_SIZE) outgrad = torch.empty_like(x) weight_grad = torch.zeros_like(weight) diff --git a/kerops/ops/conv/__init__.py b/kerops/ops/conv/__init__.py index 9c8bf22..2cd5276 100644 --- a/kerops/ops/conv/__init__.py +++ b/kerops/ops/conv/__init__.py @@ -1,2 +1,3 @@ +from .conv import Conv3d from .dwconv import DWConv from .dwconv_wgrad import DWConvWGRAD diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py new file mode 100644 index 0000000..0c41b0e --- /dev/null +++ b/kerops/ops/conv/conv.py @@ -0,0 +1,44 @@ +import torch +from triton import language as tl, next_power_of_2 + +from ...kernels.conv import _Conv_cl3d_impl_V5 +from ...settings import ConfigurableArg, configure +from ...utils import cdiv + + +@configure( + ACCTYPE='float32', + num_warps=4, + D_BLOCK=64, + CIN_BLOCK=16, +) +def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg): + channels, H, W, D = x.shape[1:] + + assert x.ndim == 5 + assert x.is_contiguous(memory_format=torch.channels_last_3d) + assert x.dtype == weight.dtype == torch.float16 + assert channels == next_power_of_2(channels) + assert list(weight.shape) == [3, 3, 3, channels, channels] # 3, 3, 3, out_channels, in_channels + assert D_BLOCK == next_power_of_2(D_BLOCK) + + ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] + output = torch.empty_like(x) + grid = (cdiv(W, 2), cdiv(H, 2), cdiv(D, D_BLOCK)) + + for unbatched_x, unbatched_y in zip(x, output): + _Conv_cl3d_impl_V5[grid]( + unbatched_x, + weight, + unbatched_y, + H, + W, + D, + D_BLOCK=D_BLOCK, + ACCTYPE=ACCTYPE, + CHANNELS=channels, + CIN_BLOCK=CIN_BLOCK, + num_warps=num_warps, + ) + + return output diff --git a/kerops/ops/conv/dwconv.py b/kerops/ops/conv/dwconv.py index 3827c19..7f9d7e3 100644 --- a/kerops/ops/conv/dwconv.py +++ b/kerops/ops/conv/dwconv.py @@ -1,10 +1,9 @@ -from math import ceil - import torch from triton import language as tl, next_power_of_2 from ...kernels.dw_conv import _DWConv_cl3d_impl from ...settings import ConfigurableArg, confexc, configure +from ...utils import cdiv @confexc(KeyError) @@ -39,9 +38,9 @@ def DWConv(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D output = torch.empty_like(x) - H_grid = ceil(H / 2) - W_grid = ceil(W / 2) - D_grid = ceil(D / D_block) + H_grid = cdiv(H, 2) + W_grid = cdiv(W, 2) + D_grid = cdiv(D, D_block) grid = (H_grid, W_grid, D_grid) for unbatched_x, unbatched_y in zip(x, output): diff --git a/kerops/ops/conv/dwconv_wgrad.py b/kerops/ops/conv/dwconv_wgrad.py index 8a0848c..b4bcd92 100644 --- a/kerops/ops/conv/dwconv_wgrad.py +++ b/kerops/ops/conv/dwconv_wgrad.py @@ -1,10 +1,9 @@ -from math import ceil - import torch from triton import language as tl, next_power_of_2 from ...kernels.dw_conv import _DWConv_wgrad_cl3d_impl from ...settings import ConfigurableArg, confexc, configure +from ...utils import cdiv @confexc(KeyError) @@ -46,9 +45,9 @@ def DWConvWGRAD( bsize, _, H, W, D = x.shape batch_stride, _, H_stride, W_stride, _ = x.stride() - H_grid = ceil(H / (2 * ILP)) - W_grid = ceil(W / 2) - D_grid = ceil(D / D_block) + H_grid = cdiv(H, 2 * ILP) + W_grid = cdiv(W, 2) + D_grid = cdiv(D, D_block) grid = (H_grid, W_grid, D_grid) grad_w = torch.zeros([bsize, H_grid * W_grid * D_grid, 3, 3, 3, channels], device=x.device, dtype=torch.float16) diff --git a/kerops/ops/linear/linear_bias_relu_linear_add.py b/kerops/ops/linear/linear_bias_relu_linear_add.py index a0f425d..93b5406 100644 --- a/kerops/ops/linear/linear_bias_relu_linear_add.py +++ b/kerops/ops/linear/linear_bias_relu_linear_add.py @@ -1,10 +1,9 @@ -from math import ceil - import torch from triton import next_power_of_2 from ...kernels.linear import _LinBReLULinAdd from ...settings import ConfigurableArg, confexc, configure +from ...utils import cdiv @confexc(KeyError) @@ -50,7 +49,7 @@ def LinBReLULinAdd( assert add_other.is_contiguous(memory_format=torch.channels_last_3d) numel_no_channels = numel // in_channels - grid_size = ceil(numel_no_channels / (D_block * ILP)) + grid_size = cdiv(numel_no_channels, D_block * ILP) output = torch.empty_like(x) diff --git a/kerops/ops/linear/linear_bias_relu_linear_backward.py b/kerops/ops/linear/linear_bias_relu_linear_backward.py index 35b70af..9be1a7e 100644 --- a/kerops/ops/linear/linear_bias_relu_linear_backward.py +++ b/kerops/ops/linear/linear_bias_relu_linear_backward.py @@ -1,10 +1,9 @@ -from math import ceil - import torch from triton import next_power_of_2 from ...kernels.linear import _LinBReLULinBackward from ...settings import ConfigurableArg, confexc, configure +from ...utils import cdiv @confexc(KeyError) @@ -43,7 +42,7 @@ def LinBReLULinBackward( assert grad.is_contiguous(memory_format=torch.channels_last_3d) numel_no_channels = numel // in_channels - grid_size = ceil(numel_no_channels / (D_block * ILP)) + grid_size = cdiv(numel_no_channels, D_block * ILP) x_grad = torch.empty_like(x) weight_up_grad = torch.zeros([grid_size, in_channels, hidden_channels], dtype=torch.float16, device='cuda') diff --git a/kerops/ops/linear/relu_linear_add.py b/kerops/ops/linear/relu_linear_add.py index 5d828c1..2cbbcf8 100644 --- a/kerops/ops/linear/relu_linear_add.py +++ b/kerops/ops/linear/relu_linear_add.py @@ -1,10 +1,9 @@ -from math import ceil - import torch from triton import next_power_of_2 from ...kernels.linear import _ReLULinearAdd from ...settings import ConfigurableArg, confexc, configure +from ...utils import cdiv @confexc(KeyError) @@ -41,7 +40,7 @@ def ReLULinearAdd( assert add_other.is_contiguous(memory_format=torch.channels_last_3d) numel_no_channels = numel // in_channels - grid_size = ceil(numel_no_channels / (D_block * ILP)) + grid_size = cdiv(numel_no_channels, D_block * ILP) output = torch.empty_like(add_other) diff --git a/kerops/ops/linear/relu_linear_backward.py b/kerops/ops/linear/relu_linear_backward.py index a8994eb..d3080fb 100644 --- a/kerops/ops/linear/relu_linear_backward.py +++ b/kerops/ops/linear/relu_linear_backward.py @@ -1,10 +1,9 @@ -from math import ceil - import torch from triton import next_power_of_2 from ...kernels.linear import _ReLULinearAddBackward from ...settings import ConfigurableArg, confexc, configure +from ...utils import cdiv @confexc(KeyError) @@ -47,7 +46,7 @@ def ReLULinearBackward( numel_no_channels = numel // out_channels - grid_size = ceil(numel_no_channels / (D_block * ILP)) + grid_size = cdiv(numel_no_channels, D_block * ILP) bsize, _, H, W, D = grad.shape x_grad = torch.empty_like(x) diff --git a/kerops/ops/quantization.py b/kerops/ops/quantization.py index d5054b5..8108a06 100644 --- a/kerops/ops/quantization.py +++ b/kerops/ops/quantization.py @@ -1,9 +1,10 @@ -from math import ceil, floor, log2 +from math import floor, log2 import torch from ..kernels.quantization import _DequantUint8Window_impl, _QuantUint8Window_impl from ..settings import ConfigurableArg, configure, get_l1_cache +from ...utils import cdiv @configure(num_warps=4, l1_cache_bytes=get_l1_cache) @@ -14,7 +15,7 @@ def QuantUint8Window(x, window, *, num_warps: ConfigurableArg, l1_cache_bytes: C BLOCK_SIZE = int(2 ** (floor(log2(BLOCK_SIZE)))) output = torch.empty_like(x, dtype=torch.uint8) - numblocks = ceil(numel / BLOCK_SIZE) + numblocks = cdiv(numel, BLOCK_SIZE) _QuantUint8Window_impl[(numblocks,)](x, output, numel, window, BLOCK_SIZE, num_warps=num_warps) return output @@ -29,7 +30,7 @@ def DequantUint8Window(x, init_dtype, window, num_warps: ConfigurableArg, l1_cac BLOCK_SIZE = min(MAX_SIZE, numel) BLOCK_SIZE = int(2 ** (floor(log2(BLOCK_SIZE)))) - numblocks = ceil(numel / BLOCK_SIZE) + numblocks = cdiv(numel, BLOCK_SIZE) _DequantUint8Window_impl[(numblocks,)](x, output, numel, window, BLOCK_SIZE, num_warps=num_warps) return output diff --git a/kerops/ops/stats.py b/kerops/ops/stats.py index cea4043..0a46799 100644 --- a/kerops/ops/stats.py +++ b/kerops/ops/stats.py @@ -1,11 +1,12 @@ from functools import reduce -from math import ceil, floor, log2 +from math import floor, log2 import torch from triton import next_power_of_2 from ..kernels.stats import _Stats_cl3d_backward_impl, _Stats_cl3d_impl from ..settings import ConfigurableArg, configure, get_l1_cache +from ...utils import cdiv @configure(l1_cache_bytes=get_l1_cache, num_warps=4) @@ -22,7 +23,7 @@ def Stats(x, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): other = min(MAX_SIZE // num_channels, numel_no_channels) other = int(2 ** (floor(log2(other)))) BLOCK_SIZE = num_channels * other - grid_size = ceil(numel / BLOCK_SIZE) + grid_size = cdiv(numel, BLOCK_SIZE) mean = torch.zeros(num_channels, dtype=torch.float32, device=x.device) sqmean = torch.zeros(num_channels, dtype=torch.float32, device=x.device) @@ -46,7 +47,7 @@ def StatsBackward(x, mean_grad, sqmean_grad, *, l1_cache_bytes: ConfigurableArg, other = min(MAX_SIZE // num_channels, numel_no_channels) other = int(2 ** (floor(log2(other)))) BLOCK_SIZE = num_channels * other - grid_size = ceil(numel / BLOCK_SIZE) + grid_size = cdiv(numel, BLOCK_SIZE) output_grad = torch.empty_like(x) diff --git a/kerops/utils.py b/kerops/utils.py index 61d70e6..337f0c4 100644 --- a/kerops/utils.py +++ b/kerops/utils.py @@ -100,3 +100,10 @@ def weight_grad_similarity(input, other, rtol_cos=1e-4, atol_cos=1e-4, rtol_len= debug_info('All stages passed') return True + + +def cdiv(a, b): + assert type(a) == int + assert type(b) == int + + return (a + b - 1) // b From ecbd7bf0f6c7ac51ead580a067fa62993c14e247 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Tue, 24 Feb 2026 11:55:42 +0100 Subject: [PATCH 03/40] fix --- kerops/ops/conv/conv.py | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py index 0c41b0e..19019d4 100644 --- a/kerops/ops/conv/conv.py +++ b/kerops/ops/conv/conv.py @@ -13,17 +13,30 @@ CIN_BLOCK=16, ) def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg): - channels, H, W, D = x.shape[1:] + assert x.device == weight.device + assert x.is_cuda assert x.ndim == 5 + assert weight.ndim == 5 + bsize, x_channels, H, W, D = x.shape + in_channels, out_channels = weight.shape[-2:] + assert x_channels == next_power_of_2(x_channels) + assert list(weight.shape) == [3, 3, 3, in_channels, out_channels] # 3, 3, 3, in_channels, out_channels + assert in_channels == x_channels + assert out_channels == next_power_of_2(out_channels) + assert x.is_contiguous(memory_format=torch.channels_last_3d) + assert weight.is_contiguous() + assert x.dtype == weight.dtype == torch.float16 - assert channels == next_power_of_2(channels) - assert list(weight.shape) == [3, 3, 3, channels, channels] # 3, 3, 3, out_channels, in_channels + assert D_BLOCK == next_power_of_2(D_BLOCK) + assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) + assert CIN_BLOCK <= in_channels + assert ACCTYPE in ('float16', 'float32') ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] - output = torch.empty_like(x) + output = torch.empty([bsize, H, W, D, out_channels], device=x.device, dtype=x.dtype, layout=x.layout).permute(0, -1, 1, 2, 3) grid = (cdiv(W, 2), cdiv(H, 2), cdiv(D, D_BLOCK)) for unbatched_x, unbatched_y in zip(x, output): @@ -36,7 +49,8 @@ def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D D, D_BLOCK=D_BLOCK, ACCTYPE=ACCTYPE, - CHANNELS=channels, + IN_CHANNELS=in_channels, + OUT_CHANNELS=out_channels, CIN_BLOCK=CIN_BLOCK, num_warps=num_warps, ) From c4c70b6d236d025234ef266117f22c708d092324 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Thu, 26 Feb 2026 17:55:39 +0100 Subject: [PATCH 04/40] conv3d configure --- kerops/ops/conv/conv.py | 56 ++++++++++++++++++++++++++++++++++++++--- 1 file changed, 52 insertions(+), 4 deletions(-) diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py index 19019d4..fcb57ea 100644 --- a/kerops/ops/conv/conv.py +++ b/kerops/ops/conv/conv.py @@ -2,15 +2,63 @@ from triton import language as tl, next_power_of_2 from ...kernels.conv import _Conv_cl3d_impl_V5 -from ...settings import ConfigurableArg, configure +from ...settings import ConfigurableArg, configure, confexc from ...utils import cdiv +@confexc(KeyError) +def num_warps(in_channels, out_channels): + return { + (16, 16): 4, + (16, 32): 4, + (32, 16): 4, + (32, 32): 2, + (32, 64): 4, + (64, 32): 1, + (64, 64): 2, + (64, 128): 4, + (128, 64): 4, + (128, 128): 4, + }[(in_channels, out_channels)] + + +@confexc(KeyError) +def d_block(in_channels, out_channels): + return { + (16, 16): 64, + (16, 32): 64, + (32, 16): 64, + (32, 32): 32, + (32, 64): 32, + (64, 32): 32, + (64, 64): 32, + (64, 128): 16, + (128, 64): 16, + (128, 128): 16, + }[(in_channels, out_channels)] + + +@confexc(KeyError) +def cin_block(in_channels, out_channels): + return { + (16, 16): 16, + (16, 32): 16, + (32, 16): 32, + (32, 32): 16, + (32, 64): 16, + (64, 32): 16, + (64, 64): 16, + (64, 128): 16, + (128, 64): 16, + (128, 128): 16, + }[(in_channels, out_channels)] + + @configure( ACCTYPE='float32', - num_warps=4, - D_BLOCK=64, - CIN_BLOCK=16, + num_warps=lambda weight: num_warps(*weight.shape[-2:]), + D_BLOCK=lambda weight: d_block(*weight.shape[-2:]), + CIN_BLOCK=lambda weight: cin_block(*weight.shape[-2:]), ) def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg): assert x.device == weight.device From 7dec7475d9763a3fa7b8a7361c39974eccbdc80f Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 27 Feb 2026 11:14:06 +0100 Subject: [PATCH 05/40] ApplyBNReLUConv3d draft --- kerops/kernels/conv.py | 124 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 123 insertions(+), 1 deletion(-) diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index 50790e2..578406d 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -128,8 +128,130 @@ def _Conv_cl3d_impl_V5( if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) > 0): acc11 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w - 1]) + omask = d_offset_shifted < D + tl.store(output_ptr + output_offset, acc00, mask=omask & ((W_cell * 2) < W) & ((H_cell * 2) < H)) + tl.store(output_ptr + output_offset + OUT_CHANNELS * D, acc01, mask=omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2) < H)) + tl.store(output_ptr + output_offset + OUT_CHANNELS * D * W, acc10, mask=omask & ((W_cell * 2) < W) & ((H_cell * 2 + 1) < H)) + tl.store(output_ptr + output_offset + OUT_CHANNELS * D * W + OUT_CHANNELS * D, acc11, mask=omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2 + 1) < H)) + + +@triton.jit +def _ApplyBNReLUConv_cl3d_impl( + input_ptr, + bn_weight_ptr, + bn_bias_ptr, + weight_ptr, + output_ptr, + H, + W, + D, + D_BLOCK: tl.constexpr, + ACCTYPE: tl.constexpr, + IN_CHANNELS: tl.constexpr, + OUT_CHANNELS: tl.constexpr, + CIN_BLOCK: tl.constexpr, +): + W_cell = tl.program_id(0) + H_cell = tl.program_id(1) + D_cell = tl.program_id(2) + + CIN_STEPS: tl.constexpr = IN_CHANNELS // CIN_BLOCK + + in_channels_offset = tl.arange(0, CIN_BLOCK) + out_channels_offset = tl.arange(0, OUT_CHANNELS) + d_offset = tl.arange(0, D_BLOCK) + d_offset_shifted = d_offset[:, None] + D_cell * D_BLOCK + + input_offset = d_offset[:, None] * IN_CHANNELS + tl.max_contiguous(tl.multiple_of(in_channels_offset, CIN_BLOCK), CIN_BLOCK)[None, :] + output_offset = d_offset[:, None] * OUT_CHANNELS + tl.max_contiguous(tl.multiple_of(out_channels_offset, OUT_CHANNELS), OUT_CHANNELS)[None, :] + weight_offset = in_channels_offset[:, None] * OUT_CHANNELS + tl.max_contiguous(tl.multiple_of(out_channels_offset, OUT_CHANNELS), OUT_CHANNELS)[None, :] + + input_ptr += D_cell * D_BLOCK * IN_CHANNELS + input_ptr += W_cell * 2 * IN_CHANNELS * D + input_ptr += H_cell * 2 * IN_CHANNELS * D * W + + output_ptr += D_cell * D_BLOCK * OUT_CHANNELS + output_ptr += W_cell * 2 * OUT_CHANNELS * D + output_ptr += H_cell * 2 * OUT_CHANNELS * D * W + + acc00 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + acc01 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + acc10 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + acc11 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + + for h_block in tl.static_range(0, 2): + for w_block in tl.static_range(0, 2): + for cin in tl.static_range(0, CIN_STEPS): + bn_weight = tl.load(bn_weight_ptr + in_channels_offset + cin * CIN_BLOCK)[None, :] + bn_bias = tl.load(bn_bias_ptr + in_channels_offset + cin * CIN_BLOCK)[None, :] + zero = tl.zeros([1], dtype=tl.float16) + + for dd in tl.static_range(-1, 2): # MB other order? + w_ptr = ( + weight_ptr + + (dd + 1) * IN_CHANNELS * OUT_CHANNELS + + w_block * IN_CHANNELS * OUT_CHANNELS * 3 + + h_block * IN_CHANNELS * OUT_CHANNELS * 9 + + cin * CIN_BLOCK * OUT_CHANNELS + ) + + weights = [ + [ + tl.load(w_ptr + weight_offset), + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 3), + ], + [ + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 9), + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 12), + ] + ] + + i_ptr = ( + input_ptr + + (h_block * 2 - 1) * IN_CHANNELS * D * W + + (w_block * 2 - 1) * IN_CHANNELS * D + + dd * IN_CHANNELS + + cin * CIN_BLOCK + ) + mask = ((d_offset_shifted + dd) < D) & ((d_offset_shifted + dd) >= 0) + + m00 = ((H_cell * 2 + h_block * 2 - 1) < H) & ((H_cell * 2 + h_block * 2 - 1) >= 0) & ((W_cell * 2 + w_block * 2 - 1) < W) & ((W_cell * 2 + w_block * 2 - 1) >= 0) + m01 = ((H_cell * 2 + h_block * 2 - 1) < H) & ((H_cell * 2 + h_block * 2 - 1) >= 0) & ((W_cell * 2 + w_block * 2) < W) & ((W_cell * 2 + w_block * 2) >= 0) + m10 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2 - 1) < W) & ((W_cell * 2 + w_block * 2 - 1) >= 0) + m11 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2) < W) & ((W_cell * 2 + w_block * 2) >= 0) + + xs = [ + [ + tl.maximum(tl.fma(tl.load(i_ptr + input_offset, mask=mask & m00, other=0.0), bn_weight, bn_bias), tl.zeros([1], dtype=tl.float16)), + tl.maximum(tl.fma(tl.load(i_ptr + input_offset + IN_CHANNELS * D, mask=mask & m01, other=0.0), bn_weight, bn_bias), tl.zeros([1], dtype=tl.float16)) + ], + [ + tl.maximum(tl.fma(tl.load(i_ptr + input_offset + IN_CHANNELS * D * W, mask=mask & m10, other=0.0), bn_weight, bn_bias), tl.zeros([1], dtype=tl.float16)), + tl.maximum(tl.fma(tl.load(i_ptr + input_offset + IN_CHANNELS * D + IN_CHANNELS * D * W, mask=mask & m11, other=0.0), bn_weight, bn_bias), tl.zeros([1], dtype=tl.float16)) + ] + ] + + for h in tl.static_range(0, 2): + for w in tl.static_range(0, 2): + # h_weight_idx = 2 * h_block + h - acc_abs_h + 1 - h_block <-- weights window shift + # <---x_h-------> ^-- +1 since weight indexed from 0 + + # acc00 + if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) < 3): + acc00 += tl.dot(xs[h][w], weights[h_block + h][w_block + w]) + + # acc01 + if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) > 0): + acc01 += tl.dot(xs[h][w], weights[h_block + h][w_block + w - 1]) + + # acc10 + if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) < 3): + acc10 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w]) + + # acc11 + if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) > 0): + acc11 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w - 1]) - # ADD BOUNDARY CHECK omask = d_offset_shifted < D tl.store(output_ptr + output_offset, acc00, mask=omask & ((W_cell * 2) < W) & ((H_cell * 2) < H)) tl.store(output_ptr + output_offset + OUT_CHANNELS * D, acc01, mask=omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2) < H)) From 5e3b1b81f080c99e57a10f873182a44922be7d29 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 27 Feb 2026 14:06:05 +0100 Subject: [PATCH 06/40] ApplyBNReLUConv3d --- kerops/kernels/conv.py | 35 +++++++++++++++------- kerops/ops/conv/__init__.py | 2 +- kerops/ops/conv/conv.py | 58 ++++++++++++++++++++++++++++++++++++- 3 files changed, 83 insertions(+), 12 deletions(-) diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index 578406d..c7bd37d 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -179,12 +179,13 @@ def _ApplyBNReLUConv_cl3d_impl( acc10 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) acc11 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + zero = tl.zeros([1], dtype=tl.float16) + for h_block in tl.static_range(0, 2): for w_block in tl.static_range(0, 2): for cin in tl.static_range(0, CIN_STEPS): bn_weight = tl.load(bn_weight_ptr + in_channels_offset + cin * CIN_BLOCK)[None, :] bn_bias = tl.load(bn_bias_ptr + in_channels_offset + cin * CIN_BLOCK)[None, :] - zero = tl.zeros([1], dtype=tl.float16) for dd in tl.static_range(-1, 2): # MB other order? w_ptr = ( @@ -222,12 +223,12 @@ def _ApplyBNReLUConv_cl3d_impl( xs = [ [ - tl.maximum(tl.fma(tl.load(i_ptr + input_offset, mask=mask & m00, other=0.0), bn_weight, bn_bias), tl.zeros([1], dtype=tl.float16)), - tl.maximum(tl.fma(tl.load(i_ptr + input_offset + IN_CHANNELS * D, mask=mask & m01, other=0.0), bn_weight, bn_bias), tl.zeros([1], dtype=tl.float16)) + tl.load(i_ptr + input_offset, mask=mask & m00), + tl.load(i_ptr + input_offset + IN_CHANNELS * D, mask=mask & m01) ], [ - tl.maximum(tl.fma(tl.load(i_ptr + input_offset + IN_CHANNELS * D * W, mask=mask & m10, other=0.0), bn_weight, bn_bias), tl.zeros([1], dtype=tl.float16)), - tl.maximum(tl.fma(tl.load(i_ptr + input_offset + IN_CHANNELS * D + IN_CHANNELS * D * W, mask=mask & m11, other=0.0), bn_weight, bn_bias), tl.zeros([1], dtype=tl.float16)) + tl.load(i_ptr + input_offset + IN_CHANNELS * D * W, mask=mask & m10), + tl.load(i_ptr + input_offset + IN_CHANNELS * D + IN_CHANNELS * D * W, mask=mask & m11) ] ] @@ -235,22 +236,36 @@ def _ApplyBNReLUConv_cl3d_impl( for w in tl.static_range(0, 2): # h_weight_idx = 2 * h_block + h - acc_abs_h + 1 - h_block <-- weights window shift # <---x_h-------> ^-- +1 since weight indexed from 0 - + + x = xs[h][w].to(tl.float32) + x = x * bn_weight + bn_bias + x = x.to(tl.float16) + x = tl.maximum(x, zero) + x = tl.where( + mask + & ((H_cell * 2 + h_block * 2 - 1 + h) < H) + & ((H_cell * 2 + h_block * 2 - 1 + h) >= 0) + & ((W_cell * 2 + w_block * 2 - 1 + w) < W) + & ((W_cell * 2 + w_block * 2 - 1 + w) >= 0), + x, + zero + ) + # acc00 if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) < 3): - acc00 += tl.dot(xs[h][w], weights[h_block + h][w_block + w]) + acc00 += tl.dot(x, weights[h_block + h][w_block + w]) # acc01 if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) > 0): - acc01 += tl.dot(xs[h][w], weights[h_block + h][w_block + w - 1]) + acc01 += tl.dot(x, weights[h_block + h][w_block + w - 1]) # acc10 if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) < 3): - acc10 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w]) + acc10 += tl.dot(x, weights[h_block + h - 1][w_block + w]) # acc11 if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) > 0): - acc11 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w - 1]) + acc11 += tl.dot(x, weights[h_block + h - 1][w_block + w - 1]) omask = d_offset_shifted < D tl.store(output_ptr + output_offset, acc00, mask=omask & ((W_cell * 2) < W) & ((H_cell * 2) < H)) diff --git a/kerops/ops/conv/__init__.py b/kerops/ops/conv/__init__.py index 2cd5276..fe8ddd9 100644 --- a/kerops/ops/conv/__init__.py +++ b/kerops/ops/conv/__init__.py @@ -1,3 +1,3 @@ -from .conv import Conv3d +from .conv import Conv3d, ApplyBNReLUConv3d from .dwconv import DWConv from .dwconv_wgrad import DWConvWGRAD diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py index fcb57ea..fe47030 100644 --- a/kerops/ops/conv/conv.py +++ b/kerops/ops/conv/conv.py @@ -1,7 +1,7 @@ import torch from triton import language as tl, next_power_of_2 -from ...kernels.conv import _Conv_cl3d_impl_V5 +from ...kernels.conv import _Conv_cl3d_impl_V5, _ApplyBNReLUConv_cl3d_impl from ...settings import ConfigurableArg, configure, confexc from ...utils import cdiv @@ -104,3 +104,59 @@ def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D ) return output + + +@configure( + ACCTYPE='float32', + num_warps=lambda weight: num_warps(*weight.shape[-2:]), + D_BLOCK=lambda weight: d_block(*weight.shape[-2:]), + CIN_BLOCK=lambda weight: cin_block(*weight.shape[-2:]), +) +def ApplyBNReLUConv3d(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg): + assert x.device == weight.device == bn_weight.device == bn_bias.device + assert x.is_cuda + + assert x.ndim == 5 + assert weight.ndim == 5 + assert bn_weight.ndim == bn_bias.ndim == 1 + bsize, x_channels, H, W, D = x.shape + in_channels, out_channels = weight.shape[-2:] + assert x_channels == next_power_of_2(x_channels) + assert list(weight.shape) == [3, 3, 3, in_channels, out_channels] # 3, 3, 3, in_channels, out_channels + assert in_channels == x_channels == bn_weight.numel() == bn_bias.numel() + assert out_channels == next_power_of_2(out_channels) + + assert x.is_contiguous(memory_format=torch.channels_last_3d) + assert weight.is_contiguous() + + assert x.dtype == weight.dtype == torch.float16 + assert bn_weight.dtype == bn_bias.dtype == torch.float32 + + assert D_BLOCK == next_power_of_2(D_BLOCK) + assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) + assert CIN_BLOCK <= in_channels + assert ACCTYPE in ('float16', 'float32') + + ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] + output = torch.empty([bsize, H, W, D, out_channels], device=x.device, dtype=x.dtype, layout=x.layout).permute(0, -1, 1, 2, 3) + grid = (cdiv(W, 2), cdiv(H, 2), cdiv(D, D_BLOCK)) + + for unbatched_x, unbatched_y in zip(x, output): + _ApplyBNReLUConv_cl3d_impl[grid]( + unbatched_x, + bn_weight, + bn_bias, + weight, + unbatched_y, + H, + W, + D, + D_BLOCK=D_BLOCK, + ACCTYPE=ACCTYPE, + IN_CHANNELS=in_channels, + OUT_CHANNELS=out_channels, + CIN_BLOCK=CIN_BLOCK, + num_warps=num_warps, + ) + + return output From 46ca76dbf6cbe69860124830172698e94efebdbc Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 27 Feb 2026 14:15:54 +0100 Subject: [PATCH 07/40] Neater style --- kerops/kernels/conv.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index c7bd37d..3c81aa2 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -237,19 +237,20 @@ def _ApplyBNReLUConv_cl3d_impl( # h_weight_idx = 2 * h_block + h - acc_abs_h + 1 - h_block <-- weights window shift # <---x_h-------> ^-- +1 since weight indexed from 0 + if h == 0 and w == 0: + valid = mask & m00 + elif h == 0 and w == 1: + valid = mask & m01 + elif h == 1 and w == 0: + valid = mask & m10 + else: + valid = mask & m11 + x = xs[h][w].to(tl.float32) x = x * bn_weight + bn_bias x = x.to(tl.float16) x = tl.maximum(x, zero) - x = tl.where( - mask - & ((H_cell * 2 + h_block * 2 - 1 + h) < H) - & ((H_cell * 2 + h_block * 2 - 1 + h) >= 0) - & ((W_cell * 2 + w_block * 2 - 1 + w) < W) - & ((W_cell * 2 + w_block * 2 - 1 + w) >= 0), - x, - zero - ) + x = tl.where(mask & valid, x, zero) # acc00 if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) < 3): From 8440b702c9d836f3d1401e738a3088b19633e8f1 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 27 Feb 2026 22:53:08 +0100 Subject: [PATCH 08/40] ApplyBNReLUConv3d -> ApplyBNReLUConv3dStats; fixes --- kerops/kernels/conv.py | 32 ++++++++++++++++++++++++++++++-- kerops/ops/addition.py | 2 +- kerops/ops/avgpool.py | 2 +- kerops/ops/bnrelu.py | 2 +- kerops/ops/conv/__init__.py | 2 +- kerops/ops/conv/conv.py | 17 ++++++++++++----- kerops/ops/quantization.py | 2 +- kerops/ops/stats.py | 2 +- 8 files changed, 48 insertions(+), 13 deletions(-) diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index 3c81aa2..35357bc 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -136,12 +136,14 @@ def _Conv_cl3d_impl_V5( @triton.jit -def _ApplyBNReLUConv_cl3d_impl( +def _ApplyBNReLUConvStats_cl3d_impl( input_ptr, bn_weight_ptr, bn_bias_ptr, weight_ptr, output_ptr, + mean_ptr, + sqmean_ptr, H, W, D, @@ -174,11 +176,17 @@ def _ApplyBNReLUConv_cl3d_impl( output_ptr += W_cell * 2 * OUT_CHANNELS * D output_ptr += H_cell * 2 * OUT_CHANNELS * D * W + mean_ptr += (W_cell + tl.num_programs(0) * H_cell + tl.num_programs(0) * tl.num_programs(1) * D_cell) * OUT_CHANNELS + sqmean_ptr += (W_cell + tl.num_programs(0) * H_cell + tl.num_programs(0) * tl.num_programs(1) * D_cell) * OUT_CHANNELS + acc00 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) acc01 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) acc10 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) acc11 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + mean = tl.zeros([OUT_CHANNELS], dtype=ACCTYPE) + sqmean = tl.zeros([OUT_CHANNELS], dtype=ACCTYPE) + zero = tl.zeros([1], dtype=tl.float16) for h_block in tl.static_range(0, 2): @@ -221,6 +229,7 @@ def _ApplyBNReLUConv_cl3d_impl( m10 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2 - 1) < W) & ((W_cell * 2 + w_block * 2 - 1) >= 0) m11 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2) < W) & ((W_cell * 2 + w_block * 2) >= 0) + # other=0.0 is NOT necessary since it is filtrated via tl.where(mask & mXX, x, zero) <-> ReLU(x) xs = [ [ tl.load(i_ptr + input_offset, mask=mask & m00), @@ -250,7 +259,7 @@ def _ApplyBNReLUConv_cl3d_impl( x = x * bn_weight + bn_bias x = x.to(tl.float16) x = tl.maximum(x, zero) - x = tl.where(mask & valid, x, zero) + x = tl.where(valid, x, zero) # acc00 if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) < 3): @@ -273,3 +282,22 @@ def _ApplyBNReLUConv_cl3d_impl( tl.store(output_ptr + output_offset + OUT_CHANNELS * D, acc01, mask=omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2) < H)) tl.store(output_ptr + output_offset + OUT_CHANNELS * D * W, acc10, mask=omask & ((W_cell * 2) < W) & ((H_cell * 2 + 1) < H)) tl.store(output_ptr + output_offset + OUT_CHANNELS * D * W + OUT_CHANNELS * D, acc11, mask=omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2 + 1) < H)) + + mean = tl.sum( + tl.where(omask & ((W_cell * 2) < W) & ((H_cell * 2) < H), acc00, 0.0) + + tl.where(omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2) < H),acc01, 0.0) + + tl.where(omask & ((W_cell * 2) < W) & ((H_cell * 2 + 1) < H), acc10, 0.0) + + tl.where(omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2 + 1) < H), acc11, 0.0), + axis=0 + ) + + sqmean = tl.sum( + tl.where(omask & ((W_cell * 2) < W) & ((H_cell * 2) < H), acc00 * acc00, 0.0) + + tl.where(omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2) < H), acc01 * acc01, 0.0) + + tl.where(omask & ((W_cell * 2) < W) & ((H_cell * 2 + 1) < H), acc10 * acc10, 0.0) + + tl.where(omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2 + 1) < H), acc11 * acc11, 0.0), + axis=0 + ) + + tl.store(mean_ptr + out_channels_offset, mean) + tl.store(sqmean_ptr + out_channels_offset, sqmean) diff --git a/kerops/ops/addition.py b/kerops/ops/addition.py index 22ac7f1..e530a18 100644 --- a/kerops/ops/addition.py +++ b/kerops/ops/addition.py @@ -6,7 +6,7 @@ from ..kernels.addition import _AddStats_cl3d_backward_impl, _AddStats_cl3d_impl from ..settings import ConfigurableArg, configure, get_l1_cache -from ...utils import cdiv +from ..utils import cdiv @configure(l1_cache_bytes=get_l1_cache, num_warps=8) diff --git a/kerops/ops/avgpool.py b/kerops/ops/avgpool.py index 1ac44b7..e222644 100644 --- a/kerops/ops/avgpool.py +++ b/kerops/ops/avgpool.py @@ -5,7 +5,7 @@ from ..kernels.avgpool import _AvgPoolCeilStats_cl3d_backward_impl, _AvgPoolCeilStats_cl3d_impl from ..settings import ConfigurableArg, configure, get_l1_cache -from ...utils import cdiv +from ..utils import cdiv @configure( diff --git a/kerops/ops/bnrelu.py b/kerops/ops/bnrelu.py index de45391..a055d70 100644 --- a/kerops/ops/bnrelu.py +++ b/kerops/ops/bnrelu.py @@ -6,7 +6,7 @@ from ..kernels.bnrelu import _ApplyBNReLU_cl3d_backward_impl, _ApplyBNReLU_cl3d_impl from ..settings import ConfigurableArg, configure, get_l1_cache -from ...utils import cdiv +from ..utils import cdiv @configure(l1_cache_bytes=get_l1_cache, num_warps=8) diff --git a/kerops/ops/conv/__init__.py b/kerops/ops/conv/__init__.py index fe8ddd9..cdc47b2 100644 --- a/kerops/ops/conv/__init__.py +++ b/kerops/ops/conv/__init__.py @@ -1,3 +1,3 @@ -from .conv import Conv3d, ApplyBNReLUConv3d +from .conv import Conv3d, ApplyBNReLUConv3dStats from .dwconv import DWConv from .dwconv_wgrad import DWConvWGRAD diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py index fe47030..15b034f 100644 --- a/kerops/ops/conv/conv.py +++ b/kerops/ops/conv/conv.py @@ -1,7 +1,8 @@ +import numpy as np import torch from triton import language as tl, next_power_of_2 -from ...kernels.conv import _Conv_cl3d_impl_V5, _ApplyBNReLUConv_cl3d_impl +from ...kernels.conv import _Conv_cl3d_impl_V5, _ApplyBNReLUConvStats_cl3d_impl from ...settings import ConfigurableArg, configure, confexc from ...utils import cdiv @@ -112,7 +113,7 @@ def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D D_BLOCK=lambda weight: d_block(*weight.shape[-2:]), CIN_BLOCK=lambda weight: cin_block(*weight.shape[-2:]), ) -def ApplyBNReLUConv3d(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg): +def ApplyBNReLUConv3dStats(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg): assert x.device == weight.device == bn_weight.device == bn_bias.device assert x.is_cuda @@ -140,14 +141,20 @@ def ApplyBNReLUConv3d(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfigurableArg ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] output = torch.empty([bsize, H, W, D, out_channels], device=x.device, dtype=x.dtype, layout=x.layout).permute(0, -1, 1, 2, 3) grid = (cdiv(W, 2), cdiv(H, 2), cdiv(D, D_BLOCK)) + mean = torch.zeros([bsize, np.prod(grid), out_channels], device=x.device, dtype=torch.float32) + sqmean = torch.zeros([bsize, np.prod(grid), out_channels], device=x.device, dtype=torch.float32) - for unbatched_x, unbatched_y in zip(x, output): - _ApplyBNReLUConv_cl3d_impl[grid]( + numel_no_channels = bsize * H * W * D + + for unbatched_x, unbatched_y, unbatched_mean, unbatched_sqmean in zip(x, output, mean, sqmean): + _ApplyBNReLUConvStats_cl3d_impl[grid]( unbatched_x, bn_weight, bn_bias, weight, unbatched_y, + unbatched_mean, + unbatched_sqmean, H, W, D, @@ -159,4 +166,4 @@ def ApplyBNReLUConv3d(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfigurableArg num_warps=num_warps, ) - return output + return output, mean.sum(dim=(0, 1)) / numel_no_channels, sqmean.sum(dim=(0, 1)) / numel_no_channels diff --git a/kerops/ops/quantization.py b/kerops/ops/quantization.py index 8108a06..903c639 100644 --- a/kerops/ops/quantization.py +++ b/kerops/ops/quantization.py @@ -4,7 +4,7 @@ from ..kernels.quantization import _DequantUint8Window_impl, _QuantUint8Window_impl from ..settings import ConfigurableArg, configure, get_l1_cache -from ...utils import cdiv +from ..utils import cdiv @configure(num_warps=4, l1_cache_bytes=get_l1_cache) diff --git a/kerops/ops/stats.py b/kerops/ops/stats.py index 0a46799..69e14a1 100644 --- a/kerops/ops/stats.py +++ b/kerops/ops/stats.py @@ -6,7 +6,7 @@ from ..kernels.stats import _Stats_cl3d_backward_impl, _Stats_cl3d_impl from ..settings import ConfigurableArg, configure, get_l1_cache -from ...utils import cdiv +from ..utils import cdiv @configure(l1_cache_bytes=get_l1_cache, num_warps=4) From 578b0b9534e133e0ff8d85f368416681d946abad Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Sat, 28 Feb 2026 00:34:14 +0100 Subject: [PATCH 09/40] rename test --- tests/ops/{test_conv.py => test_dwconv.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/ops/{test_conv.py => test_dwconv.py} (100%) diff --git a/tests/ops/test_conv.py b/tests/ops/test_dwconv.py similarity index 100% rename from tests/ops/test_conv.py rename to tests/ops/test_dwconv.py From 4f438572a410d57c1261a223ca992a1c8f38a81a Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Sun, 1 Mar 2026 12:27:05 +0100 Subject: [PATCH 10/40] Conv3d tests --- tests/conftest.py | 10 ++++++++++ tests/ops/test_conv.py | 38 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+) create mode 100644 tests/ops/test_conv.py diff --git a/tests/conftest.py b/tests/conftest.py index 60a2840..f68b91f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -29,3 +29,13 @@ def other_2(request): @pytest.fixture(params=[2, 3, 13, 16, 32, 37, 53, 64]) def other_3(request): return request.param + + +@pytest.fixture(params=[16, 32, 64, 128]) +def conv_in_channels(request): + return request.param + + +@pytest.fixture(params=[16, 32, 64, 128]) +def conv_out_channels(request): + return request.param diff --git a/tests/ops/test_conv.py b/tests/ops/test_conv.py new file mode 100644 index 0000000..f779c76 --- /dev/null +++ b/tests/ops/test_conv.py @@ -0,0 +1,38 @@ +import math + +import torch +from torch import nn +from torch.nn import functional as F + +from kerops.ops.conv import Conv3d +from kerops.utils import allclose_two_stage + + +def test_conv(bsize, conv_in_channels, conv_out_channels, other_1, other_2, other_3): + if not (conv_in_channels == conv_out_channels or conv_in_channels * 2 == conv_out_channels or conv_in_channels == 2 * conv_out_channels): + return + + torch.manual_seed(322) + x = torch.randn(bsize, conv_in_channels, other_1, other_2, other_3, device='cuda', dtype=torch.float16).to( + memory_format=torch.channels_last_3d + ) + + weight = torch.empty(conv_out_channels, conv_in_channels, 3, 3, 3, device='cuda', dtype=torch.float16) + nn.init.kaiming_uniform_(weight, a=math.sqrt(5)) + + weight_kernel = weight.permute(2, 3, 4, 1, 0).contiguous() + + out_standard = F.conv3d(x, weight, None, padding=(1, 1, 1), stride=(1, 1, 1)) + out_check = Conv3d(x, weight_kernel) + + assert allclose_two_stage( + out_standard, + out_check, + rtol_strict=1e-4, + atol_strict=1e-3, + rtol_narrow=1e-3, + atol_narrow=2e-3, + debug_info='print', + ) + + torch.cuda.empty_cache() From 7ad8c9a9ee59378d4c0357cb89831e9a9c1e3d69 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Thu, 12 Mar 2026 16:04:47 +0100 Subject: [PATCH 11/40] notebook, profiling, bump triton, perf adjust, batching support --- .gitignore | 4 + kerops/kernels/conv.py | 7 +- kerops/ops/conv/conv.py | 35 ++-- notebooks/playground.ipynb | 316 +++++++++++++++++++++++++++++++++++++ requirements.txt | 2 +- 5 files changed, 344 insertions(+), 20 deletions(-) create mode 100644 notebooks/playground.ipynb diff --git a/.gitignore b/.gitignore index f4f8bbe..6ec4a5f 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,7 @@ pip-delete-this-directory.txt # Jupyter Notebook .ipynb_checkpoints/ + +# Profiling +*.ncu-rep +PROF/ diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index 35357bc..a5064b8 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -33,7 +33,10 @@ def _Conv_cl3d_impl_V5( ): W_cell = tl.program_id(0) H_cell = tl.program_id(1) - D_cell = tl.program_id(2) + BD_cell = tl.program_id(2) + + B_cell = BD_cell // tl.cdiv(D, D_BLOCK) + D_cell = BD_cell % tl.cdiv(D, D_BLOCK) CIN_STEPS: tl.constexpr = IN_CHANNELS // CIN_BLOCK @@ -49,10 +52,12 @@ def _Conv_cl3d_impl_V5( input_ptr += D_cell * D_BLOCK * IN_CHANNELS input_ptr += W_cell * 2 * IN_CHANNELS * D input_ptr += H_cell * 2 * IN_CHANNELS * D * W + input_ptr += B_cell * H * W * D * IN_CHANNELS output_ptr += D_cell * D_BLOCK * OUT_CHANNELS output_ptr += W_cell * 2 * OUT_CHANNELS * D output_ptr += H_cell * 2 * OUT_CHANNELS * D * W + output_ptr += B_cell * H * W * D * OUT_CHANNELS acc00 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) acc01 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py index 15b034f..4b5a3af 100644 --- a/kerops/ops/conv/conv.py +++ b/kerops/ops/conv/conv.py @@ -16,7 +16,7 @@ def num_warps(in_channels, out_channels): (32, 32): 2, (32, 64): 4, (64, 32): 1, - (64, 64): 2, + (64, 64): 4, (64, 128): 4, (128, 64): 4, (128, 128): 4, @@ -86,23 +86,22 @@ def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] output = torch.empty([bsize, H, W, D, out_channels], device=x.device, dtype=x.dtype, layout=x.layout).permute(0, -1, 1, 2, 3) - grid = (cdiv(W, 2), cdiv(H, 2), cdiv(D, D_BLOCK)) - - for unbatched_x, unbatched_y in zip(x, output): - _Conv_cl3d_impl_V5[grid]( - unbatched_x, - weight, - unbatched_y, - H, - W, - D, - D_BLOCK=D_BLOCK, - ACCTYPE=ACCTYPE, - IN_CHANNELS=in_channels, - OUT_CHANNELS=out_channels, - CIN_BLOCK=CIN_BLOCK, - num_warps=num_warps, - ) + grid = (cdiv(W, 2), cdiv(H, 2), cdiv(D, D_BLOCK) * bsize) + + _Conv_cl3d_impl_V5[grid]( + x, + weight, + output, + H, + W, + D, + D_BLOCK=D_BLOCK, + ACCTYPE=ACCTYPE, + IN_CHANNELS=in_channels, + OUT_CHANNELS=out_channels, + CIN_BLOCK=CIN_BLOCK, + num_warps=num_warps, + ) return output diff --git a/notebooks/playground.ipynb b/notebooks/playground.ipynb new file mode 100644 index 0000000..59a5f91 --- /dev/null +++ b/notebooks/playground.ipynb @@ -0,0 +1,316 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "3600ef2a-552d-4798-98cb-1c0d2ae9da70", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "from torch import nn\n", + "from kerops.ops.conv import Conv3d" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "78508bd3-37b2-4684-88a2-857614673b8b", + "metadata": {}, + "outputs": [], + "source": [ + "import itertools\n", + "from tqdm.notebook import tqdm\n", + "from time import perf_counter, sleep\n", + "import numpy as np\n", + "\n", + "\n", + "def mean_std_percentile(x, lo=20, hi=80):\n", + " x = np.asarray(x, dtype=np.float32)\n", + "\n", + " p_lo, p_hi = np.percentile(x, [lo, hi])\n", + "\n", + " mask = (x >= p_lo) & (x <= p_hi)\n", + " x_mid = x[mask]\n", + "\n", + " return x_mid.mean(), x_mid.std()\n", + "\n", + "\n", + "def bench(func, *args, warmup=25, sleep_ms=100, n_iters=50, q_show=10, quantiles=(20, 80), **specset):\n", + " results = []\n", + "\n", + " keys = list(specset.keys())\n", + " values = list(specset.values())\n", + " configs = list(itertools.product(*values))\n", + "\n", + " for config in tqdm(configs, desc=\"Benchmark configs\"):\n", + " if sleep_ms is not None:\n", + " sleep(sleep_ms / 1000)\n", + " kwargs = dict(zip(keys, config))\n", + "\n", + " try:\n", + " func(*args, **kwargs)\n", + " torch.cuda.synchronize()\n", + " except Exception as e:\n", + " pass\n", + "\n", + " for _ in range(warmup):\n", + " func(*args, **kwargs)\n", + " torch.cuda.synchronize()\n", + "\n", + " times_ms = []\n", + "\n", + " for _ in range(n_iters):\n", + " start = perf_counter()\n", + " func(*args, **kwargs)\n", + " torch.cuda.synchronize()\n", + " end = perf_counter()\n", + " times_ms.append((end - start) * 1e3)\n", + "\n", + " mean, std = mean_std_percentile(times_ms, *quantiles)\n", + " results.append({\n", + " \"spec\": kwargs,\n", + " \"mean_ms\": float(mean),\n", + " \"std_ms\": float(std),\n", + " })\n", + "\n", + " best_result = min(results, key=lambda x: x[\"mean_ms\"])\n", + " best_mean = best_result['mean_ms']\n", + " best_std = best_result['std_ms']\n", + " best_spec = best_result['spec']\n", + " print(f'Best spec - {best_mean:.3f}+-{best_std:.3f}ms {best_spec}')\n", + "\n", + " good_results = []\n", + " for result in results:\n", + " if best_mean * (1 + q_show / 100) >= result[\"mean_ms\"] and result['spec'] != best_spec:\n", + " good_results.append(result)\n", + "\n", + " if good_results:\n", + " print(30 * '-')\n", + " print(\"Other good specs:\")\n", + "\n", + " for result in good_results:\n", + " print(f'{result['mean_ms']:.3f}+-{result['std_ms']:.3f}ms {result['spec']}')\n", + " else:\n", + " print(f'Other specs have a time difference of more than {q_show}%')" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "47bf9383-f4a8-44a7-9037-b79c285074e3", + "metadata": {}, + "outputs": [], + "source": [ + "CIN = 64\n", + "COUT = 64\n", + "S = 64\n", + "\n", + "x = torch.randn(2, CIN, S, S, S, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", + "w = torch.randn(3, 3, 3, CIN, COUT, device='cuda', dtype=torch.float16)\n", + "conv = nn.Conv3d(CIN, COUT, 3, padding=1, bias=False, device='cuda')" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "ba72112e-f726-4582-a5f0-1ff6d356d9b4", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1.83 ms ± 65 μs per loop (mean ± std. dev. of 5 runs, 5 loops each)\n" + ] + } + ], + "source": [ + "%%timeit -r 5 -n 5\n", + "with torch.inference_mode(), torch.amp.autocast('cuda'):\n", + " conv(x)\n", + "\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "3f417833-9005-4b2b-9ab3-78cdef7389be", + "metadata": {}, + "outputs": [], + "source": [ + "Conv3d(x, w)\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "10ac3075-1799-44e3-80f6-5e6b81f70322", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1.88 ms ± 94.1 μs per loop (mean ± std. dev. of 5 runs, 5 loops each)\n" + ] + } + ], + "source": [ + "%%timeit -r 5 -n 5\n", + "Conv3d(x, w)\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "86210699-1c69-4d30-9718-f25a57320435", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "The slowest run took 6430.09 times longer than the fastest. This could mean that an intermediate result is being cached.\n", + "2.3 s ± 4.59 s per loop (mean ± std. dev. of 5 runs, 5 loops each)\n" + ] + } + ], + "source": [ + "%%timeit -r 5 -n 5\n", + "Conv3d(x, w, COUT_BLOCK=32, num_warps=2)\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "id": "ba01cad5-c171-470b-ad72-422d4de2bfc5", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1.88 ms ± 59.8 μs per loop (mean ± std. dev. of 5 runs, 5 loops each)\n" + ] + } + ], + "source": [ + "%%timeit -r 5 -n 5\n", + "Conv3d(x, w, COUT_BLOCK=32, num_warps=2)\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "999d773f-2848-4641-9d96-81ebfb6b806c", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "a48876bd6e9643efb55d926eb03a8324", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Benchmark configs: 0%| | 0/18 [00:00=3.6.0 torch From 19e7da7f5a76ea240f1291f5083643d031733a78 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Mon, 16 Mar 2026 15:24:39 +0100 Subject: [PATCH 12/40] V6 Conv3d --- kerops/kernels/conv.py | 198 +++++++++++++++++++++++++------------ kerops/ops/conv/conv.py | 10 +- notebooks/playground.ipynb | 151 ++++++++++++++++++---------- 3 files changed, 240 insertions(+), 119 deletions(-) diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index a5064b8..67b13b7 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -6,6 +6,7 @@ - Conv3dV3: `tl.max_contiguous(tl.multiple_of(channels_offset, CHANNELS), CHANNELS)` - Conv3dV4: Input channels BLOCKCING - Conv3dV5: output blocking by HW: 2x2 tile +- Conv3dV6: loading order choice and weight-major alorithm choice No impact: - `x @ w` and `w @ x` orientation via ORDER parameter @@ -15,10 +16,12 @@ - `x: [D_BLOCK, CHANNELS] -> x: [NEAR, D_BLOCK, CHANNELS]` where NEAR=4 represents neighbours; `w: [4, CHANNELS_IN, CHANNELS_OUT]` - Output channels BLOCKCING (why?) - Output channels BLOCKCING via expanding grid + - TMA (why???) + - tl.swizzle2d and a X-major tile with sizes 2 and 4 """ @triton.jit -def _Conv_cl3d_impl_V5( +def _Conv_cl3d_impl_V6( input_ptr, weight_ptr, output_ptr, @@ -30,6 +33,8 @@ def _Conv_cl3d_impl_V5( IN_CHANNELS: tl.constexpr, OUT_CHANNELS: tl.constexpr, CIN_BLOCK: tl.constexpr, + WEIGHT_MAJOR: tl.constexpr, + LOAD_WEIGHT_FIRST: tl.constexpr, ): W_cell = tl.program_id(0) H_cell = tl.program_id(1) @@ -64,74 +69,137 @@ def _Conv_cl3d_impl_V5( acc10 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) acc11 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) - for h_block in tl.static_range(0, 2): - for w_block in tl.static_range(0, 2): - for cin in tl.static_range(0, CIN_STEPS): - for dd in tl.static_range(-1, 2): # MB other order? - w_ptr = ( - weight_ptr - + (dd + 1) * IN_CHANNELS * OUT_CHANNELS - + w_block * IN_CHANNELS * OUT_CHANNELS * 3 - + h_block * IN_CHANNELS * OUT_CHANNELS * 9 - + cin * CIN_BLOCK * OUT_CHANNELS - ) - - weights = [ - [ - tl.load(w_ptr + weight_offset), - tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 3), - ], - [ - tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 9), - tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 12), + if WEIGHT_MAJOR: + for weight_h in tl.static_range(0, 3): + for weight_w in tl.static_range(0, 3): + for cin in tl.static_range(0, CIN_STEPS): + for weight_d in tl.static_range(0, 3): + w_ptr = ( + weight_ptr + + weight_d * IN_CHANNELS * OUT_CHANNELS + + weight_w * IN_CHANNELS * OUT_CHANNELS * 3 + + weight_h * IN_CHANNELS * OUT_CHANNELS * 9 + + cin * CIN_BLOCK * OUT_CHANNELS + ) + + if LOAD_WEIGHT_FIRST: + weight = tl.load(w_ptr + weight_offset) + + i_ptr = ( + input_ptr + + (weight_h - 1) * IN_CHANNELS * D * W + + (weight_w - 1) * IN_CHANNELS * D + + (weight_d - 1) * IN_CHANNELS + + cin * CIN_BLOCK + ) + + mask = ((d_offset_shifted + weight_d - 1) < D) & ((d_offset_shifted + weight_d - 1) >= 0) + + m00 = ((H_cell * 2 + weight_h - 1) < H) & ((H_cell * 2 + weight_h - 1) >= 0) & ((W_cell * 2 + weight_w - 1) < W) & ((W_cell * 2 + weight_w - 1) >= 0) + m01 = ((H_cell * 2 + weight_h - 1) < H) & ((H_cell * 2 + weight_h - 1) >= 0) & ((W_cell * 2 + weight_w) < W) & ((W_cell * 2 + weight_w) >= 0) + m10 = ((H_cell * 2 + weight_h) < H) & ((H_cell * 2 + weight_h) >= 0) & ((W_cell * 2 + weight_w - 1) < W) & ((W_cell * 2 + weight_w - 1) >= 0) + m11 = ((H_cell * 2 + weight_h) < H) & ((H_cell * 2 + weight_h) >= 0) & ((W_cell * 2 + weight_w) < W) & ((W_cell * 2 + weight_w) >= 0) + + xs = [ + [ + tl.load(i_ptr + input_offset, mask=mask & m00, other=0.0), + tl.load(i_ptr + input_offset + IN_CHANNELS * D, mask=mask & m01, other=0.0) + ], + [ + tl.load(i_ptr + input_offset + IN_CHANNELS * D * W, mask=mask & m10, other=0.0), + tl.load(i_ptr + input_offset + IN_CHANNELS * D + IN_CHANNELS * D * W, mask=mask & m11, other=0.0) + ] ] - ] - - i_ptr = ( - input_ptr - + (h_block * 2 - 1) * IN_CHANNELS * D * W - + (w_block * 2 - 1) * IN_CHANNELS * D - + dd * IN_CHANNELS - + cin * CIN_BLOCK - ) - mask = ((d_offset_shifted + dd) < D) & ((d_offset_shifted + dd) >= 0) - m00 = ((H_cell * 2 + h_block * 2 - 1) < H) & ((H_cell * 2 + h_block * 2 - 1) >= 0) & ((W_cell * 2 + w_block * 2 - 1) < W) & ((W_cell * 2 + w_block * 2 - 1) >= 0) - m01 = ((H_cell * 2 + h_block * 2 - 1) < H) & ((H_cell * 2 + h_block * 2 - 1) >= 0) & ((W_cell * 2 + w_block * 2) < W) & ((W_cell * 2 + w_block * 2) >= 0) - m10 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2 - 1) < W) & ((W_cell * 2 + w_block * 2 - 1) >= 0) - m11 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2) < W) & ((W_cell * 2 + w_block * 2) >= 0) - - xs = [ - [ - tl.load(i_ptr + input_offset, mask=mask & m00, other=0.0), - tl.load(i_ptr + input_offset + IN_CHANNELS * D, mask=mask & m01, other=0.0) - ], - [ - tl.load(i_ptr + input_offset + IN_CHANNELS * D * W, mask=mask & m10, other=0.0), - tl.load(i_ptr + input_offset + IN_CHANNELS * D + IN_CHANNELS * D * W, mask=mask & m11, other=0.0) + if not LOAD_WEIGHT_FIRST: + weight = tl.load(w_ptr + weight_offset) + + acc00 += tl.dot(xs[0][0], weight) + acc01 += tl.dot(xs[0][1], weight) + acc10 += tl.dot(xs[1][0], weight) + acc11 += tl.dot(xs[1][1], weight) + else: + for h_block in tl.static_range(0, 2): + for w_block in tl.static_range(0, 2): + for cin in tl.static_range(0, CIN_STEPS): + for dd in tl.static_range(-1, 2): # MB other order? + w_ptr = ( + weight_ptr + + (dd + 1) * IN_CHANNELS * OUT_CHANNELS + + w_block * IN_CHANNELS * OUT_CHANNELS * 3 + + h_block * IN_CHANNELS * OUT_CHANNELS * 9 + + cin * CIN_BLOCK * OUT_CHANNELS + ) + + if LOAD_WEIGHT_FIRST: + weights = [ + [ + tl.load(w_ptr + weight_offset), + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 3), + ], + [ + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 9), + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 12), + ] + ] + + i_ptr = ( + input_ptr + + (h_block * 2 - 1) * IN_CHANNELS * D * W + + (w_block * 2 - 1) * IN_CHANNELS * D + + dd * IN_CHANNELS + + cin * CIN_BLOCK + ) + mask = ((d_offset_shifted + dd) < D) & ((d_offset_shifted + dd) >= 0) + + m00 = ((H_cell * 2 + h_block * 2 - 1) < H) & ((H_cell * 2 + h_block * 2 - 1) >= 0) & ((W_cell * 2 + w_block * 2 - 1) < W) & ((W_cell * 2 + w_block * 2 - 1) >= 0) + m01 = ((H_cell * 2 + h_block * 2 - 1) < H) & ((H_cell * 2 + h_block * 2 - 1) >= 0) & ((W_cell * 2 + w_block * 2) < W) & ((W_cell * 2 + w_block * 2) >= 0) + m10 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2 - 1) < W) & ((W_cell * 2 + w_block * 2 - 1) >= 0) + m11 = ((H_cell * 2 + h_block * 2) < H) & ((H_cell * 2 + h_block * 2) >= 0) & ((W_cell * 2 + w_block * 2) < W) & ((W_cell * 2 + w_block * 2) >= 0) + + xs = [ + [ + tl.load(i_ptr + input_offset, mask=mask & m00, other=0.0), + tl.load(i_ptr + input_offset + IN_CHANNELS * D, mask=mask & m01, other=0.0) + ], + [ + tl.load(i_ptr + input_offset + IN_CHANNELS * D * W, mask=mask & m10, other=0.0), + tl.load(i_ptr + input_offset + IN_CHANNELS * D + IN_CHANNELS * D * W, mask=mask & m11, other=0.0) + ] ] - ] - for h in tl.static_range(0, 2): - for w in tl.static_range(0, 2): - # h_weight_idx = 2 * h_block + h - acc_abs_h + 1 - h_block <-- weights window shift - # <---x_h-------> ^-- +1 since weight indexed from 0 - - # acc00 - if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) < 3): - acc00 += tl.dot(xs[h][w], weights[h_block + h][w_block + w]) - - # acc01 - if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) > 0): - acc01 += tl.dot(xs[h][w], weights[h_block + h][w_block + w - 1]) - - # acc10 - if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) < 3): - acc10 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w]) - - # acc11 - if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) > 0): - acc11 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w - 1]) + if not LOAD_WEIGHT_FIRST: + weights = [ + [ + tl.load(w_ptr + weight_offset), + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 3), + ], + [ + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 9), + tl.load(w_ptr + weight_offset + IN_CHANNELS * OUT_CHANNELS * 12), + ] + ] + + for h in tl.static_range(0, 2): + for w in tl.static_range(0, 2): + # h_weight_idx = 2 * h_block + h - acc_abs_h + 1 - h_block <-- weights window shift + # <---x_h-------> ^-- +1 since weight indexed from 0 + + # acc00 + if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) < 3): + acc00 += tl.dot(xs[h][w], weights[h_block + h][w_block + w]) + + # acc01 + if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) > 0): + acc01 += tl.dot(xs[h][w], weights[h_block + h][w_block + w - 1]) + + # acc10 + if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) < 3): + acc10 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w]) + + # acc11 + if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) > 0): + acc11 += tl.dot(xs[h][w], weights[h_block + h - 1][w_block + w - 1]) omask = d_offset_shifted < D tl.store(output_ptr + output_offset, acc00, mask=omask & ((W_cell * 2) < W) & ((H_cell * 2) < H)) diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py index 4b5a3af..877a670 100644 --- a/kerops/ops/conv/conv.py +++ b/kerops/ops/conv/conv.py @@ -2,7 +2,7 @@ import torch from triton import language as tl, next_power_of_2 -from ...kernels.conv import _Conv_cl3d_impl_V5, _ApplyBNReLUConvStats_cl3d_impl +from ...kernels.conv import _Conv_cl3d_impl_V6, _ApplyBNReLUConvStats_cl3d_impl from ...settings import ConfigurableArg, configure, confexc from ...utils import cdiv @@ -60,8 +60,10 @@ def cin_block(in_channels, out_channels): num_warps=lambda weight: num_warps(*weight.shape[-2:]), D_BLOCK=lambda weight: d_block(*weight.shape[-2:]), CIN_BLOCK=lambda weight: cin_block(*weight.shape[-2:]), + LOAD_WEIGHT_FIRST=True, + WEIGHT_MAJOR=False, ) -def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg): +def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg, LOAD_WEIGHT_FIRST: ConfigurableArg, WEIGHT_MAJOR: ConfigurableArg): assert x.device == weight.device assert x.is_cuda @@ -88,7 +90,7 @@ def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D output = torch.empty([bsize, H, W, D, out_channels], device=x.device, dtype=x.dtype, layout=x.layout).permute(0, -1, 1, 2, 3) grid = (cdiv(W, 2), cdiv(H, 2), cdiv(D, D_BLOCK) * bsize) - _Conv_cl3d_impl_V5[grid]( + _Conv_cl3d_impl_V6[grid]( x, weight, output, @@ -100,6 +102,8 @@ def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D IN_CHANNELS=in_channels, OUT_CHANNELS=out_channels, CIN_BLOCK=CIN_BLOCK, + LOAD_WEIGHT_FIRST=LOAD_WEIGHT_FIRST, + WEIGHT_MAJOR=WEIGHT_MAJOR, num_warps=num_warps, ) diff --git a/notebooks/playground.ipynb b/notebooks/playground.ipynb index 59a5f91..60fb843 100644 --- a/notebooks/playground.ipynb +++ b/notebooks/playground.ipynb @@ -97,23 +97,24 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 2, "id": "47bf9383-f4a8-44a7-9037-b79c285074e3", "metadata": {}, "outputs": [], "source": [ "CIN = 64\n", "COUT = 64\n", - "S = 64\n", + "S = 96\n", + "D = 96\n", "\n", - "x = torch.randn(2, CIN, S, S, S, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", + "x = torch.randn(2, CIN, S, S, D, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", "w = torch.randn(3, 3, 3, CIN, COUT, device='cuda', dtype=torch.float16)\n", "conv = nn.Conv3d(CIN, COUT, 3, padding=1, bias=False, device='cuda')" ] }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 13, "id": "ba72112e-f726-4582-a5f0-1ff6d356d9b4", "metadata": {}, "outputs": [ @@ -121,12 +122,12 @@ "name": "stdout", "output_type": "stream", "text": [ - "1.83 ms ± 65 μs per loop (mean ± std. dev. of 5 runs, 5 loops each)\n" + "5.96 ms ± 68 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" ] } ], "source": [ - "%%timeit -r 5 -n 5\n", + "%%timeit -r 10 -n 10\n", "with torch.inference_mode(), torch.amp.autocast('cuda'):\n", " conv(x)\n", "\n", @@ -135,7 +136,7 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 4, "id": "3f417833-9005-4b2b-9ab3-78cdef7389be", "metadata": {}, "outputs": [], @@ -146,80 +147,101 @@ }, { "cell_type": "code", - "execution_count": 17, - "id": "10ac3075-1799-44e3-80f6-5e6b81f70322", + "execution_count": 14, + "id": "f7731c96-8bb7-4ece-ab40-d438beeba82e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "1.88 ms ± 94.1 μs per loop (mean ± std. dev. of 5 runs, 5 loops each)\n" + "6 ms ± 81.2 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" ] } ], "source": [ - "%%timeit -r 5 -n 5\n", + "%%timeit -r 10 -n 10\n", "Conv3d(x, w)\n", "torch.cuda.synchronize()" ] }, { "cell_type": "code", - "execution_count": 22, - "id": "86210699-1c69-4d30-9718-f25a57320435", + "execution_count": 6, + "id": "925fbd7c-e860-4f80-9c5c-84c886175c62", + "metadata": {}, + "outputs": [], + "source": [ + "Conv3d(x, w, **{'num_warps': 2, 'D_BLOCK': 32, 'CIN_BLOCK': 32, 'LOAD_WEIGHT_FIRST': True, 'WEIGHT_MAJOR': True})\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "fc10543a-982d-48fa-bfbb-0e2872b55968", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "The slowest run took 6430.09 times longer than the fastest. This could mean that an intermediate result is being cached.\n", - "2.3 s ± 4.59 s per loop (mean ± std. dev. of 5 runs, 5 loops each)\n" + "5.86 ms ± 57.5 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" ] } ], "source": [ - "%%timeit -r 5 -n 5\n", - "Conv3d(x, w, COUT_BLOCK=32, num_warps=2)\n", + "%%timeit -r 10 -n 10\n", + "Conv3d(x, w, **{'num_warps': 2, 'D_BLOCK': 32, 'CIN_BLOCK': 32, 'LOAD_WEIGHT_FIRST': True, 'WEIGHT_MAJOR': True})\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "6a6246f6-be7b-4683-ae51-7aa22ee620e3", + "metadata": {}, + "outputs": [], + "source": [ + "Conv3d(x, w, **{'num_warps': 2, 'D_BLOCK': 16, 'CIN_BLOCK': 16, 'LOAD_WEIGHT_FIRST': False, 'WEIGHT_MAJOR': True})\n", "torch.cuda.synchronize()" ] }, { "cell_type": "code", - "execution_count": 28, - "id": "ba01cad5-c171-470b-ad72-422d4de2bfc5", + "execution_count": 12, + "id": "df11394a-9430-4c40-b64d-a1999f31bad3", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "1.88 ms ± 59.8 μs per loop (mean ± std. dev. of 5 runs, 5 loops each)\n" + "5.93 ms ± 70.1 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" ] } ], "source": [ - "%%timeit -r 5 -n 5\n", - "Conv3d(x, w, COUT_BLOCK=32, num_warps=2)\n", + "%%timeit -r 10 -n 10\n", + "Conv3d(x, w, **{'num_warps': 2, 'D_BLOCK': 16, 'CIN_BLOCK': 16, 'LOAD_WEIGHT_FIRST': False, 'WEIGHT_MAJOR': True})\n", "torch.cuda.synchronize()" ] }, { "cell_type": "code", - "execution_count": 6, - "id": "999d773f-2848-4641-9d96-81ebfb6b806c", + "execution_count": 4, + "id": "cd845ca5-cdb0-4fa4-9f7b-987e95152aca", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "a48876bd6e9643efb55d926eb03a8324", + "model_id": "58bf24ace46648b0ad4199ed9dc082bd", "version_major": 2, "version_minor": 0 }, "text/plain": [ - "Benchmark configs: 0%| | 0/18 [00:00 Date: Mon, 16 Mar 2026 21:09:43 +0100 Subject: [PATCH 13/40] rename --- kerops/ops/addition.py | 6 +++--- kerops/ops/avgpool.py | 8 ++++---- kerops/ops/bnrelu.py | 6 +++--- kerops/ops/conv/conv.py | 6 +++--- kerops/ops/conv/dwconv.py | 4 ++-- kerops/ops/conv/dwconv_wgrad.py | 4 ++-- kerops/ops/linear/linear_bias_relu_linear_add.py | 8 ++++---- kerops/ops/linear/linear_bias_relu_linear_backward.py | 8 ++++---- kerops/ops/linear/relu_linear_add.py | 8 ++++---- kerops/ops/linear/relu_linear_backward.py | 8 ++++---- kerops/ops/quantization.py | 6 +++--- kerops/ops/stats.py | 6 +++--- kerops/settings/__init__.py | 2 +- kerops/settings/utils.py | 8 ++++---- notebooks/playground.ipynb | 2 +- 15 files changed, 45 insertions(+), 45 deletions(-) diff --git a/kerops/ops/addition.py b/kerops/ops/addition.py index e530a18..e6b29ba 100644 --- a/kerops/ops/addition.py +++ b/kerops/ops/addition.py @@ -5,12 +5,12 @@ from triton import next_power_of_2 from ..kernels.addition import _AddStats_cl3d_backward_impl, _AddStats_cl3d_impl -from ..settings import ConfigurableArg, configure, get_l1_cache +from ..settings import ConfArg, configure, get_l1_cache from ..utils import cdiv @configure(l1_cache_bytes=get_l1_cache, num_warps=8) -def AddStats(x, y, inplace=False, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +def AddStats(x, y, inplace=False, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() assert x.shape == y.shape @@ -52,7 +52,7 @@ def AddStats(x, y, inplace=False, *, l1_cache_bytes: ConfigurableArg, num_warps: @configure(l1_cache_bytes=get_l1_cache, num_warps=8) def AddStatsBackward( - add_grad, mean_grad, sqmean_grad, add_result, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg + add_grad, mean_grad, sqmean_grad, add_result, *, l1_cache_bytes: ConfArg, num_warps: ConfArg ): num_channels = add_grad.shape[1] numel = add_grad.numel() diff --git a/kerops/ops/avgpool.py b/kerops/ops/avgpool.py index e222644..32c08ef 100644 --- a/kerops/ops/avgpool.py +++ b/kerops/ops/avgpool.py @@ -4,7 +4,7 @@ from triton import next_power_of_2 from ..kernels.avgpool import _AvgPoolCeilStats_cl3d_backward_impl, _AvgPoolCeilStats_cl3d_impl -from ..settings import ConfigurableArg, configure, get_l1_cache +from ..settings import ConfArg, configure, get_l1_cache from ..utils import cdiv @@ -12,7 +12,7 @@ l1_cache_bytes=get_l1_cache, num_warps=2, ) -def AvgPoolCeilStats(x, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +def AvgPoolCeilStats(x, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] input_d = x.shape[-1] MAX_SIZE = l1_cache_bytes // x.element_size() # 32768 for fp16 @@ -68,8 +68,8 @@ def AvgPoolCeilStatsBackward( output, outgrad_shape, *, - l1_cache_bytes: ConfigurableArg, - num_warps: ConfigurableArg, + l1_cache_bytes: ConfArg, + num_warps: ConfArg, ): MAX_SIZE = l1_cache_bytes // inpgrad.element_size() # 32768 for fp16 bsize, num_channels, h_outgrad, w_outgrad, d_outgrad = outgrad_shape diff --git a/kerops/ops/bnrelu.py b/kerops/ops/bnrelu.py index a055d70..a39d4b1 100644 --- a/kerops/ops/bnrelu.py +++ b/kerops/ops/bnrelu.py @@ -5,12 +5,12 @@ from triton import next_power_of_2 from ..kernels.bnrelu import _ApplyBNReLU_cl3d_backward_impl, _ApplyBNReLU_cl3d_impl -from ..settings import ConfigurableArg, configure, get_l1_cache +from ..settings import ConfArg, configure, get_l1_cache from ..utils import cdiv @configure(l1_cache_bytes=get_l1_cache, num_warps=8) -def ApplyBNReLU(x, weight, bias, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +def ApplyBNReLU(x, weight, bias, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() assert x.ndim == 5 @@ -43,7 +43,7 @@ def ApplyBNReLU(x, weight, bias, *, l1_cache_bytes: ConfigurableArg, num_warps: @configure(l1_cache_bytes=get_l1_cache, num_warps=8) -def ApplyBNReLUBackward(x, weight, bias, grad, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +def ApplyBNReLUBackward(x, weight, bias, grad, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() assert x.ndim == 5 diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py index 877a670..17e3c07 100644 --- a/kerops/ops/conv/conv.py +++ b/kerops/ops/conv/conv.py @@ -3,7 +3,7 @@ from triton import language as tl, next_power_of_2 from ...kernels.conv import _Conv_cl3d_impl_V6, _ApplyBNReLUConvStats_cl3d_impl -from ...settings import ConfigurableArg, configure, confexc +from ...settings import ConfArg, configure, confexc from ...utils import cdiv @@ -63,7 +63,7 @@ def cin_block(in_channels, out_channels): LOAD_WEIGHT_FIRST=True, WEIGHT_MAJOR=False, ) -def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg, LOAD_WEIGHT_FIRST: ConfigurableArg, WEIGHT_MAJOR: ConfigurableArg): +def Conv3d(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_BLOCK: ConfArg, CIN_BLOCK: ConfArg, LOAD_WEIGHT_FIRST: ConfArg, WEIGHT_MAJOR: ConfArg): assert x.device == weight.device assert x.is_cuda @@ -116,7 +116,7 @@ def Conv3d(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D D_BLOCK=lambda weight: d_block(*weight.shape[-2:]), CIN_BLOCK=lambda weight: cin_block(*weight.shape[-2:]), ) -def ApplyBNReLUConv3dStats(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_BLOCK: ConfigurableArg, CIN_BLOCK: ConfigurableArg): +def ApplyBNReLUConv3dStats(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_BLOCK: ConfArg, CIN_BLOCK: ConfArg): assert x.device == weight.device == bn_weight.device == bn_bias.device assert x.is_cuda diff --git a/kerops/ops/conv/dwconv.py b/kerops/ops/conv/dwconv.py index 7f9d7e3..9585057 100644 --- a/kerops/ops/conv/dwconv.py +++ b/kerops/ops/conv/dwconv.py @@ -2,7 +2,7 @@ from triton import language as tl, next_power_of_2 from ...kernels.dw_conv import _DWConv_cl3d_impl -from ...settings import ConfigurableArg, confexc, configure +from ...settings import ConfArg, confexc, configure from ...utils import cdiv @@ -21,7 +21,7 @@ def dblock(channels): num_warps=lambda x: warps(x.shape[1]), D_block=lambda x: dblock(x.shape[1]), ) -def DWConv(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_block: ConfigurableArg): +def DWConv(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_block: ConfArg): channels = x.shape[1] assert x.ndim == 5 diff --git a/kerops/ops/conv/dwconv_wgrad.py b/kerops/ops/conv/dwconv_wgrad.py index b4bcd92..604aa17 100644 --- a/kerops/ops/conv/dwconv_wgrad.py +++ b/kerops/ops/conv/dwconv_wgrad.py @@ -2,7 +2,7 @@ from triton import language as tl, next_power_of_2 from ...kernels.dw_conv import _DWConv_wgrad_cl3d_impl -from ...settings import ConfigurableArg, confexc, configure +from ...settings import ConfArg, confexc, configure from ...utils import cdiv @@ -28,7 +28,7 @@ def ilp(channels): ILP=lambda x: ilp(x.shape[1]), ) def DWConvWGRAD( - x, grad, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D_block: ConfigurableArg, ILP: ConfigurableArg + x, grad, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_block: ConfArg, ILP: ConfArg ): channels = x.shape[1] diff --git a/kerops/ops/linear/linear_bias_relu_linear_add.py b/kerops/ops/linear/linear_bias_relu_linear_add.py index 93b5406..7586e61 100644 --- a/kerops/ops/linear/linear_bias_relu_linear_add.py +++ b/kerops/ops/linear/linear_bias_relu_linear_add.py @@ -2,7 +2,7 @@ from triton import next_power_of_2 from ...kernels.linear import _LinBReLULinAdd -from ...settings import ConfigurableArg, confexc, configure +from ...settings import ConfArg, confexc, configure from ...utils import cdiv @@ -30,9 +30,9 @@ def LinBReLULinAdd( bias, add_other, *, - num_warps: ConfigurableArg, - D_block: ConfigurableArg, - ILP: ConfigurableArg, + num_warps: ConfArg, + D_block: ConfArg, + ILP: ConfArg, ): in_channels = x.shape[1] hidden_channels = 2 * in_channels diff --git a/kerops/ops/linear/linear_bias_relu_linear_backward.py b/kerops/ops/linear/linear_bias_relu_linear_backward.py index 9be1a7e..14022fc 100644 --- a/kerops/ops/linear/linear_bias_relu_linear_backward.py +++ b/kerops/ops/linear/linear_bias_relu_linear_backward.py @@ -2,7 +2,7 @@ from triton import next_power_of_2 from ...kernels.linear import _LinBReLULinBackward -from ...settings import ConfigurableArg, confexc, configure +from ...settings import ConfArg, confexc, configure from ...utils import cdiv @@ -23,9 +23,9 @@ def LinBReLULinBackward( weight_down, bias, *, - num_warps: ConfigurableArg, - D_block: ConfigurableArg, - ILP: ConfigurableArg, + num_warps: ConfArg, + D_block: ConfArg, + ILP: ConfArg, ): in_channels = x.shape[1] hidden_channels = 2 * in_channels diff --git a/kerops/ops/linear/relu_linear_add.py b/kerops/ops/linear/relu_linear_add.py index 2cbbcf8..63c99c8 100644 --- a/kerops/ops/linear/relu_linear_add.py +++ b/kerops/ops/linear/relu_linear_add.py @@ -2,7 +2,7 @@ from triton import next_power_of_2 from ...kernels.linear import _ReLULinearAdd -from ...settings import ConfigurableArg, confexc, configure +from ...settings import ConfArg, confexc, configure from ...utils import cdiv @@ -21,9 +21,9 @@ def ReLULinearAdd( weight, add_other, *, - num_warps: ConfigurableArg, - D_block: ConfigurableArg, - ILP: ConfigurableArg, + num_warps: ConfArg, + D_block: ConfArg, + ILP: ConfArg, ): in_channels = x.shape[1] out_channels = add_other.shape[1] diff --git a/kerops/ops/linear/relu_linear_backward.py b/kerops/ops/linear/relu_linear_backward.py index d3080fb..7ea907d 100644 --- a/kerops/ops/linear/relu_linear_backward.py +++ b/kerops/ops/linear/relu_linear_backward.py @@ -2,7 +2,7 @@ from triton import next_power_of_2 from ...kernels.linear import _ReLULinearAddBackward -from ...settings import ConfigurableArg, confexc, configure +from ...settings import ConfArg, confexc, configure from ...utils import cdiv @@ -26,9 +26,9 @@ def ReLULinearBackward( grad, weight, *, - num_warps: ConfigurableArg, - D_block: ConfigurableArg, - ILP: ConfigurableArg, + num_warps: ConfArg, + D_block: ConfArg, + ILP: ConfArg, ): in_channels = x.shape[1] out_channels = grad.shape[1] diff --git a/kerops/ops/quantization.py b/kerops/ops/quantization.py index 903c639..5d87658 100644 --- a/kerops/ops/quantization.py +++ b/kerops/ops/quantization.py @@ -3,12 +3,12 @@ import torch from ..kernels.quantization import _DequantUint8Window_impl, _QuantUint8Window_impl -from ..settings import ConfigurableArg, configure, get_l1_cache +from ..settings import ConfArg, configure, get_l1_cache from ..utils import cdiv @configure(num_warps=4, l1_cache_bytes=get_l1_cache) -def QuantUint8Window(x, window, *, num_warps: ConfigurableArg, l1_cache_bytes: ConfigurableArg): +def QuantUint8Window(x, window, *, num_warps: ConfArg, l1_cache_bytes: ConfArg): numel = x.numel() MAX_SIZE = l1_cache_bytes // (2 * x.element_size()) BLOCK_SIZE = min(MAX_SIZE, numel) @@ -22,7 +22,7 @@ def QuantUint8Window(x, window, *, num_warps: ConfigurableArg, l1_cache_bytes: C @configure(num_warps=4, l1_cache_bytes=get_l1_cache) -def DequantUint8Window(x, init_dtype, window, num_warps: ConfigurableArg, l1_cache_bytes: ConfigurableArg): +def DequantUint8Window(x, init_dtype, window, num_warps: ConfArg, l1_cache_bytes: ConfArg): numel = x.numel() output = torch.empty_like(x, dtype=init_dtype) diff --git a/kerops/ops/stats.py b/kerops/ops/stats.py index 69e14a1..232f2aa 100644 --- a/kerops/ops/stats.py +++ b/kerops/ops/stats.py @@ -5,12 +5,12 @@ from triton import next_power_of_2 from ..kernels.stats import _Stats_cl3d_backward_impl, _Stats_cl3d_impl -from ..settings import ConfigurableArg, configure, get_l1_cache +from ..settings import ConfArg, configure, get_l1_cache from ..utils import cdiv @configure(l1_cache_bytes=get_l1_cache, num_warps=4) -def Stats(x, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +def Stats(x, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() assert x.ndim == 5 @@ -33,7 +33,7 @@ def Stats(x, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): @configure(l1_cache_bytes=get_l1_cache, num_warps=4) -def StatsBackward(x, mean_grad, sqmean_grad, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +def StatsBackward(x, mean_grad, sqmean_grad, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() assert x.ndim == 5 diff --git a/kerops/settings/__init__.py b/kerops/settings/__init__.py index 1974121..c74a353 100644 --- a/kerops/settings/__init__.py +++ b/kerops/settings/__init__.py @@ -1,3 +1,3 @@ from .hardware_conf import get_l1_cache -from .utils import ConfigurableArg, CongiguratorError +from .utils import ConfArg, CongiguratorError from .wrapper import confexc, configure diff --git a/kerops/settings/utils.py b/kerops/settings/utils.py index 668cbe6..5855e34 100644 --- a/kerops/settings/utils.py +++ b/kerops/settings/utils.py @@ -1,7 +1,7 @@ from inspect import Parameter -class ConfigurableArg: +class ConfArg: pass @@ -11,14 +11,14 @@ class CongiguratorError(Exception): def validate_signature(signature): for param in signature.parameters.values(): - if param.annotation is ConfigurableArg and param.kind is not Parameter.KEYWORD_ONLY: + if param.annotation is ConfArg and param.kind is not Parameter.KEYWORD_ONLY: raise RuntimeError(f'ConfigurableArg must be keyword-only - {param.name}') - elif param.annotation is not ConfigurableArg and param.kind is Parameter.KEYWORD_ONLY: + elif param.annotation is not ConfArg and param.kind is Parameter.KEYWORD_ONLY: raise RuntimeError(f'non-ConfigurableArg must not be keyword-only - {param.name}') def get_config_args(signature): - return [param.name for param in signature.parameters.values() if param.annotation is ConfigurableArg] + return [param.name for param in signature.parameters.values() if param.annotation is ConfArg] def get_standard_args(signature): diff --git a/notebooks/playground.ipynb b/notebooks/playground.ipynb index 60fb843..9dde9e5 100644 --- a/notebooks/playground.ipynb +++ b/notebooks/playground.ipynb @@ -335,7 +335,7 @@ { "cell_type": "code", "execution_count": null, - "id": "69b41523-9e6d-4725-8da6-485b40d626a0", + "id": "d28d18e2-1ac4-4a11-bf12-6e64397080f9", "metadata": {}, "outputs": [], "source": [] From 85eeae5e88ec7be428f8966a796d8f42ee661ada Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Thu, 19 Mar 2026 15:42:45 +0100 Subject: [PATCH 14/40] Conv3dWgrad --- kerops/kernels/conv.py | 123 +++++++ kerops/ops/conv/__init__.py | 1 + kerops/ops/conv/conv_wgrad.py | 130 +++++++ notebooks/playground.ipynb | 641 +++++++++++++++++++++++++--------- tests/ops/test_conv.py | 40 ++- 5 files changed, 773 insertions(+), 162 deletions(-) create mode 100644 kerops/ops/conv/conv_wgrad.py diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index 67b13b7..edfa615 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -208,6 +208,129 @@ def _Conv_cl3d_impl_V6( tl.store(output_ptr + output_offset + OUT_CHANNELS * D * W + OUT_CHANNELS * D, acc11, mask=omask & ((W_cell * 2 + 1) < W) & ((H_cell * 2 + 1) < H)) +""" +- Conv3dV2: grad and x tiling added +""" + +@triton.jit +def _Conv_wgrad_cl3d_impl_V2( + grad_ptr, + input_ptr, + weight_grad_ptr, + H, + W, + D, + num_buffers, + ACCTYPE: tl.constexpr, + D_BLOCK: tl.constexpr, + IN_CHANNELS: tl.constexpr, + OUT_CHANNELS: tl.constexpr, + CIN_BLOCK: tl.constexpr, + COUT_BLOCK: tl.constexpr, +): + WCOUT_pid = tl.program_id(0) + H_pid = tl.program_id(1) + BD_pid = tl.program_id(2) + + W_pid = WCOUT_pid // tl.cdiv(OUT_CHANNELS, COUT_BLOCK) + COUT_pid = WCOUT_pid % tl.cdiv(OUT_CHANNELS, COUT_BLOCK) + + B_pid = BD_pid // tl.cdiv(D, D_BLOCK) + D_pid = BD_pid % tl.cdiv(D, D_BLOCK) + + CIN_STEPS: tl.constexpr = IN_CHANNELS // CIN_BLOCK + + linear = WCOUT_pid + H_pid * tl.num_programs(0) + BD_pid * tl.num_programs(0) * tl.num_programs(1) + buffer_idx = linear % num_buffers + + in_channels_offset = tl.arange(0, CIN_BLOCK) + out_channels_offset = tl.arange(0, COUT_BLOCK) + d_offset = tl.arange(0, D_BLOCK) + + grad_offset = d_offset[:, None] * OUT_CHANNELS + out_channels_offset[None, :] + input_offset = d_offset[None, :] * IN_CHANNELS + in_channels_offset[:, None] + weight_grad_offset = in_channels_offset[:, None] * OUT_CHANNELS + out_channels_offset[None, :] + + grad_ptr += COUT_pid * COUT_BLOCK + grad_ptr += D_pid * D_BLOCK * OUT_CHANNELS + grad_ptr += W_pid * 2 * D * OUT_CHANNELS + grad_ptr += H_pid * 2 * D * W * OUT_CHANNELS + grad_ptr += B_pid * H * W * D * OUT_CHANNELS + + input_ptr += D_pid * D_BLOCK * IN_CHANNELS + input_ptr += W_pid * 2 * D * IN_CHANNELS + input_ptr += H_pid * 2 * D * W * IN_CHANNELS + input_ptr += B_pid * H * W * D * IN_CHANNELS + + weight_grad_ptr += COUT_pid * COUT_BLOCK + + g01 = ((W_pid * 2 + 1) < W) + g10 = ((H_pid * 2 + 1) < H) + g11 = ((W_pid * 2 + 1) < W) & ((H_pid * 2 + 1) < H) + + gmask = (d_offset < (D - D_pid * D_BLOCK)) + gmask = gmask[:, None] + + grads = [ + [ + tl.load(grad_ptr + grad_offset, other=0, mask=gmask), + tl.load(grad_ptr + grad_offset + OUT_CHANNELS * D, other=0, mask=gmask & g01), + ], + [ + tl.load(grad_ptr + grad_offset + OUT_CHANNELS * D * W, other=0, mask=gmask & g10), + tl.load(grad_ptr + grad_offset + OUT_CHANNELS * D * W + OUT_CHANNELS * D, other=0, mask=gmask & g11) + ] + ] + + for h in tl.static_range(-1, 2): + for w in tl.static_range(-1, 2): + for d in tl.static_range(-1, 2): + for cin in tl.static_range(0, CIN_STEPS): + wgrad = tl.zeros([CIN_BLOCK, COUT_BLOCK], dtype=ACCTYPE) + + x_ptr = ( + input_ptr + + h * IN_CHANNELS * D * W + + w * IN_CHANNELS * D + + d * IN_CHANNELS + + cin * CIN_BLOCK + ) + + xmask = (d_offset < (D - D_pid * D_BLOCK - d)) & (d_offset >= (- D_pid * D_BLOCK - d)) + xmask = xmask[None, :] + + x00 = ((H_pid * 2 + h) < H) & ((H_pid * 2 + h) >= 0) & ((W_pid * 2 + w) < W) & ((W_pid * 2 + w) >= 0) + x01 = ((H_pid * 2 + h) < H) & ((H_pid * 2 + h) >= 0) & ((W_pid * 2 + w + 1) < W) & ((W_pid * 2 + w + 1) >= 0) + x10 = ((H_pid * 2 + h + 1) < H) & ((H_pid * 2 + h + 1) >= 0) & ((W_pid * 2 + w) < W) & ((W_pid * 2 + w) >= 0) + x11 = ((H_pid * 2 + h + 1) < H) & ((H_pid * 2 + h + 1) >= 0) & ((W_pid * 2 + w + 1) < W) & ((W_pid * 2 + w + 1) >= 0) + + xs = [ + [ + tl.load(x_ptr + input_offset, other=0, mask=xmask & x00), + tl.load(x_ptr + input_offset + IN_CHANNELS * D, other=0, mask=xmask & x01), + ], + [ + tl.load(x_ptr + input_offset + IN_CHANNELS * D * W, other=0, mask=xmask & x10), + tl.load(x_ptr + input_offset + IN_CHANNELS * D * W + IN_CHANNELS * D, other=0, mask=xmask & x11), + ] + ] + + for kh in tl.static_range(0, 2): + for kw in tl.static_range(0, 2): + wgrad += tl.dot(xs[kh][kw], grads[kh][kw]) + + + w_ptr = ( + weight_grad_ptr + + buffer_idx * IN_CHANNELS * OUT_CHANNELS * 3 * 3 * 3 + + (h + 1) * IN_CHANNELS * OUT_CHANNELS * 3 * 3 + + (w + 1) * IN_CHANNELS * OUT_CHANNELS * 3 + + (d + 1) * IN_CHANNELS * OUT_CHANNELS + + cin * CIN_BLOCK * OUT_CHANNELS + ) + tl.atomic_add(w_ptr + weight_grad_offset, wgrad, sem='relaxed') + + @triton.jit def _ApplyBNReLUConvStats_cl3d_impl( input_ptr, diff --git a/kerops/ops/conv/__init__.py b/kerops/ops/conv/__init__.py index cdc47b2..82b82b9 100644 --- a/kerops/ops/conv/__init__.py +++ b/kerops/ops/conv/__init__.py @@ -1,3 +1,4 @@ from .conv import Conv3d, ApplyBNReLUConv3dStats +from .conv_wgrad import Conv3dWgrad from .dwconv import DWConv from .dwconv_wgrad import DWConvWGRAD diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py new file mode 100644 index 0000000..a280003 --- /dev/null +++ b/kerops/ops/conv/conv_wgrad.py @@ -0,0 +1,130 @@ +import torch +from triton import language as tl, next_power_of_2 + +from ...kernels.conv import _Conv_wgrad_cl3d_impl_V2 +from ...settings import ConfArg, configure, confexc +from ...utils import cdiv + + +@confexc(KeyError) +def num_warps(in_channels, out_channels): + # BE AWARE of swap_grad_with_input - it permutes grad with x + return { + (16, 16): 1, + (16, 32): 2, + (32, 16): 2, + (32, 32): 2, + (32, 64): 4, + (64, 32): 4, + (64, 64): 2, + (64, 128): 4, + (128, 64): 4, + (128, 128): 4, + }[(in_channels, out_channels)] + + +@confexc(KeyError) +def CIN_BLOCK(in_channels, out_channels): + # BE AWARE of swap_grad_with_input - it permutes grad with x + return { + (16, 16): 16, + (16, 32): 16, + (32, 16): 16, + (32, 32): 32, + (32, 64): 16, + (64, 32): 16, + (64, 64): 32, + (64, 128): 32, + (128, 64): 32, + (128, 128): 16, + }[(in_channels, out_channels)] + + +@confexc(KeyError) +def COUT_BLOCK(in_channels, out_channels): + # BE AWARE of swap_grad_with_input - it permutes grad with x + return { + (16, 16): 16, + (16, 32): 32, + (32, 16): 32, + (32, 32): 32, + (32, 64): 64, + (64, 32): 64, + (64, 64): 64, + (64, 128): 128, + (128, 64): 128, + (128, 128): 128, + }[(in_channels, out_channels)] + + +# TODO: looks like this trick can be fixed with a better algorithm. +def swap_grad_with_input(in_channels, out_channels): + return out_channels < in_channels + + +@configure( + ACCTYPE='float32', + num_warps=lambda grad, x: num_warps(x.shape[1], grad.shape[1]), + REDUCTION_FACTOR=32, + CIN_BLOCK=lambda grad, x: CIN_BLOCK(x.shape[1], grad.shape[1]), + COUT_BLOCK=lambda grad, x: COUT_BLOCK(x.shape[1], grad.shape[1]), + D_BLOCK=32, + SWAP_GRAD_WITH_INPUT=lambda grad, x: swap_grad_with_input(x.shape[1], grad.shape[1]) +) +def Conv3dWgrad(grad, x, *, D_BLOCK: ConfArg, ACCTYPE: ConfArg, num_warps: ConfArg, REDUCTION_FACTOR: ConfArg, CIN_BLOCK: ConfArg, COUT_BLOCK: ConfArg, SWAP_GRAD_WITH_INPUT: ConfArg): + if SWAP_GRAD_WITH_INPUT: + grad, x = x, grad + + assert x.device == grad.device + assert x.is_cuda + + assert x.ndim == grad.ndim == 5 + xbsize, in_channels, xH, xW, xD = x.shape + gbsize, out_channels, gH, gW, gD = grad.shape + assert in_channels == next_power_of_2(in_channels) + assert out_channels == next_power_of_2(out_channels) + assert [xbsize, xH, xW, xD] == [gbsize, gH, gW, gD] + + assert x.is_contiguous(memory_format=torch.channels_last_3d) + assert grad.is_contiguous(memory_format=torch.channels_last_3d) + + assert x.dtype == grad.dtype == torch.float16 + + assert D_BLOCK == next_power_of_2(D_BLOCK) + assert ACCTYPE in ('float16', 'float32') + assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) + assert CIN_BLOCK <= in_channels + assert COUT_BLOCK == next_power_of_2(COUT_BLOCK) + assert COUT_BLOCK <= out_channels + assert isinstance(REDUCTION_FACTOR, int) + + ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] + num_buffers = cdiv(xW, REDUCTION_FACTOR * 2) * cdiv(xH, REDUCTION_FACTOR * 2) * cdiv(xD, D_BLOCK * REDUCTION_FACTOR) * xbsize + weight_grad = torch.zeros([num_buffers, 3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float32) + grid = (cdiv(xW, 2) * cdiv(out_channels, COUT_BLOCK), cdiv(xH, 2), cdiv(xD, D_BLOCK) * xbsize) + + _Conv_wgrad_cl3d_impl_V2[grid]( + grad, + x, + weight_grad, + xH, + xW, + xD, + num_buffers, + IN_CHANNELS=in_channels, + OUT_CHANNELS=out_channels, + ACCTYPE=ACCTYPE, + D_BLOCK=D_BLOCK, + CIN_BLOCK=CIN_BLOCK, + COUT_BLOCK=COUT_BLOCK, + num_warps=num_warps, + ) + + weight_grad = weight_grad.sum(dim=0).to(torch.float16) + + if SWAP_GRAD_WITH_INPUT: + weight_grad = weight_grad.permute(0, 1, 2, 4, 3) + weight_grad = torch.flip(weight_grad, dims=(0, 1, 2)) + weight_grad = weight_grad.contiguous() + + return weight_grad diff --git a/notebooks/playground.ipynb b/notebooks/playground.ipynb index 9dde9e5..80fd15b 100644 --- a/notebooks/playground.ipynb +++ b/notebooks/playground.ipynb @@ -9,7 +9,7 @@ "source": [ "import torch\n", "from torch import nn\n", - "from kerops.ops.conv import Conv3d" + "from kerops.ops.conv import Conv3d, Conv3dWgrad" ] }, { @@ -96,246 +96,567 @@ ] }, { - "cell_type": "code", - "execution_count": 2, - "id": "47bf9383-f4a8-44a7-9037-b79c285074e3", - "metadata": {}, - "outputs": [], + "cell_type": "markdown", + "id": "432e294a-c563-4dad-ad5c-b03093d2a61f", + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, "source": [ - "CIN = 64\n", - "COUT = 64\n", - "S = 96\n", - "D = 96\n", - "\n", - "x = torch.randn(2, CIN, S, S, D, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", - "w = torch.randn(3, 3, 3, CIN, COUT, device='cuda', dtype=torch.float16)\n", - "conv = nn.Conv3d(CIN, COUT, 3, padding=1, bias=False, device='cuda')" + "### Conv3dWgradV1" ] }, { "cell_type": "code", - "execution_count": 13, - "id": "ba72112e-f726-4582-a5f0-1ff6d356d9b4", + "execution_count": 69, + "id": "4bc6cd74-9594-47d6-a6f7-ed5ee36373de", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "5.96 ms ± 68 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" - ] - } - ], + "outputs": [], "source": [ - "%%timeit -r 10 -n 10\n", - "with torch.inference_mode(), torch.amp.autocast('cuda'):\n", - " conv(x)\n", + "import triton\n", + "from triton import language as tl, next_power_of_2\n", + "\n", + "\n", + "@triton.jit\n", + "def _Conv_wgrad_cl3d_impl_V1(\n", + " grad_ptr,\n", + " input_ptr,\n", + " weight_grad_ptr,\n", + " H,\n", + " W,\n", + " D,\n", + " num_buffers,\n", + " ACCTYPE: tl.constexpr,\n", + " D_BLOCK: tl.constexpr,\n", + " IN_CHANNELS: tl.constexpr,\n", + " OUT_CHANNELS: tl.constexpr,\n", + " CIN_BLOCK: tl.constexpr,\n", + " COUT_BLOCK: tl.constexpr,\n", + "):\n", + " W_pid = tl.program_id(0)\n", + " H_pid = tl.program_id(1)\n", + " D_pid = tl.program_id(2)\n", + "\n", + " CIN_STEPS: tl.constexpr = IN_CHANNELS // CIN_BLOCK\n", + " COUT_STEPS: tl.constexpr = OUT_CHANNELS // COUT_BLOCK\n", "\n", - "torch.cuda.synchronize()" + " linear = W_pid + H_pid * tl.num_programs(0) + D_pid * tl.num_programs(0) * tl.num_programs(1)\n", + " buffer_idx = linear % num_buffers\n", + " \n", + " in_channels_offset = tl.arange(0, CIN_BLOCK)\n", + " out_channels_offset = tl.arange(0, COUT_BLOCK)\n", + " d_offset = tl.arange(0, D_BLOCK)\n", + "\n", + " grad_offset = d_offset[:, None] * OUT_CHANNELS + out_channels_offset[None, :]\n", + " input_offset = d_offset[None, :] * IN_CHANNELS + in_channels_offset[:, None]\n", + " weight_grad_offset = in_channels_offset[:, None] * OUT_CHANNELS + out_channels_offset[None, :]\n", + "\n", + " grad_ptr += D_pid * D_BLOCK * OUT_CHANNELS\n", + " grad_ptr += W_pid * D * OUT_CHANNELS\n", + " grad_ptr += H_pid * D * W * OUT_CHANNELS\n", + "\n", + " input_ptr += D_pid * D_BLOCK * IN_CHANNELS\n", + " input_ptr += W_pid * D * IN_CHANNELS\n", + " input_ptr += H_pid * D * W * IN_CHANNELS\n", + "\n", + " for cout in tl.static_range(0, COUT_STEPS):\n", + " grad = tl.load(grad_ptr + cout * COUT_BLOCK + grad_offset, other=0, mask=(d_offset < (D - D_pid * D_BLOCK))[:, None])\n", + " \n", + " for cin in tl.static_range(0, CIN_STEPS):\n", + " for h in tl.static_range(-1, 2):\n", + " for w in tl.static_range(-1, 2):\n", + " for d in tl.static_range(-1, 2):\n", + " x_ptr = (\n", + " input_ptr\n", + " + h * IN_CHANNELS * D * W\n", + " + w * IN_CHANNELS * D\n", + " + d * IN_CHANNELS\n", + " + cin * CIN_BLOCK\n", + " )\n", + " \n", + " mask = (d_offset < (D - D_pid * D_BLOCK - d))[None, :] & (d_offset >= (- D_pid * D_BLOCK - d))[None, :] & ((H_pid + h) < H) & ((H_pid + h) >= 0) & ((W_pid + w) < W) & ((W_pid + w) >= 0)\n", + " \n", + " x = tl.load(x_ptr + input_offset, other=0, mask=mask)\n", + " \n", + " wgrad = tl.dot(x, grad, out_dtype=ACCTYPE)\n", + " \n", + " w_ptr = (\n", + " weight_grad_ptr\n", + " + buffer_idx * IN_CHANNELS * OUT_CHANNELS * 3 * 3 * 3\n", + " + (h + 1) * IN_CHANNELS * OUT_CHANNELS * 3 * 3\n", + " + (w + 1) * IN_CHANNELS * OUT_CHANNELS * 3\n", + " + (d + 1) * IN_CHANNELS * OUT_CHANNELS\n", + " + cin * CIN_BLOCK * OUT_CHANNELS\n", + " + cout * COUT_BLOCK\n", + " )\n", + " tl.atomic_add(w_ptr + weight_grad_offset, wgrad, sem='relaxed')" ] }, { "cell_type": "code", "execution_count": 4, - "id": "3f417833-9005-4b2b-9ab3-78cdef7389be", + "id": "378154ee-095c-4e04-b314-db73ad4f0b7f", "metadata": {}, "outputs": [], "source": [ - "Conv3d(x, w)\n", - "torch.cuda.synchronize()" + "from kerops.utils import cdiv\n", + "\n", + "\n", + "def Conv3dWgradV1(grad, x, D_BLOCK, ACCTYPE, num_warps, REDUCTION_FACTOR, CIN_BLOCK, COUT_BLOCK):\n", + " assert x.device == grad.device\n", + " assert x.is_cuda\n", + "\n", + " assert x.ndim == grad.ndim == 5\n", + " xbsize, in_channels, xH, xW, xD = x.shape\n", + " gbsize, out_channels, gH, gW, gD = grad.shape\n", + " assert in_channels == next_power_of_2(in_channels)\n", + " assert out_channels == next_power_of_2(out_channels)\n", + " assert [xbsize, xH, xW, xD] == [gbsize, gH, gW, gD]\n", + "\n", + " assert xbsize == 1 # TODO: batched version\n", + "\n", + " assert x.is_contiguous(memory_format=torch.channels_last_3d)\n", + " assert grad.is_contiguous(memory_format=torch.channels_last_3d)\n", + "\n", + " assert x.dtype == grad.dtype == torch.float16\n", + "\n", + " assert D_BLOCK == next_power_of_2(D_BLOCK)\n", + " assert ACCTYPE in ('float16', 'float32')\n", + " assert CIN_BLOCK == next_power_of_2(CIN_BLOCK)\n", + " assert CIN_BLOCK <= in_channels\n", + " assert COUT_BLOCK == next_power_of_2(COUT_BLOCK)\n", + " assert COUT_BLOCK <= out_channels\n", + " assert isinstance(REDUCTION_FACTOR, int) or REDUCTION_FACTOR == None\n", + "\n", + " ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE]\n", + " num_buffers = 1 if REDUCTION_FACTOR is None else cdiv(xW, REDUCTION_FACTOR) * cdiv(xH, REDUCTION_FACTOR) * cdiv(D, D_BLOCK * REDUCTION_FACTOR)\n", + " weight_grad = torch.zeros([num_buffers, 3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float16)\n", + " grid = (xW, xH, cdiv(xD, D_BLOCK))\n", + "\n", + " _Conv_wgrad_cl3d_impl_V1[grid](\n", + " grad,\n", + " x,\n", + " weight_grad,\n", + " xH,\n", + " xW,\n", + " xD,\n", + " num_buffers,\n", + " IN_CHANNELS=in_channels,\n", + " OUT_CHANNELS=out_channels,\n", + " ACCTYPE=ACCTYPE,\n", + " D_BLOCK=D_BLOCK,\n", + " CIN_BLOCK=CIN_BLOCK,\n", + " COUT_BLOCK=COUT_BLOCK,\n", + " num_warps=num_warps,\n", + " )\n", + "\n", + " return weight_grad.sum(dim=0)" ] }, { - "cell_type": "code", - "execution_count": 14, - "id": "f7731c96-8bb7-4ece-ab40-d438beeba82e", + "cell_type": "markdown", + "id": "2b12f8ca-7a33-471e-b6f8-0d868e355b9a", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "6 ms ± 81.2 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" - ] - } - ], "source": [ - "%%timeit -r 10 -n 10\n", - "Conv3d(x, w)\n", - "torch.cuda.synchronize()" + "### Conv3dWgradV2" ] }, { "cell_type": "code", - "execution_count": 6, - "id": "925fbd7c-e860-4f80-9c5c-84c886175c62", - "metadata": {}, + "execution_count": 150, + "id": "ec8a8fa1-6afe-4a25-9ca9-54b21816edcb", + "metadata": { + "jupyter": { + "source_hidden": true + } + }, "outputs": [], "source": [ - "Conv3d(x, w, **{'num_warps': 2, 'D_BLOCK': 32, 'CIN_BLOCK': 32, 'LOAD_WEIGHT_FIRST': True, 'WEIGHT_MAJOR': True})\n", - "torch.cuda.synchronize()" + "import triton\n", + "from triton import language as tl, next_power_of_2\n", + "\n", + "\n", + "@triton.jit\n", + "def _Conv_wgrad_cl3d_impl_V2(\n", + " grad_ptr,\n", + " input_ptr,\n", + " weight_grad_ptr,\n", + " H,\n", + " W,\n", + " D,\n", + " num_buffers,\n", + " ACCTYPE: tl.constexpr,\n", + " D_BLOCK: tl.constexpr,\n", + " IN_CHANNELS: tl.constexpr,\n", + " OUT_CHANNELS: tl.constexpr,\n", + " CIN_BLOCK: tl.constexpr,\n", + " COUT_BLOCK: tl.constexpr,\n", + "):\n", + " WCOUT_pid = tl.program_id(0)\n", + " H_pid = tl.program_id(1)\n", + " BD_pid = tl.program_id(2)\n", + "\n", + " W_pid = WCOUT_pid // tl.cdiv(OUT_CHANNELS, COUT_BLOCK)\n", + " COUT_pid = WCOUT_pid % tl.cdiv(OUT_CHANNELS, COUT_BLOCK)\n", + "\n", + " B_pid = BD_pid // tl.cdiv(D, D_BLOCK)\n", + " D_pid = BD_pid % tl.cdiv(D, D_BLOCK)\n", + "\n", + " CIN_STEPS: tl.constexpr = IN_CHANNELS // CIN_BLOCK\n", + "\n", + " linear = W_pid + H_pid * tl.num_programs(0) + BD_pid * tl.num_programs(0) * tl.num_programs(1)\n", + " buffer_idx = linear % num_buffers\n", + " \n", + " in_channels_offset = tl.arange(0, CIN_BLOCK)\n", + " out_channels_offset = tl.arange(0, COUT_BLOCK)\n", + " d_offset = tl.arange(0, D_BLOCK)\n", + "\n", + " grad_offset = d_offset[:, None] * OUT_CHANNELS + out_channels_offset[None, :]\n", + " input_offset = d_offset[None, :] * IN_CHANNELS + in_channels_offset[:, None]\n", + " weight_grad_offset = in_channels_offset[:, None] * OUT_CHANNELS + out_channels_offset[None, :]\n", + "\n", + " grad_ptr += COUT_pid * COUT_BLOCK\n", + " grad_ptr += D_pid * D_BLOCK * OUT_CHANNELS\n", + " grad_ptr += W_pid * 2 * D * OUT_CHANNELS\n", + " grad_ptr += H_pid * 2 * D * W * OUT_CHANNELS\n", + " grad_ptr += B_pid * H * W * D * OUT_CHANNELS\n", + "\n", + " input_ptr += D_pid * D_BLOCK * IN_CHANNELS\n", + " input_ptr += W_pid * 2 * D * IN_CHANNELS\n", + " input_ptr += H_pid * 2 * D * W * IN_CHANNELS\n", + " input_ptr += B_pid * H * W * D * IN_CHANNELS\n", + "\n", + " weight_grad_ptr += COUT_pid * COUT_BLOCK\n", + "\n", + " g01 = ((W_pid * 2 + 1) < W)\n", + " g10 = ((H_pid * 2 + 1) < H)\n", + " g11 = ((W_pid * 2 + 1) < W) & ((H_pid * 2 + 1) < H)\n", + "\n", + " gmask = (d_offset < (D - D_pid * D_BLOCK))\n", + " gmask = gmask[:, None]\n", + "\n", + " grads = [\n", + " [\n", + " tl.load(grad_ptr + grad_offset, other=0, mask=gmask),\n", + " tl.load(grad_ptr + grad_offset + OUT_CHANNELS * D, other=0, mask=gmask & g01),\n", + " ],\n", + " [\n", + " tl.load(grad_ptr + grad_offset + OUT_CHANNELS * D * W, other=0, mask=gmask & g10),\n", + " tl.load(grad_ptr + grad_offset + OUT_CHANNELS * D * W + OUT_CHANNELS * D, other=0, mask=gmask & g11)\n", + " ]\n", + " ]\n", + "\n", + " for h in tl.static_range(-1, 2):\n", + " for w in tl.static_range(-1, 2):\n", + " for d in tl.static_range(-1, 2):\n", + " for cin in tl.static_range(0, CIN_STEPS):\n", + " wgrad = tl.zeros([CIN_BLOCK, COUT_BLOCK], dtype=ACCTYPE)\n", + " \n", + " x_ptr = (\n", + " input_ptr\n", + " + h * IN_CHANNELS * D * W\n", + " + w * IN_CHANNELS * D\n", + " + d * IN_CHANNELS\n", + " + cin * CIN_BLOCK\n", + " )\n", + " \n", + " xmask = (d_offset < (D - D_pid * D_BLOCK - d)) & (d_offset >= (- D_pid * D_BLOCK - d))\n", + " xmask = xmask[None, :]\n", + "\n", + " x00 = ((H_pid * 2 + h) < H) & ((H_pid * 2 + h) >= 0) & ((W_pid * 2 + w) < W) & ((W_pid * 2 + w) >= 0)\n", + " x01 = ((H_pid * 2 + h) < H) & ((H_pid * 2 + h) >= 0) & ((W_pid * 2 + w + 1) < W) & ((W_pid * 2 + w + 1) >= 0)\n", + " x10 = ((H_pid * 2 + h + 1) < H) & ((H_pid * 2 + h + 1) >= 0) & ((W_pid * 2 + w) < W) & ((W_pid * 2 + w) >= 0)\n", + " x11 = ((H_pid * 2 + h + 1) < H) & ((H_pid * 2 + h + 1) >= 0) & ((W_pid * 2 + w + 1) < W) & ((W_pid * 2 + w + 1) >= 0)\n", + "\n", + " xs = [\n", + " [\n", + " tl.load(x_ptr + input_offset, other=0, mask=xmask & x00),\n", + " tl.load(x_ptr + input_offset + IN_CHANNELS * D, other=0, mask=xmask & x01),\n", + " ],\n", + " [\n", + " tl.load(x_ptr + input_offset + IN_CHANNELS * D * W, other=0, mask=xmask & x10),\n", + " tl.load(x_ptr + input_offset + IN_CHANNELS * D * W + IN_CHANNELS * D, other=0, mask=xmask & x11),\n", + " ]\n", + " ]\n", + " \n", + " for kh in tl.static_range(0, 2):\n", + " for kw in tl.static_range(0, 2):\n", + " wgrad += tl.dot(xs[kh][kw], grads[kh][kw])\n", + "\n", + " \n", + " w_ptr = (\n", + " weight_grad_ptr\n", + " + buffer_idx * IN_CHANNELS * OUT_CHANNELS * 3 * 3 * 3\n", + " + (h + 1) * IN_CHANNELS * OUT_CHANNELS * 3 * 3\n", + " + (w + 1) * IN_CHANNELS * OUT_CHANNELS * 3\n", + " + (d + 1) * IN_CHANNELS * OUT_CHANNELS\n", + " + cin * CIN_BLOCK * OUT_CHANNELS\n", + " )\n", + " tl.atomic_add(w_ptr + weight_grad_offset, wgrad, sem='relaxed')" ] }, { "cell_type": "code", - "execution_count": 7, - "id": "fc10543a-982d-48fa-bfbb-0e2872b55968", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "5.86 ms ± 57.5 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" - ] + "execution_count": 151, + "id": "d28d18e2-1ac4-4a11-bf12-6e64397080f9", + "metadata": { + "jupyter": { + "source_hidden": true } - ], - "source": [ - "%%timeit -r 10 -n 10\n", - "Conv3d(x, w, **{'num_warps': 2, 'D_BLOCK': 32, 'CIN_BLOCK': 32, 'LOAD_WEIGHT_FIRST': True, 'WEIGHT_MAJOR': True})\n", - "torch.cuda.synchronize()" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "6a6246f6-be7b-4683-ae51-7aa22ee620e3", - "metadata": {}, + }, "outputs": [], "source": [ - "Conv3d(x, w, **{'num_warps': 2, 'D_BLOCK': 16, 'CIN_BLOCK': 16, 'LOAD_WEIGHT_FIRST': False, 'WEIGHT_MAJOR': True})\n", - "torch.cuda.synchronize()" + "from kerops.utils import cdiv\n", + "\n", + "\n", + "def Conv3dWgradV2(grad, x, D_BLOCK, ACCTYPE, num_warps, REDUCTION_FACTOR, CIN_BLOCK, COUT_BLOCK):\n", + " assert x.device == grad.device\n", + " assert x.is_cuda\n", + "\n", + " assert x.ndim == grad.ndim == 5\n", + " xbsize, in_channels, xH, xW, xD = x.shape\n", + " gbsize, out_channels, gH, gW, gD = grad.shape\n", + " assert in_channels == next_power_of_2(in_channels)\n", + " assert out_channels == next_power_of_2(out_channels)\n", + " assert [xbsize, xH, xW, xD] == [gbsize, gH, gW, gD]\n", + "\n", + " assert x.is_contiguous(memory_format=torch.channels_last_3d)\n", + " assert grad.is_contiguous(memory_format=torch.channels_last_3d)\n", + "\n", + " assert x.dtype == grad.dtype == torch.float16\n", + "\n", + " assert D_BLOCK == next_power_of_2(D_BLOCK)\n", + " assert ACCTYPE in ('float16', 'float32')\n", + " assert CIN_BLOCK == next_power_of_2(CIN_BLOCK)\n", + " assert CIN_BLOCK <= in_channels\n", + " assert isinstance(REDUCTION_FACTOR, int)\n", + "\n", + " ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE]\n", + " num_buffers = cdiv(xW, REDUCTION_FACTOR * 2) * cdiv(xH, REDUCTION_FACTOR * 2) * cdiv(xD, D_BLOCK * REDUCTION_FACTOR) * xbsize\n", + " weight_grad = torch.zeros([num_buffers, 3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float32)\n", + " grid = (cdiv(xW, 2) * cdiv(out_channels, COUT_BLOCK), cdiv(xH, 2), cdiv(xD, D_BLOCK) * xbsize)\n", + "\n", + " _Conv_wgrad_cl3d_impl_V2[grid](\n", + " grad,\n", + " x,\n", + " weight_grad,\n", + " xH,\n", + " xW,\n", + " xD,\n", + " num_buffers,\n", + " IN_CHANNELS=in_channels,\n", + " OUT_CHANNELS=out_channels,\n", + " ACCTYPE=ACCTYPE,\n", + " D_BLOCK=D_BLOCK,\n", + " CIN_BLOCK=CIN_BLOCK,\n", + " COUT_BLOCK=COUT_BLOCK,\n", + " num_warps=num_warps,\n", + " )\n", + "\n", + " return weight_grad.sum(dim=0).to(torch.float16)" ] }, { "cell_type": "code", - "execution_count": 12, - "id": "df11394a-9430-4c40-b64d-a1999f31bad3", + "execution_count": 3, + "id": "4e6da0ab-c024-495f-92e6-d17143eb928a", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "5.93 ms ± 70.1 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" + "All stages passed\n" ] - } - ], - "source": [ - "%%timeit -r 10 -n 10\n", - "Conv3d(x, w, **{'num_warps': 2, 'D_BLOCK': 16, 'CIN_BLOCK': 16, 'LOAD_WEIGHT_FIRST': False, 'WEIGHT_MAJOR': True})\n", - "torch.cuda.synchronize()" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "cd845ca5-cdb0-4fa4-9f7b-987e95152aca", - "metadata": {}, - "outputs": [ + }, { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "58bf24ace46648b0ad4199ed9dc082bd", - "version_major": 2, - "version_minor": 0 - }, "text/plain": [ - "Benchmark configs: 0%| | 0/48 [00:00 Date: Thu, 19 Mar 2026 21:48:07 +0100 Subject: [PATCH 15/40] Conv3d perf tuning --- kerops/ops/conv/conv.py | 30 ++++++++++++---- notebooks/playground.ipynb | 70 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 91 insertions(+), 9 deletions(-) diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py index 17e3c07..bb0e68a 100644 --- a/kerops/ops/conv/conv.py +++ b/kerops/ops/conv/conv.py @@ -14,11 +14,11 @@ def num_warps(in_channels, out_channels): (16, 32): 4, (32, 16): 4, (32, 32): 2, - (32, 64): 4, - (64, 32): 1, + (32, 64): 2, + (64, 32): 2, (64, 64): 4, - (64, 128): 4, - (128, 64): 4, + (64, 128): 2, + (128, 64): 2, (128, 128): 4, }[(in_channels, out_channels)] @@ -44,24 +44,40 @@ def cin_block(in_channels, out_channels): return { (16, 16): 16, (16, 32): 16, - (32, 16): 32, + (32, 16): 16, (32, 32): 16, (32, 64): 16, (64, 32): 16, - (64, 64): 16, + (64, 64): 32, (64, 128): 16, (128, 64): 16, (128, 128): 16, }[(in_channels, out_channels)] +@confexc(KeyError) +def weight_major(in_channels, out_channels): + return { + (16, 16): False, + (16, 32): True, + (32, 16): False, + (32, 32): False, + (32, 64): True, + (64, 32): False, + (64, 64): True, + (64, 128): True, + (128, 64): True, + (128, 128): True, + }[(in_channels, out_channels)] + + @configure( ACCTYPE='float32', num_warps=lambda weight: num_warps(*weight.shape[-2:]), D_BLOCK=lambda weight: d_block(*weight.shape[-2:]), CIN_BLOCK=lambda weight: cin_block(*weight.shape[-2:]), LOAD_WEIGHT_FIRST=True, - WEIGHT_MAJOR=False, + WEIGHT_MAJOR=lambda weight: weight_major(*weight.shape[-2:]), ) def Conv3d(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_BLOCK: ConfArg, CIN_BLOCK: ConfArg, LOAD_WEIGHT_FIRST: ConfArg, WEIGHT_MAJOR: ConfArg): assert x.device == weight.device diff --git a/notebooks/playground.ipynb b/notebooks/playground.ipynb index 80fd15b..2737c5e 100644 --- a/notebooks/playground.ipynb +++ b/notebooks/playground.ipynb @@ -20,9 +20,12 @@ "outputs": [], "source": [ "import itertools\n", + "import json\n", "from tqdm.notebook import tqdm\n", "from time import perf_counter, sleep\n", + "\n", "import numpy as np\n", + "from joblib import Parallel, delayed\n", "\n", "\n", "def mean_std_percentile(x, lo=20, hi=80):\n", @@ -36,13 +39,20 @@ " return x_mid.mean(), x_mid.std()\n", "\n", "\n", - "def bench(func, *args, warmup=25, sleep_ms=100, n_iters=50, q_show=10, quantiles=(20, 80), **specset):\n", + "def bench(func, *args, n_jobs_precompile=4, warmup=25, sleep_ms=100, n_iters=50, q_show=10, quantiles=(20, 80), savefile=None, **specset):\n", " results = []\n", "\n", " keys = list(specset.keys())\n", " values = list(specset.values())\n", " configs = list(itertools.product(*values))\n", "\n", + " def precompile_call(config):\n", + " kwargs = dict(zip(keys, config))\n", + " func(*args, **kwargs)\n", + "\n", + " n_jobs_precompile = min(n_jobs_precompile, len(configs))\n", + " Parallel(n_jobs=n_jobs_precompile, backend='threading')(delayed(precompile_call)(config) for config in tqdm(configs, desc=\"Precompiling\"))\n", + "\n", " for config in tqdm(configs, desc=\"Benchmark configs\"):\n", " if sleep_ms is not None:\n", " sleep(sleep_ms / 1000)\n", @@ -92,7 +102,63 @@ " for result in good_results:\n", " print(f'{result['mean_ms']:.3f}+-{result['std_ms']:.3f}ms {result['spec']}')\n", " else:\n", - " print(f'Other specs have a time difference of more than {q_show}%')" + " print(f'Other specs have a time difference of more than {q_show}%')\n", + "\n", + " if savefile is not None:\n", + " with open(savefile, 'w') as f:\n", + " json.dump({'best': best_result, 'good_results': good_results}, f)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "05308eb4-225e-496b-b785-397de9dd7c83", + "metadata": {}, + "outputs": [], + "source": [ + "from functools import partial\n", + "from kerops.ops.conv.conv import num_warps, d_block, cin_block\n", + "\n", + "\n", + "cudnn_conv = partial(nn.functional.conv3d, stride=1, padding=1)\n", + "\n", + "\n", + "channels = [16, 32, 64, 128]\n", + "\n", + "\n", + "def get_size(CIN, COUT):\n", + " if CIN <= 32 and COUT <= 32:\n", + " return 128\n", + " elif CIN <= 64 and COUT <= 64:\n", + " return 96\n", + " else:\n", + " return 64\n", + "\n", + "\n", + "for CIN in channels:\n", + " for COUT in channels:\n", + " if CIN == COUT or 2 * CIN == COUT or CIN == 2 * COUT:\n", + " print(f'{CIN=} {COUT=}')\n", + "\n", + " BS = 1\n", + " S = get_size(CIN, COUT)\n", + " D = get_size(CIN, COUT)\n", + "\n", + " x = torch.randn(BS, CIN, S, S, D, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", + " grad = torch.randn(BS, COUT, S, S, D, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", + " w = torch.randn(3, 3, 3, CIN, COUT, device='cuda', dtype=torch.float16)\n", + " weight = w.permute(-1, -2, 0, 1, 2).contiguous()\n", + "\n", + " print('default')\n", + " bench(Conv3d, x, w)\n", + " print('cudnn')\n", + " bench(cudnn_conv, x, weight, None)\n", + " print('bench')\n", + " CIN_BLOCKS = [16, 32] if CIN >= 32 else [16]\n", + " D_BLOCKS = [16, 32, 64] if (CIN <= 32 and COUT <= 32) else [16, 32]\n", + " bench(Conv3d, x, w, savefile=f'Conv3dV6_{CIN}_{COUT}.json', num_warps=[1, 2, 4], D_BLOCK=D_BLOCKS, CIN_BLOCK=CIN_BLOCKS, LOAD_WEIGHT_FIRST=[True, False], WEIGHT_MAJOR=[False, True])\n", + " \n", + " print(30 * '-')" ] }, { From a5d617e1dcc1bdd6f84c923cc23524746db3554b Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Wed, 25 Mar 2026 15:27:30 +0300 Subject: [PATCH 16/40] Revision of ConfiguredFunction and configuration --- kerops/settings/__init__.py | 5 +- kerops/settings/hardware_conf.py | 11 -- kerops/settings/kernel_config.py | 229 +++++++++++++++++++++++++++++++ kerops/settings/utils.py | 34 ----- kerops/settings/wrapper.py | 141 +++++++++---------- 5 files changed, 295 insertions(+), 125 deletions(-) delete mode 100644 kerops/settings/hardware_conf.py create mode 100644 kerops/settings/kernel_config.py delete mode 100644 kerops/settings/utils.py diff --git a/kerops/settings/__init__.py b/kerops/settings/__init__.py index c74a353..f7e4913 100644 --- a/kerops/settings/__init__.py +++ b/kerops/settings/__init__.py @@ -1,3 +1,2 @@ -from .hardware_conf import get_l1_cache -from .utils import ConfArg, CongiguratorError -from .wrapper import confexc, configure +from .wrapper import ConfiguredFunction, ConfArg +from .kernel_config import StaticKernelConfig, RuleKernelConfig, TableKernelConfig diff --git a/kerops/settings/hardware_conf.py b/kerops/settings/hardware_conf.py deleted file mode 100644 index e9eba7e..0000000 --- a/kerops/settings/hardware_conf.py +++ /dev/null @@ -1,11 +0,0 @@ -L1_CACHE_BYTES = 65536 - - -def get_l1_cache(): - global L1_CACHE_BYTES - return L1_CACHE_BYTES - - -def set_l1_cache(new_cache): - global L1_CACHE_BYTES - L1_CACHE_BYTES = new_cache diff --git a/kerops/settings/kernel_config.py b/kerops/settings/kernel_config.py new file mode 100644 index 0000000..6bfeed5 --- /dev/null +++ b/kerops/settings/kernel_config.py @@ -0,0 +1,229 @@ +import tomllib +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Callable, Any +from inspect import signature as get_signature +from warnings import warn + +from torch import Tensor +from torch.cuda import get_device_name + + +class KernelConfigBase(ABC): + @property + @abstractmethod + def arg_names(self) -> list[str]: + ... + + @property + @abstractmethod + def confarg_names(self) -> list[str]: + ... + + @abstractmethod + def __call__(self, *input_args) -> dict[str, Any]: + ... + + @abstractmethod + def __repr__(self): + ... + + def can_be_configured(self, *input_args): + try: + self(*input_args) + return True + except Exception: + return False + + +class StaticKernelConfig(KernelConfigBase): + def __init__(self, name: str | None = None, **params): + self.params = params + self.arg_names = [] + self.confarg_names = list(params.keys()) + self.name = name + + def can_be_configured(self): + return True + + def __call__(self): + return self.params + + def __repr__(self): + prefix = f'{self.name}: ' if self.name else '' + input_s = ', '.join(self.arg_names) + conf_s = ', '.join(self.confarg_names) + return ( + f'{prefix}StaticKernelConfig({input_s})' + f' -> ({conf_s})' + ) + + +class RuleKernelConfig(KernelConfigBase): + def __init__( + self, + args_to_problem_sizes: Callable, + problem_size_names: list[str], + rule: Callable, + confarg_names: list[str], + name: str | None = None + ): + sig = get_signature(args_to_problem_sizes) + + self.arg_names = [param.name for param in sig.parameters.values()] + self.problem_size_names = problem_size_names + self.args_to_problem_sizes = args_to_problem_sizes + + sig = get_signature(rule) + + if any(param.name != problem_size_name for param, problem_size_name in zip(sig.parameters.values(), problem_size_names, strict=True)): + raise ValueError(f'Rule required {list(sig.parameters.values())} as input, but provided {problem_size_names}') + + self.confarg_names = confarg_names + self.rule = rule + self.name = name + + def __call__(self, *input_args): + problem_sizes = self.args_to_problem_sizes(*input_args) + + if not isinstance(problem_sizes, tuple): + problem_sizes = (problem_sizes,) + + return self.rule(*problem_sizes) + + def __repr__(self): + prefix = f'{self.name}: ' if self.name else '' + input_s = ', '.join(self.arg_names) + inter_s = ', '.join(self.problem_size_names) + conf_s = ', '.join(self.confarg_names) + return ( + f'{prefix}RuleKernelConfig({input_s})' + f' -[{inter_s}]-> ({conf_s})' + ) + + +class TableKernelConfig(KernelConfigBase): + def __init__( + self, + problem_size_names: list[str], + confarg_names: list[str], + args_to_problem_sizes: Callable, + fallback_device: str | None = None, + toml_path: str | None = None, + name: str | None = None + ): + sig = get_signature(args_to_problem_sizes) + + self.arg_names = [param.name for param in sig.parameters.values()] + self.problem_size_names = problem_size_names + self.confarg_names = confarg_names + self.args_to_problem_sizes = args_to_problem_sizes + + # { gpu_name: { problem_sizes: configured_args_dict } } + self._configs: dict[str, dict[tuple, dict[str, Any]]] = {} + self.fallback_device = fallback_device + self.name = name + + if toml_path is not None: + self.load_toml(toml_path) + + def load_toml(self, path: str | Path): + path = Path(path) + with open(path, 'rb') as f: + raw = tomllib.load(f) + + meta = raw.get('meta') + if meta is None: + raise ValueError('TOML must contain a [meta] section') + + toml_problem_size_names = meta.get('problem_size_names') + if toml_problem_size_names is None: + raise ValueError('[meta] must contain problem_size_names') + + if list(toml_problem_size_names) != self.problem_size_names: + raise ValueError( + f'problem_size_names mismatch: ' + f'toml={toml_problem_size_names}, config={self.problem_size_names}' + ) + + for gpu_name, gpu_data in raw.items(): + if gpu_name == 'meta': + continue + + configs_raw = gpu_data.get('configs') + if not isinstance(configs_raw, list): + raise ValueError(f'[{gpu_name}] must contain an array of configs') + + parsed: dict[tuple, dict[str, Any]] = {} + + for entry in configs_raw: + if 'problem_sizes' not in entry: + raise ValueError(f'[{gpu_name}] config entry missing "problem_sizes" field: {entry}') + + problem_sizes = tuple(entry['problem_sizes']) + if len(problem_sizes) != len(self.problem_size_names): + raise ValueError( + f'[{gpu_name}] problem_sizes tuple length {len(problem_sizes)} ' + f'!= problem_size_names length {len(self.problem_size_names)}' + ) + + configured = {k: v for k, v in entry.items() if k != 'problem_sizes'} + missing = set(self.confarg_names) - configured.keys() + extra = configured.keys() - set(self.confarg_names) + + if missing: + raise ValueError(f'[{gpu_name}] config entry missing keys: {missing}') + if extra: + raise ValueError(f'[{gpu_name}] config entry has unexpected keys: {extra}') + + parsed[problem_sizes] = configured + + self._configs[gpu_name] = parsed + + if self.fallback_device is not None and self.fallback_device not in self._configs: + raise ValueError( + f'fallback_device={self.fallback_device!r} not found in loaded configs: {list(self._configs)}' + ) + + def get_device_name_from_args(self, *input_args): + device_indices = {arg.device.index for arg in input_args if isinstance (arg, Tensor) and arg.device.type == 'cuda'} + + if len(device_indices) == 1: + return get_device_name(device_indices.pop()) + elif len(device_indices) == 0: + raise RuntimeError('Cannot configure due to non-cuda args') + else: + raise RuntimeError(f'Expected all tensors to be on the same GPU, got CUDA-devices:{device_indices}') + + def __call__(self, *input_args) -> dict[str, Any]: + device = self.get_device_name_from_args(*input_args) + + problem_sizes = self.args_to_problem_sizes(*input_args) + + if not isinstance(problem_sizes, tuple): + problem_sizes = (problem_sizes,) + + if device not in self._configs: + if self.fallback_device is None: + raise RuntimeError("Fallback device has not been set") + + warn(f'{self} is not configured for {device}, fallback to {self.fallback_device}', stacklevel=2) + device_configs = self._configs[self.fallback_device] + device = self.fallback_device + else: + device_configs = self._configs[device] + + if problem_sizes not in device_configs: + raise KeyError(f'No config for problem_sizes {problem_sizes} on device {device}') + + return device_configs[problem_sizes] + + def __repr__(self): + prefix = f'{self.name}: ' if self.name else '' + input_s = ', '.join(self.arg_names) + inter_s = ', '.join(self.problem_size_names) + conf_s = ', '.join(self.confarg_names) + return ( + f'{prefix}TableKernelConfig({input_s})' + f' -[{inter_s}]-> ({conf_s})' + ) diff --git a/kerops/settings/utils.py b/kerops/settings/utils.py deleted file mode 100644 index 5855e34..0000000 --- a/kerops/settings/utils.py +++ /dev/null @@ -1,34 +0,0 @@ -from inspect import Parameter - - -class ConfArg: - pass - - -class CongiguratorError(Exception): - pass - - -def validate_signature(signature): - for param in signature.parameters.values(): - if param.annotation is ConfArg and param.kind is not Parameter.KEYWORD_ONLY: - raise RuntimeError(f'ConfigurableArg must be keyword-only - {param.name}') - elif param.annotation is not ConfArg and param.kind is Parameter.KEYWORD_ONLY: - raise RuntimeError(f'non-ConfigurableArg must not be keyword-only - {param.name}') - - -def get_config_args(signature): - return [param.name for param in signature.parameters.values() if param.annotation is ConfArg] - - -def get_standard_args(signature): - return [ - param.name - for param in signature.parameters.values() - if param.kind is Parameter.POSITIONAL_ONLY or param.kind is Parameter.POSITIONAL_OR_KEYWORD - ] - - -def configs_match(configurable_args, configurators_names): - if set(configurable_args) != set(configurators_names): - raise RuntimeError(f'Configuration mismatch, {configurable_args=}, {configurators_names=}') diff --git a/kerops/settings/wrapper.py b/kerops/settings/wrapper.py index 914b97e..9e2cf98 100644 --- a/kerops/settings/wrapper.py +++ b/kerops/settings/wrapper.py @@ -1,107 +1,94 @@ -import inspect +from inspect import Parameter, signature as get_signature from functools import wraps from typing import Callable -from .utils import CongiguratorError, configs_match, get_config_args, get_standard_args, validate_signature +from .kernel_config import KernelConfig -class EmptyKwarg: +class ConfArg: pass -class ConfiguredFunction: - def __init__(self, origin_function, signature, configurable_args, usual_args, **configurators): - self.origin_function = origin_function - self.signature = signature - self.configurable_args = configurable_args - self.usual_args = usual_args - self.configurators = configurators - - def __repr__(self): - def format_configurator(configurator): - if isinstance(configurator, Callable): - params = ', '.join(inspect.signature(configurator).parameters) - return f'Configurator({params})' - return str(configurator) - - configurators_repr = '\n'.join( - f'{confarg}: {format_configurator(configurator)}' for confarg, configurator in self.configurators.items() - ) +EmptyKwarg = object() - return f'{self.origin_function.__name__}{self.signature}\n{configurators_repr}' - @staticmethod - def configurator_call(args, configurator, usual_args): - if isinstance(configurator, Callable): - conf_sign = inspect.signature(configurator) - - # take argnames from configurator, map args with respect to origin function's argnames - conf_args = [args[usual_args.index(param.name)] for param in conf_sign.parameters.values()] - - return configurator(*conf_args) - else: - return configurator +class ConfiguredFunction: + def __init__(self, function: Callable, kernel_config: KernelConfig): + self.function = function - def __call__(self, *args, **kwargs): - tmp_kwargs = {**{arg: EmptyKwarg for arg in self.configurable_args}, **kwargs} + signature = get_signature(function) + for param in signature.parameters.values(): + if param.kind is Parameter.VAR_POSITIONAL: + raise TypeError(f'VAR_POSITIONAL (*args) is not supported - {param.name}') - bind = self.signature.bind(*args, **tmp_kwargs) - bind.apply_defaults() + elif param.annotation is ConfArg: + if param.kind is not Parameter.KEYWORD_ONLY: + raise TypeError(f'ConfArg must be keyword-only - {param.name}') + + if param.default is not param.empty: + raise TypeError(f'ConfArg must not have default value - {param.name}:{param.default}') - configured_kwargs = { - k: ( - self.configurator_call(bind.args, self.configurators[k], self.usual_args) - if input_v is EmptyKwarg - else input_v - ) - for k, input_v in bind.kwargs.items() - } + elif param.annotation is not ConfArg and param.kind is Parameter.KEYWORD_ONLY: + raise TypeError(f'non-ConfArg must not be keyword-only - {param.name}') - return self.origin_function(*bind.args, **configured_kwargs) + self.signature = signature - def can_be_configured(self, **kwargs): - try: - for configurator in self.configurators.values(): - if isinstance(configurator, Callable): - confargs = inspect.signature(configurator).parameters - _ = configurator(**{arg: kwargs[arg] for arg in confargs}) + self.confargs = [param.name for param in self.signature.parameters.values() if param.annotation is ConfArg] + self.usual_args = [param.name for param in self.signature.parameters.values() if param.annotation is not ConfArg] - except CongiguratorError: - return False + self.register_kernel_config(kernel_config) - return True + def register_kernel_config(self, kernel_config: KernelConfig): + configured_arg_names = kernel_config.confarg_names + input_arg_names = kernel_config.arg_names - def reconfigure(self, **new_configurators): - configs_match(self.configurable_args, new_configurators.keys()) - self.configurators = new_configurators + if set(self.confargs) != set(configured_arg_names): + raise ValueError( + f'Configuration mismatch, confargs={self.confargs}, configured_arg_names={configured_arg_names}' + ) + for arg in input_arg_names: + if arg not in self.usual_args: + raise ValueError(f"{kernel_config.__class__.__name__} expects unknown arg: {arg}") -def configure(**configurators): - def wrapper(function): - signature = inspect.signature(function) + self.kernel_config = kernel_config + self.kernel_input_arg_indices = [self.usual_args.index(arg) for arg in input_arg_names] - validate_signature(signature) + def kernel_config_call(self, usual_args): + return self.kernel_config(*(usual_args[idx] for idx in self.kernel_input_arg_indices)) - configurable_args = get_config_args(signature) + def __call__(self, *args, **kwargs): + # all confargs and, maybe, some usual args passed as keyword-argument + full_kwargs = {arg: EmptyKwarg for arg in self.confargs} + full_kwargs.update(kwargs) - usual_args = get_standard_args(signature) + bind = self.signature.bind(*args, **full_kwargs) + bind.apply_defaults() - configs_match(configurable_args, configurators.keys()) + # after binding kwargs consists of ConfArgs ONLY, and kwargs.keys() == self.confargs + # args does no contain ant ConfArg + args, kwargs = bind.args, bind.kwargs - return wraps(function)(ConfiguredFunction(function, signature, configurable_args, usual_args, **configurators)) + # if all ConfArgs are overriden there is no reason to call kernel_config + if any(v is EmptyKwarg for v in kwargs.values()): + # configuration - ConfArg is configured in priority order for overriden values from kwargs + configured_kwargs = self.kernel_config_call(bind.args) - return wrapper + configured_kwargs = { + k: configured_kwargs[k] if v is EmptyKwarg else v + for k, v in kwargs.items() + } + else: + configured_kwargs = kwargs + return self.function(*args, **configured_kwargs) -def confexc(*exceptions): - def wrapper(configurator): - @wraps(configurator) - def wrapped(*args, **kwargs): - try: - return configurator(*args, **kwargs) - except exceptions as e: - raise CongiguratorError(str(e)) + @classmethod + def configure(cls, kernel_config: KernelConfig): + def wrapper(function: Callable): + return wraps(function)(cls(function, kernel_config)) - return wrapped + return wrapper - return wrapper + def __repr__(self): + return f'{self.function.__name__}{self.signature}\n{self.kernel_config}' From eff09762d83427f9a83fa59397fc3d3b799ef4cf Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Thu, 2 Apr 2026 17:22:08 +0200 Subject: [PATCH 17/40] autotuner --- kerops/settings/autotune.py | 153 ++++++++ kerops/settings/kernel_config.py | 25 +- kerops/settings/utils.py | 13 + kerops/settings/wrapper.py | 8 +- notebooks/config_builder.ipynb | 377 ++++++++++++++++++ notebooks/playground.ipynb | 643 +++++++++++++++++++++++++++++-- 6 files changed, 1151 insertions(+), 68 deletions(-) create mode 100644 kerops/settings/autotune.py create mode 100644 kerops/settings/utils.py create mode 100644 notebooks/config_builder.ipynb diff --git a/kerops/settings/autotune.py b/kerops/settings/autotune.py new file mode 100644 index 0000000..6efd244 --- /dev/null +++ b/kerops/settings/autotune.py @@ -0,0 +1,153 @@ +import itertools +import torch +from pathlib import Path +from typing import Callable +from tqdm.notebook import tqdm +from time import perf_counter, sleep +from warnings import warn + +import numpy as np +from joblib import Parallel, delayed + +from .utils import get_device_name_from_args + + +def mean_std_percentile(x, lo=20, hi=80): + x = np.asarray(x, dtype=np.float32) + p_lo, p_hi = np.percentile(x, [lo, hi]) + mask = (x >= p_lo) & (x <= p_hi) + x_mid = x[mask] + return x_mid.mean(), x_mid.std() + + +def bench_single( + func, + args, + keys, + configs, + warmup, + sleep_ms, + n_iters, + quantiles, +): + results = [] + device = get_device_name_from_args(*args) + + for config in tqdm(configs, desc="Benchmark configs", leave=False): + kwargs = dict(zip(keys, config)) + + if sleep_ms is not None: + sleep(sleep_ms / 1000) + + try: + func(*args, **kwargs) + torch.cuda.synchronize() + except Exception: + results.append({"spec": kwargs, "mean_ms": float("inf"), "std_ms": 0.0}) + continue + + for _ in range(warmup): + func(*args, **kwargs) + torch.cuda.synchronize() + + times_ms = [] + for _ in range(n_iters): + start = perf_counter() + func(*args, **kwargs) + torch.cuda.synchronize() + end = perf_counter() + times_ms.append((end - start) * 1e3) + + mean, std = mean_std_percentile(times_ms, *quantiles) + results.append({"spec": kwargs, "mean_ms": float(mean), "std_ms": float(std), "device": device}) + + return results + + +def _build_toml( + device: str, + problem_size_names: list[str], + entries: list[dict], # [{"problem_sizes": [...], **kernel_params}] +) -> str: + lines = [] + + lines.append("[meta]") + lines.append(f'problem_size_names = {problem_size_names}') + lines.append("") + + lines.append(f'["{device}"]') + for entry in entries: + lines.append(f' [["{device}".configs]]') + ps = list(entry["problem_sizes"]) + lines.append(f" problem_sizes = {ps}") + for k, v in entry["kernel_params"].items(): + val = str(v).lower() if isinstance(v, bool) else v + lines.append(f" {k} = {val}") + lines.append("") + + return "\n".join(lines) + + +def autotune( + func, + generate_inputs: Callable[[dict], tuple], + problem_sizes: list[dict], + pruning_rule: Callable[[dict, dict], bool] | None = None, + toml_path: str | Path | None = None, + n_jobs_precompile: int = 4, + warmup: int = 25, + sleep_ms: int = 100, + n_iters: int = 50, + quantiles: tuple = (20, 80), + **specset, +): + keys = list(specset.keys()) + values = list(specset.values()) + configs = list(itertools.product(*values)) + + problem_size_names = list(next(iter(problem_sizes)).keys()) + assert all(set(problem_size_names) == set(problem_size.keys()) for problem_size in problem_sizes) + + def precompile_call(*args, config): + kwargs = dict(zip(keys, config)) + try: + func(*args, **kwargs) + except Exception: + pass + + n_jobs = min(n_jobs_precompile, len(configs)) + toml_entries = [] + devices = [] + + for problem_size in tqdm(problem_sizes, desc="Problem sizes", leave=False): + args = generate_inputs(problem_size) + ps_values = list(problem_size.values()) + pruned_configs = [config for config in configs if pruning_rule is None or pruning_rule(problem_size, dict(zip(keys, config)))] + + Parallel(n_jobs=n_jobs, backend='threading')( + delayed(precompile_call)(*args, config=config) + for config in tqdm(pruned_configs, desc="Precompiling", leave=False) + ) + + results = bench_single(func, args, keys, pruned_configs, warmup, sleep_ms, n_iters, quantiles) + + valid = [r for r in results if r["mean_ms"] != float("inf")] + if not valid: + warn(f"No valid configs for problem_size={problem_size}, skipping", stacklevel=2) + continue + + best = min(valid, key=lambda r: r["mean_ms"]) + + toml_entries.append({ + "problem_sizes": ps_values, + "kernel_params": best["spec"], + }) + + devices.append(best["device"]) + + devices_found = set(device for device in devices) + assert len(devices_found) == 1, f"Expected single-device autotune, got devices {devices_found}" + + if toml_path is not None: + toml_str = _build_toml(devices_found.pop(), problem_size_names, toml_entries) + Path(toml_path).write_text(toml_str, encoding="utf-8") diff --git a/kerops/settings/kernel_config.py b/kerops/settings/kernel_config.py index 6bfeed5..d6f10f5 100644 --- a/kerops/settings/kernel_config.py +++ b/kerops/settings/kernel_config.py @@ -5,21 +5,10 @@ from inspect import signature as get_signature from warnings import warn -from torch import Tensor -from torch.cuda import get_device_name +from .utils import get_device_name_from_args class KernelConfigBase(ABC): - @property - @abstractmethod - def arg_names(self) -> list[str]: - ... - - @property - @abstractmethod - def confarg_names(self) -> list[str]: - ... - @abstractmethod def __call__(self, *input_args) -> dict[str, Any]: ... @@ -185,18 +174,8 @@ def load_toml(self, path: str | Path): f'fallback_device={self.fallback_device!r} not found in loaded configs: {list(self._configs)}' ) - def get_device_name_from_args(self, *input_args): - device_indices = {arg.device.index for arg in input_args if isinstance (arg, Tensor) and arg.device.type == 'cuda'} - - if len(device_indices) == 1: - return get_device_name(device_indices.pop()) - elif len(device_indices) == 0: - raise RuntimeError('Cannot configure due to non-cuda args') - else: - raise RuntimeError(f'Expected all tensors to be on the same GPU, got CUDA-devices:{device_indices}') - def __call__(self, *input_args) -> dict[str, Any]: - device = self.get_device_name_from_args(*input_args) + device = get_device_name_from_args(*input_args) problem_sizes = self.args_to_problem_sizes(*input_args) diff --git a/kerops/settings/utils.py b/kerops/settings/utils.py new file mode 100644 index 0000000..a23b71c --- /dev/null +++ b/kerops/settings/utils.py @@ -0,0 +1,13 @@ +from torch import Tensor +from torch.cuda import get_device_name + + +def get_device_name_from_args(*args): + device_indices = {arg.device.index for arg in args if isinstance (arg, Tensor) and arg.device.type == 'cuda'} + + if len(device_indices) == 1: + return get_device_name(device_indices.pop()) + elif len(device_indices) == 0: + raise RuntimeError('Cannot configure due to non-cuda args') + else: + raise RuntimeError(f'Expected all tensors to be on the same GPU, got CUDA-devices:{device_indices}') diff --git a/kerops/settings/wrapper.py b/kerops/settings/wrapper.py index 9e2cf98..9bb3c7b 100644 --- a/kerops/settings/wrapper.py +++ b/kerops/settings/wrapper.py @@ -2,7 +2,7 @@ from functools import wraps from typing import Callable -from .kernel_config import KernelConfig +from .kernel_config import KernelConfigBase class ConfArg: @@ -13,7 +13,7 @@ class ConfArg: class ConfiguredFunction: - def __init__(self, function: Callable, kernel_config: KernelConfig): + def __init__(self, function: Callable, kernel_config: KernelConfigBase): self.function = function signature = get_signature(function) @@ -38,7 +38,7 @@ def __init__(self, function: Callable, kernel_config: KernelConfig): self.register_kernel_config(kernel_config) - def register_kernel_config(self, kernel_config: KernelConfig): + def register_kernel_config(self, kernel_config: KernelConfigBase): configured_arg_names = kernel_config.confarg_names input_arg_names = kernel_config.arg_names @@ -84,7 +84,7 @@ def __call__(self, *args, **kwargs): return self.function(*args, **configured_kwargs) @classmethod - def configure(cls, kernel_config: KernelConfig): + def configure(cls, kernel_config: KernelConfigBase): def wrapper(function: Callable): return wraps(function)(cls(function, kernel_config)) diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb new file mode 100644 index 0000000..0f56501 --- /dev/null +++ b/notebooks/config_builder.ipynb @@ -0,0 +1,377 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "edf736df-e0c0-43c9-8e7b-c24a3f94596e", + "metadata": {}, + "outputs": [], + "source": [ + "from kerops.settings.autotune import autotune\n", + "import torch" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "96fb3e4a-cf21-4e19-8e46-74075a363611", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "4b3106d2430b458a885903bd064f0911", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Problem sizes: 0%| | 0/10 [00:00 Date: Thu, 2 Apr 2026 20:14:43 +0200 Subject: [PATCH 18/40] Conv configs --- .gitignore | 2 + kerops/ops/assets/BNReLUConv3d.toml | 63 ++++ kerops/ops/assets/Conv3d.toml | 83 +++++ kerops/ops/assets/__init__.py | 3 + kerops/ops/conv/conv.py | 195 +++++++----- kerops/settings/__init__.py | 1 + kerops/settings/autotune.py | 2 +- notebooks/config_builder.ipynb | 478 ++++++++++++++++++++++++++++ 8 files changed, 743 insertions(+), 84 deletions(-) create mode 100644 kerops/ops/assets/BNReLUConv3d.toml create mode 100644 kerops/ops/assets/Conv3d.toml create mode 100644 kerops/ops/assets/__init__.py diff --git a/.gitignore b/.gitignore index 6ec4a5f..d5faeb9 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,5 @@ pip-delete-this-directory.txt # Profiling *.ncu-rep PROF/ + +.vscode diff --git a/kerops/ops/assets/BNReLUConv3d.toml b/kerops/ops/assets/BNReLUConv3d.toml new file mode 100644 index 0000000..7d499a8 --- /dev/null +++ b/kerops/ops/assets/BNReLUConv3d.toml @@ -0,0 +1,63 @@ +[meta] +problem_size_names = ['in_channels', 'out_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 16] + num_warps = 4 + D_BLOCK = 64 + CIN_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 32] + num_warps = 2 + D_BLOCK = 32 + CIN_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 16] + num_warps = 4 + D_BLOCK = 64 + CIN_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 32] + num_warps = 2 + D_BLOCK = 32 + CIN_BLOCK = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 64] + num_warps = 2 + D_BLOCK = 32 + CIN_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 32] + num_warps = 2 + D_BLOCK = 32 + CIN_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 64] + num_warps = 2 + D_BLOCK = 32 + CIN_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 128] + num_warps = 2 + D_BLOCK = 16 + CIN_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 64] + num_warps = 4 + D_BLOCK = 16 + CIN_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 128] + num_warps = 2 + D_BLOCK = 16 + CIN_BLOCK = 16 diff --git a/kerops/ops/assets/Conv3d.toml b/kerops/ops/assets/Conv3d.toml new file mode 100644 index 0000000..0d7afcd --- /dev/null +++ b/kerops/ops/assets/Conv3d.toml @@ -0,0 +1,83 @@ +[meta] +problem_size_names = ['in_channels', 'out_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 16] + num_warps = 2 + D_BLOCK = 32 + CIN_BLOCK = 16 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 32] + num_warps = 4 + D_BLOCK = 64 + CIN_BLOCK = 16 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 16] + num_warps = 4 + D_BLOCK = 64 + CIN_BLOCK = 16 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 32] + num_warps = 2 + D_BLOCK = 32 + CIN_BLOCK = 16 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 64] + num_warps = 2 + D_BLOCK = 16 + CIN_BLOCK = 16 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 32] + num_warps = 2 + D_BLOCK = 32 + CIN_BLOCK = 16 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 64] + num_warps = 2 + D_BLOCK = 16 + CIN_BLOCK = 64 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 128] + num_warps = 2 + D_BLOCK = 16 + CIN_BLOCK = 32 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 64] + num_warps = 2 + D_BLOCK = 16 + CIN_BLOCK = 64 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 128] + num_warps = 4 + D_BLOCK = 16 + CIN_BLOCK = 16 + LOAD_WEIGHT_FIRST = true + WEIGHT_MAJOR = true diff --git a/kerops/ops/assets/__init__.py b/kerops/ops/assets/__init__.py new file mode 100644 index 0000000..7eb5302 --- /dev/null +++ b/kerops/ops/assets/__init__.py @@ -0,0 +1,3 @@ +from pathlib import Path + +ASSETS_ROOT = Path(__file__).absolute().parent diff --git a/kerops/ops/conv/conv.py b/kerops/ops/conv/conv.py index bb0e68a..882be37 100644 --- a/kerops/ops/conv/conv.py +++ b/kerops/ops/conv/conv.py @@ -2,84 +2,22 @@ import torch from triton import language as tl, next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.conv import _Conv_cl3d_impl_V6, _ApplyBNReLUConvStats_cl3d_impl -from ...settings import ConfArg, configure, confexc +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv -@confexc(KeyError) -def num_warps(in_channels, out_channels): - return { - (16, 16): 4, - (16, 32): 4, - (32, 16): 4, - (32, 32): 2, - (32, 64): 2, - (64, 32): 2, - (64, 64): 4, - (64, 128): 2, - (128, 64): 2, - (128, 128): 4, - }[(in_channels, out_channels)] - - -@confexc(KeyError) -def d_block(in_channels, out_channels): - return { - (16, 16): 64, - (16, 32): 64, - (32, 16): 64, - (32, 32): 32, - (32, 64): 32, - (64, 32): 32, - (64, 64): 32, - (64, 128): 16, - (128, 64): 16, - (128, 128): 16, - }[(in_channels, out_channels)] - - -@confexc(KeyError) -def cin_block(in_channels, out_channels): - return { - (16, 16): 16, - (16, 32): 16, - (32, 16): 16, - (32, 32): 16, - (32, 64): 16, - (64, 32): 16, - (64, 64): 32, - (64, 128): 16, - (128, 64): 16, - (128, 128): 16, - }[(in_channels, out_channels)] - - -@confexc(KeyError) -def weight_major(in_channels, out_channels): - return { - (16, 16): False, - (16, 32): True, - (32, 16): False, - (32, 32): False, - (32, 64): True, - (64, 32): False, - (64, 64): True, - (64, 128): True, - (128, 64): True, - (128, 128): True, - }[(in_channels, out_channels)] - - -@configure( - ACCTYPE='float32', - num_warps=lambda weight: num_warps(*weight.shape[-2:]), - D_BLOCK=lambda weight: d_block(*weight.shape[-2:]), - CIN_BLOCK=lambda weight: cin_block(*weight.shape[-2:]), - LOAD_WEIGHT_FIRST=True, - WEIGHT_MAJOR=lambda weight: weight_major(*weight.shape[-2:]), +conv3d_config = TableKernelConfig( + problem_size_names=['in_channels', 'out_channels'], + confarg_names=['num_warps', 'D_BLOCK', 'CIN_BLOCK', 'WEIGHT_MAJOR', 'LOAD_WEIGHT_FIRST'], + args_to_problem_sizes=lambda weight: tuple(weight.shape[-2:]), + toml_path=ASSETS_ROOT / 'Conv3d.toml' ) -def Conv3d(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_BLOCK: ConfArg, CIN_BLOCK: ConfArg, LOAD_WEIGHT_FIRST: ConfArg, WEIGHT_MAJOR: ConfArg): + + +@ConfiguredFunction.configure(conv3d_config) +def Conv3d(x, weight, *, num_warps: ConfArg, D_BLOCK: ConfArg, CIN_BLOCK: ConfArg, LOAD_WEIGHT_FIRST: ConfArg, WEIGHT_MAJOR: ConfArg): assert x.device == weight.device assert x.is_cuda @@ -100,9 +38,8 @@ def Conv3d(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_BLOCK: ConfArg, assert D_BLOCK == next_power_of_2(D_BLOCK) assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) assert CIN_BLOCK <= in_channels - assert ACCTYPE in ('float16', 'float32') - ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] + ACCTYPE = tl.float32 output = torch.empty([bsize, H, W, D, out_channels], device=x.device, dtype=x.dtype, layout=x.layout).permute(0, -1, 1, 2, 3) grid = (cdiv(W, 2), cdiv(H, 2), cdiv(D, D_BLOCK) * bsize) @@ -126,13 +63,16 @@ def Conv3d(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_BLOCK: ConfArg, return output -@configure( - ACCTYPE='float32', - num_warps=lambda weight: num_warps(*weight.shape[-2:]), - D_BLOCK=lambda weight: d_block(*weight.shape[-2:]), - CIN_BLOCK=lambda weight: cin_block(*weight.shape[-2:]), +bnreluconv3d_config = TableKernelConfig( + problem_size_names=['in_channels', 'out_channels'], + confarg_names=['D_BLOCK', 'num_warps', 'CIN_BLOCK'], + args_to_problem_sizes=lambda weight: tuple(weight.shape[-2:]), + toml_path=ASSETS_ROOT / 'BNReLUConv3d.toml' ) -def ApplyBNReLUConv3dStats(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_BLOCK: ConfArg, CIN_BLOCK: ConfArg): + + +@ConfiguredFunction.configure(bnreluconv3d_config) +def ApplyBNReLUConv3dStats(x, bn_weight, bn_bias, weight, *, num_warps: ConfArg, D_BLOCK: ConfArg, CIN_BLOCK: ConfArg): assert x.device == weight.device == bn_weight.device == bn_bias.device assert x.is_cuda @@ -155,9 +95,8 @@ def ApplyBNReLUConv3dStats(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfArg, n assert D_BLOCK == next_power_of_2(D_BLOCK) assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) assert CIN_BLOCK <= in_channels - assert ACCTYPE in ('float16', 'float32') - ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] + ACCTYPE = tl.float32 output = torch.empty([bsize, H, W, D, out_channels], device=x.device, dtype=x.dtype, layout=x.layout).permute(0, -1, 1, 2, 3) grid = (cdiv(W, 2), cdiv(H, 2), cdiv(D, D_BLOCK)) mean = torch.zeros([bsize, np.prod(grid), out_channels], device=x.device, dtype=torch.float32) @@ -186,3 +125,93 @@ def ApplyBNReLUConv3dStats(x, bn_weight, bn_bias, weight, *, ACCTYPE: ConfArg, n ) return output, mean.sum(dim=(0, 1)) / numel_no_channels, sqmean.sum(dim=(0, 1)) / numel_no_channels + + +def generate_inputs_conv(problem_sizes, device='cuda'): + in_channels, out_channels = problem_sizes['in_channels'], problem_sizes['out_channels'] + + if in_channels <= 32 and out_channels <= 32: + base = 128 + elif in_channels <= 64 and out_channels <= 64: + base = 96 + else: + base = 64 + + x = torch.randn(1, in_channels, base, base, base, device=device, dtype=torch.float16).to(memory_format=torch.channels_last_3d) + w = torch.randn(3, 3, 3, in_channels, out_channels, device=device, dtype=torch.float16) + + return x, w + + +def generate_inputs_bnreluconv(problem_sizes, device='cuda'): + x, w = generate_inputs_conv(problem_sizes, device) + + in_channels = problem_sizes['in_channels'] + bn_weight = torch.randn(in_channels, device=device, dtype=torch.float32) + bn_bias = torch.randn(in_channels, device=device, dtype=torch.float32) + + return x, bn_weight, bn_bias, w + + +def pruning_rule(problem_size, named_config): + D_BLOCK = named_config['D_BLOCK'] + CIN_BLOCK = named_config['CIN_BLOCK'] + + in_channels, out_channels = problem_size['in_channels'], problem_size['out_channels'] + + if in_channels >= 32 and out_channels >= 32 and D_BLOCK > 32: + return False + + if (in_channels >= 128 or out_channels >= 128) and D_BLOCK > 16: + return False + + if CIN_BLOCK > in_channels: + return False + + return True + + +def autotune_conv(toml_path, **autotune_kwargs): + channels = [2 ** i for i in range(4, 8)] + problem_sizes = [ + {'in_channels': cin, 'out_channels': cout} + for cin in channels + for cout in channels + if (cin == 2 * cout) or (cin * 2 == cout) or (cin == cout) + ] + + autotune( + getattr(Conv3d, 'function', Conv3d), + generate_inputs_conv, + problem_sizes, + pruning_rule, + toml_path, + **autotune_kwargs, + num_warps=[2, 4], + D_BLOCK=[16, 32, 64], + CIN_BLOCK=[16, 32, 64], + LOAD_WEIGHT_FIRST=[True, False], + WEIGHT_MAJOR=[True, False] + ) + + +def autotune_bnreluconv(toml_path, **autotune_kwargs): + channels = [2 ** i for i in range(4, 8)] + problem_sizes = [ + {'in_channels': cin, 'out_channels': cout} + for cin in channels + for cout in channels + if (cin == 2 * cout) or (cin * 2 == cout) or (cin == cout) + ] + + autotune( + getattr(ApplyBNReLUConv3dStats, 'function', ApplyBNReLUConv3dStats), + generate_inputs_bnreluconv, + problem_sizes, + pruning_rule, + toml_path, + **autotune_kwargs, + num_warps=[2, 4], + D_BLOCK=[16, 32, 64], + CIN_BLOCK=[16, 32, 64], + ) diff --git a/kerops/settings/__init__.py b/kerops/settings/__init__.py index f7e4913..f54d629 100644 --- a/kerops/settings/__init__.py +++ b/kerops/settings/__init__.py @@ -1,2 +1,3 @@ +from .autotune import autotune from .wrapper import ConfiguredFunction, ConfArg from .kernel_config import StaticKernelConfig, RuleKernelConfig, TableKernelConfig diff --git a/kerops/settings/autotune.py b/kerops/settings/autotune.py index 6efd244..1e09fde 100644 --- a/kerops/settings/autotune.py +++ b/kerops/settings/autotune.py @@ -97,7 +97,7 @@ def autotune( n_jobs_precompile: int = 4, warmup: int = 25, sleep_ms: int = 100, - n_iters: int = 50, + n_iters: int = 100, quantiles: tuple = (20, 80), **specset, ): diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb index 0f56501..a1d2c6b 100644 --- a/notebooks/config_builder.ipynb +++ b/notebooks/config_builder.ipynb @@ -11,6 +11,51 @@ "import torch" ] }, + { + "cell_type": "code", + "execution_count": 2, + "id": "c34b5a83-ee1e-4c03-b4da-3700561abcda", + "metadata": {}, + "outputs": [], + "source": [ + "def generate_inputs(problem_sizes):\n", + " in_channels, out_channels = problem_sizes['in_channels'], problem_sizes['out_channels']\n", + "\n", + " if in_channels <= 32 and out_channels <= 32:\n", + " base = 128\n", + " elif in_channels <= 64 and out_channels <= 64:\n", + " base = 96\n", + " else:\n", + " base = 64\n", + "\n", + " x = torch.randn(1, in_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", + " w = torch.randn(3, 3, 3, in_channels, out_channels, device='cuda', dtype=torch.float16)\n", + "\n", + " return x, w\n", + "\n", + "problem_size_names = ['in_channels', 'out_channels']\n", + "\n", + "channels = [2 ** i for i in range(4, 8)]\n", + "problem_sizes = [{'in_channels': cin, 'out_channels': cout} for cin in channels for cout in channels if (cin == 2 * cout) or (cin * 2 == cout) or (cin == cout)]\n", + "\n", + "def pruning_rule(problem_size, named_config):\n", + " D_BLOCK = named_config['D_BLOCK']\n", + " CIN_BLOCK = named_config['CIN_BLOCK']\n", + "\n", + " in_channels, out_channels = problem_size['in_channels'], problem_size['out_channels']\n", + "\n", + " if in_channels >= 32 and out_channels >= 32 and D_BLOCK > 32:\n", + " return False\n", + "\n", + " if (in_channels >= 128 or out_channels >= 128) and D_BLOCK > 16:\n", + " return False\n", + "\n", + " if CIN_BLOCK > in_channels:\n", + " return False\n", + "\n", + " return True" + ] + }, { "cell_type": "code", "execution_count": 4, @@ -332,6 +377,314 @@ ")" ] }, + { + "cell_type": "code", + "execution_count": 1, + "id": "7a2b9df2-d284-408f-b660-92bc5d2d2fe1", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "a2eb9a8b47d34263be6845a343711f42", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Problem sizes: 0%| | 0/10 [00:00 Date: Fri, 3 Apr 2026 01:00:04 +0200 Subject: [PATCH 19/40] Conv3dWgrad config --- kerops/ops/assets/Conv3dWgrad.toml | 93 ++++ kerops/ops/conv/conv_wgrad.py | 154 +++--- notebooks/config_builder.ipynb | 800 +---------------------------- 3 files changed, 195 insertions(+), 852 deletions(-) create mode 100644 kerops/ops/assets/Conv3dWgrad.toml diff --git a/kerops/ops/assets/Conv3dWgrad.toml b/kerops/ops/assets/Conv3dWgrad.toml new file mode 100644 index 0000000..4f0e09f --- /dev/null +++ b/kerops/ops/assets/Conv3dWgrad.toml @@ -0,0 +1,93 @@ +[meta] +problem_size_names = ['in_channels', 'out_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 16] + num_warps = 1 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 16 + COUT_BLOCK = 16 + SWAP_GRAD_WITH_INPUT = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 32] + num_warps = 1 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 16 + COUT_BLOCK = 32 + SWAP_GRAD_WITH_INPUT = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 16] + num_warps = 1 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 16 + COUT_BLOCK = 32 + SWAP_GRAD_WITH_INPUT = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 32] + num_warps = 2 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 32 + COUT_BLOCK = 32 + SWAP_GRAD_WITH_INPUT = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 64] + num_warps = 2 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 16 + COUT_BLOCK = 64 + SWAP_GRAD_WITH_INPUT = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 32] + num_warps = 2 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 16 + COUT_BLOCK = 64 + SWAP_GRAD_WITH_INPUT = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 64] + num_warps = 2 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 32 + COUT_BLOCK = 64 + SWAP_GRAD_WITH_INPUT = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 128] + num_warps = 4 + D_BLOCK = 16 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 64 + COUT_BLOCK = 64 + SWAP_GRAD_WITH_INPUT = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 64] + num_warps = 4 + D_BLOCK = 16 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 64 + COUT_BLOCK = 64 + SWAP_GRAD_WITH_INPUT = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 128] + num_warps = 4 + D_BLOCK = 16 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 64 + COUT_BLOCK = 64 + SWAP_GRAD_WITH_INPUT = false diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index a280003..804644f 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -1,77 +1,32 @@ import torch from triton import language as tl, next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.conv import _Conv_wgrad_cl3d_impl_V2 -from ...settings import ConfArg, configure, confexc +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv -@confexc(KeyError) -def num_warps(in_channels, out_channels): - # BE AWARE of swap_grad_with_input - it permutes grad with x - return { - (16, 16): 1, - (16, 32): 2, - (32, 16): 2, - (32, 32): 2, - (32, 64): 4, - (64, 32): 4, - (64, 64): 2, - (64, 128): 4, - (128, 64): 4, - (128, 128): 4, - }[(in_channels, out_channels)] - - -@confexc(KeyError) -def CIN_BLOCK(in_channels, out_channels): - # BE AWARE of swap_grad_with_input - it permutes grad with x - return { - (16, 16): 16, - (16, 32): 16, - (32, 16): 16, - (32, 32): 32, - (32, 64): 16, - (64, 32): 16, - (64, 64): 32, - (64, 128): 32, - (128, 64): 32, - (128, 128): 16, - }[(in_channels, out_channels)] - - -@confexc(KeyError) -def COUT_BLOCK(in_channels, out_channels): - # BE AWARE of swap_grad_with_input - it permutes grad with x - return { - (16, 16): 16, - (16, 32): 32, - (32, 16): 32, - (32, 32): 32, - (32, 64): 64, - (64, 32): 64, - (64, 64): 64, - (64, 128): 128, - (128, 64): 128, - (128, 128): 128, - }[(in_channels, out_channels)] - - -# TODO: looks like this trick can be fixed with a better algorithm. -def swap_grad_with_input(in_channels, out_channels): - return out_channels < in_channels - - -@configure( - ACCTYPE='float32', - num_warps=lambda grad, x: num_warps(x.shape[1], grad.shape[1]), - REDUCTION_FACTOR=32, - CIN_BLOCK=lambda grad, x: CIN_BLOCK(x.shape[1], grad.shape[1]), - COUT_BLOCK=lambda grad, x: COUT_BLOCK(x.shape[1], grad.shape[1]), - D_BLOCK=32, - SWAP_GRAD_WITH_INPUT=lambda grad, x: swap_grad_with_input(x.shape[1], grad.shape[1]) +conv3d_wgrad_config = TableKernelConfig( + problem_size_names=['in_channels', 'out_channels'], + confarg_names=['num_warps', 'D_BLOCK', 'REDUCTION_FACTOR', 'CIN_BLOCK', 'COUT_BLOCK', 'SWAP_GRAD_WITH_INPUT'], + args_to_problem_sizes=lambda grad, x: (x.shape[1], grad.shape[1]), + toml_path=ASSETS_ROOT / 'Conv3dWgrad.toml' ) -def Conv3dWgrad(grad, x, *, D_BLOCK: ConfArg, ACCTYPE: ConfArg, num_warps: ConfArg, REDUCTION_FACTOR: ConfArg, CIN_BLOCK: ConfArg, COUT_BLOCK: ConfArg, SWAP_GRAD_WITH_INPUT: ConfArg): + + +@ConfiguredFunction.configure(conv3d_wgrad_config) +def Conv3dWgrad( + grad, + x, + *, + num_warps: ConfArg, + D_BLOCK: ConfArg, + REDUCTION_FACTOR: ConfArg, + CIN_BLOCK: ConfArg, + COUT_BLOCK: ConfArg, + SWAP_GRAD_WITH_INPUT: ConfArg +): if SWAP_GRAD_WITH_INPUT: grad, x = x, grad @@ -91,14 +46,13 @@ def Conv3dWgrad(grad, x, *, D_BLOCK: ConfArg, ACCTYPE: ConfArg, num_warps: ConfA assert x.dtype == grad.dtype == torch.float16 assert D_BLOCK == next_power_of_2(D_BLOCK) - assert ACCTYPE in ('float16', 'float32') assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) assert CIN_BLOCK <= in_channels assert COUT_BLOCK == next_power_of_2(COUT_BLOCK) assert COUT_BLOCK <= out_channels assert isinstance(REDUCTION_FACTOR, int) - ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] + ACCTYPE = tl.float32 num_buffers = cdiv(xW, REDUCTION_FACTOR * 2) * cdiv(xH, REDUCTION_FACTOR * 2) * cdiv(xD, D_BLOCK * REDUCTION_FACTOR) * xbsize weight_grad = torch.zeros([num_buffers, 3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float32) grid = (cdiv(xW, 2) * cdiv(out_channels, COUT_BLOCK), cdiv(xH, 2), cdiv(xD, D_BLOCK) * xbsize) @@ -128,3 +82,67 @@ def Conv3dWgrad(grad, x, *, D_BLOCK: ConfArg, ACCTYPE: ConfArg, num_warps: ConfA weight_grad = weight_grad.contiguous() return weight_grad + + +def pruning_rule(problem_size, named_config): + D_BLOCK = named_config['D_BLOCK'] + CIN_BLOCK = named_config['CIN_BLOCK'] + COUT_BLOCK = named_config['COUT_BLOCK'] + SWAP_GRAD_WITH_INPUT = named_config['SWAP_GRAD_WITH_INPUT'] + + in_channels, out_channels = problem_size['in_channels'], problem_size['out_channels'] + + if in_channels >= 32 and out_channels >= 32 and D_BLOCK > 32: + return False + + if (in_channels >= 128 or out_channels >= 128) and D_BLOCK > 16: + return False + + if (CIN_BLOCK > in_channels and not SWAP_GRAD_WITH_INPUT) or (CIN_BLOCK > out_channels and SWAP_GRAD_WITH_INPUT): + return False + + if (COUT_BLOCK > out_channels and not SWAP_GRAD_WITH_INPUT) or (COUT_BLOCK > in_channels and SWAP_GRAD_WITH_INPUT): + return False + + return True + + +def generate_inputs_conv_wgrad(problem_sizes): + in_channels, out_channels = problem_sizes['in_channels'], problem_sizes['out_channels'] + + if in_channels <= 32 and out_channels <= 32: + base = 128 + elif in_channels <= 64 and out_channels <= 64: + base = 96 + else: + base = 64 + + x = torch.randn(1, in_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + grad = torch.randn(1, out_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + + return grad, x + + +def autotune_conv_wgrad(toml_path, **autotune_kwargs): + channels = [2 ** i for i in range(4, 8)] + problem_sizes = [ + {'in_channels': cin, 'out_channels': cout} + for cin in channels + for cout in channels + if (cin == 2 * cout) or (cin * 2 == cout) or (cin == cout) + ] + + autotune( + getattr(Conv3dWgrad, 'function', Conv3dWgrad), + generate_inputs_conv_wgrad, + problem_sizes, + pruning_rule, + toml_path, + **autotune_kwargs, + num_warps=[1, 2, 4], + D_BLOCK=[16, 32], + REDUCTION_FACTOR=[32], + CIN_BLOCK=[16, 32, 64], + COUT_BLOCK=[16, 32, 64], + SWAP_GRAD_WITH_INPUT=[False, True] + ) diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb index a1d2c6b..44d3978 100644 --- a/notebooks/config_builder.ipynb +++ b/notebooks/config_builder.ipynb @@ -13,822 +13,54 @@ }, { "cell_type": "code", - "execution_count": 2, - "id": "c34b5a83-ee1e-4c03-b4da-3700561abcda", + "execution_count": null, + "id": "6e35524d-0c68-4fd8-859c-fb581cbd96f8", "metadata": {}, "outputs": [], "source": [ - "def generate_inputs(problem_sizes):\n", - " in_channels, out_channels = problem_sizes['in_channels'], problem_sizes['out_channels']\n", - "\n", - " if in_channels <= 32 and out_channels <= 32:\n", - " base = 128\n", - " elif in_channels <= 64 and out_channels <= 64:\n", - " base = 96\n", - " else:\n", - " base = 64\n", - "\n", - " x = torch.randn(1, in_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", - " w = torch.randn(3, 3, 3, in_channels, out_channels, device='cuda', dtype=torch.float16)\n", - "\n", - " return x, w\n", - "\n", - "problem_size_names = ['in_channels', 'out_channels']\n", - "\n", - "channels = [2 ** i for i in range(4, 8)]\n", - "problem_sizes = [{'in_channels': cin, 'out_channels': cout} for cin in channels for cout in channels if (cin == 2 * cout) or (cin * 2 == cout) or (cin == cout)]\n", - "\n", - "def pruning_rule(problem_size, named_config):\n", - " D_BLOCK = named_config['D_BLOCK']\n", - " CIN_BLOCK = named_config['CIN_BLOCK']\n", - "\n", - " in_channels, out_channels = problem_size['in_channels'], problem_size['out_channels']\n", - "\n", - " if in_channels >= 32 and out_channels >= 32 and D_BLOCK > 32:\n", - " return False\n", - "\n", - " if (in_channels >= 128 or out_channels >= 128) and D_BLOCK > 16:\n", - " return False\n", - "\n", - " if CIN_BLOCK > in_channels:\n", - " return False\n", - "\n", - " return True" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "96fb3e4a-cf21-4e19-8e46-74075a363611", - "metadata": {}, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "4b3106d2430b458a885903bd064f0911", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Problem sizes: 0%| | 0/10 [00:00 Date: Fri, 3 Apr 2026 12:28:45 +0200 Subject: [PATCH 20/40] DWConv config --- kerops/kernels/dw_conv.py | 14 +-- kerops/ops/assets/DWConv.toml | 28 +++++ kerops/ops/conv/dwconv.py | 79 +++++++++---- notebooks/config_builder.ipynb | 203 ++++++++++++++++++++++++++++----- 4 files changed, 271 insertions(+), 53 deletions(-) create mode 100644 kerops/ops/assets/DWConv.toml diff --git a/kerops/kernels/dw_conv.py b/kerops/kernels/dw_conv.py index 4bcf056..a74a1ea 100644 --- a/kerops/kernels/dw_conv.py +++ b/kerops/kernels/dw_conv.py @@ -30,8 +30,8 @@ def _DWConv_cl3d_impl( offset = d_offset[:, None, None] * channels + channels_offset[None, None, :] + near_offset[None, :, None] * channels mask = d_offset[:, None, None] + near_offset[None, :, None] < D - D_block * D_cell - mask = mask and (d_offset[:, None, None] + near_offset[None, :, None] >= 0 - D_block * D_cell) - mask = mask and (near_offset[None, :, None] != 2) + mask = mask & (d_offset[:, None, None] + near_offset[None, :, None] >= 0 - D_block * D_cell) + mask = mask & (near_offset[None, :, None] != 2) weight_offset = channels_offset[None, None, :] + tl.arange(0, 4)[None, :, None] * channels weight_mask = tl.arange(0, 4)[None, :, None] != 3 @@ -64,7 +64,7 @@ def _DWConv_cl3d_impl( load_next = (2 * H_cell + i < H and 2 * H_cell + i >= 0) and (2 * W_cell + j < W and 2 * W_cell + j >= 0) tmp_input_ptr = input_ptr + (2 * H_cell + i) * H_stride + (2 * W_cell + j) * W_stride - x = tl.load(tmp_input_ptr + offset, mask=(load_all or load_next) and mask) + x = tl.load(tmp_input_ptr + offset, mask=(load_all or load_next) & mask) for k in tl.static_range(0, 16): if k == 0: @@ -127,19 +127,19 @@ def _DWConv_cl3d_impl( load_next = (2 * H_cell + i < H and 2 * H_cell + i >= 0) and (2 * W_cell + j < W and 2 * W_cell + j >= 0) tmp_input_ptr = input_ptr + (2 * H_cell + i) * H_stride + (2 * W_cell + j) * W_stride - x = tl.load(tmp_input_ptr + offset, mask=(load_all or load_next) and mask) + x = tl.load(tmp_input_ptr + offset, mask=(load_all or load_next) & mask) tmp_output_ptr = output_ptr + (2 * H_cell) * H_stride + (2 * W_cell) * W_stride tl.store(tmp_output_ptr + out_offset, h0_w0, mask=out_mask) tmp_output_ptr = output_ptr + (2 * H_cell) * H_stride + (2 * W_cell + 1) * W_stride - tl.store(tmp_output_ptr + out_offset, h0_w1, mask=out_mask and W1_store) + tl.store(tmp_output_ptr + out_offset, h0_w1, mask=out_mask & W1_store) tmp_output_ptr = output_ptr + (2 * H_cell + 1) * H_stride + (2 * W_cell) * W_stride - tl.store(tmp_output_ptr + out_offset, h1_w0, mask=out_mask and H1_store) + tl.store(tmp_output_ptr + out_offset, h1_w0, mask=out_mask & H1_store) tmp_output_ptr = output_ptr + (2 * H_cell + 1) * H_stride + (2 * W_cell + 1) * W_stride - tl.store(tmp_output_ptr + out_offset, h1_w1, mask=out_mask and (H1_store and W1_store)) + tl.store(tmp_output_ptr + out_offset, h1_w1, mask=out_mask & (H1_store and W1_store)) # TODO: single kernel for both grad_X and grad_W diff --git a/kerops/ops/assets/DWConv.toml b/kerops/ops/assets/DWConv.toml new file mode 100644 index 0000000..b359064 --- /dev/null +++ b/kerops/ops/assets/DWConv.toml @@ -0,0 +1,28 @@ +[meta] +problem_size_names = ['channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [8] + num_warps = 2 + D_BLOCK = 64 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16] + num_warps = 4 + D_BLOCK = 64 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32] + num_warps = 2 + D_BLOCK = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64] + num_warps = 4 + D_BLOCK = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128] + num_warps = 2 + D_BLOCK = 8 diff --git a/kerops/ops/conv/dwconv.py b/kerops/ops/conv/dwconv.py index 9585057..aa51926 100644 --- a/kerops/ops/conv/dwconv.py +++ b/kerops/ops/conv/dwconv.py @@ -1,27 +1,22 @@ import torch from triton import language as tl, next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.dw_conv import _DWConv_cl3d_impl -from ...settings import ConfArg, confexc, configure +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv -@confexc(KeyError) -def warps(channels): - return {8: 1, 16: 2, 32: 2, 64: 2, 128: 4}[channels] - - -@confexc(KeyError) -def dblock(channels): - return {8: 32, 16: 32, 32: 16, 64: 8, 128: 8}[channels] +dwconv_config = TableKernelConfig( + problem_size_names=['channels'], + confarg_names=['num_warps', 'D_BLOCK'], + args_to_problem_sizes=lambda weight: (weight.shape[-1], ), + toml_path=ASSETS_ROOT / 'DWConv.toml' +) -@configure( - ACCTYPE='float32', - num_warps=lambda x: warps(x.shape[1]), - D_block=lambda x: dblock(x.shape[1]), -) -def DWConv(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_block: ConfArg): +@ConfiguredFunction.configure(dwconv_config) +def DWConv(x, weight, *, num_warps: ConfArg, D_BLOCK: ConfArg): channels = x.shape[1] assert x.ndim == 5 @@ -29,10 +24,9 @@ def DWConv(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_block: ConfArg) assert x.dtype == weight.dtype == torch.float16 assert channels == next_power_of_2(channels) assert list(weight.shape) == [3, 3, 3, channels] - assert D_block == next_power_of_2(D_block) - - ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] + assert D_BLOCK == next_power_of_2(D_BLOCK) + ACCTYPE = tl.float32 bsize, _, H, W, D = x.shape batch_stride, _, H_stride, W_stride, _ = x.stride() @@ -40,7 +34,7 @@ def DWConv(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_block: ConfArg) H_grid = cdiv(H, 2) W_grid = cdiv(W, 2) - D_grid = cdiv(D, D_block) + D_grid = cdiv(D, D_BLOCK) grid = (H_grid, W_grid, D_grid) for unbatched_x, unbatched_y in zip(x, output): @@ -55,8 +49,53 @@ def DWConv(x, weight, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_block: ConfArg) W_stride, ACCTYPE, channels, - D_block, + D_BLOCK, num_warps=num_warps, ) return output + + +def generate_inputs_dwconv(problem_sizes, device='cuda'): + channels = problem_sizes['channels'] + + if channels <= 32: + base = 256 + elif channels <= 64: + base = 192 + else: + base = 128 + + x = torch.randn(1, channels, base, base, base, device=device, dtype=torch.float16).to(memory_format=torch.channels_last_3d) + w = torch.randn(3, 3, 3, channels, device=device, dtype=torch.float16) + + return x, w + + +def pruning_rule(problem_size, named_config): + D_BLOCK = named_config['D_BLOCK'] + + channels = problem_size['channels'] + + if channels >= 32 and D_BLOCK > 32: + return False + + if channels >= 128 and D_BLOCK > 16: + return False + + return True + + +def autotune_dwconv(toml_path, **autotune_kwargs): + problem_sizes = [{'channels': channels} for channels in [2 ** i for i in range(3, 8)]] + + autotune( + getattr(DWConv, 'function', DWConv), + generate_inputs_dwconv, + problem_sizes, + pruning_rule, + toml_path, + **autotune_kwargs, + num_warps=[1, 2, 4], + D_BLOCK=[8, 16, 32, 64], + ) diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb index 44d3978..43eaad6 100644 --- a/notebooks/config_builder.ipynb +++ b/notebooks/config_builder.ipynb @@ -2,65 +2,216 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, - "id": "edf736df-e0c0-43c9-8e7b-c24a3f94596e", - "metadata": {}, - "outputs": [], - "source": [ - "from kerops.settings.autotune import autotune\n", - "import torch" - ] - }, - { - "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "6e35524d-0c68-4fd8-859c-fb581cbd96f8", "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "109b9736c921497db5c3eaabfc6f7480", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Problem sizes: 0%| | 0/5 [00:00 Date: Fri, 3 Apr 2026 14:48:34 +0200 Subject: [PATCH 21/40] DWConvWGRAD config --- kerops/kernels/dw_conv.py | 12 +- kerops/ops/assets/DWConvWGRAD.toml | 33 +++++ kerops/ops/conv/dwconv_wgrad.py | 86 +++++++++---- notebooks/config_builder.ipynb | 193 ++++------------------------- 4 files changed, 120 insertions(+), 204 deletions(-) create mode 100644 kerops/ops/assets/DWConvWGRAD.toml diff --git a/kerops/kernels/dw_conv.py b/kerops/kernels/dw_conv.py index a74a1ea..f2a6865 100644 --- a/kerops/kernels/dw_conv.py +++ b/kerops/kernels/dw_conv.py @@ -176,8 +176,8 @@ def _DWConv_wgrad_cl3d_impl( offset = d_offset[None, None, :] * channels + channels_offset[None, :, None] + near_offset[:, None, None] * channels mask = d_offset[None, None, :] + near_offset[:, None, None] < D - D_block * D_cell - mask = mask and (d_offset[None, None, :] + near_offset[:, None, None] >= 0 - D_block * D_cell) - mask = mask and (near_offset[:, None, None] != 2) + mask = mask & (d_offset[None, None, :] + near_offset[:, None, None] >= 0 - D_block * D_cell) + mask = mask & (near_offset[:, None, None] != 2) grad_offset = d_offset[None, :] * channels + channels_offset[:, None] grad_mask = d_offset[None, :] < D - D_block * D_cell @@ -201,16 +201,16 @@ def _DWConv_wgrad_cl3d_impl( W1_load = 2 * W_cell + 1 < W tmp_input_ptr = input_ptr + 2 * H_cell * H_stride + 2 * W_cell * W_stride - x_h0_w0 = tl.load(tmp_input_ptr + offset, mask=mask and H0_load) + x_h0_w0 = tl.load(tmp_input_ptr + offset, mask=mask & H0_load) tmp_input_ptr = input_ptr + (2 * H_cell + 1) * H_stride + 2 * W_cell * W_stride - x_h1_w0 = tl.load(tmp_input_ptr + offset, mask=mask and H1_load) + x_h1_w0 = tl.load(tmp_input_ptr + offset, mask=mask & H1_load) tmp_input_ptr = input_ptr + 2 * H_cell * H_stride + (2 * W_cell + 1) * W_stride - x_h0_w1 = tl.load(tmp_input_ptr + offset, mask=mask and (W1_load and H0_load)) + x_h0_w1 = tl.load(tmp_input_ptr + offset, mask=mask & (W1_load & H0_load)) tmp_input_ptr = input_ptr + (2 * H_cell + 1) * H_stride + (2 * W_cell + 1) * W_stride - x_h1_w1 = tl.load(tmp_input_ptr + offset, mask=mask and (W1_load and H1_load)) + x_h1_w1 = tl.load(tmp_input_ptr + offset, mask=mask & (W1_load & H1_load)) for k in tl.static_range(0, 16): i = (k % 4) - 1 diff --git a/kerops/ops/assets/DWConvWGRAD.toml b/kerops/ops/assets/DWConvWGRAD.toml new file mode 100644 index 0000000..ed829ed --- /dev/null +++ b/kerops/ops/assets/DWConvWGRAD.toml @@ -0,0 +1,33 @@ +[meta] +problem_size_names = ['channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [8] + num_warps = 1 + D_BLOCK = 64 + ILP = 1 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16] + num_warps = 1 + D_BLOCK = 32 + ILP = 2 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32] + num_warps = 1 + D_BLOCK = 16 + ILP = 2 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64] + num_warps = 1 + D_BLOCK = 8 + ILP = 4 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128] + num_warps = 2 + D_BLOCK = 8 + ILP = 4 diff --git a/kerops/ops/conv/dwconv_wgrad.py b/kerops/ops/conv/dwconv_wgrad.py index 604aa17..d7de08a 100644 --- a/kerops/ops/conv/dwconv_wgrad.py +++ b/kerops/ops/conv/dwconv_wgrad.py @@ -1,34 +1,23 @@ import torch from triton import language as tl, next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.dw_conv import _DWConv_wgrad_cl3d_impl -from ...settings import ConfArg, confexc, configure +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv -@confexc(KeyError) -def warps(channels): - return {8: 1, 16: 1, 32: 1, 64: 1, 128: 2}[channels] - - -@confexc(KeyError) -def dblock(channels): - return {8: 32, 16: 32, 32: 16, 64: 8, 128: 8}[channels] - - -@confexc(KeyError) -def ilp(channels): - return {8: 1, 16: 1, 32: 2, 64: 3, 128: 3}[channels] +dwconv_wgrad_config = TableKernelConfig( + problem_size_names=['channels'], + confarg_names=['num_warps', 'D_BLOCK', 'ILP'], + args_to_problem_sizes=lambda x: (x.shape[1], ), + toml_path=ASSETS_ROOT / 'DWConvWGRAD.toml' +) -@configure( - ACCTYPE='float32', - num_warps=lambda x: warps(x.shape[1]), - D_block=lambda x: dblock(x.shape[1]), - ILP=lambda x: ilp(x.shape[1]), -) +@ConfiguredFunction.configure(dwconv_wgrad_config) def DWConvWGRAD( - x, grad, *, ACCTYPE: ConfArg, num_warps: ConfArg, D_block: ConfArg, ILP: ConfArg + x, grad, *, num_warps: ConfArg, D_BLOCK: ConfArg, ILP: ConfArg ): channels = x.shape[1] @@ -38,16 +27,15 @@ def DWConvWGRAD( assert grad.is_contiguous(memory_format=torch.channels_last_3d) assert x.dtype == grad.dtype == torch.float16 assert channels == next_power_of_2(channels) - assert D_block == next_power_of_2(D_block) - - ACCTYPE = {'float32': tl.float32, 'float16': tl.float16}[ACCTYPE] + assert D_BLOCK == next_power_of_2(D_BLOCK) + ACCTYPE =tl.float32 bsize, _, H, W, D = x.shape batch_stride, _, H_stride, W_stride, _ = x.stride() H_grid = cdiv(H, 2 * ILP) W_grid = cdiv(W, 2) - D_grid = cdiv(D, D_block) + D_grid = cdiv(D, D_BLOCK) grid = (H_grid, W_grid, D_grid) grad_w = torch.zeros([bsize, H_grid * W_grid * D_grid, 3, 3, 3, channels], device=x.device, dtype=torch.float16) @@ -65,7 +53,7 @@ def DWConvWGRAD( W_stride, ACCTYPE, channels, - D_block, + D_BLOCK, WD_grid, D_grid, H_grid, @@ -76,3 +64,49 @@ def DWConvWGRAD( grad_w = grad_w.sum(dim=(0, 1)) return grad_w + + +def generate_inputs_dwconv_wgrad(problem_sizes, device='cuda'): + channels = problem_sizes['channels'] + + if channels <= 32: + base = 256 + elif channels <= 64: + base = 192 + else: + base = 128 + + x = torch.randn(1, channels, base, base, base, device=device, dtype=torch.float16).to(memory_format=torch.channels_last_3d) + grad = torch.randn(1, channels, base, base, base, device=device, dtype=torch.float16).to(memory_format=torch.channels_last_3d) + + return x, grad + + +def pruning_rule(problem_size, named_config): + D_BLOCK = named_config['D_BLOCK'] + + channels = problem_size['channels'] + + if channels >= 32 and D_BLOCK > 32: + return False + + if channels >= 128 and D_BLOCK > 16: + return False + + return True + + +def autotune_dwconv_wgrad(toml_path, **autotune_kwargs): + problem_sizes = [{'channels': channels} for channels in [2 ** i for i in range(3, 8)]] + + autotune( + getattr(DWConvWGRAD, 'function', DWConvWGRAD), + generate_inputs_dwconv_wgrad, + problem_sizes, + pruning_rule, + toml_path, + **autotune_kwargs, + num_warps=[1, 2, 4], + D_BLOCK=[8, 16, 32, 64], + ILP=[1, 2, 3, 4] + ) diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb index 43eaad6..d7cc994 100644 --- a/notebooks/config_builder.ipynb +++ b/notebooks/config_builder.ipynb @@ -2,192 +2,41 @@ "cells": [ { "cell_type": "code", - "execution_count": 2, + "execution_count": null, + "id": "ead17e8f-7f50-4f5e-b0ee-6e541d3b3605", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "from kerops.ops.conv.dwconv_wgrad import DWConvWGRAD, autotune_dwconv_wgrad, generate_inputs_dwconv_wgrad\n", + "from kerops.ops.assets import ASSETS_ROOT" + ] + }, + { + "cell_type": "code", + "execution_count": null, "id": "6e35524d-0c68-4fd8-859c-fb581cbd96f8", "metadata": {}, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "109b9736c921497db5c3eaabfc6f7480", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Problem sizes: 0%| | 0/5 [00:00 Date: Fri, 3 Apr 2026 14:54:02 +0200 Subject: [PATCH 22/40] Add config --- kerops/ops/addition.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/kerops/ops/addition.py b/kerops/ops/addition.py index e6b29ba..52a1f3d 100644 --- a/kerops/ops/addition.py +++ b/kerops/ops/addition.py @@ -5,11 +5,17 @@ from triton import next_power_of_2 from ..kernels.addition import _AddStats_cl3d_backward_impl, _AddStats_cl3d_impl -from ..settings import ConfArg, configure, get_l1_cache +from ..settings import ConfArg, StaticKernelConfig, ConfiguredFunction from ..utils import cdiv -@configure(l1_cache_bytes=get_l1_cache, num_warps=8) +add_config = StaticKernelConfig( + l1_cache_bytes=65536, + num_warps=8 +) + + +@ConfiguredFunction.configure(add_config) def AddStats(x, y, inplace=False, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() @@ -50,10 +56,8 @@ def AddStats(x, y, inplace=False, *, l1_cache_bytes: ConfArg, num_warps: ConfArg return output, mean, sqmean -@configure(l1_cache_bytes=get_l1_cache, num_warps=8) -def AddStatsBackward( - add_grad, mean_grad, sqmean_grad, add_result, *, l1_cache_bytes: ConfArg, num_warps: ConfArg -): +@ConfiguredFunction.configure(add_config) +def AddStatsBackward(add_grad, mean_grad, sqmean_grad, add_result, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = add_grad.shape[1] numel = add_grad.numel() assert add_result.shape == add_grad.shape From f7263174bec26dcaf2e4dba14a8ad4ef44454773 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 3 Apr 2026 14:56:23 +0200 Subject: [PATCH 23/40] Avgpool config --- kerops/ops/avgpool.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/kerops/ops/avgpool.py b/kerops/ops/avgpool.py index 32c08ef..a54c6de 100644 --- a/kerops/ops/avgpool.py +++ b/kerops/ops/avgpool.py @@ -4,14 +4,17 @@ from triton import next_power_of_2 from ..kernels.avgpool import _AvgPoolCeilStats_cl3d_backward_impl, _AvgPoolCeilStats_cl3d_impl -from ..settings import ConfArg, configure, get_l1_cache +from ..settings import ConfArg, StaticKernelConfig, ConfiguredFunction from ..utils import cdiv -@configure( - l1_cache_bytes=get_l1_cache, - num_warps=2, +avgpool_config = StaticKernelConfig( + l1_cache_bytes=65536, + num_warps=2 ) + + +@ConfiguredFunction.configure(avgpool_config) def AvgPoolCeilStats(x, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] input_d = x.shape[-1] @@ -60,7 +63,13 @@ def AvgPoolCeilStats(x, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): return output, mean, sqmean -@configure(l1_cache_bytes=get_l1_cache, num_warps=4) +avgpool_backward_config = StaticKernelConfig( + l1_cache_bytes=65536, + num_warps=4 +) + + +@ConfiguredFunction.configure(avgpool_backward_config) def AvgPoolCeilStatsBackward( inpgrad, meangrad, From 7f6a2cfccc80f436c05b081f0a54b52d5c7dbc49 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 3 Apr 2026 14:57:21 +0200 Subject: [PATCH 24/40] BNReLU config --- kerops/ops/bnrelu.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/kerops/ops/bnrelu.py b/kerops/ops/bnrelu.py index a39d4b1..1701491 100644 --- a/kerops/ops/bnrelu.py +++ b/kerops/ops/bnrelu.py @@ -5,11 +5,17 @@ from triton import next_power_of_2 from ..kernels.bnrelu import _ApplyBNReLU_cl3d_backward_impl, _ApplyBNReLU_cl3d_impl -from ..settings import ConfArg, configure, get_l1_cache +from ..settings import ConfArg, StaticKernelConfig, ConfiguredFunction from ..utils import cdiv -@configure(l1_cache_bytes=get_l1_cache, num_warps=8) +bnrelu_config = StaticKernelConfig( + l1_cache_bytes=65536, + num_warps=8 +) + + +@ConfiguredFunction.configure(bnrelu_config) def ApplyBNReLU(x, weight, bias, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() @@ -42,7 +48,7 @@ def ApplyBNReLU(x, weight, bias, *, l1_cache_bytes: ConfArg, num_warps: ConfArg) return output -@configure(l1_cache_bytes=get_l1_cache, num_warps=8) +@ConfiguredFunction.configure(bnrelu_config) def ApplyBNReLUBackward(x, weight, bias, grad, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() From 7c8ab6c2a0928853f8524a449636956fedeed88d Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 3 Apr 2026 14:58:22 +0200 Subject: [PATCH 25/40] Quant config --- kerops/ops/quantization.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/kerops/ops/quantization.py b/kerops/ops/quantization.py index 5d87658..62fa55f 100644 --- a/kerops/ops/quantization.py +++ b/kerops/ops/quantization.py @@ -3,12 +3,18 @@ import torch from ..kernels.quantization import _DequantUint8Window_impl, _QuantUint8Window_impl -from ..settings import ConfArg, configure, get_l1_cache +from ..settings import ConfArg, StaticKernelConfig, ConfiguredFunction from ..utils import cdiv -@configure(num_warps=4, l1_cache_bytes=get_l1_cache) -def QuantUint8Window(x, window, *, num_warps: ConfArg, l1_cache_bytes: ConfArg): +quant_config = StaticKernelConfig( + l1_cache_bytes=65536, + num_warps=4 +) + + +@ConfiguredFunction.configure(quant_config) +def QuantUint8Window(x, window, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): numel = x.numel() MAX_SIZE = l1_cache_bytes // (2 * x.element_size()) BLOCK_SIZE = min(MAX_SIZE, numel) @@ -21,8 +27,8 @@ def QuantUint8Window(x, window, *, num_warps: ConfArg, l1_cache_bytes: ConfArg): return output -@configure(num_warps=4, l1_cache_bytes=get_l1_cache) -def DequantUint8Window(x, init_dtype, window, num_warps: ConfArg, l1_cache_bytes: ConfArg): +@ConfiguredFunction.configure(quant_config) +def DequantUint8Window(x, init_dtype, window, l1_cache_bytes: ConfArg, num_warps: ConfArg): numel = x.numel() output = torch.empty_like(x, dtype=init_dtype) From 3213cc3ee83311c285b4cb3e1d265f616c5fac5b Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 3 Apr 2026 14:59:42 +0200 Subject: [PATCH 26/40] Stats config --- kerops/ops/stats.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/kerops/ops/stats.py b/kerops/ops/stats.py index 232f2aa..8ccb538 100644 --- a/kerops/ops/stats.py +++ b/kerops/ops/stats.py @@ -5,11 +5,17 @@ from triton import next_power_of_2 from ..kernels.stats import _Stats_cl3d_backward_impl, _Stats_cl3d_impl -from ..settings import ConfArg, configure, get_l1_cache +from ..settings import ConfArg, StaticKernelConfig, ConfiguredFunction from ..utils import cdiv -@configure(l1_cache_bytes=get_l1_cache, num_warps=4) +stats_config = StaticKernelConfig( + l1_cache_bytes=65536, + num_warps=4 +) + + +@ConfiguredFunction.configure(stats_config) def Stats(x, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() @@ -32,7 +38,7 @@ def Stats(x, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): return mean, sqmean -@configure(l1_cache_bytes=get_l1_cache, num_warps=4) +@ConfiguredFunction.configure(stats_config) def StatsBackward(x, mean_grad, sqmean_grad, *, l1_cache_bytes: ConfArg, num_warps: ConfArg): num_channels = x.shape[1] numel = x.numel() From ea7e4c406a9b3d5311e240fd3207df3a064b38d8 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 3 Apr 2026 16:54:21 +0200 Subject: [PATCH 27/40] test fix --- tests/ops/test_conv.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/ops/test_conv.py b/tests/ops/test_conv.py index 1b4ab71..224392a 100644 --- a/tests/ops/test_conv.py +++ b/tests/ops/test_conv.py @@ -4,10 +4,11 @@ from torch import nn from torch.nn import functional as F -from kerops.ops.conv import Conv3dWgrad +from kerops.ops.conv import Conv3d, Conv3dWgrad from kerops.utils import allclose_two_stage, weight_grad_similarity +# TODO: small shapes causes recompilations for some reasons regardless of configuration def test_conv(bsize, conv_in_channels, conv_out_channels, other_1, other_2, other_3): if not (conv_in_channels == conv_out_channels or conv_in_channels * 2 == conv_out_channels or conv_in_channels == 2 * conv_out_channels): return From 236035d9b3e7fb051dc4db75741ccd2375057736 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 3 Apr 2026 18:55:11 +0200 Subject: [PATCH 28/40] LinBNReLULinAdd config --- kerops/ops/assets/LinBReLULinAdd.toml | 21 +++ .../ops/linear/linear_bias_relu_linear_add.py | 78 ++++++--- kerops/settings/autotune.py | 5 +- notebooks/config_builder.ipynb | 148 ++++++++++++++++-- 4 files changed, 217 insertions(+), 35 deletions(-) create mode 100644 kerops/ops/assets/LinBReLULinAdd.toml diff --git a/kerops/ops/assets/LinBReLULinAdd.toml b/kerops/ops/assets/LinBReLULinAdd.toml new file mode 100644 index 0000000..6d52e72 --- /dev/null +++ b/kerops/ops/assets/LinBReLULinAdd.toml @@ -0,0 +1,21 @@ +[meta] +problem_size_names = ['in_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16] + num_warps = 4 + D_BLOCK = 32 + ILP = 1 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32] + num_warps = 2 + D_BLOCK = 16 + ILP = 1 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64] + num_warps = 4 + D_BLOCK = 32 + ILP = 2 diff --git a/kerops/ops/linear/linear_bias_relu_linear_add.py b/kerops/ops/linear/linear_bias_relu_linear_add.py index 7586e61..7d048c7 100644 --- a/kerops/ops/linear/linear_bias_relu_linear_add.py +++ b/kerops/ops/linear/linear_bias_relu_linear_add.py @@ -1,28 +1,21 @@ import torch from triton import next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.linear import _LinBReLULinAdd -from ...settings import ConfArg, confexc, configure +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv -@confexc(KeyError) -def dblock(channels): - # '64: 32' removed, fix bwd func - return {16: 32, 32: 16}[channels] - - -@confexc(KeyError) -def ilp(channels): - # '64: 8' removed, fix bwd func - return {16: 1, 32: 2}[channels] +lin_bn_relu_lin_add_config = TableKernelConfig( + problem_size_names=['in_channels'], + confarg_names=['num_warps', 'D_BLOCK', 'ILP'], + args_to_problem_sizes=lambda x: (x.shape[1], ), + toml_path=ASSETS_ROOT / 'LinBReLULinAdd.toml' +) -@configure( - num_warps=4, - D_block=lambda x: dblock(x.shape[1]), - ILP=lambda x: ilp(x.shape[1]), -) +@ConfiguredFunction.configure(lin_bn_relu_lin_add_config) def LinBReLULinAdd( x, weight_up, @@ -31,7 +24,7 @@ def LinBReLULinAdd( add_other, *, num_warps: ConfArg, - D_block: ConfArg, + D_BLOCK: ConfArg, ILP: ConfArg, ): in_channels = x.shape[1] @@ -49,7 +42,7 @@ def LinBReLULinAdd( assert add_other.is_contiguous(memory_format=torch.channels_last_3d) numel_no_channels = numel // in_channels - grid_size = cdiv(numel_no_channels, D_block * ILP) + grid_size = cdiv(numel_no_channels, D_BLOCK * ILP) output = torch.empty_like(x) @@ -63,9 +56,56 @@ def LinBReLULinAdd( numel_no_channels, in_channels, hidden_channels, - D_block, + D_BLOCK, ILP, num_warps=num_warps, ) return output + + +def pruning_rule(problem_size, named_config): + D_BLOCK = named_config['D_BLOCK'] + + in_channels = problem_size['in_channels'] + + if in_channels >= 32 and D_BLOCK > 32: + return False + + return True + + +def generate_inputs_lin_bn_relu_lin_add(problem_sizes): + in_channels = problem_sizes['in_channels'] + + if in_channels <= 32: + base = 128 + elif in_channels <= 64: + base = 96 + else: + base = 64 + + x = torch.randn(1, in_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + weight_up = torch.randn(in_channels, 2 * in_channels, device='cuda', dtype=torch.float16) + weight_down = torch.randn(2 * in_channels, in_channels, device='cuda', dtype=torch.float16) + bias = torch.randn(2 * in_channels, device='cuda', dtype=torch.float16) + add_other = torch.randn(1, in_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + + return x, weight_up, weight_down, bias, add_other + + +def autotune_lin_bn_relu_lin_add(toml_path, **autotune_kwargs): + channels = [16, 32, 64] + problem_sizes = [{'in_channels': cin} for cin in channels] + + autotune( + getattr(LinBReLULinAdd, 'function', LinBReLULinAdd), + generate_inputs_lin_bn_relu_lin_add, + problem_sizes, + pruning_rule, + toml_path, + **autotune_kwargs, + num_warps=[1, 2, 4], + D_BLOCK=[16, 32, 64], + ILP=[1, 2, 3, 4] + ) diff --git a/kerops/settings/autotune.py b/kerops/settings/autotune.py index 1e09fde..2998e03 100644 --- a/kerops/settings/autotune.py +++ b/kerops/settings/autotune.py @@ -110,10 +110,7 @@ def autotune( def precompile_call(*args, config): kwargs = dict(zip(keys, config)) - try: - func(*args, **kwargs) - except Exception: - pass + func(*args, **kwargs) n_jobs = min(n_jobs_precompile, len(configs)) toml_entries = [] diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb index d7cc994..0fb57f4 100644 --- a/notebooks/config_builder.ipynb +++ b/notebooks/config_builder.ipynb @@ -2,41 +2,141 @@ "cells": [ { "cell_type": "code", - "execution_count": null, + "execution_count": 18, "id": "ead17e8f-7f50-4f5e-b0ee-6e541d3b3605", "metadata": {}, "outputs": [], "source": [ "import torch\n", - "from kerops.ops.conv.dwconv_wgrad import DWConvWGRAD, autotune_dwconv_wgrad, generate_inputs_dwconv_wgrad\n", + "from torch.nn import functional as F\n", + "from kerops.ops.linear.linear_bias_relu_linear_add import LinBReLULinAdd, autotune_lin_bn_relu_lin_add, generate_inputs_lin_bn_relu_lin_add\n", "from kerops.ops.assets import ASSETS_ROOT" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "6e35524d-0c68-4fd8-859c-fb581cbd96f8", "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "42f06acaca8c48b99bcce155288611ab", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Problem sizes: 0%| | 0/3 [00:00 Date: Fri, 3 Apr 2026 19:09:18 +0200 Subject: [PATCH 29/40] LinBReLULinBackward config --- kerops/ops/assets/LinBReLULinBackward.toml | 21 +++++ .../linear_bias_relu_linear_backward.py | 72 ++++++++++++--- notebooks/config_builder.ipynb | 87 ++++++++++++++----- 3 files changed, 146 insertions(+), 34 deletions(-) create mode 100644 kerops/ops/assets/LinBReLULinBackward.toml diff --git a/kerops/ops/assets/LinBReLULinBackward.toml b/kerops/ops/assets/LinBReLULinBackward.toml new file mode 100644 index 0000000..a4f2262 --- /dev/null +++ b/kerops/ops/assets/LinBReLULinBackward.toml @@ -0,0 +1,21 @@ +[meta] +problem_size_names = ['in_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16] + num_warps = 2 + D_BLOCK = 64 + ILP = 4 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32] + num_warps = 2 + D_BLOCK = 32 + ILP = 4 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64] + num_warps = 4 + D_BLOCK = 32 + ILP = 1 diff --git a/kerops/ops/linear/linear_bias_relu_linear_backward.py b/kerops/ops/linear/linear_bias_relu_linear_backward.py index 14022fc..9170329 100644 --- a/kerops/ops/linear/linear_bias_relu_linear_backward.py +++ b/kerops/ops/linear/linear_bias_relu_linear_backward.py @@ -1,21 +1,22 @@ import torch from triton import next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.linear import _LinBReLULinBackward -from ...settings import ConfArg, confexc, configure +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv -@confexc(KeyError) -def ilp(channels): - return {16: 8, 32: 9}[channels] +lin_bn_relu_lin_backward_config = TableKernelConfig( + problem_size_names=['in_channels'], + confarg_names=['num_warps', 'D_BLOCK', 'ILP'], + args_to_problem_sizes=lambda x: (x.shape[1], ), + toml_path=ASSETS_ROOT / 'LinBReLULinBackward.toml' +) -@configure( - num_warps=2, - D_block=32, - ILP=lambda x: ilp(x.shape[1]), -) +# TODO: channels=64 bad performance +@ConfiguredFunction.configure(lin_bn_relu_lin_backward_config) def LinBReLULinBackward( x, grad, @@ -24,7 +25,7 @@ def LinBReLULinBackward( bias, *, num_warps: ConfArg, - D_block: ConfArg, + D_BLOCK: ConfArg, ILP: ConfArg, ): in_channels = x.shape[1] @@ -42,7 +43,7 @@ def LinBReLULinBackward( assert grad.is_contiguous(memory_format=torch.channels_last_3d) numel_no_channels = numel // in_channels - grid_size = cdiv(numel_no_channels, D_block * ILP) + grid_size = cdiv(numel_no_channels, D_BLOCK * ILP) x_grad = torch.empty_like(x) weight_up_grad = torch.zeros([grid_size, in_channels, hidden_channels], dtype=torch.float16, device='cuda') @@ -62,9 +63,56 @@ def LinBReLULinBackward( numel_no_channels, in_channels, hidden_channels, - D_block, + D_BLOCK, ILP, num_warps=num_warps, ) return x_grad, weight_up_grad.sum(dim=0), weight_down_grad.sum(dim=0), bias_grad.sum(dim=0) + + +def pruning_rule(problem_size, named_config): + D_BLOCK = named_config['D_BLOCK'] + + in_channels = problem_size['in_channels'] + + if in_channels >= 32 and D_BLOCK > 32: + return False + + return True + + +def generate_inputs_lin_bn_relu_lin_backward(problem_sizes): + in_channels = problem_sizes['in_channels'] + + if in_channels <= 32: + base = 128 + elif in_channels <= 64: + base = 96 + else: + base = 64 + + x = torch.randn(1, in_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + grad = torch.randn(1, in_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + weight_up = torch.randn(in_channels, 2 * in_channels, device='cuda', dtype=torch.float16) + weight_down = torch.randn(2 * in_channels, in_channels, device='cuda', dtype=torch.float16) + bias = torch.randn(2 * in_channels, device='cuda', dtype=torch.float16) + + return x, grad, weight_up, weight_down, bias + + +def autotune_lin_bn_relu_lin_backward(toml_path, **autotune_kwargs): + channels = [16, 32, 64] + problem_sizes = [{'in_channels': cin} for cin in channels] + + autotune( + getattr(LinBReLULinBackward, 'function', LinBReLULinBackward), + generate_inputs_lin_bn_relu_lin_backward, + problem_sizes, + pruning_rule, + toml_path, + **autotune_kwargs, + num_warps=[1, 2, 4], + D_BLOCK=[16, 32, 64], + ILP=[1, 2, 3, 4] + ) diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb index 0fb57f4..6700176 100644 --- a/notebooks/config_builder.ipynb +++ b/notebooks/config_builder.ipynb @@ -2,27 +2,27 @@ "cells": [ { "cell_type": "code", - "execution_count": 18, + "execution_count": 1, "id": "ead17e8f-7f50-4f5e-b0ee-6e541d3b3605", "metadata": {}, "outputs": [], "source": [ "import torch\n", "from torch.nn import functional as F\n", - "from kerops.ops.linear.linear_bias_relu_linear_add import LinBReLULinAdd, autotune_lin_bn_relu_lin_add, generate_inputs_lin_bn_relu_lin_add\n", + "from kerops.ops.linear.linear_bias_relu_linear_backward import LinBReLULinBackward, autotune_lin_bn_relu_lin_backward, generate_inputs_lin_bn_relu_lin_backward\n", "from kerops.ops.assets import ASSETS_ROOT" ] }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 3, "id": "6e35524d-0c68-4fd8-859c-fb581cbd96f8", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "42f06acaca8c48b99bcce155288611ab", + "model_id": "63ce1f349116411cb5919edcfa954931", "version_major": 2, "version_minor": 0 }, @@ -36,7 +36,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "5dc7230045ac4f6a8aa26e79f2f31fe0", + "model_id": "04dc5f5f82fe414ebc8826a56e3e537d", "version_major": 2, "version_minor": 0 }, @@ -50,7 +50,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "594da6e14684416b80876bbe59e903e3", + "model_id": "d94099afed9e4127a9b3dfb9e4774e44", "version_major": 2, "version_minor": 0 }, @@ -64,7 +64,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "30eb0406aff849348327934ba3a0f825", + "model_id": "3c0da22266a04ca2a0a814eabe63f8e9", "version_major": 2, "version_minor": 0 }, @@ -78,7 +78,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "be5782ae668d4acd911cfffac378f435", + "model_id": "216e23dccb5c49408b8e8699519227bd", "version_major": 2, "version_minor": 0 }, @@ -92,7 +92,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "65250453a646454eab48590b68e8ba96", + "model_id": "6760a126ca1d40b2948046538ee4129b", "version_major": 2, "version_minor": 0 }, @@ -106,7 +106,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "0826f5473a59415588cbbbfd7d304fd9", + "model_id": "cb84a90eff86404a813592daf76b3b93", "version_major": 2, "version_minor": 0 }, @@ -119,24 +119,24 @@ } ], "source": [ - "autotune_lin_bn_relu_lin_add(ASSETS_ROOT / 'LinBReLULinAdd.toml', n_jobs_precompile=4)" + "autotune_lin_bn_relu_lin_backward(ASSETS_ROOT / 'LinBReLULinBackward.toml', n_jobs_precompile=4)" ] }, { "cell_type": "code", - "execution_count": 70, + "execution_count": 37, "id": "df9f647e-f1d1-4185-a017-22bde26a7577", "metadata": {}, "outputs": [], "source": [ "channels = 64\n", "\n", - "x, weight_up, weight_down, bias, add_other = generate_inputs_lin_bn_relu_lin_add({'in_channels': channels})" + "x, grad, weight_up, weight_down, bias = generate_inputs_lin_bn_relu_lin_backward({'in_channels': channels})" ] }, { "cell_type": "code", - "execution_count": 72, + "execution_count": 39, "id": "ca511514-ec6d-4a20-b54f-53b1bc4fa0c0", "metadata": {}, "outputs": [ @@ -144,19 +144,63 @@ "name": "stdout", "output_type": "stream", "text": [ - "636 μs ± 43.6 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" + "5.3 ms ± 79.6 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" ] } ], "source": [ "%%timeit -r 10 -n 10\n", - "LinBReLULinAdd(x, weight_up, weight_down, bias, add_other)\n", + "LinBReLULinBackward(x, grad, weight_up, weight_down, bias)\n", "torch.cuda.synchronize()" ] }, { "cell_type": "code", - "execution_count": 74, + "execution_count": 33, + "id": "6d41b707-c363-4d10-aaf6-86ffdd5a6ddd", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Sequential(\n", + " (0): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))\n", + " (1): ReLU(inplace=True)\n", + " (2): Conv3d(128, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)\n", + ")" + ] + }, + "execution_count": 33, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from torch import nn\n", + "\n", + "\n", + "model = nn.Sequential(\n", + " nn.Conv3d(channels, channels * 2, kernel_size=1, bias=True),\n", + " nn.ReLU(inplace=True),\n", + " nn.Conv3d(2 * channels, channels, kernel_size=1, bias=False)\n", + ").cuda()\n", + "model.half()" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "id": "c634d98b-4b39-4b8e-bd40-70ee85712601", + "metadata": {}, + "outputs": [], + "source": [ + "with torch.amp.autocast('cuda'):\n", + " output = model(x)" + ] + }, + { + "cell_type": "code", + "execution_count": 36, "id": "3376c641-96ab-48ad-be94-ba566cdb981e", "metadata": {}, "outputs": [ @@ -164,17 +208,16 @@ "name": "stdout", "output_type": "stream", "text": [ - "2.62 ms ± 89.6 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" + "3.29 ms ± 73 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" ] } ], "source": [ "%%timeit -r 10 -n 10\n", + "model.zero_grad()\n", + "\n", "with torch.amp.autocast('cuda'), torch.inference_mode():\n", - " o = F.linear(x.permute(0, 2, 3, 4, 1), weight_up.T, bias)\n", - " o = F.relu(o)\n", - " o = F.linear(o, weight_down.T, None)\n", - " o += add_other.permute(0, 2, 3, 4, 1)\n", + " output.backward(grad, retain_graph=True)\n", "torch.cuda.synchronize()" ] }, From 2b1676ec55b403f41d625bb0aac747e0c1cc7ffb Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 3 Apr 2026 21:22:39 +0200 Subject: [PATCH 30/40] LinBReLULinAdd config --- kerops/ops/assets/ReLULinAdd.toml | 27 ++++++ kerops/ops/linear/relu_linear_add.py | 61 ++++++++++--- notebooks/config_builder.ipynb | 128 +++++++++++++-------------- 3 files changed, 136 insertions(+), 80 deletions(-) create mode 100644 kerops/ops/assets/ReLULinAdd.toml diff --git a/kerops/ops/assets/ReLULinAdd.toml b/kerops/ops/assets/ReLULinAdd.toml new file mode 100644 index 0000000..03e2594 --- /dev/null +++ b/kerops/ops/assets/ReLULinAdd.toml @@ -0,0 +1,27 @@ +[meta] +problem_size_names = ['in_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16] + num_warps = 4 + D_BLOCK = 64 + ILP = 1 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32] + num_warps = 2 + D_BLOCK = 16 + ILP = 1 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64] + num_warps = 4 + D_BLOCK = 16 + ILP = 1 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128] + num_warps = 4 + D_BLOCK = 16 + ILP = 1 diff --git a/kerops/ops/linear/relu_linear_add.py b/kerops/ops/linear/relu_linear_add.py index 63c99c8..4bdb958 100644 --- a/kerops/ops/linear/relu_linear_add.py +++ b/kerops/ops/linear/relu_linear_add.py @@ -1,28 +1,28 @@ import torch from triton import next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.linear import _ReLULinearAdd -from ...settings import ConfArg, confexc, configure +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv -@confexc(KeyError) -def ilp(in_channels): - return {16: 8, 32: 8, 64: 4, 128: 4}[in_channels] +relu_lin_add_config = TableKernelConfig( + problem_size_names=['in_channels'], + confarg_names=['num_warps', 'D_BLOCK', 'ILP'], + args_to_problem_sizes=lambda x: (x.shape[1], ), + toml_path=ASSETS_ROOT / 'ReLULinAdd.toml' +) -@configure( - num_warps=4, - D_block=16, - ILP=lambda x: ilp(x.shape[1]), -) +@ConfiguredFunction.configure(relu_lin_add_config) def ReLULinearAdd( x, weight, add_other, *, num_warps: ConfArg, - D_block: ConfArg, + D_BLOCK: ConfArg, ILP: ConfArg, ): in_channels = x.shape[1] @@ -34,18 +34,53 @@ def ReLULinearAdd( assert x.shape[0] == add_other.shape[0] assert in_channels == next_power_of_2(in_channels) assert out_channels == next_power_of_2(out_channels) - assert list(weight.shape) == [in_channels, out_channels] + assert in_channels == out_channels * 2 + assert list(weight.shape) == [in_channels, out_channels], ([in_channels, out_channels], weight.shape) assert x.dtype == weight.dtype == add_other.dtype == torch.float16 assert x.is_contiguous(memory_format=torch.channels_last_3d) assert add_other.is_contiguous(memory_format=torch.channels_last_3d) numel_no_channels = numel // in_channels - grid_size = cdiv(numel_no_channels, D_block * ILP) + grid_size = cdiv(numel_no_channels, D_BLOCK * ILP) output = torch.empty_like(add_other) _ReLULinearAdd[(grid_size,)]( - x, weight, add_other, output, numel_no_channels, in_channels, out_channels, D_block, ILP, num_warps=num_warps + x, weight, add_other, output, numel_no_channels, in_channels, out_channels, D_BLOCK, ILP, num_warps=num_warps ) return output + + +def generate_inputs_relu_lin_add(problem_sizes): + in_channels = problem_sizes['in_channels'] + + if in_channels <= 32: + base = 128 + elif in_channels <= 64: + base = 96 + else: + base = 64 + + x = torch.randn(1, in_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + weight = torch.randn(in_channels, in_channels // 2, device='cuda', dtype=torch.float16) + add_other = torch.randn(1, in_channels // 2, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + + return x, weight, add_other + + +def autotune_relu_lin_add(toml_path, **autotune_kwargs): + channels = [16, 32, 64, 128] + problem_sizes = [{'in_channels': cin} for cin in channels] + + autotune( + getattr(ReLULinearAdd, 'function', ReLULinearAdd), + generate_inputs_relu_lin_add, + problem_sizes, + pruning_rule=None, + toml_path=toml_path, + **autotune_kwargs, + num_warps=[1, 2, 4], + D_BLOCK=[16, 32, 64], + ILP=[1, 2, 3, 4, 8, 16] + ) diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb index 6700176..d6121c7 100644 --- a/notebooks/config_builder.ipynb +++ b/notebooks/config_builder.ipynb @@ -9,25 +9,25 @@ "source": [ "import torch\n", "from torch.nn import functional as F\n", - "from kerops.ops.linear.linear_bias_relu_linear_backward import LinBReLULinBackward, autotune_lin_bn_relu_lin_backward, generate_inputs_lin_bn_relu_lin_backward\n", + "from kerops.ops.linear.relu_linear_add import ReLULinearAdd, autotune_relu_lin_add, generate_inputs_relu_lin_add\n", "from kerops.ops.assets import ASSETS_ROOT" ] }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 2, "id": "6e35524d-0c68-4fd8-859c-fb581cbd96f8", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "63ce1f349116411cb5919edcfa954931", + "model_id": "8f59215d899f4360927bb80e4b3cd56a", "version_major": 2, "version_minor": 0 }, "text/plain": [ - "Problem sizes: 0%| | 0/3 [00:00 Date: Fri, 3 Apr 2026 21:52:34 +0200 Subject: [PATCH 31/40] ReLULinearBackward config --- kerops/kernels/linear.py | 3 +- kerops/ops/assets/ReLULinBackward.toml | 21 ++++++ kerops/ops/linear/relu_linear_backward.py | 68 ++++++++++++----- notebooks/config_builder.ipynb | 92 +++++++++-------------- 4 files changed, 106 insertions(+), 78 deletions(-) create mode 100644 kerops/ops/assets/ReLULinBackward.toml diff --git a/kerops/kernels/linear.py b/kerops/kernels/linear.py index bbb9245..1ff09cf 100644 --- a/kerops/kernels/linear.py +++ b/kerops/kernels/linear.py @@ -45,8 +45,9 @@ def _ReLULinearAdd( add_ptr += out_channels * D_block +# TODO: suboptimal @triton.jit -def _ReLULinearAddBackward( +def _ReLULinearBackward( input_ptr, grad_ptr, input_grad_ptr, diff --git a/kerops/ops/assets/ReLULinBackward.toml b/kerops/ops/assets/ReLULinBackward.toml new file mode 100644 index 0000000..08c0f43 --- /dev/null +++ b/kerops/ops/assets/ReLULinBackward.toml @@ -0,0 +1,21 @@ +[meta] +problem_size_names = ['in_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32] + num_warps = 2 + D_BLOCK = 64 + ILP = 8 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64] + num_warps = 4 + D_BLOCK = 64 + ILP = 8 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128] + num_warps = 8 + D_BLOCK = 64 + ILP = 16 diff --git a/kerops/ops/linear/relu_linear_backward.py b/kerops/ops/linear/relu_linear_backward.py index 7ea907d..8c1de2e 100644 --- a/kerops/ops/linear/relu_linear_backward.py +++ b/kerops/ops/linear/relu_linear_backward.py @@ -1,33 +1,28 @@ import torch from triton import next_power_of_2 -from ...kernels.linear import _ReLULinearAddBackward -from ...settings import ConfArg, confexc, configure +from ..assets import ASSETS_ROOT +from ...kernels.linear import _ReLULinearBackward +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv -@confexc(KeyError) -def warps(in_channels): - return {16: 4, 32: 8, 64: 8, 128: 8}[in_channels] - - -@confexc(KeyError) -def dblock(in_channels): - return {16: 16, 32: 32, 64: 32, 128: 32}[in_channels] +relu_lin_backward_config = TableKernelConfig( + problem_size_names=['in_channels'], + confarg_names=['num_warps', 'D_BLOCK', 'ILP'], + args_to_problem_sizes=lambda x: (x.shape[1], ), + toml_path=ASSETS_ROOT / 'ReLULinBackward.toml' +) -@configure( - num_warps=lambda x: warps(x.shape[1]), - D_block=lambda x: dblock(x.shape[1]), - ILP=16, -) +@ConfiguredFunction.configure(relu_lin_backward_config) def ReLULinearBackward( x, grad, weight, *, num_warps: ConfArg, - D_block: ConfArg, + D_BLOCK: ConfArg, ILP: ConfArg, ): in_channels = x.shape[1] @@ -39,6 +34,7 @@ def ReLULinearBackward( assert grad.shape[0] == x.shape[0] assert in_channels == next_power_of_2(in_channels) assert out_channels == next_power_of_2(out_channels) + assert in_channels == 2 * out_channels assert list(weight.shape) == [in_channels, out_channels] assert x.dtype == grad.dtype == weight.dtype == torch.float16 assert x.is_contiguous(memory_format=torch.channels_last_3d) @@ -46,13 +42,13 @@ def ReLULinearBackward( numel_no_channels = numel // out_channels - grid_size = cdiv(numel_no_channels, D_block * ILP) + grid_size = cdiv(numel_no_channels, D_BLOCK * ILP) bsize, _, H, W, D = grad.shape x_grad = torch.empty_like(x) weight_grad = torch.zeros([grid_size, in_channels, out_channels], dtype=torch.float16, device='cuda') - _ReLULinearAddBackward[(grid_size,)]( + _ReLULinearBackward[(grid_size,)]( x, grad, x_grad, @@ -61,9 +57,43 @@ def ReLULinearBackward( numel_no_channels, in_channels, out_channels, - D_block, + D_BLOCK, ILP, num_warps=num_warps, ) return x_grad, weight_grad.sum(dim=0) + + +def generate_inputs_relu_lin_backward(problem_sizes): + in_channels = problem_sizes['in_channels'] + + if in_channels <= 32: + base = 128 + elif in_channels <= 64: + base = 96 + else: + base = 64 + + x = torch.randn(1, in_channels, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + grad = torch.randn(1, in_channels // 2, base, base, base, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d) + weight = torch.randn(in_channels, in_channels // 2, device='cuda', dtype=torch.float16) + + return x, grad, weight + + +def autotune_relu_lin_backward(toml_path, **autotune_kwargs): + channels = [32, 64, 128] + problem_sizes = [{'in_channels': cin} for cin in channels] + + autotune( + getattr(ReLULinearBackward, 'function', ReLULinearBackward), + generate_inputs_relu_lin_backward, + problem_sizes, + pruning_rule=None, + toml_path=toml_path, + **autotune_kwargs, + num_warps=[1, 2, 4, 8], + D_BLOCK=[16, 32, 64], + ILP=[1, 2, 3, 4, 8, 16] + ) diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb index d6121c7..ee62027 100644 --- a/notebooks/config_builder.ipynb +++ b/notebooks/config_builder.ipynb @@ -2,14 +2,15 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": 20, "id": "ead17e8f-7f50-4f5e-b0ee-6e541d3b3605", "metadata": {}, "outputs": [], "source": [ "import torch\n", + "from torch import nn\n", "from torch.nn import functional as F\n", - "from kerops.ops.linear.relu_linear_add import ReLULinearAdd, autotune_relu_lin_add, generate_inputs_relu_lin_add\n", + "from kerops.ops.linear.relu_linear_backward import ReLULinearBackward, autotune_relu_lin_backward, generate_inputs_relu_lin_backward\n", "from kerops.ops.assets import ASSETS_ROOT" ] }, @@ -22,12 +23,12 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "8f59215d899f4360927bb80e4b3cd56a", + "model_id": "869514dd7ec64fb2aa9df82d9adef393", "version_major": 2, "version_minor": 0 }, "text/plain": [ - "Problem sizes: 0%| | 0/4 [00:00 Date: Mon, 6 Apr 2026 19:51:00 +0200 Subject: [PATCH 32/40] small fix --- kerops/ops/assets/Conv3dWgrad.toml | 6 +++--- kerops/ops/conv/conv_wgrad.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/kerops/ops/assets/Conv3dWgrad.toml b/kerops/ops/assets/Conv3dWgrad.toml index 4f0e09f..5c321be 100644 --- a/kerops/ops/assets/Conv3dWgrad.toml +++ b/kerops/ops/assets/Conv3dWgrad.toml @@ -71,7 +71,7 @@ problem_size_names = ['in_channels', 'out_channels'] D_BLOCK = 16 REDUCTION_FACTOR = 32 CIN_BLOCK = 64 - COUT_BLOCK = 64 + COUT_BLOCK = 128 SWAP_GRAD_WITH_INPUT = false [["NVIDIA GeForce RTX 5070".configs]] @@ -80,7 +80,7 @@ problem_size_names = ['in_channels', 'out_channels'] D_BLOCK = 16 REDUCTION_FACTOR = 32 CIN_BLOCK = 64 - COUT_BLOCK = 64 + COUT_BLOCK = 128 SWAP_GRAD_WITH_INPUT = true [["NVIDIA GeForce RTX 5070".configs]] @@ -89,5 +89,5 @@ problem_size_names = ['in_channels', 'out_channels'] D_BLOCK = 16 REDUCTION_FACTOR = 32 CIN_BLOCK = 64 - COUT_BLOCK = 64 + COUT_BLOCK = 128 SWAP_GRAD_WITH_INPUT = false diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index 804644f..56c6447 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -142,7 +142,7 @@ def autotune_conv_wgrad(toml_path, **autotune_kwargs): num_warps=[1, 2, 4], D_BLOCK=[16, 32], REDUCTION_FACTOR=[32], - CIN_BLOCK=[16, 32, 64], - COUT_BLOCK=[16, 32, 64], + CIN_BLOCK=[2 ** i for i in range(4, 8)], + COUT_BLOCK=[2 ** i for i in range(4, 8)], SWAP_GRAD_WITH_INPUT=[False, True] ) From 52780b9fb93cad966a6d7d7c874cb45112379688 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Tue, 7 Apr 2026 13:43:15 +0200 Subject: [PATCH 33/40] add comparator to autotuner --- kerops/ops/conv/conv_wgrad.py | 25 +- kerops/settings/autotune.py | 43 ++- notebooks/Conv3dWgrad.ipynb | 488 ++++++++++++++++++++++++++++++++++ 3 files changed, 551 insertions(+), 5 deletions(-) create mode 100644 notebooks/Conv3dWgrad.ipynb diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index 56c6447..c3fae85 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -123,6 +123,26 @@ def generate_inputs_conv_wgrad(problem_sizes): return grad, x +def comparator(grad, x): + cout = grad.shape[1] + cin = x.shape[1] + dummy_weight = torch.empty(cout, cin, 3, 3, 3, device='cuda', dtype=torch.float16) + + torch.ops.aten.convolution_backward( + grad, + x, + dummy_weight, + [0], # bias_sizes + [1, 1, 1], # stride + [1, 1, 1], # padding + [1, 1, 1], # dilation + False, # transposed + [0, 0, 0], # output padding + 1, # groups + [False, True, False], # output_mask - grad_inpt, grad_weight, grad_bias + ) + + def autotune_conv_wgrad(toml_path, **autotune_kwargs): channels = [2 ** i for i in range(4, 8)] problem_sizes = [ @@ -138,11 +158,12 @@ def autotune_conv_wgrad(toml_path, **autotune_kwargs): problem_sizes, pruning_rule, toml_path, + comparator=comparator, **autotune_kwargs, num_warps=[1, 2, 4], D_BLOCK=[16, 32], REDUCTION_FACTOR=[32], - CIN_BLOCK=[2 ** i for i in range(4, 8)], - COUT_BLOCK=[2 ** i for i in range(4, 8)], + CIN_BLOCK=channels, + COUT_BLOCK=channels, SWAP_GRAD_WITH_INPUT=[False, True] ) diff --git a/kerops/settings/autotune.py b/kerops/settings/autotune.py index 2998e03..216cd74 100644 --- a/kerops/settings/autotune.py +++ b/kerops/settings/autotune.py @@ -1,4 +1,5 @@ import itertools +import traceback as tb import torch from pathlib import Path from typing import Callable @@ -64,6 +65,27 @@ def bench_single( return results +def compare(best_entry, comparator, args, n_iters, quantiles): + best_ms = best_entry['mean_ms'] + + start = perf_counter() + end = perf_counter() + + times_ms = [] + for _ in range(n_iters): + start = perf_counter() + comparator(*args) + torch.cuda.synchronize() + end = perf_counter() + times_ms.append((end - start) * 1e3) + + mean_comparator, _ = mean_std_percentile(times_ms, *quantiles) + + ratio = mean_comparator / best_ms + + return ratio + + def _build_toml( device: str, problem_size_names: list[str], @@ -99,6 +121,7 @@ def autotune( sleep_ms: int = 100, n_iters: int = 100, quantiles: tuple = (20, 80), + comparator: Callable | None = None, **specset, ): keys = list(specset.keys()) @@ -109,8 +132,13 @@ def autotune( assert all(set(problem_size_names) == set(problem_size.keys()) for problem_size in problem_sizes) def precompile_call(*args, config): - kwargs = dict(zip(keys, config)) - func(*args, **kwargs) + try: + kwargs = dict(zip(keys, config)) + func(*args, **kwargs) + except Exception as e: + return ''.join(tb.format_exception(e)) + + return None n_jobs = min(n_jobs_precompile, len(configs)) toml_entries = [] @@ -121,11 +149,16 @@ def precompile_call(*args, config): ps_values = list(problem_size.values()) pruned_configs = [config for config in configs if pruning_rule is None or pruning_rule(problem_size, dict(zip(keys, config)))] - Parallel(n_jobs=n_jobs, backend='threading')( + precompile_statuses = Parallel(n_jobs=n_jobs, backend='threading')( delayed(precompile_call)(*args, config=config) for config in tqdm(pruned_configs, desc="Precompiling", leave=False) ) + if not any(status is None for status in precompile_statuses): + example_tb = next(status for status in precompile_statuses if status is not None) + + raise RuntimeError(f'Precompilation failed - all configs cause exception.\nExample:\n{example_tb}') + results = bench_single(func, args, keys, pruned_configs, warmup, sleep_ms, n_iters, quantiles) valid = [r for r in results if r["mean_ms"] != float("inf")] @@ -135,6 +168,10 @@ def precompile_call(*args, config): best = min(valid, key=lambda r: r["mean_ms"]) + if comparator is not None: + ratio = compare(best, comparator, args, n_iters, quantiles) + print(f'{problem_size=} best ratio - {ratio:.3f} (bigger is better)') + toml_entries.append({ "problem_sizes": ps_values, "kernel_params": best["spec"], diff --git a/notebooks/Conv3dWgrad.ipynb b/notebooks/Conv3dWgrad.ipynb new file mode 100644 index 0000000..f2533ba --- /dev/null +++ b/notebooks/Conv3dWgrad.ipynb @@ -0,0 +1,488 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 3, + "id": "4d9e0870-6d1a-484c-ab0e-84ca2a72f9e8", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "from kerops.ops.conv.conv_wgrad import Conv3dWgrad, conv3d_wgrad_config, autotune_conv_wgrad\n", + "\n", + "Conv3dWgrad.register_kernel_config(conv3d_wgrad_config)\n", + "\n", + "cin = 128\n", + "cout = 128\n", + "S = D = 64\n", + "x = torch.randn(1, cin, S, S, D, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", + "grad = torch.randn(1, cout, S, S, D, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "339e0403-64ef-46e1-87d7-c3c164a79ec6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "5.03 ms ± 54.2 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" + ] + } + ], + "source": [ + "%%timeit -r 10 -n 10\n", + "Conv3dWgrad(grad, x)\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "79789d16-7bb7-44ae-bf48-b65c711ee62a", + "metadata": {}, + "outputs": [], + "source": [ + "weight = torch.randn(cout, cin, 3, 3, 3, device='cuda', dtype=torch.float16)" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "04dc085a-9952-4d19-a2f7-b8a298707ec4", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3.71 ms ± 67.9 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" + ] + } + ], + "source": [ + "%%timeit -r 10 -n 10\n", + "torch.ops.aten.convolution_backward(\n", + " grad,\n", + " x,\n", + " weight,\n", + " [0], # bias_sizes\n", + " [1, 1, 1], # stride\n", + " [1, 1, 1], # padding\n", + " [1, 1, 1], # dilation\n", + " False, # transposed\n", + " [0, 0, 0], # output padding\n", + " 1, # groups!\n", + " [False, True, False], # output_mask - grad_inpt, grad_weight, grad_bias\n", + ")\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "ff720fb6-fb1f-4786-ba90-bdd268c4ed46", + "metadata": {}, + "outputs": [], + "source": [ + "import triton\n", + "from triton import language as tl, next_power_of_2\n", + "\n", + "\n", + "@triton.jit\n", + "def make_offset(h_str, w_str, d_str, H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr):\n", + " d_off = tl.arange(0, D_BLOCK)\n", + " w_off = tl.arange(0, W_BLOCK)\n", + " h_off = tl.arange(0, H_BLOCK)\n", + "\n", + " offset = h_off[:, None, None] * h_str + w_off[None, :, None] * w_str + d_off[None, None, :] * d_str\n", + " offset = offset.reshape((H_BLOCK * W_BLOCK * D_BLOCK))\n", + "\n", + " return offset\n", + "\n", + "\n", + "@triton.jit\n", + "def make_mask(curr_h, curr_w, curr_d, H, W, D, H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr):\n", + " mask_d = ((tl.arange(0, D_BLOCK) + curr_d) >= 0) & ((tl.arange(0, D_BLOCK) + curr_d) < D)\n", + " mask_w = ((tl.arange(0, W_BLOCK) + curr_w) >= 0) & ((tl.arange(0, W_BLOCK) + curr_w) < W)\n", + " mask_h = ((tl.arange(0, H_BLOCK) + curr_h) >= 0) & ((tl.arange(0, H_BLOCK) + curr_h) < H)\n", + "\n", + " mask = mask_h[:, None, None] & mask_w[None, :, None] & mask_d[None, None, :]\n", + " mask = mask.reshape((H_BLOCK * W_BLOCK * D_BLOCK))\n", + "\n", + " return mask\n", + "\n", + "\n", + "@triton.jit\n", + "def _Conv_wgrad_cl3d_impl(\n", + " grad_ptr,\n", + " input_ptr,\n", + " weight_grad_ptr,\n", + " H,\n", + " W,\n", + " D,\n", + " ACCTYPE: tl.constexpr,\n", + " H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr,\n", + " IN_CHANNELS: tl.constexpr, OUT_CHANNELS: tl.constexpr,\n", + " CIN_BLOCK: tl.constexpr, COUT_BLOCK: tl.constexpr,\n", + " SPLIT_K: tl.constexpr,\n", + "):\n", + " khwd_pid = tl.program_id(0)\n", + " cin_pid = tl.program_id(1)\n", + " cout_pid = tl.program_id(2)\n", + "\n", + " k_pid = khwd_pid // 27\n", + " hwd_pid = khwd_pid % 27\n", + "\n", + " block_d = hwd_pid % 3\n", + " hwd_pid = hwd_pid // 3\n", + " block_w = hwd_pid % 3\n", + " block_h = hwd_pid // 3\n", + "\n", + " cin_offset = tl.arange(0, CIN_BLOCK)\n", + " cout_offset = tl.arange(0, COUT_BLOCK)\n", + " geom_offset = make_offset(W * D, D, 1, H_BLOCK, W_BLOCK, D_BLOCK)\n", + "\n", + " grad_offset = geom_offset[:, None] * OUT_CHANNELS + cout_offset[None, :]\n", + " input_offset = geom_offset[None, :] * IN_CHANNELS + cin_offset[:, None]\n", + " weight_grad_offset = cin_offset[:, None] * OUT_CHANNELS + cout_offset[None, :]\n", + "\n", + " grad_ptr += cout_pid * COUT_BLOCK\n", + " grad_ptr += k_pid * H_BLOCK * OUT_CHANNELS * D * W\n", + " input_ptr += cin_pid * CIN_BLOCK\n", + " input_ptr += (block_d - 1) * IN_CHANNELS\n", + " input_ptr += k_pid * H_BLOCK * IN_CHANNELS * D * W\n", + " input_ptr += (block_w - 1) * IN_CHANNELS * D\n", + " input_ptr += (block_h - 1) * IN_CHANNELS * D * W\n", + " weight_grad_ptr += cin_pid * CIN_BLOCK * OUT_CHANNELS + cout_pid * COUT_BLOCK\n", + " weight_grad_ptr += block_d * IN_CHANNELS * OUT_CHANNELS\n", + " weight_grad_ptr += block_w * IN_CHANNELS * OUT_CHANNELS * 3\n", + " weight_grad_ptr += block_h * IN_CHANNELS * OUT_CHANNELS * 9\n", + "\n", + " weight_grad = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=ACCTYPE)\n", + "\n", + " for grad_h in range(0, tl.cdiv(H, H_BLOCK * SPLIT_K)):\n", + " for grad_w in range(0, tl.cdiv(W, W_BLOCK)):\n", + " for grad_d in range(0, tl.cdiv(D, D_BLOCK)):\n", + " grad_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK, grad_w * W_BLOCK, grad_d * D_BLOCK, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", + " grad_iter_ptr = (\n", + " grad_ptr\n", + " + grad_h * H_BLOCK * W * D * OUT_CHANNELS * SPLIT_K\n", + " + grad_w * W_BLOCK * D * OUT_CHANNELS\n", + " + grad_d * D_BLOCK * OUT_CHANNELS\n", + " )\n", + " grad = tl.load(grad_iter_ptr + grad_offset, mask=grad_mask[:, None], other=0)\n", + "\n", + " x_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK + block_d - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", + " x_iter_ptr = (\n", + " input_ptr\n", + " + grad_h * H_BLOCK * W * D * IN_CHANNELS * SPLIT_K\n", + " + grad_w * W_BLOCK * D * IN_CHANNELS\n", + " + grad_d * D_BLOCK * IN_CHANNELS\n", + " )\n", + " x = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0)\n", + "\n", + " weight_grad += tl.dot(x, grad)\n", + "\n", + " tl.atomic_add(weight_grad_ptr + weight_grad_offset, weight_grad, sem='relaxed')" + ] + }, + { + "cell_type": "code", + "execution_count": 63, + "id": "a4f627e6-f633-41aa-9dbd-910e0a126e4d", + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "@triton.jit\n", + "def _Conv_wgrad_cl3d_impl_prefetch(\n", + " grad_ptr,\n", + " input_ptr,\n", + " weight_grad_ptr,\n", + " H,\n", + " W,\n", + " D,\n", + " ACCTYPE: tl.constexpr,\n", + " H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr,\n", + " IN_CHANNELS: tl.constexpr, OUT_CHANNELS: tl.constexpr,\n", + " CIN_BLOCK: tl.constexpr, COUT_BLOCK: tl.constexpr,\n", + "):\n", + " hwd_pid = tl.program_id(0)\n", + " cin_pid = tl.program_id(1)\n", + " cout_pid = tl.program_id(2)\n", + "\n", + " #hwd_pid = tl.program_id(0)\n", + " #cin_cout_pid = tl.program_id(1)\n", + "\n", + " block_d = hwd_pid % 3\n", + " hwd_pid = hwd_pid // 3\n", + " block_w = hwd_pid % 3\n", + " block_h = hwd_pid // 3\n", + "\n", + " #cin_pid = cin_cout_pid % tl.cdiv(IN_CHANNELS, CIN_BLOCK)\n", + " #cout_pid = cin_cout_pid // tl.cdiv(IN_CHANNELS, CIN_BLOCK)\n", + "\n", + " cin_offset = tl.arange(0, CIN_BLOCK)\n", + " cout_offset = tl.arange(0, COUT_BLOCK)\n", + " geom_offset = make_offset(W * D, D, 1, H_BLOCK, W_BLOCK, D_BLOCK)\n", + "\n", + " grad_offset = geom_offset[:, None] * OUT_CHANNELS + cout_offset[None, :]\n", + " input_offset = geom_offset[None, :] * IN_CHANNELS + cin_offset[:, None]\n", + " weight_grad_offset = cin_offset[:, None] * OUT_CHANNELS + cout_offset[None, :]\n", + "\n", + " grad_ptr += cout_pid * COUT_BLOCK\n", + " input_ptr += cin_pid * CIN_BLOCK\n", + " input_ptr += (block_d - 1) * IN_CHANNELS\n", + " input_ptr += (block_w - 1) * IN_CHANNELS * D\n", + " input_ptr += (block_h - 1) * IN_CHANNELS * D * W\n", + " weight_grad_ptr += cin_pid * CIN_BLOCK * OUT_CHANNELS + cout_pid * COUT_BLOCK\n", + " weight_grad_ptr += block_d * IN_CHANNELS * OUT_CHANNELS\n", + " weight_grad_ptr += block_w * IN_CHANNELS * OUT_CHANNELS * 3\n", + " weight_grad_ptr += block_h * IN_CHANNELS * OUT_CHANNELS * 9\n", + "\n", + " weight_grad = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=tl.float32)\n", + " \n", + " grad_mask = make_mask(0, 0, 0, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", + " grad_next = tl.load(grad_ptr + grad_offset, mask=grad_mask[:, None], other=0)\n", + "\n", + " x_mask = make_mask(block_h - 1, block_w - 1, block_d - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", + " x_next = tl.load(input_ptr + input_offset, mask=x_mask[None, :], other=0)\n", + "\n", + " h_bound = tl.cdiv(H, H_BLOCK)\n", + " w_bound = tl.cdiv(W, W_BLOCK)\n", + " d_bound = tl.cdiv(D, D_BLOCK)\n", + "\n", + " for grad_h in range(0, h_bound):\n", + " for grad_w in range(0, w_bound):\n", + " for grad_d in range(0, d_bound):\n", + " grad = grad_next\n", + " x = x_next\n", + "\n", + " if ((grad_h + 1) != h_bound) and ((grad_w + 1) != w_bound) and ((grad_d + 1) != d_bound):\n", + " grad_mask = make_mask(grad_h * H_BLOCK, grad_w * W_BLOCK, grad_d * D_BLOCK, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", + " grad_iter_ptr = (\n", + " grad_ptr\n", + " + grad_h * H_BLOCK * W * D * OUT_CHANNELS\n", + " + grad_w * W_BLOCK * D * OUT_CHANNELS\n", + " + grad_d * D_BLOCK * OUT_CHANNELS\n", + " )\n", + " grad_next = tl.load(grad_iter_ptr + grad_offset, mask=grad_mask[:, None], other=0)\n", + " \n", + " x_mask = make_mask(grad_h * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK + block_d - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", + " x_iter_ptr = (\n", + " input_ptr\n", + " + grad_h * H_BLOCK * W * D * IN_CHANNELS\n", + " + grad_w * W_BLOCK * D * IN_CHANNELS\n", + " + grad_d * D_BLOCK * IN_CHANNELS\n", + " )\n", + " x_next = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0)\n", + "\n", + " weight_grad += tl.dot(x, grad)\n", + "\n", + " tl.store(weight_grad_ptr + weight_grad_offset, weight_grad)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "67fae11d-4fb5-4aef-bcc0-34548a10d9ae", + "metadata": {}, + "outputs": [], + "source": [ + "from kerops.settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction\n", + "from kerops.utils import cdiv\n", + "\n", + "\n", + "def Conv3dWgrad_plain(\n", + " grad,\n", + " x,\n", + " *,\n", + " num_warps: ConfArg,\n", + " H_BLOCK: ConfArg,\n", + " W_BLOCK: ConfArg,\n", + " D_BLOCK: ConfArg,\n", + " CIN_BLOCK: ConfArg, \n", + " COUT_BLOCK: ConfArg,\n", + " SPLIT_K: ConfArg,\n", + "):\n", + " assert x.device == grad.device\n", + " assert x.is_cuda\n", + "\n", + " assert x.ndim == grad.ndim == 5\n", + " xbsize, in_channels, xH, xW, xD = x.shape\n", + " gbsize, out_channels, gH, gW, gD = grad.shape\n", + " assert in_channels == next_power_of_2(in_channels)\n", + " assert out_channels == next_power_of_2(out_channels)\n", + " assert [xbsize, xH, xW, xD] == [gbsize, gH, gW, gD]\n", + "\n", + " assert xbsize == gbsize == 1\n", + "\n", + " assert x.is_contiguous(memory_format=torch.channels_last_3d)\n", + " assert grad.is_contiguous(memory_format=torch.channels_last_3d)\n", + "\n", + " assert x.dtype == grad.dtype == torch.float16\n", + "\n", + " assert D_BLOCK == next_power_of_2(D_BLOCK)\n", + " assert W_BLOCK == next_power_of_2(W_BLOCK)\n", + " assert D_BLOCK == next_power_of_2(D_BLOCK)\n", + "\n", + " ACCTYPE = tl.float32\n", + " weight_grad = torch.zeros([3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float32)\n", + "\n", + " grid = (\n", + " 27 * SPLIT_K,\n", + " cdiv(in_channels, CIN_BLOCK),\n", + " cdiv(out_channels, COUT_BLOCK)\n", + " )\n", + "\n", + " _Conv_wgrad_cl3d_impl[grid](\n", + " grad,\n", + " x,\n", + " weight_grad,\n", + " xH,\n", + " xW,\n", + " xD,\n", + " ACCTYPE=ACCTYPE,\n", + " H_BLOCK=H_BLOCK, W_BLOCK=W_BLOCK, D_BLOCK=D_BLOCK,\n", + " IN_CHANNELS=in_channels, OUT_CHANNELS=out_channels,\n", + " CIN_BLOCK=CIN_BLOCK, COUT_BLOCK=COUT_BLOCK,\n", + " SPLIT_K=SPLIT_K,\n", + " num_warps=num_warps,\n", + " )\n", + " \n", + " weight_grad = weight_grad.to(torch.float16)\n", + "\n", + " return weight_grad" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "3af06690-a21b-46a9-bef8-89e0cbf96b89", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "torch.Size([1, 128, 64, 64, 64])\n", + "torch.Size([1, 128, 64, 64, 64])\n" + ] + } + ], + "source": [ + "print(grad.shape)\n", + "print(x.shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "2bc19bff-2358-4e05-a48b-4e8854086c12", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3.52 ms ± 86.1 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" + ] + } + ], + "source": [ + "%%timeit -r 10 -n 10\n", + "Conv3dWgrad_plain(grad, x, num_warps=4, H_BLOCK=4, W_BLOCK=4, D_BLOCK=4, CIN_BLOCK=64, COUT_BLOCK=64, SPLIT_K=4)\n", + "torch.cuda.synchronize()" + ] + }, + { + "cell_type": "code", + "execution_count": 76, + "id": "51713754-efba-49c4-bfa0-30de870cd325", + "metadata": {}, + "outputs": [], + "source": [ + "out_ = Conv3dWgrad_plain(grad, x, num_warps=4, H_BLOCK=4, W_BLOCK=4, D_BLOCK=4, CIN_BLOCK=64, COUT_BLOCK=64, SPLIT_K=4)\n", + "out = Conv3dWgrad(grad, x)" + ] + }, + { + "cell_type": "code", + "execution_count": 104, + "id": "3ba679a8-b6e9-4e45-94f8-cb81e6e1959d", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor(8., device='cuda:0', dtype=torch.float16)" + ] + }, + "execution_count": 104, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "torch.abs(out_ - out).max()" + ] + }, + { + "cell_type": "code", + "execution_count": 105, + "id": "ca52ffef-3655-4539-8886-a8efec678c63", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor(6480., device='cuda:0', dtype=torch.float16)" + ] + }, + "execution_count": 105, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "torch.abs(out_).max()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8e5f1a53-50de-4230-a931-98d96a805469", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 0a55970a2b6d6378fcee49f59283c0a225294190 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Tue, 7 Apr 2026 14:32:41 +0200 Subject: [PATCH 34/40] SplitK Conv3dWgrad --- kerops/kernels/conv.py | 98 +++++++++++++++++++++++++++++ kerops/ops/conv/conv_wgrad.py | 112 +++++++++++++++++++++++++++++++++- 2 files changed, 209 insertions(+), 1 deletion(-) diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index edfa615..90cc223 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -331,6 +331,104 @@ def _Conv_wgrad_cl3d_impl_V2( tl.atomic_add(w_ptr + weight_grad_offset, wgrad, sem='relaxed') +@triton.jit +def make_offset(h_str, w_str, d_str, H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr): + d_off = tl.arange(0, D_BLOCK) + w_off = tl.arange(0, W_BLOCK) + h_off = tl.arange(0, H_BLOCK) + + offset = h_off[:, None, None] * h_str + w_off[None, :, None] * w_str + d_off[None, None, :] * d_str + offset = offset.reshape((H_BLOCK * W_BLOCK * D_BLOCK)) + + return offset + + +@triton.jit +def make_mask(curr_h, curr_w, curr_d, H, W, D, H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr): + mask_d = ((tl.arange(0, D_BLOCK) + curr_d) >= 0) & ((tl.arange(0, D_BLOCK) + curr_d) < D) + mask_w = ((tl.arange(0, W_BLOCK) + curr_w) >= 0) & ((tl.arange(0, W_BLOCK) + curr_w) < W) + mask_h = ((tl.arange(0, H_BLOCK) + curr_h) >= 0) & ((tl.arange(0, H_BLOCK) + curr_h) < H) + + mask = mask_h[:, None, None] & mask_w[None, :, None] & mask_d[None, None, :] + mask = mask.reshape((H_BLOCK * W_BLOCK * D_BLOCK)) + + return mask + + +@triton.jit +def _Conv_wgrad_cl3d_splitKonH_impl( + grad_ptr, + input_ptr, + weight_grad_ptr, + H, + W, + D, + ACCTYPE: tl.constexpr, + H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr, + IN_CHANNELS: tl.constexpr, OUT_CHANNELS: tl.constexpr, + CIN_BLOCK: tl.constexpr, COUT_BLOCK: tl.constexpr, + SPLIT_K: tl.constexpr, +): + khwd_pid = tl.program_id(0) + cin_pid = tl.program_id(1) + cout_pid = tl.program_id(2) + + k_pid = khwd_pid // 27 + hwd_pid = khwd_pid % 27 + + block_d = hwd_pid % 3 + hwd_pid = hwd_pid // 3 + block_w = hwd_pid % 3 + block_h = hwd_pid // 3 + + cin_offset = tl.arange(0, CIN_BLOCK) + cout_offset = tl.arange(0, COUT_BLOCK) + geom_offset = make_offset(W * D, D, 1, H_BLOCK, W_BLOCK, D_BLOCK) + + grad_offset = geom_offset[:, None] * OUT_CHANNELS + cout_offset[None, :] + input_offset = geom_offset[None, :] * IN_CHANNELS + cin_offset[:, None] + weight_grad_offset = cin_offset[:, None] * OUT_CHANNELS + cout_offset[None, :] + + grad_ptr += cout_pid * COUT_BLOCK + grad_ptr += k_pid * H_BLOCK * OUT_CHANNELS * D * W + input_ptr += cin_pid * CIN_BLOCK + input_ptr += (block_d - 1) * IN_CHANNELS + input_ptr += k_pid * H_BLOCK * IN_CHANNELS * D * W + input_ptr += (block_w - 1) * IN_CHANNELS * D + input_ptr += (block_h - 1) * IN_CHANNELS * D * W + weight_grad_ptr += cin_pid * CIN_BLOCK * OUT_CHANNELS + cout_pid * COUT_BLOCK + weight_grad_ptr += block_d * IN_CHANNELS * OUT_CHANNELS + weight_grad_ptr += block_w * IN_CHANNELS * OUT_CHANNELS * 3 + weight_grad_ptr += block_h * IN_CHANNELS * OUT_CHANNELS * 9 + + weight_grad = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=ACCTYPE) + + for grad_h in range(0, tl.cdiv(H, H_BLOCK * SPLIT_K)): + for grad_w in range(0, tl.cdiv(W, W_BLOCK)): + for grad_d in range(0, tl.cdiv(D, D_BLOCK)): + grad_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK, grad_w * W_BLOCK, grad_d * D_BLOCK, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) + grad_iter_ptr = ( + grad_ptr + + grad_h * H_BLOCK * W * D * OUT_CHANNELS * SPLIT_K + + grad_w * W_BLOCK * D * OUT_CHANNELS + + grad_d * D_BLOCK * OUT_CHANNELS + ) + grad = tl.load(grad_iter_ptr + grad_offset, mask=grad_mask[:, None], other=0) + + x_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK + block_d - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) + x_iter_ptr = ( + input_ptr + + grad_h * H_BLOCK * W * D * IN_CHANNELS * SPLIT_K + + grad_w * W_BLOCK * D * IN_CHANNELS + + grad_d * D_BLOCK * IN_CHANNELS + ) + x = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0) + + weight_grad += tl.dot(x, grad) + + tl.atomic_add(weight_grad_ptr + weight_grad_offset, weight_grad, sem='relaxed') + + @triton.jit def _ApplyBNReLUConvStats_cl3d_impl( input_ptr, diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index c3fae85..e4c2eee 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -2,7 +2,7 @@ from triton import language as tl, next_power_of_2 from ..assets import ASSETS_ROOT -from ...kernels.conv import _Conv_wgrad_cl3d_impl_V2 +from ...kernels.conv import _Conv_wgrad_cl3d_impl_V2, _Conv_wgrad_cl3d_splitKonH_impl from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv @@ -84,6 +84,70 @@ def Conv3dWgrad( return weight_grad +def Conv3dWgrad_splitKonH( + grad, + x, + *, + num_warps: ConfArg, + H_BLOCK: ConfArg, + W_BLOCK: ConfArg, + D_BLOCK: ConfArg, + CIN_BLOCK: ConfArg, + COUT_BLOCK: ConfArg, + SPLIT_K: ConfArg, +): + assert x.device == grad.device + assert x.is_cuda + + assert x.ndim == grad.ndim == 5 + xbsize, in_channels, xH, xW, xD = x.shape + gbsize, out_channels, gH, gW, gD = grad.shape + assert in_channels == next_power_of_2(in_channels) + assert out_channels == next_power_of_2(out_channels) + assert [xbsize, xH, xW, xD] == [gbsize, gH, gW, gD] + + assert xbsize == gbsize == 1 + + assert x.is_contiguous(memory_format=torch.channels_last_3d) + assert grad.is_contiguous(memory_format=torch.channels_last_3d) + + assert x.dtype == grad.dtype == torch.float16 + + assert H_BLOCK == next_power_of_2(H_BLOCK) + assert W_BLOCK == next_power_of_2(W_BLOCK) + assert D_BLOCK == next_power_of_2(D_BLOCK) + assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) + assert CIN_BLOCK <= in_channels + assert COUT_BLOCK == next_power_of_2(COUT_BLOCK) + assert COUT_BLOCK <= out_channels + + ACCTYPE = tl.float32 + weight_grad = torch.zeros([3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float32) + + grid = ( + 27 * SPLIT_K, + cdiv(in_channels, CIN_BLOCK), + cdiv(out_channels, COUT_BLOCK) + ) + + _Conv_wgrad_cl3d_splitKonH_impl[grid]( + grad, + x, + weight_grad, + xH, xW, xD, + ACCTYPE=ACCTYPE, + H_BLOCK=H_BLOCK, W_BLOCK=W_BLOCK, D_BLOCK=D_BLOCK, + IN_CHANNELS=in_channels, OUT_CHANNELS=out_channels, + CIN_BLOCK=CIN_BLOCK, COUT_BLOCK=COUT_BLOCK, + SPLIT_K=SPLIT_K, + num_warps=num_warps, + ) + + weight_grad = weight_grad.to(torch.float16) + + return weight_grad + + def pruning_rule(problem_size, named_config): D_BLOCK = named_config['D_BLOCK'] CIN_BLOCK = named_config['CIN_BLOCK'] @@ -107,6 +171,25 @@ def pruning_rule(problem_size, named_config): return True +# TODO: non-equal block sizes? +def pruning_rule_splitKonH(problem_size, named_config): + H_BLOCK = named_config['H_BLOCK'] + W_BLOCK = named_config['W_BLOCK'] + D_BLOCK = named_config['D_BLOCK'] + CIN_BLOCK = named_config['CIN_BLOCK'] + COUT_BLOCK = named_config['COUT_BLOCK'] + + in_channels, out_channels = problem_size['in_channels'], problem_size['out_channels'] + + if in_channels < CIN_BLOCK or out_channels < COUT_BLOCK: + return False + + if H_BLOCK != W_BLOCK or H_BLOCK != D_BLOCK or W_BLOCK != D_BLOCK: + return False + + return True + + def generate_inputs_conv_wgrad(problem_sizes): in_channels, out_channels = problem_sizes['in_channels'], problem_sizes['out_channels'] @@ -167,3 +250,30 @@ def autotune_conv_wgrad(toml_path, **autotune_kwargs): COUT_BLOCK=channels, SWAP_GRAD_WITH_INPUT=[False, True] ) + + +def autotune_conv_wgrad_splitKonH(toml_path, **autotune_kwargs): + channels = [2 ** i for i in range(4, 8)] + problem_sizes = [ + {'in_channels': cin, 'out_channels': cout} + for cin in channels + for cout in channels + if (cin == 2 * cout) or (cin * 2 == cout) or (cin == cout) + ] + + autotune( + getattr(Conv3dWgrad_splitKonH, 'function', Conv3dWgrad_splitKonH), + generate_inputs_conv_wgrad, + problem_sizes, + pruning_rule_splitKonH, + toml_path, + comparator=comparator, + **autotune_kwargs, + num_warps=[1, 2, 4], + H_BLOCK=[2, 4, 8, 16], + W_BLOCK=[2, 4, 8, 16], + D_BLOCK=[2, 4, 8, 16], + CIN_BLOCK=channels, + COUT_BLOCK=channels, + SPLIT_K=[1, 2, 4, 8, 16] + ) From 3b66f9c936adee08795c411d7d4dd0e610313188 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 10 Apr 2026 13:11:50 +0200 Subject: [PATCH 35/40] Optimize Conv3dWgradSplitKonH, configure, remove SWAP_GRAD_WITH_INPUT from Conv3dWgrad --- kerops/kernels/conv.py | 47 ++++++------ kerops/ops/assets/Conv3dWgrad.toml | 22 ++---- kerops/ops/assets/Conv3dWgradSplitk.toml | 93 ++++++++++++++++++++++++ kerops/ops/conv/conv_wgrad.py | 55 +++++++------- 4 files changed, 149 insertions(+), 68 deletions(-) create mode 100644 kerops/ops/assets/Conv3dWgradSplitk.toml diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index 90cc223..a7f492c 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -365,9 +365,9 @@ def _Conv_wgrad_cl3d_splitKonH_impl( D, ACCTYPE: tl.constexpr, H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr, - IN_CHANNELS: tl.constexpr, OUT_CHANNELS: tl.constexpr, + IN_CHANNELS, OUT_CHANNELS, CIN_BLOCK: tl.constexpr, COUT_BLOCK: tl.constexpr, - SPLIT_K: tl.constexpr, + SPLIT_K, ): khwd_pid = tl.program_id(0) cin_pid = tl.program_id(1) @@ -404,27 +404,28 @@ def _Conv_wgrad_cl3d_splitKonH_impl( weight_grad = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=ACCTYPE) for grad_h in range(0, tl.cdiv(H, H_BLOCK * SPLIT_K)): - for grad_w in range(0, tl.cdiv(W, W_BLOCK)): - for grad_d in range(0, tl.cdiv(D, D_BLOCK)): - grad_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK, grad_w * W_BLOCK, grad_d * D_BLOCK, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) - grad_iter_ptr = ( - grad_ptr - + grad_h * H_BLOCK * W * D * OUT_CHANNELS * SPLIT_K - + grad_w * W_BLOCK * D * OUT_CHANNELS - + grad_d * D_BLOCK * OUT_CHANNELS - ) - grad = tl.load(grad_iter_ptr + grad_offset, mask=grad_mask[:, None], other=0) - - x_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK + block_d - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) - x_iter_ptr = ( - input_ptr - + grad_h * H_BLOCK * W * D * IN_CHANNELS * SPLIT_K - + grad_w * W_BLOCK * D * IN_CHANNELS - + grad_d * D_BLOCK * IN_CHANNELS - ) - x = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0) - - weight_grad += tl.dot(x, grad) + if grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK < H: + for grad_w in range(0, tl.cdiv(W, W_BLOCK)): + for grad_d in range(0, tl.cdiv(D, D_BLOCK)): + grad_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK, grad_w * W_BLOCK, grad_d * D_BLOCK, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) + grad_iter_ptr = ( + grad_ptr + + grad_h * H_BLOCK * W * D * OUT_CHANNELS * SPLIT_K + + grad_w * W_BLOCK * D * OUT_CHANNELS + + grad_d * D_BLOCK * OUT_CHANNELS + ) + grad = tl.load(grad_iter_ptr + grad_offset, mask=grad_mask[:, None], other=0) + + x_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK + block_d - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) + x_iter_ptr = ( + input_ptr + + grad_h * H_BLOCK * W * D * IN_CHANNELS * SPLIT_K + + grad_w * W_BLOCK * D * IN_CHANNELS + + grad_d * D_BLOCK * IN_CHANNELS + ) + x = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0) + + weight_grad += tl.dot(x, grad) tl.atomic_add(weight_grad_ptr + weight_grad_offset, weight_grad, sem='relaxed') diff --git a/kerops/ops/assets/Conv3dWgrad.toml b/kerops/ops/assets/Conv3dWgrad.toml index 5c321be..ed9717b 100644 --- a/kerops/ops/assets/Conv3dWgrad.toml +++ b/kerops/ops/assets/Conv3dWgrad.toml @@ -9,7 +9,6 @@ problem_size_names = ['in_channels', 'out_channels'] REDUCTION_FACTOR = 32 CIN_BLOCK = 16 COUT_BLOCK = 16 - SWAP_GRAD_WITH_INPUT = false [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [16, 32] @@ -18,16 +17,14 @@ problem_size_names = ['in_channels', 'out_channels'] REDUCTION_FACTOR = 32 CIN_BLOCK = 16 COUT_BLOCK = 32 - SWAP_GRAD_WITH_INPUT = false [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [32, 16] num_warps = 1 D_BLOCK = 32 REDUCTION_FACTOR = 32 - CIN_BLOCK = 16 - COUT_BLOCK = 32 - SWAP_GRAD_WITH_INPUT = true + CIN_BLOCK = 32 + COUT_BLOCK = 16 [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [32, 32] @@ -36,7 +33,6 @@ problem_size_names = ['in_channels', 'out_channels'] REDUCTION_FACTOR = 32 CIN_BLOCK = 32 COUT_BLOCK = 32 - SWAP_GRAD_WITH_INPUT = false [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [32, 64] @@ -45,16 +41,14 @@ problem_size_names = ['in_channels', 'out_channels'] REDUCTION_FACTOR = 32 CIN_BLOCK = 16 COUT_BLOCK = 64 - SWAP_GRAD_WITH_INPUT = false [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [64, 32] - num_warps = 2 + num_warps = 4 D_BLOCK = 32 REDUCTION_FACTOR = 32 - CIN_BLOCK = 16 - COUT_BLOCK = 64 - SWAP_GRAD_WITH_INPUT = true + CIN_BLOCK = 64 + COUT_BLOCK = 32 [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [64, 64] @@ -63,7 +57,6 @@ problem_size_names = ['in_channels', 'out_channels'] REDUCTION_FACTOR = 32 CIN_BLOCK = 32 COUT_BLOCK = 64 - SWAP_GRAD_WITH_INPUT = false [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [64, 128] @@ -72,7 +65,6 @@ problem_size_names = ['in_channels', 'out_channels'] REDUCTION_FACTOR = 32 CIN_BLOCK = 64 COUT_BLOCK = 128 - SWAP_GRAD_WITH_INPUT = false [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [128, 64] @@ -80,8 +72,7 @@ problem_size_names = ['in_channels', 'out_channels'] D_BLOCK = 16 REDUCTION_FACTOR = 32 CIN_BLOCK = 64 - COUT_BLOCK = 128 - SWAP_GRAD_WITH_INPUT = true + COUT_BLOCK = 64 [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [128, 128] @@ -90,4 +81,3 @@ problem_size_names = ['in_channels', 'out_channels'] REDUCTION_FACTOR = 32 CIN_BLOCK = 64 COUT_BLOCK = 128 - SWAP_GRAD_WITH_INPUT = false diff --git a/kerops/ops/assets/Conv3dWgradSplitk.toml b/kerops/ops/assets/Conv3dWgradSplitk.toml new file mode 100644 index 0000000..290eab9 --- /dev/null +++ b/kerops/ops/assets/Conv3dWgradSplitk.toml @@ -0,0 +1,93 @@ +[meta] +problem_size_names = ['in_channels', 'out_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 16] + num_warps = 1 + H_BLOCK = 4 + WD_BLOCK = 4 + CIN_BLOCK = 16 + COUT_BLOCK = 16 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 32] + num_warps = 2 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 16 + COUT_BLOCK = 32 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 16] + num_warps = 2 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 32 + COUT_BLOCK = 16 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 32] + num_warps = 4 + H_BLOCK = 4 + WD_BLOCK = 4 + CIN_BLOCK = 32 + COUT_BLOCK = 32 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 64] + num_warps = 4 + H_BLOCK = 4 + WD_BLOCK = 4 + CIN_BLOCK = 32 + COUT_BLOCK = 64 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 32] + num_warps = 2 + H_BLOCK = 4 + WD_BLOCK = 4 + CIN_BLOCK = 64 + COUT_BLOCK = 32 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 64] + num_warps = 4 + H_BLOCK = 8 + WD_BLOCK = 2 + CIN_BLOCK = 64 + COUT_BLOCK = 64 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 128] + num_warps = 4 + H_BLOCK = 4 + WD_BLOCK = 4 + CIN_BLOCK = 64 + COUT_BLOCK = 128 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 64] + num_warps = 4 + H_BLOCK = 4 + WD_BLOCK = 4 + CIN_BLOCK = 128 + COUT_BLOCK = 64 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 128] + num_warps = 2 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 64 + COUT_BLOCK = 128 + SPLIT_K = 16 diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index e4c2eee..2af1f4c 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -9,7 +9,7 @@ conv3d_wgrad_config = TableKernelConfig( problem_size_names=['in_channels', 'out_channels'], - confarg_names=['num_warps', 'D_BLOCK', 'REDUCTION_FACTOR', 'CIN_BLOCK', 'COUT_BLOCK', 'SWAP_GRAD_WITH_INPUT'], + confarg_names=['num_warps', 'D_BLOCK', 'REDUCTION_FACTOR', 'CIN_BLOCK', 'COUT_BLOCK'], args_to_problem_sizes=lambda grad, x: (x.shape[1], grad.shape[1]), toml_path=ASSETS_ROOT / 'Conv3dWgrad.toml' ) @@ -25,11 +25,7 @@ def Conv3dWgrad( REDUCTION_FACTOR: ConfArg, CIN_BLOCK: ConfArg, COUT_BLOCK: ConfArg, - SWAP_GRAD_WITH_INPUT: ConfArg ): - if SWAP_GRAD_WITH_INPUT: - grad, x = x, grad - assert x.device == grad.device assert x.is_cuda @@ -76,22 +72,25 @@ def Conv3dWgrad( weight_grad = weight_grad.sum(dim=0).to(torch.float16) - if SWAP_GRAD_WITH_INPUT: - weight_grad = weight_grad.permute(0, 1, 2, 4, 3) - weight_grad = torch.flip(weight_grad, dims=(0, 1, 2)) - weight_grad = weight_grad.contiguous() - return weight_grad +conv3d_wgrad_splitk_config = TableKernelConfig( + problem_size_names=['in_channels', 'out_channels'], + confarg_names=['num_warps', 'H_BLOCK', 'WD_BLOCK', 'CIN_BLOCK', 'COUT_BLOCK', 'SPLIT_K'], + args_to_problem_sizes=lambda grad, x: (x.shape[1], grad.shape[1]), + toml_path=ASSETS_ROOT / 'Conv3dWgradSplitk.toml' +) + + +@ConfiguredFunction.configure(conv3d_wgrad_splitk_config) def Conv3dWgrad_splitKonH( grad, x, *, num_warps: ConfArg, H_BLOCK: ConfArg, - W_BLOCK: ConfArg, - D_BLOCK: ConfArg, + WD_BLOCK: ConfArg, CIN_BLOCK: ConfArg, COUT_BLOCK: ConfArg, SPLIT_K: ConfArg, @@ -114,8 +113,7 @@ def Conv3dWgrad_splitKonH( assert x.dtype == grad.dtype == torch.float16 assert H_BLOCK == next_power_of_2(H_BLOCK) - assert W_BLOCK == next_power_of_2(W_BLOCK) - assert D_BLOCK == next_power_of_2(D_BLOCK) + assert WD_BLOCK == next_power_of_2(WD_BLOCK) assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) assert CIN_BLOCK <= in_channels assert COUT_BLOCK == next_power_of_2(COUT_BLOCK) @@ -136,7 +134,7 @@ def Conv3dWgrad_splitKonH( weight_grad, xH, xW, xD, ACCTYPE=ACCTYPE, - H_BLOCK=H_BLOCK, W_BLOCK=W_BLOCK, D_BLOCK=D_BLOCK, + H_BLOCK=H_BLOCK, W_BLOCK=WD_BLOCK, D_BLOCK=WD_BLOCK, IN_CHANNELS=in_channels, OUT_CHANNELS=out_channels, CIN_BLOCK=CIN_BLOCK, COUT_BLOCK=COUT_BLOCK, SPLIT_K=SPLIT_K, @@ -152,7 +150,6 @@ def pruning_rule(problem_size, named_config): D_BLOCK = named_config['D_BLOCK'] CIN_BLOCK = named_config['CIN_BLOCK'] COUT_BLOCK = named_config['COUT_BLOCK'] - SWAP_GRAD_WITH_INPUT = named_config['SWAP_GRAD_WITH_INPUT'] in_channels, out_channels = problem_size['in_channels'], problem_size['out_channels'] @@ -162,20 +159,15 @@ def pruning_rule(problem_size, named_config): if (in_channels >= 128 or out_channels >= 128) and D_BLOCK > 16: return False - if (CIN_BLOCK > in_channels and not SWAP_GRAD_WITH_INPUT) or (CIN_BLOCK > out_channels and SWAP_GRAD_WITH_INPUT): - return False - - if (COUT_BLOCK > out_channels and not SWAP_GRAD_WITH_INPUT) or (COUT_BLOCK > in_channels and SWAP_GRAD_WITH_INPUT): + if CIN_BLOCK > in_channels or COUT_BLOCK > out_channels: return False return True -# TODO: non-equal block sizes? def pruning_rule_splitKonH(problem_size, named_config): H_BLOCK = named_config['H_BLOCK'] - W_BLOCK = named_config['W_BLOCK'] - D_BLOCK = named_config['D_BLOCK'] + WD_BLOCK = named_config['WD_BLOCK'] CIN_BLOCK = named_config['CIN_BLOCK'] COUT_BLOCK = named_config['COUT_BLOCK'] @@ -183,8 +175,14 @@ def pruning_rule_splitKonH(problem_size, named_config): if in_channels < CIN_BLOCK or out_channels < COUT_BLOCK: return False - - if H_BLOCK != W_BLOCK or H_BLOCK != D_BLOCK or W_BLOCK != D_BLOCK: + + if in_channels // CIN_BLOCK > 4 or out_channels // COUT_BLOCK > 4: + return False + + if H_BLOCK * WD_BLOCK * WD_BLOCK < 16: + return False + + if H_BLOCK * WD_BLOCK * WD_BLOCK > 64: return False return True @@ -248,7 +246,6 @@ def autotune_conv_wgrad(toml_path, **autotune_kwargs): REDUCTION_FACTOR=[32], CIN_BLOCK=channels, COUT_BLOCK=channels, - SWAP_GRAD_WITH_INPUT=[False, True] ) @@ -271,9 +268,9 @@ def autotune_conv_wgrad_splitKonH(toml_path, **autotune_kwargs): **autotune_kwargs, num_warps=[1, 2, 4], H_BLOCK=[2, 4, 8, 16], - W_BLOCK=[2, 4, 8, 16], - D_BLOCK=[2, 4, 8, 16], + WD_BLOCK=[2, 4, 8, 16], CIN_BLOCK=channels, COUT_BLOCK=channels, - SPLIT_K=[1, 2, 4, 8, 16] + SPLIT_K=[4, 8, 16], + n_iters=85 ) From 74a3ca181f7fb3fc754b1a9c002c12e1512bed1c Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 10 Apr 2026 13:14:34 +0200 Subject: [PATCH 36/40] rename Conv3dWgrad -> Conv3dWgrad_grad_based --- .../{Conv3dWgrad.toml => Conv3dWgrad_grad_based.toml} | 0 kerops/ops/conv/__init__.py | 2 +- kerops/ops/conv/conv_wgrad.py | 8 ++++---- 3 files changed, 5 insertions(+), 5 deletions(-) rename kerops/ops/assets/{Conv3dWgrad.toml => Conv3dWgrad_grad_based.toml} (100%) diff --git a/kerops/ops/assets/Conv3dWgrad.toml b/kerops/ops/assets/Conv3dWgrad_grad_based.toml similarity index 100% rename from kerops/ops/assets/Conv3dWgrad.toml rename to kerops/ops/assets/Conv3dWgrad_grad_based.toml diff --git a/kerops/ops/conv/__init__.py b/kerops/ops/conv/__init__.py index 82b82b9..795402d 100644 --- a/kerops/ops/conv/__init__.py +++ b/kerops/ops/conv/__init__.py @@ -1,4 +1,4 @@ from .conv import Conv3d, ApplyBNReLUConv3dStats -from .conv_wgrad import Conv3dWgrad +from .conv_wgrad import Conv3dWgrad_grad_based from .dwconv import DWConv from .dwconv_wgrad import DWConvWGRAD diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index 2af1f4c..eebd36d 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -11,12 +11,12 @@ problem_size_names=['in_channels', 'out_channels'], confarg_names=['num_warps', 'D_BLOCK', 'REDUCTION_FACTOR', 'CIN_BLOCK', 'COUT_BLOCK'], args_to_problem_sizes=lambda grad, x: (x.shape[1], grad.shape[1]), - toml_path=ASSETS_ROOT / 'Conv3dWgrad.toml' + toml_path=ASSETS_ROOT / 'Conv3dWgrad_grad_based.toml' ) @ConfiguredFunction.configure(conv3d_wgrad_config) -def Conv3dWgrad( +def Conv3dWgrad_grad_based( grad, x, *, @@ -224,7 +224,7 @@ def comparator(grad, x): ) -def autotune_conv_wgrad(toml_path, **autotune_kwargs): +def autotune_conv_wgrad_grad_based(toml_path, **autotune_kwargs): channels = [2 ** i for i in range(4, 8)] problem_sizes = [ {'in_channels': cin, 'out_channels': cout} @@ -234,7 +234,7 @@ def autotune_conv_wgrad(toml_path, **autotune_kwargs): ] autotune( - getattr(Conv3dWgrad, 'function', Conv3dWgrad), + getattr(Conv3dWgrad_grad_based, 'function', Conv3dWgrad_grad_based), generate_inputs_conv_wgrad, problem_sizes, pruning_rule, From 3898acd70bc97af242038b42a2a4929fe702cc5f Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 10 Apr 2026 13:18:00 +0200 Subject: [PATCH 37/40] Standartize Conv3dWgrad splitk --- kerops/kernels/conv.py | 2 +- ...v3dWgradSplitk.toml => Conv3dWgrad_splitk.toml} | 0 kerops/ops/conv/conv_wgrad.py | 14 +++++++------- 3 files changed, 8 insertions(+), 8 deletions(-) rename kerops/ops/assets/{Conv3dWgradSplitk.toml => Conv3dWgrad_splitk.toml} (100%) diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index a7f492c..029b1cd 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -356,7 +356,7 @@ def make_mask(curr_h, curr_w, curr_d, H, W, D, H_BLOCK: tl.constexpr, W_BLOCK: t @triton.jit -def _Conv_wgrad_cl3d_splitKonH_impl( +def _Conv_wgrad_cl3d_splitk_impl( grad_ptr, input_ptr, weight_grad_ptr, diff --git a/kerops/ops/assets/Conv3dWgradSplitk.toml b/kerops/ops/assets/Conv3dWgrad_splitk.toml similarity index 100% rename from kerops/ops/assets/Conv3dWgradSplitk.toml rename to kerops/ops/assets/Conv3dWgrad_splitk.toml diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index eebd36d..9405fe1 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -2,12 +2,12 @@ from triton import language as tl, next_power_of_2 from ..assets import ASSETS_ROOT -from ...kernels.conv import _Conv_wgrad_cl3d_impl_V2, _Conv_wgrad_cl3d_splitKonH_impl +from ...kernels.conv import _Conv_wgrad_cl3d_impl_V2, _Conv_wgrad_cl3d_splitk_impl from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv -conv3d_wgrad_config = TableKernelConfig( +conv3d_wgrad_grad_based_config = TableKernelConfig( problem_size_names=['in_channels', 'out_channels'], confarg_names=['num_warps', 'D_BLOCK', 'REDUCTION_FACTOR', 'CIN_BLOCK', 'COUT_BLOCK'], args_to_problem_sizes=lambda grad, x: (x.shape[1], grad.shape[1]), @@ -15,7 +15,7 @@ ) -@ConfiguredFunction.configure(conv3d_wgrad_config) +@ConfiguredFunction.configure(conv3d_wgrad_grad_based_config) def Conv3dWgrad_grad_based( grad, x, @@ -79,12 +79,12 @@ def Conv3dWgrad_grad_based( problem_size_names=['in_channels', 'out_channels'], confarg_names=['num_warps', 'H_BLOCK', 'WD_BLOCK', 'CIN_BLOCK', 'COUT_BLOCK', 'SPLIT_K'], args_to_problem_sizes=lambda grad, x: (x.shape[1], grad.shape[1]), - toml_path=ASSETS_ROOT / 'Conv3dWgradSplitk.toml' + toml_path=ASSETS_ROOT / 'Conv3dWgrad_splitk.toml' ) @ConfiguredFunction.configure(conv3d_wgrad_splitk_config) -def Conv3dWgrad_splitKonH( +def Conv3dWgrad_splitk( grad, x, *, @@ -128,7 +128,7 @@ def Conv3dWgrad_splitKonH( cdiv(out_channels, COUT_BLOCK) ) - _Conv_wgrad_cl3d_splitKonH_impl[grid]( + _Conv_wgrad_cl3d_splitk_impl[grid]( grad, x, weight_grad, @@ -259,7 +259,7 @@ def autotune_conv_wgrad_splitKonH(toml_path, **autotune_kwargs): ] autotune( - getattr(Conv3dWgrad_splitKonH, 'function', Conv3dWgrad_splitKonH), + getattr(Conv3dWgrad_splitk, 'function', Conv3dWgrad_splitk), generate_inputs_conv_wgrad, problem_sizes, pruning_rule_splitKonH, From 76697762a6792b42e35d2334bfb9453725375ea1 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 10 Apr 2026 13:29:13 +0200 Subject: [PATCH 38/40] ConvWgrad as combination of different kernels --- kerops/ops/assets/Conv3dWgrad.toml | 53 ++++ kerops/ops/conv/conv_wgrad.py | 55 ++++ notebooks/Conv3dWgrad.ipynb | 488 ----------------------------- 3 files changed, 108 insertions(+), 488 deletions(-) create mode 100644 kerops/ops/assets/Conv3dWgrad.toml delete mode 100644 notebooks/Conv3dWgrad.ipynb diff --git a/kerops/ops/assets/Conv3dWgrad.toml b/kerops/ops/assets/Conv3dWgrad.toml new file mode 100644 index 0000000..c3cfe45 --- /dev/null +++ b/kerops/ops/assets/Conv3dWgrad.toml @@ -0,0 +1,53 @@ +[meta] +problem_size_names = ['in_channels', 'out_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 16] + IMPL_ID = 0 + SWAP_GRAD_X = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 32] + IMPL_ID = 0 + SWAP_GRAD_X = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 16] + IMPL_ID = 0 + SWAP_GRAD_X = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 32] + IMPL_ID = 0 + SWAP_GRAD_X = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 64] + IMPL_ID = 0 + SWAP_GRAD_X = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 32] + IMPL_ID = 0 + SWAP_GRAD_X = true + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 64] + IMPL_ID = 1 + SWAP_GRAD_X = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 128] + IMPL_ID = 1 + SWAP_GRAD_X = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 64] + IMPL_ID = 1 + SWAP_GRAD_X = false + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 128] + IMPL_ID = 1 + SWAP_GRAD_X = false diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index 9405fe1..7b7d72c 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -146,6 +146,39 @@ def Conv3dWgrad_splitk( return weight_grad +conv3d_wgrad = TableKernelConfig( + problem_size_names=['in_channels', 'out_channels'], + confarg_names=['IMPL_ID', 'SWAP_GRAD_X'], + args_to_problem_sizes=lambda grad, x: (x.shape[1], grad.shape[1]), + toml_path=ASSETS_ROOT / 'Conv3dWgrad.toml' +) + + +@ConfiguredFunction.configure(conv3d_wgrad) +def Conv3dWgrad( + grad, + x, + *, + IMPL_ID: ConfArg, + SWAP_GRAD_X: ConfArg, +): + if SWAP_GRAD_X: + grad, x = x, grad + + match IMPL_ID: + case 0: + weight_grad = Conv3dWgrad_grad_based(grad, x) + case 1: + weight_grad = Conv3dWgrad_splitk(grad, x) + + if SWAP_GRAD_X: + weight_grad = weight_grad.permute(0, 1, 2, 4, 3) + weight_grad = torch.flip(weight_grad, dims=(0, 1, 2)) + weight_grad = weight_grad.contiguous() + + return weight_grad + + def pruning_rule(problem_size, named_config): D_BLOCK = named_config['D_BLOCK'] CIN_BLOCK = named_config['CIN_BLOCK'] @@ -274,3 +307,25 @@ def autotune_conv_wgrad_splitKonH(toml_path, **autotune_kwargs): SPLIT_K=[4, 8, 16], n_iters=85 ) + + +def autotune_conv_wgrad(toml_path, **autotune_kwargs): + channels = [2 ** i for i in range(4, 8)] + problem_sizes = [ + {'in_channels': cin, 'out_channels': cout} + for cin in channels + for cout in channels + if (cin == 2 * cout) or (cin * 2 == cout) or (cin == cout) + ] + + autotune( + getattr(Conv3dWgrad, 'function', Conv3dWgrad), + generate_inputs_conv_wgrad, + problem_sizes, + pruning_rule=None, + toml_path=toml_path, + comparator=comparator, + **autotune_kwargs, + IMPL_ID=[0, 1], + SWAP_GRAD_X=[False, True] + ) diff --git a/notebooks/Conv3dWgrad.ipynb b/notebooks/Conv3dWgrad.ipynb deleted file mode 100644 index f2533ba..0000000 --- a/notebooks/Conv3dWgrad.ipynb +++ /dev/null @@ -1,488 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 3, - "id": "4d9e0870-6d1a-484c-ab0e-84ca2a72f9e8", - "metadata": {}, - "outputs": [], - "source": [ - "import torch\n", - "from kerops.ops.conv.conv_wgrad import Conv3dWgrad, conv3d_wgrad_config, autotune_conv_wgrad\n", - "\n", - "Conv3dWgrad.register_kernel_config(conv3d_wgrad_config)\n", - "\n", - "cin = 128\n", - "cout = 128\n", - "S = D = 64\n", - "x = torch.randn(1, cin, S, S, D, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)\n", - "grad = torch.randn(1, cout, S, S, D, device='cuda', dtype=torch.float16).to(memory_format=torch.channels_last_3d)" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "339e0403-64ef-46e1-87d7-c3c164a79ec6", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "5.03 ms ± 54.2 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" - ] - } - ], - "source": [ - "%%timeit -r 10 -n 10\n", - "Conv3dWgrad(grad, x)\n", - "torch.cuda.synchronize()" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "79789d16-7bb7-44ae-bf48-b65c711ee62a", - "metadata": {}, - "outputs": [], - "source": [ - "weight = torch.randn(cout, cin, 3, 3, 3, device='cuda', dtype=torch.float16)" - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "id": "04dc085a-9952-4d19-a2f7-b8a298707ec4", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "3.71 ms ± 67.9 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" - ] - } - ], - "source": [ - "%%timeit -r 10 -n 10\n", - "torch.ops.aten.convolution_backward(\n", - " grad,\n", - " x,\n", - " weight,\n", - " [0], # bias_sizes\n", - " [1, 1, 1], # stride\n", - " [1, 1, 1], # padding\n", - " [1, 1, 1], # dilation\n", - " False, # transposed\n", - " [0, 0, 0], # output padding\n", - " 1, # groups!\n", - " [False, True, False], # output_mask - grad_inpt, grad_weight, grad_bias\n", - ")\n", - "torch.cuda.synchronize()" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "id": "ff720fb6-fb1f-4786-ba90-bdd268c4ed46", - "metadata": {}, - "outputs": [], - "source": [ - "import triton\n", - "from triton import language as tl, next_power_of_2\n", - "\n", - "\n", - "@triton.jit\n", - "def make_offset(h_str, w_str, d_str, H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr):\n", - " d_off = tl.arange(0, D_BLOCK)\n", - " w_off = tl.arange(0, W_BLOCK)\n", - " h_off = tl.arange(0, H_BLOCK)\n", - "\n", - " offset = h_off[:, None, None] * h_str + w_off[None, :, None] * w_str + d_off[None, None, :] * d_str\n", - " offset = offset.reshape((H_BLOCK * W_BLOCK * D_BLOCK))\n", - "\n", - " return offset\n", - "\n", - "\n", - "@triton.jit\n", - "def make_mask(curr_h, curr_w, curr_d, H, W, D, H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr):\n", - " mask_d = ((tl.arange(0, D_BLOCK) + curr_d) >= 0) & ((tl.arange(0, D_BLOCK) + curr_d) < D)\n", - " mask_w = ((tl.arange(0, W_BLOCK) + curr_w) >= 0) & ((tl.arange(0, W_BLOCK) + curr_w) < W)\n", - " mask_h = ((tl.arange(0, H_BLOCK) + curr_h) >= 0) & ((tl.arange(0, H_BLOCK) + curr_h) < H)\n", - "\n", - " mask = mask_h[:, None, None] & mask_w[None, :, None] & mask_d[None, None, :]\n", - " mask = mask.reshape((H_BLOCK * W_BLOCK * D_BLOCK))\n", - "\n", - " return mask\n", - "\n", - "\n", - "@triton.jit\n", - "def _Conv_wgrad_cl3d_impl(\n", - " grad_ptr,\n", - " input_ptr,\n", - " weight_grad_ptr,\n", - " H,\n", - " W,\n", - " D,\n", - " ACCTYPE: tl.constexpr,\n", - " H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr,\n", - " IN_CHANNELS: tl.constexpr, OUT_CHANNELS: tl.constexpr,\n", - " CIN_BLOCK: tl.constexpr, COUT_BLOCK: tl.constexpr,\n", - " SPLIT_K: tl.constexpr,\n", - "):\n", - " khwd_pid = tl.program_id(0)\n", - " cin_pid = tl.program_id(1)\n", - " cout_pid = tl.program_id(2)\n", - "\n", - " k_pid = khwd_pid // 27\n", - " hwd_pid = khwd_pid % 27\n", - "\n", - " block_d = hwd_pid % 3\n", - " hwd_pid = hwd_pid // 3\n", - " block_w = hwd_pid % 3\n", - " block_h = hwd_pid // 3\n", - "\n", - " cin_offset = tl.arange(0, CIN_BLOCK)\n", - " cout_offset = tl.arange(0, COUT_BLOCK)\n", - " geom_offset = make_offset(W * D, D, 1, H_BLOCK, W_BLOCK, D_BLOCK)\n", - "\n", - " grad_offset = geom_offset[:, None] * OUT_CHANNELS + cout_offset[None, :]\n", - " input_offset = geom_offset[None, :] * IN_CHANNELS + cin_offset[:, None]\n", - " weight_grad_offset = cin_offset[:, None] * OUT_CHANNELS + cout_offset[None, :]\n", - "\n", - " grad_ptr += cout_pid * COUT_BLOCK\n", - " grad_ptr += k_pid * H_BLOCK * OUT_CHANNELS * D * W\n", - " input_ptr += cin_pid * CIN_BLOCK\n", - " input_ptr += (block_d - 1) * IN_CHANNELS\n", - " input_ptr += k_pid * H_BLOCK * IN_CHANNELS * D * W\n", - " input_ptr += (block_w - 1) * IN_CHANNELS * D\n", - " input_ptr += (block_h - 1) * IN_CHANNELS * D * W\n", - " weight_grad_ptr += cin_pid * CIN_BLOCK * OUT_CHANNELS + cout_pid * COUT_BLOCK\n", - " weight_grad_ptr += block_d * IN_CHANNELS * OUT_CHANNELS\n", - " weight_grad_ptr += block_w * IN_CHANNELS * OUT_CHANNELS * 3\n", - " weight_grad_ptr += block_h * IN_CHANNELS * OUT_CHANNELS * 9\n", - "\n", - " weight_grad = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=ACCTYPE)\n", - "\n", - " for grad_h in range(0, tl.cdiv(H, H_BLOCK * SPLIT_K)):\n", - " for grad_w in range(0, tl.cdiv(W, W_BLOCK)):\n", - " for grad_d in range(0, tl.cdiv(D, D_BLOCK)):\n", - " grad_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK, grad_w * W_BLOCK, grad_d * D_BLOCK, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", - " grad_iter_ptr = (\n", - " grad_ptr\n", - " + grad_h * H_BLOCK * W * D * OUT_CHANNELS * SPLIT_K\n", - " + grad_w * W_BLOCK * D * OUT_CHANNELS\n", - " + grad_d * D_BLOCK * OUT_CHANNELS\n", - " )\n", - " grad = tl.load(grad_iter_ptr + grad_offset, mask=grad_mask[:, None], other=0)\n", - "\n", - " x_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK + block_d - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", - " x_iter_ptr = (\n", - " input_ptr\n", - " + grad_h * H_BLOCK * W * D * IN_CHANNELS * SPLIT_K\n", - " + grad_w * W_BLOCK * D * IN_CHANNELS\n", - " + grad_d * D_BLOCK * IN_CHANNELS\n", - " )\n", - " x = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0)\n", - "\n", - " weight_grad += tl.dot(x, grad)\n", - "\n", - " tl.atomic_add(weight_grad_ptr + weight_grad_offset, weight_grad, sem='relaxed')" - ] - }, - { - "cell_type": "code", - "execution_count": 63, - "id": "a4f627e6-f633-41aa-9dbd-910e0a126e4d", - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "@triton.jit\n", - "def _Conv_wgrad_cl3d_impl_prefetch(\n", - " grad_ptr,\n", - " input_ptr,\n", - " weight_grad_ptr,\n", - " H,\n", - " W,\n", - " D,\n", - " ACCTYPE: tl.constexpr,\n", - " H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr,\n", - " IN_CHANNELS: tl.constexpr, OUT_CHANNELS: tl.constexpr,\n", - " CIN_BLOCK: tl.constexpr, COUT_BLOCK: tl.constexpr,\n", - "):\n", - " hwd_pid = tl.program_id(0)\n", - " cin_pid = tl.program_id(1)\n", - " cout_pid = tl.program_id(2)\n", - "\n", - " #hwd_pid = tl.program_id(0)\n", - " #cin_cout_pid = tl.program_id(1)\n", - "\n", - " block_d = hwd_pid % 3\n", - " hwd_pid = hwd_pid // 3\n", - " block_w = hwd_pid % 3\n", - " block_h = hwd_pid // 3\n", - "\n", - " #cin_pid = cin_cout_pid % tl.cdiv(IN_CHANNELS, CIN_BLOCK)\n", - " #cout_pid = cin_cout_pid // tl.cdiv(IN_CHANNELS, CIN_BLOCK)\n", - "\n", - " cin_offset = tl.arange(0, CIN_BLOCK)\n", - " cout_offset = tl.arange(0, COUT_BLOCK)\n", - " geom_offset = make_offset(W * D, D, 1, H_BLOCK, W_BLOCK, D_BLOCK)\n", - "\n", - " grad_offset = geom_offset[:, None] * OUT_CHANNELS + cout_offset[None, :]\n", - " input_offset = geom_offset[None, :] * IN_CHANNELS + cin_offset[:, None]\n", - " weight_grad_offset = cin_offset[:, None] * OUT_CHANNELS + cout_offset[None, :]\n", - "\n", - " grad_ptr += cout_pid * COUT_BLOCK\n", - " input_ptr += cin_pid * CIN_BLOCK\n", - " input_ptr += (block_d - 1) * IN_CHANNELS\n", - " input_ptr += (block_w - 1) * IN_CHANNELS * D\n", - " input_ptr += (block_h - 1) * IN_CHANNELS * D * W\n", - " weight_grad_ptr += cin_pid * CIN_BLOCK * OUT_CHANNELS + cout_pid * COUT_BLOCK\n", - " weight_grad_ptr += block_d * IN_CHANNELS * OUT_CHANNELS\n", - " weight_grad_ptr += block_w * IN_CHANNELS * OUT_CHANNELS * 3\n", - " weight_grad_ptr += block_h * IN_CHANNELS * OUT_CHANNELS * 9\n", - "\n", - " weight_grad = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=tl.float32)\n", - " \n", - " grad_mask = make_mask(0, 0, 0, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", - " grad_next = tl.load(grad_ptr + grad_offset, mask=grad_mask[:, None], other=0)\n", - "\n", - " x_mask = make_mask(block_h - 1, block_w - 1, block_d - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", - " x_next = tl.load(input_ptr + input_offset, mask=x_mask[None, :], other=0)\n", - "\n", - " h_bound = tl.cdiv(H, H_BLOCK)\n", - " w_bound = tl.cdiv(W, W_BLOCK)\n", - " d_bound = tl.cdiv(D, D_BLOCK)\n", - "\n", - " for grad_h in range(0, h_bound):\n", - " for grad_w in range(0, w_bound):\n", - " for grad_d in range(0, d_bound):\n", - " grad = grad_next\n", - " x = x_next\n", - "\n", - " if ((grad_h + 1) != h_bound) and ((grad_w + 1) != w_bound) and ((grad_d + 1) != d_bound):\n", - " grad_mask = make_mask(grad_h * H_BLOCK, grad_w * W_BLOCK, grad_d * D_BLOCK, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", - " grad_iter_ptr = (\n", - " grad_ptr\n", - " + grad_h * H_BLOCK * W * D * OUT_CHANNELS\n", - " + grad_w * W_BLOCK * D * OUT_CHANNELS\n", - " + grad_d * D_BLOCK * OUT_CHANNELS\n", - " )\n", - " grad_next = tl.load(grad_iter_ptr + grad_offset, mask=grad_mask[:, None], other=0)\n", - " \n", - " x_mask = make_mask(grad_h * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK + block_d - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK)\n", - " x_iter_ptr = (\n", - " input_ptr\n", - " + grad_h * H_BLOCK * W * D * IN_CHANNELS\n", - " + grad_w * W_BLOCK * D * IN_CHANNELS\n", - " + grad_d * D_BLOCK * IN_CHANNELS\n", - " )\n", - " x_next = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0)\n", - "\n", - " weight_grad += tl.dot(x, grad)\n", - "\n", - " tl.store(weight_grad_ptr + weight_grad_offset, weight_grad)" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "67fae11d-4fb5-4aef-bcc0-34548a10d9ae", - "metadata": {}, - "outputs": [], - "source": [ - "from kerops.settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction\n", - "from kerops.utils import cdiv\n", - "\n", - "\n", - "def Conv3dWgrad_plain(\n", - " grad,\n", - " x,\n", - " *,\n", - " num_warps: ConfArg,\n", - " H_BLOCK: ConfArg,\n", - " W_BLOCK: ConfArg,\n", - " D_BLOCK: ConfArg,\n", - " CIN_BLOCK: ConfArg, \n", - " COUT_BLOCK: ConfArg,\n", - " SPLIT_K: ConfArg,\n", - "):\n", - " assert x.device == grad.device\n", - " assert x.is_cuda\n", - "\n", - " assert x.ndim == grad.ndim == 5\n", - " xbsize, in_channels, xH, xW, xD = x.shape\n", - " gbsize, out_channels, gH, gW, gD = grad.shape\n", - " assert in_channels == next_power_of_2(in_channels)\n", - " assert out_channels == next_power_of_2(out_channels)\n", - " assert [xbsize, xH, xW, xD] == [gbsize, gH, gW, gD]\n", - "\n", - " assert xbsize == gbsize == 1\n", - "\n", - " assert x.is_contiguous(memory_format=torch.channels_last_3d)\n", - " assert grad.is_contiguous(memory_format=torch.channels_last_3d)\n", - "\n", - " assert x.dtype == grad.dtype == torch.float16\n", - "\n", - " assert D_BLOCK == next_power_of_2(D_BLOCK)\n", - " assert W_BLOCK == next_power_of_2(W_BLOCK)\n", - " assert D_BLOCK == next_power_of_2(D_BLOCK)\n", - "\n", - " ACCTYPE = tl.float32\n", - " weight_grad = torch.zeros([3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float32)\n", - "\n", - " grid = (\n", - " 27 * SPLIT_K,\n", - " cdiv(in_channels, CIN_BLOCK),\n", - " cdiv(out_channels, COUT_BLOCK)\n", - " )\n", - "\n", - " _Conv_wgrad_cl3d_impl[grid](\n", - " grad,\n", - " x,\n", - " weight_grad,\n", - " xH,\n", - " xW,\n", - " xD,\n", - " ACCTYPE=ACCTYPE,\n", - " H_BLOCK=H_BLOCK, W_BLOCK=W_BLOCK, D_BLOCK=D_BLOCK,\n", - " IN_CHANNELS=in_channels, OUT_CHANNELS=out_channels,\n", - " CIN_BLOCK=CIN_BLOCK, COUT_BLOCK=COUT_BLOCK,\n", - " SPLIT_K=SPLIT_K,\n", - " num_warps=num_warps,\n", - " )\n", - " \n", - " weight_grad = weight_grad.to(torch.float16)\n", - "\n", - " return weight_grad" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "3af06690-a21b-46a9-bef8-89e0cbf96b89", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "torch.Size([1, 128, 64, 64, 64])\n", - "torch.Size([1, 128, 64, 64, 64])\n" - ] - } - ], - "source": [ - "print(grad.shape)\n", - "print(x.shape)" - ] - }, - { - "cell_type": "code", - "execution_count": 16, - "id": "2bc19bff-2358-4e05-a48b-4e8854086c12", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "3.52 ms ± 86.1 μs per loop (mean ± std. dev. of 10 runs, 10 loops each)\n" - ] - } - ], - "source": [ - "%%timeit -r 10 -n 10\n", - "Conv3dWgrad_plain(grad, x, num_warps=4, H_BLOCK=4, W_BLOCK=4, D_BLOCK=4, CIN_BLOCK=64, COUT_BLOCK=64, SPLIT_K=4)\n", - "torch.cuda.synchronize()" - ] - }, - { - "cell_type": "code", - "execution_count": 76, - "id": "51713754-efba-49c4-bfa0-30de870cd325", - "metadata": {}, - "outputs": [], - "source": [ - "out_ = Conv3dWgrad_plain(grad, x, num_warps=4, H_BLOCK=4, W_BLOCK=4, D_BLOCK=4, CIN_BLOCK=64, COUT_BLOCK=64, SPLIT_K=4)\n", - "out = Conv3dWgrad(grad, x)" - ] - }, - { - "cell_type": "code", - "execution_count": 104, - "id": "3ba679a8-b6e9-4e45-94f8-cb81e6e1959d", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "tensor(8., device='cuda:0', dtype=torch.float16)" - ] - }, - "execution_count": 104, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "torch.abs(out_ - out).max()" - ] - }, - { - "cell_type": "code", - "execution_count": 105, - "id": "ca52ffef-3655-4539-8886-a8efec678c63", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "tensor(6480., device='cuda:0', dtype=torch.float16)" - ] - }, - "execution_count": 105, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "torch.abs(out_).max()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8e5f1a53-50de-4230-a931-98d96a805469", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.13.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} From 355e8538399bcb99fe6c0245fb3179abbe3ace0c Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Fri, 10 Apr 2026 15:57:27 +0200 Subject: [PATCH 39/40] Batch parallelism --- kerops/kernels/conv.py | 10 ++++++++-- kerops/ops/conv/__init__.py | 2 +- kerops/ops/conv/conv_wgrad.py | 7 ++++--- 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index 029b1cd..13b016d 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -371,7 +371,7 @@ def _Conv_wgrad_cl3d_splitk_impl( ): khwd_pid = tl.program_id(0) cin_pid = tl.program_id(1) - cout_pid = tl.program_id(2) + cout_batch_pid = tl.program_id(2) k_pid = khwd_pid // 27 hwd_pid = khwd_pid % 27 @@ -381,6 +381,9 @@ def _Conv_wgrad_cl3d_splitk_impl( block_w = hwd_pid % 3 block_h = hwd_pid // 3 + cout_pid = cout_batch_pid % tl.cdiv(OUT_CHANNELS, COUT_BLOCK) + batch_pid = cout_batch_pid // tl.cdiv(OUT_CHANNELS, COUT_BLOCK) + cin_offset = tl.arange(0, CIN_BLOCK) cout_offset = tl.arange(0, COUT_BLOCK) geom_offset = make_offset(W * D, D, 1, H_BLOCK, W_BLOCK, D_BLOCK) @@ -390,16 +393,19 @@ def _Conv_wgrad_cl3d_splitk_impl( weight_grad_offset = cin_offset[:, None] * OUT_CHANNELS + cout_offset[None, :] grad_ptr += cout_pid * COUT_BLOCK - grad_ptr += k_pid * H_BLOCK * OUT_CHANNELS * D * W + grad_ptr += k_pid * H_BLOCK * OUT_CHANNELS * D * W + grad_ptr += batch_pid * H * W * D * OUT_CHANNELS input_ptr += cin_pid * CIN_BLOCK input_ptr += (block_d - 1) * IN_CHANNELS input_ptr += k_pid * H_BLOCK * IN_CHANNELS * D * W input_ptr += (block_w - 1) * IN_CHANNELS * D input_ptr += (block_h - 1) * IN_CHANNELS * D * W + input_ptr += batch_pid * H * W * D * IN_CHANNELS weight_grad_ptr += cin_pid * CIN_BLOCK * OUT_CHANNELS + cout_pid * COUT_BLOCK weight_grad_ptr += block_d * IN_CHANNELS * OUT_CHANNELS weight_grad_ptr += block_w * IN_CHANNELS * OUT_CHANNELS * 3 weight_grad_ptr += block_h * IN_CHANNELS * OUT_CHANNELS * 9 + weight_grad_ptr += batch_pid * IN_CHANNELS * OUT_CHANNELS * 27 weight_grad = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=ACCTYPE) diff --git a/kerops/ops/conv/__init__.py b/kerops/ops/conv/__init__.py index 795402d..82b82b9 100644 --- a/kerops/ops/conv/__init__.py +++ b/kerops/ops/conv/__init__.py @@ -1,4 +1,4 @@ from .conv import Conv3d, ApplyBNReLUConv3dStats -from .conv_wgrad import Conv3dWgrad_grad_based +from .conv_wgrad import Conv3dWgrad from .dwconv import DWConv from .dwconv_wgrad import DWConvWGRAD diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index 7b7d72c..236035f 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -105,7 +105,7 @@ def Conv3dWgrad_splitk( assert out_channels == next_power_of_2(out_channels) assert [xbsize, xH, xW, xD] == [gbsize, gH, gW, gD] - assert xbsize == gbsize == 1 + assert xbsize == gbsize assert x.is_contiguous(memory_format=torch.channels_last_3d) assert grad.is_contiguous(memory_format=torch.channels_last_3d) @@ -120,12 +120,12 @@ def Conv3dWgrad_splitk( assert COUT_BLOCK <= out_channels ACCTYPE = tl.float32 - weight_grad = torch.zeros([3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float32) + weight_grad = torch.zeros([xbsize, 3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float32) grid = ( 27 * SPLIT_K, cdiv(in_channels, CIN_BLOCK), - cdiv(out_channels, COUT_BLOCK) + cdiv(out_channels, COUT_BLOCK) * xbsize ) _Conv_wgrad_cl3d_splitk_impl[grid]( @@ -141,6 +141,7 @@ def Conv3dWgrad_splitk( num_warps=num_warps, ) + weight_grad = torch.sum(weight_grad, dim=0) weight_grad = weight_grad.to(torch.float16) return weight_grad From 81bbf954fa54103e08c41c35e74912bc29ea1308 Mon Sep 17 00:00:00 2001 From: AnihilatorGun Date: Sat, 11 Apr 2026 13:03:17 +0200 Subject: [PATCH 40/40] new Conv3dWgrad implementation + config --- kerops/kernels/conv.py | 103 +++++++++++++++ kerops/ops/assets/Conv3dWgrad.toml | 8 +- .../ops/assets/Conv3dWgrad_splitk_3acc.toml | 93 ++++++++++++++ kerops/ops/conv/conv_wgrad.py | 118 ++++++++++++++++-- 4 files changed, 310 insertions(+), 12 deletions(-) create mode 100644 kerops/ops/assets/Conv3dWgrad_splitk_3acc.toml diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py index 13b016d..a327cb5 100644 --- a/kerops/kernels/conv.py +++ b/kerops/kernels/conv.py @@ -436,6 +436,109 @@ def _Conv_wgrad_cl3d_splitk_impl( tl.atomic_add(weight_grad_ptr + weight_grad_offset, weight_grad, sem='relaxed') +@triton.jit +def _Conv_wgrad_cl3d_splitk_3acc_impl( + grad_ptr, + input_ptr, + weight_grad_ptr, + H, + W, + D, + ACCTYPE: tl.constexpr, + H_BLOCK: tl.constexpr, W_BLOCK: tl.constexpr, D_BLOCK: tl.constexpr, + IN_CHANNELS, OUT_CHANNELS, + CIN_BLOCK: tl.constexpr, COUT_BLOCK: tl.constexpr, + SPLIT_K, +): + khw_pid = tl.program_id(0) + cin_pid = tl.program_id(1) + cout_batch_pid = tl.program_id(2) + + k_pid = khw_pid // 9 + hw_pid = khw_pid % 9 + + block_w = hw_pid % 3 + block_h = hw_pid // 3 + + cout_pid = cout_batch_pid % tl.cdiv(OUT_CHANNELS, COUT_BLOCK) + batch_pid = cout_batch_pid // tl.cdiv(OUT_CHANNELS, COUT_BLOCK) + + cin_offset = tl.arange(0, CIN_BLOCK) + cout_offset = tl.arange(0, COUT_BLOCK) + geom_offset = make_offset(W * D, D, 1, H_BLOCK, W_BLOCK, D_BLOCK) + + grad_offset = geom_offset[:, None] * OUT_CHANNELS + cout_offset[None, :] + input_offset = geom_offset[None, :] * IN_CHANNELS + cin_offset[:, None] + weight_grad_offset = cin_offset[:, None] * OUT_CHANNELS + cout_offset[None, :] + + grad_ptr += cout_pid * COUT_BLOCK + grad_ptr += k_pid * H_BLOCK * OUT_CHANNELS * D * W + grad_ptr += batch_pid * H * W * D * OUT_CHANNELS + input_ptr += cin_pid * CIN_BLOCK + input_ptr += k_pid * H_BLOCK * IN_CHANNELS * D * W + input_ptr += (block_w - 1) * IN_CHANNELS * D + input_ptr += (block_h - 1) * IN_CHANNELS * D * W + input_ptr += batch_pid * H * W * D * IN_CHANNELS + weight_grad_ptr += cin_pid * CIN_BLOCK * OUT_CHANNELS + cout_pid * COUT_BLOCK + weight_grad_ptr += block_w * IN_CHANNELS * OUT_CHANNELS * 3 + weight_grad_ptr += block_h * IN_CHANNELS * OUT_CHANNELS * 9 + weight_grad_ptr += batch_pid * IN_CHANNELS * OUT_CHANNELS * 27 + + weight_grad0 = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=ACCTYPE) + weight_grad1 = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=ACCTYPE) + weight_grad2 = tl.zeros((CIN_BLOCK, COUT_BLOCK), dtype=ACCTYPE) + + for grad_h in range(0, tl.cdiv(H, H_BLOCK * SPLIT_K)): + if grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK < H: + for grad_w in range(0, tl.cdiv(W, W_BLOCK)): + for grad_d in range(0, tl.cdiv(D, D_BLOCK)): + grad_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK, grad_w * W_BLOCK, grad_d * D_BLOCK, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) + grad_iter_ptr = ( + grad_ptr + + grad_h * H_BLOCK * W * D * OUT_CHANNELS * SPLIT_K + + grad_w * W_BLOCK * D * OUT_CHANNELS + + grad_d * D_BLOCK * OUT_CHANNELS + ) + grad = tl.load(grad_iter_ptr + grad_offset, mask=grad_mask[:, None], other=0) + + x_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK - 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) + x_iter_ptr = ( + input_ptr + + grad_h * H_BLOCK * W * D * IN_CHANNELS * SPLIT_K + + grad_w * W_BLOCK * D * IN_CHANNELS + + grad_d * D_BLOCK * IN_CHANNELS - IN_CHANNELS + ) + x = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0) + + weight_grad0 += tl.dot(x, grad) + + x_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) + x_iter_ptr = ( + input_ptr + + grad_h * H_BLOCK * W * D * IN_CHANNELS * SPLIT_K + + grad_w * W_BLOCK * D * IN_CHANNELS + + grad_d * D_BLOCK * IN_CHANNELS + ) + x = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0) + + weight_grad1 += tl.dot(x, grad) + + x_mask = make_mask(grad_h * H_BLOCK * SPLIT_K + k_pid * H_BLOCK + block_h - 1, grad_w * W_BLOCK + block_w - 1, grad_d * D_BLOCK + 1, H, W, D, H_BLOCK, W_BLOCK, D_BLOCK) + x_iter_ptr = ( + input_ptr + + grad_h * H_BLOCK * W * D * IN_CHANNELS * SPLIT_K + + grad_w * W_BLOCK * D * IN_CHANNELS + + grad_d * D_BLOCK * IN_CHANNELS + IN_CHANNELS + ) + x = tl.load(x_iter_ptr + input_offset, mask=x_mask[None, :], other=0) + + weight_grad2 += tl.dot(x, grad) + + tl.atomic_add(weight_grad_ptr + weight_grad_offset, weight_grad0, sem='relaxed') + tl.atomic_add(weight_grad_ptr + IN_CHANNELS * OUT_CHANNELS + weight_grad_offset, weight_grad1, sem='relaxed') + tl.atomic_add(weight_grad_ptr + 2 * IN_CHANNELS * OUT_CHANNELS + weight_grad_offset, weight_grad2, sem='relaxed') + + @triton.jit def _ApplyBNReLUConvStats_cl3d_impl( input_ptr, diff --git a/kerops/ops/assets/Conv3dWgrad.toml b/kerops/ops/assets/Conv3dWgrad.toml index c3cfe45..2447f36 100644 --- a/kerops/ops/assets/Conv3dWgrad.toml +++ b/kerops/ops/assets/Conv3dWgrad.toml @@ -19,22 +19,22 @@ problem_size_names = ['in_channels', 'out_channels'] [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [32, 32] - IMPL_ID = 0 + IMPL_ID = 2 SWAP_GRAD_X = false [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [32, 64] - IMPL_ID = 0 + IMPL_ID = 2 SWAP_GRAD_X = false [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [64, 32] - IMPL_ID = 0 + IMPL_ID = 2 SWAP_GRAD_X = true [["NVIDIA GeForce RTX 5070".configs]] problem_sizes = [64, 64] - IMPL_ID = 1 + IMPL_ID = 2 SWAP_GRAD_X = false [["NVIDIA GeForce RTX 5070".configs]] diff --git a/kerops/ops/assets/Conv3dWgrad_splitk_3acc.toml b/kerops/ops/assets/Conv3dWgrad_splitk_3acc.toml new file mode 100644 index 0000000..519578a --- /dev/null +++ b/kerops/ops/assets/Conv3dWgrad_splitk_3acc.toml @@ -0,0 +1,93 @@ +[meta] +problem_size_names = ['in_channels', 'out_channels'] + +["NVIDIA GeForce RTX 5070"] + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 16] + num_warps = 1 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 16 + COUT_BLOCK = 16 + SPLIT_K = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 32] + num_warps = 2 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 16 + COUT_BLOCK = 32 + SPLIT_K = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 16] + num_warps = 2 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 32 + COUT_BLOCK = 16 + SPLIT_K = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 32] + num_warps = 2 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 32 + COUT_BLOCK = 32 + SPLIT_K = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 64] + num_warps = 4 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 32 + COUT_BLOCK = 64 + SPLIT_K = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 32] + num_warps = 4 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 64 + COUT_BLOCK = 32 + SPLIT_K = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 64] + num_warps = 2 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 32 + COUT_BLOCK = 64 + SPLIT_K = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 128] + num_warps = 2 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 16 + COUT_BLOCK = 128 + SPLIT_K = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 64] + num_warps = 2 + H_BLOCK = 2 + WD_BLOCK = 4 + CIN_BLOCK = 32 + COUT_BLOCK = 64 + SPLIT_K = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 128] + num_warps = 4 + H_BLOCK = 4 + WD_BLOCK = 2 + CIN_BLOCK = 128 + COUT_BLOCK = 64 + SPLIT_K = 16 diff --git a/kerops/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py index 236035f..ca92734 100644 --- a/kerops/ops/conv/conv_wgrad.py +++ b/kerops/ops/conv/conv_wgrad.py @@ -2,7 +2,7 @@ from triton import language as tl, next_power_of_2 from ..assets import ASSETS_ROOT -from ...kernels.conv import _Conv_wgrad_cl3d_impl_V2, _Conv_wgrad_cl3d_splitk_impl +from ...kernels.conv import _Conv_wgrad_cl3d_impl_V2, _Conv_wgrad_cl3d_splitk_impl, _Conv_wgrad_cl3d_splitk_3acc_impl from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction from ...utils import cdiv @@ -147,7 +147,79 @@ def Conv3dWgrad_splitk( return weight_grad -conv3d_wgrad = TableKernelConfig( +conv3d_wgrad_splitk_3acc_config = TableKernelConfig( + problem_size_names=['in_channels', 'out_channels'], + confarg_names=['num_warps', 'H_BLOCK', 'WD_BLOCK', 'CIN_BLOCK', 'COUT_BLOCK', 'SPLIT_K'], + args_to_problem_sizes=lambda grad, x: (x.shape[1], grad.shape[1]), + toml_path=ASSETS_ROOT / 'Conv3dWgrad_splitk_3acc.toml' +) + + +@ConfiguredFunction.configure(conv3d_wgrad_splitk_3acc_config) +def Conv3dWgrad_splitk_3acc( + grad, + x, + *, + num_warps: ConfArg, + H_BLOCK: ConfArg, + WD_BLOCK: ConfArg, + CIN_BLOCK: ConfArg, + COUT_BLOCK: ConfArg, + SPLIT_K: ConfArg, +): + assert x.device == grad.device + assert x.is_cuda + + assert x.ndim == grad.ndim == 5 + xbsize, in_channels, xH, xW, xD = x.shape + gbsize, out_channels, gH, gW, gD = grad.shape + assert in_channels == next_power_of_2(in_channels) + assert out_channels == next_power_of_2(out_channels) + assert [xbsize, xH, xW, xD] == [gbsize, gH, gW, gD] + + assert xbsize == gbsize + + assert x.is_contiguous(memory_format=torch.channels_last_3d) + assert grad.is_contiguous(memory_format=torch.channels_last_3d) + + assert x.dtype == grad.dtype == torch.float16 + + assert H_BLOCK == next_power_of_2(H_BLOCK) + assert WD_BLOCK == next_power_of_2(WD_BLOCK) + assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) + assert CIN_BLOCK <= in_channels + assert COUT_BLOCK == next_power_of_2(COUT_BLOCK) + assert COUT_BLOCK <= out_channels + + ACCTYPE = tl.float32 + weight_grad = torch.zeros([xbsize, 3, 3, 3, in_channels, out_channels], device=x.device, dtype=torch.float32) + + grid = ( + 9 * SPLIT_K, + cdiv(in_channels, CIN_BLOCK), + cdiv(out_channels, COUT_BLOCK) * xbsize + ) + + _Conv_wgrad_cl3d_splitk_3acc_impl[grid]( + grad, + x, + weight_grad, + xH, xW, xD, + ACCTYPE=ACCTYPE, + H_BLOCK=H_BLOCK, W_BLOCK=WD_BLOCK, D_BLOCK=WD_BLOCK, + IN_CHANNELS=in_channels, OUT_CHANNELS=out_channels, + CIN_BLOCK=CIN_BLOCK, COUT_BLOCK=COUT_BLOCK, + SPLIT_K=SPLIT_K, + num_warps=num_warps, + ) + + weight_grad = torch.sum(weight_grad, dim=0) + weight_grad = weight_grad.to(torch.float16) + + return weight_grad + + +conv3d_wgrad_config = TableKernelConfig( problem_size_names=['in_channels', 'out_channels'], confarg_names=['IMPL_ID', 'SWAP_GRAD_X'], args_to_problem_sizes=lambda grad, x: (x.shape[1], grad.shape[1]), @@ -155,7 +227,7 @@ def Conv3dWgrad_splitk( ) -@ConfiguredFunction.configure(conv3d_wgrad) +@ConfiguredFunction.configure(conv3d_wgrad_config) def Conv3dWgrad( grad, x, @@ -171,6 +243,8 @@ def Conv3dWgrad( weight_grad = Conv3dWgrad_grad_based(grad, x) case 1: weight_grad = Conv3dWgrad_splitk(grad, x) + case 2: + weight_grad = Conv3dWgrad_splitk_3acc(grad, x) if SWAP_GRAD_X: weight_grad = weight_grad.permute(0, 1, 2, 4, 3) @@ -199,7 +273,7 @@ def pruning_rule(problem_size, named_config): return True -def pruning_rule_splitKonH(problem_size, named_config): +def pruning_rule_splitk(problem_size, named_config): H_BLOCK = named_config['H_BLOCK'] WD_BLOCK = named_config['WD_BLOCK'] CIN_BLOCK = named_config['CIN_BLOCK'] @@ -283,7 +357,7 @@ def autotune_conv_wgrad_grad_based(toml_path, **autotune_kwargs): ) -def autotune_conv_wgrad_splitKonH(toml_path, **autotune_kwargs): +def autotune_conv_wgrad_splitk(toml_path, **autotune_kwargs): channels = [2 ** i for i in range(4, 8)] problem_sizes = [ {'in_channels': cin, 'out_channels': cout} @@ -296,7 +370,7 @@ def autotune_conv_wgrad_splitKonH(toml_path, **autotune_kwargs): getattr(Conv3dWgrad_splitk, 'function', Conv3dWgrad_splitk), generate_inputs_conv_wgrad, problem_sizes, - pruning_rule_splitKonH, + pruning_rule_splitk, toml_path, comparator=comparator, **autotune_kwargs, @@ -310,6 +384,33 @@ def autotune_conv_wgrad_splitKonH(toml_path, **autotune_kwargs): ) +def autotune_conv_wgrad_splitk_3acc(toml_path, **autotune_kwargs): + channels = [2 ** i for i in range(4, 8)] + problem_sizes = [ + {'in_channels': cin, 'out_channels': cout} + for cin in channels + for cout in channels + if (cin == 2 * cout) or (cin * 2 == cout) or (cin == cout) + ] + + autotune( + getattr(Conv3dWgrad_splitk_3acc, 'function', Conv3dWgrad_splitk_3acc), + generate_inputs_conv_wgrad, + problem_sizes, + pruning_rule_splitk, + toml_path, + comparator=comparator, + **autotune_kwargs, + num_warps=[1, 2, 4], + H_BLOCK=[2, 4, 8, 16], + WD_BLOCK=[2, 4, 8, 16], + CIN_BLOCK=channels, + COUT_BLOCK=channels, + SPLIT_K=[8, 16, 32], + n_iters=85 + ) + + def autotune_conv_wgrad(toml_path, **autotune_kwargs): channels = [2 ** i for i in range(4, 8)] problem_sizes = [ @@ -327,6 +428,7 @@ def autotune_conv_wgrad(toml_path, **autotune_kwargs): toml_path=toml_path, comparator=comparator, **autotune_kwargs, - IMPL_ID=[0, 1], - SWAP_GRAD_X=[False, True] + IMPL_ID=[0, 1, 2], + SWAP_GRAD_X=[False, True], + n_iters=200 )