Conv3d, new configurator

.gitignore

@@ -22,3 +22,9 @@ pip-delete-this-directory.txt
 
 # Jupyter Notebook
 .ipynb_checkpoints/
+
+# Profiling
+*.ncu-rep
+PROF/
+
+.vscode

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
@@ -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

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,

kerops/ops/addition.py

@@ -1,15 +1,22 @@
 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 ..settings import ConfArg, StaticKernelConfig, ConfiguredFunction
+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):
+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()
     assert x.shape == y.shape
@@ -24,7 +31,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
@@ -49,10 +56,8 @@ def AddStats(x, y, inplace=False, *, l1_cache_bytes: ConfigurableArg, num_warps:
     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: ConfigurableArg, num_warps: ConfigurableArg
-):
+@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
@@ -68,7 +73,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)
 

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

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

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 = 2
+  SWAP_GRAD_X = false
+
+  [["NVIDIA GeForce RTX 5070".configs]]
+  problem_sizes = [32, 64]
+  IMPL_ID = 2
+  SWAP_GRAD_X = false
+
+  [["NVIDIA GeForce RTX 5070".configs]]
+  problem_sizes = [64, 32]
+  IMPL_ID = 2
+  SWAP_GRAD_X = true
+
+  [["NVIDIA GeForce RTX 5070".configs]]
+  problem_sizes = [64, 64]
+  IMPL_ID = 2
+  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