diff --git a/.gitignore b/.gitignore index f4f8bbe..d5faeb9 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,9 @@ pip-delete-this-directory.txt # Jupyter Notebook .ipynb_checkpoints/ + +# Profiling +*.ncu-rep +PROF/ + +.vscode diff --git a/kerops/kernels/conv.py b/kerops/kernels/conv.py new file mode 100644 index 0000000..a327cb5 --- /dev/null +++ b/kerops/kernels/conv.py @@ -0,0 +1,707 @@ +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 +- Conv3dV6: loading order choice and weight-major alorithm choice + +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 + - TMA (why???) + - tl.swizzle2d and a X-major tile with sizes 2 and 4 +""" + +@triton.jit +def _Conv_cl3d_impl_V6( + 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, + WEIGHT_MAJOR: tl.constexpr, + LOAD_WEIGHT_FIRST: tl.constexpr, +): + W_cell = tl.program_id(0) + H_cell = tl.program_id(1) + 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 + + 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 + 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) + acc10 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + acc11 = tl.zeros([D_BLOCK, OUT_CHANNELS], dtype=ACCTYPE) + + 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) + ] + ] + + 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) + ] + ] + + 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)) + 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)) + + +""" +- 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 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_splitk_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, +): + khwd_pid = tl.program_id(0) + cin_pid = tl.program_id(1) + cout_batch_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 + + 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 += (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) + + 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 + 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 _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, + bn_weight_ptr, + bn_bias_ptr, + weight_ptr, + output_ptr, + mean_ptr, + sqmean_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 + + 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): + 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, :] + + 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) + + # 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), + tl.load(i_ptr + input_offset + IN_CHANNELS * D, mask=mask & m01) + ], + [ + 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) + ] + ] + + 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 + + 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(valid, x, zero) + + # acc00 + if ((h_block * 2 + h) < 3) & ((w_block * 2 + w) < 3): + 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(x, weights[h_block + h][w_block + w - 1]) + + # acc10 + if ((h_block * 2 + h) > 0) & ((w_block * 2 + w) < 3): + 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(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)) + 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/kernels/dw_conv.py b/kerops/kernels/dw_conv.py index 4bcf056..f2a6865 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 @@ -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/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/addition.py b/kerops/ops/addition.py index 532d1b4..52a1f3d 100644 --- a/kerops/ops/addition.py +++ b/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) 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/Conv3dWgrad.toml b/kerops/ops/assets/Conv3dWgrad.toml new file mode 100644 index 0000000..2447f36 --- /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 = 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 diff --git a/kerops/ops/assets/Conv3dWgrad_grad_based.toml b/kerops/ops/assets/Conv3dWgrad_grad_based.toml new file mode 100644 index 0000000..ed9717b --- /dev/null +++ b/kerops/ops/assets/Conv3dWgrad_grad_based.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 = 1 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 16 + COUT_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [16, 32] + num_warps = 1 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 16 + COUT_BLOCK = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 16] + num_warps = 1 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 32 + COUT_BLOCK = 16 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 32] + num_warps = 2 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 32 + COUT_BLOCK = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [32, 64] + num_warps = 2 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 16 + COUT_BLOCK = 64 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 32] + num_warps = 4 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 64 + COUT_BLOCK = 32 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 64] + num_warps = 2 + D_BLOCK = 32 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 32 + COUT_BLOCK = 64 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [64, 128] + num_warps = 4 + D_BLOCK = 16 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 64 + COUT_BLOCK = 128 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 64] + num_warps = 4 + D_BLOCK = 16 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 64 + COUT_BLOCK = 64 + + [["NVIDIA GeForce RTX 5070".configs]] + problem_sizes = [128, 128] + num_warps = 4 + D_BLOCK = 16 + REDUCTION_FACTOR = 32 + CIN_BLOCK = 64 + COUT_BLOCK = 128 diff --git a/kerops/ops/assets/Conv3dWgrad_splitk.toml b/kerops/ops/assets/Conv3dWgrad_splitk.toml new file mode 100644 index 0000000..290eab9 --- /dev/null +++ b/kerops/ops/assets/Conv3dWgrad_splitk.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/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/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/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/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/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/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/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/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/avgpool.py b/kerops/ops/avgpool.py index 19df6e9..a54c6de 100644 --- a/kerops/ops/avgpool.py +++ b/kerops/ops/avgpool.py @@ -1,18 +1,21 @@ 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 ..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 ) -def AvgPoolCeilStats(x, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): + + +@ConfiguredFunction.configure(avgpool_config) +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 @@ -26,7 +29,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 @@ -60,7 +63,13 @@ def AvgPoolCeilStats(x, *, l1_cache_bytes: ConfigurableArg, num_warps: Configura 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, @@ -68,8 +77,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 f51d14a..1701491 100644 --- a/kerops/ops/bnrelu.py +++ b/kerops/ops/bnrelu.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.bnrelu import _ApplyBNReLU_cl3d_backward_impl, _ApplyBNReLU_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 ApplyBNReLU(x, weight, bias, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +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() assert x.ndim == 5 @@ -23,7 +30,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) @@ -41,8 +48,8 @@ def ApplyBNReLU(x, weight, bias, *, l1_cache_bytes: ConfigurableArg, num_warps: return output -@configure(l1_cache_bytes=get_l1_cache, num_warps=8) -def ApplyBNReLUBackward(x, weight, bias, grad, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +@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() assert x.ndim == 5 @@ -59,7 +66,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..82b82b9 100644 --- a/kerops/ops/conv/__init__.py +++ b/kerops/ops/conv/__init__.py @@ -1,2 +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.py b/kerops/ops/conv/conv.py new file mode 100644 index 0000000..882be37 --- /dev/null +++ b/kerops/ops/conv/conv.py @@ -0,0 +1,217 @@ +import numpy as np +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 autotune, ConfArg, TableKernelConfig, ConfiguredFunction +from ...utils import cdiv + + +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' +) + + +@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 + + 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 D_BLOCK == next_power_of_2(D_BLOCK) + assert CIN_BLOCK == next_power_of_2(CIN_BLOCK) + assert CIN_BLOCK <= in_channels + + 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) + + _Conv_cl3d_impl_V6[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, + LOAD_WEIGHT_FIRST=LOAD_WEIGHT_FIRST, + WEIGHT_MAJOR=WEIGHT_MAJOR, + num_warps=num_warps, + ) + + return output + + +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' +) + + +@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 + + 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 + + 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) + sqmean = torch.zeros([bsize, np.prod(grid), out_channels], device=x.device, dtype=torch.float32) + + 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, + D_BLOCK=D_BLOCK, + ACCTYPE=ACCTYPE, + IN_CHANNELS=in_channels, + OUT_CHANNELS=out_channels, + CIN_BLOCK=CIN_BLOCK, + num_warps=num_warps, + ) + + 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/ops/conv/conv_wgrad.py b/kerops/ops/conv/conv_wgrad.py new file mode 100644 index 0000000..ca92734 --- /dev/null +++ b/kerops/ops/conv/conv_wgrad.py @@ -0,0 +1,434 @@ +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, _Conv_wgrad_cl3d_splitk_impl, _Conv_wgrad_cl3d_splitk_3acc_impl +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction +from ...utils import cdiv + + +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]), + toml_path=ASSETS_ROOT / 'Conv3dWgrad_grad_based.toml' +) + + +@ConfiguredFunction.configure(conv3d_wgrad_grad_based_config) +def Conv3dWgrad_grad_based( + grad, + x, + *, + num_warps: ConfArg, + D_BLOCK: ConfArg, + REDUCTION_FACTOR: ConfArg, + CIN_BLOCK: ConfArg, + COUT_BLOCK: 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 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 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 = 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) + + _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) + + 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 / 'Conv3dWgrad_splitk.toml' +) + + +@ConfiguredFunction.configure(conv3d_wgrad_splitk_config) +def Conv3dWgrad_splitk( + 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 = ( + 27 * SPLIT_K, + cdiv(in_channels, CIN_BLOCK), + cdiv(out_channels, COUT_BLOCK) * xbsize + ) + + _Conv_wgrad_cl3d_splitk_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_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]), + toml_path=ASSETS_ROOT / 'Conv3dWgrad.toml' +) + + +@ConfiguredFunction.configure(conv3d_wgrad_config) +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) + case 2: + weight_grad = Conv3dWgrad_splitk_3acc(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'] + COUT_BLOCK = named_config['COUT_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 or COUT_BLOCK > out_channels: + return False + + return True + + +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'] + 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 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 + + +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 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_grad_based(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_grad_based, 'function', Conv3dWgrad_grad_based), + generate_inputs_conv_wgrad, + problem_sizes, + pruning_rule, + toml_path, + comparator=comparator, + **autotune_kwargs, + num_warps=[1, 2, 4], + D_BLOCK=[16, 32], + REDUCTION_FACTOR=[32], + CIN_BLOCK=channels, + COUT_BLOCK=channels, + ) + + +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} + for cin in channels + for cout in channels + if (cin == 2 * cout) or (cin * 2 == cout) or (cin == cout) + ] + + autotune( + getattr(Conv3dWgrad_splitk, 'function', Conv3dWgrad_splitk), + 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=[4, 8, 16], + n_iters=85 + ) + + +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 = [ + {'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, 2], + SWAP_GRAD_X=[False, True], + n_iters=200 + ) diff --git a/kerops/ops/conv/dwconv.py b/kerops/ops/conv/dwconv.py index 3827c19..aa51926 100644 --- a/kerops/ops/conv/dwconv.py +++ b/kerops/ops/conv/dwconv.py @@ -1,28 +1,22 @@ -from math import ceil - 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 ConfigurableArg, confexc, configure - - -@confexc(KeyError) -def warps(channels): - return {8: 1, 16: 2, 32: 2, 64: 2, 128: 4}[channels] +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction +from ...utils import cdiv -@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: ConfigurableArg, num_warps: ConfigurableArg, D_block: ConfigurableArg): +@ConfiguredFunction.configure(dwconv_config) +def DWConv(x, weight, *, num_warps: ConfArg, D_BLOCK: ConfArg): channels = x.shape[1] assert x.ndim == 5 @@ -30,18 +24,17 @@ def DWConv(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D 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() 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): @@ -56,8 +49,53 @@ def DWConv(x, weight, *, ACCTYPE: ConfigurableArg, num_warps: ConfigurableArg, D 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/kerops/ops/conv/dwconv_wgrad.py b/kerops/ops/conv/dwconv_wgrad.py index 8a0848c..d7de08a 100644 --- a/kerops/ops/conv/dwconv_wgrad.py +++ b/kerops/ops/conv/dwconv_wgrad.py @@ -1,35 +1,23 @@ -from math import ceil - 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 ConfigurableArg, confexc, configure - - -@confexc(KeyError) -def warps(channels): - return {8: 1, 16: 1, 32: 1, 64: 1, 128: 2}[channels] +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction +from ...utils import cdiv -@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: ConfigurableArg, num_warps: ConfigurableArg, D_block: ConfigurableArg, ILP: ConfigurableArg + x, grad, *, num_warps: ConfArg, D_BLOCK: ConfArg, ILP: ConfArg ): channels = x.shape[1] @@ -39,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 = 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) @@ -66,7 +53,7 @@ def DWConvWGRAD( W_stride, ACCTYPE, channels, - D_block, + D_BLOCK, WD_grid, D_grid, H_grid, @@ -77,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/kerops/ops/linear/linear_bias_relu_linear_add.py b/kerops/ops/linear/linear_bias_relu_linear_add.py index a0f425d..7d048c7 100644 --- a/kerops/ops/linear/linear_bias_relu_linear_add.py +++ b/kerops/ops/linear/linear_bias_relu_linear_add.py @@ -1,29 +1,21 @@ -from math import ceil - import torch from triton import next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.linear import _LinBReLULinAdd -from ...settings import ConfigurableArg, confexc, configure - - -@confexc(KeyError) -def dblock(channels): - # '64: 32' removed, fix bwd func - return {16: 32, 32: 16}[channels] +from ...settings import autotune, ConfArg, TableKernelConfig, ConfiguredFunction +from ...utils import cdiv -@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,9 +23,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 @@ -50,7 +42,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) @@ -64,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/ops/linear/linear_bias_relu_linear_backward.py b/kerops/ops/linear/linear_bias_relu_linear_backward.py index 35b70af..9170329 100644 --- a/kerops/ops/linear/linear_bias_relu_linear_backward.py +++ b/kerops/ops/linear/linear_bias_relu_linear_backward.py @@ -1,22 +1,22 @@ -from math import ceil - import torch from triton import next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.linear import _LinBReLULinBackward -from ...settings import ConfigurableArg, 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,9 +24,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 @@ -43,7 +43,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') @@ -63,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/kerops/ops/linear/relu_linear_add.py b/kerops/ops/linear/relu_linear_add.py index 5d828c1..4bdb958 100644 --- a/kerops/ops/linear/relu_linear_add.py +++ b/kerops/ops/linear/relu_linear_add.py @@ -1,30 +1,29 @@ -from math import ceil - import torch from triton import next_power_of_2 +from ..assets import ASSETS_ROOT from ...kernels.linear import _ReLULinearAdd -from ...settings import ConfigurableArg, 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: 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] @@ -35,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 = ceil(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/kerops/ops/linear/relu_linear_backward.py b/kerops/ops/linear/relu_linear_backward.py index a8994eb..8c1de2e 100644 --- a/kerops/ops/linear/relu_linear_backward.py +++ b/kerops/ops/linear/relu_linear_backward.py @@ -1,35 +1,29 @@ -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 ..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: ConfigurableArg, - D_block: ConfigurableArg, - ILP: ConfigurableArg, + num_warps: ConfArg, + D_BLOCK: ConfArg, + ILP: ConfArg, ): in_channels = x.shape[1] out_channels = grad.shape[1] @@ -40,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) @@ -47,13 +42,13 @@ 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) 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, @@ -62,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/kerops/ops/quantization.py b/kerops/ops/quantization.py index d5054b5..62fa55f 100644 --- a/kerops/ops/quantization.py +++ b/kerops/ops/quantization.py @@ -1,27 +1,34 @@ -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 ..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: ConfigurableArg, l1_cache_bytes: ConfigurableArg): +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) 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 -@configure(num_warps=4, l1_cache_bytes=get_l1_cache) -def DequantUint8Window(x, init_dtype, window, num_warps: ConfigurableArg, l1_cache_bytes: ConfigurableArg): +@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) @@ -29,7 +36,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..8ccb538 100644 --- a/kerops/ops/stats.py +++ b/kerops/ops/stats.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.stats import _Stats_cl3d_backward_impl, _Stats_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=4) -def Stats(x, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +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() assert x.ndim == 5 @@ -22,7 +29,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) @@ -31,8 +38,8 @@ def Stats(x, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): return mean, sqmean -@configure(l1_cache_bytes=get_l1_cache, num_warps=4) -def StatsBackward(x, mean_grad, sqmean_grad, *, l1_cache_bytes: ConfigurableArg, num_warps: ConfigurableArg): +@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() assert x.ndim == 5 @@ -46,7 +53,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/settings/__init__.py b/kerops/settings/__init__.py index 1974121..f54d629 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 .wrapper import confexc, configure +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 new file mode 100644 index 0000000..216cd74 --- /dev/null +++ b/kerops/settings/autotune.py @@ -0,0 +1,187 @@ +import itertools +import traceback as tb +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 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], + 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 = 100, + quantiles: tuple = (20, 80), + comparator: Callable | None = None, + **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): + 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 = [] + 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)))] + + 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")] + 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"]) + + 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"], + }) + + 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/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..d6f10f5 --- /dev/null +++ b/kerops/settings/kernel_config.py @@ -0,0 +1,208 @@ +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 .utils import get_device_name_from_args + + +class KernelConfigBase(ABC): + @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 __call__(self, *input_args) -> dict[str, Any]: + device = 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 index 668cbe6..a23b71c 100644 --- a/kerops/settings/utils.py +++ b/kerops/settings/utils.py @@ -1,34 +1,13 @@ -from inspect import Parameter +from torch import Tensor +from torch.cuda import get_device_name -class ConfigurableArg: - pass +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'} - -class CongiguratorError(Exception): - pass - - -def validate_signature(signature): - for param in signature.parameters.values(): - if param.annotation is ConfigurableArg 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: - 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] - - -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=}') + 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 914b97e..9bb3c7b 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 KernelConfigBase -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: KernelConfigBase): + 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: KernelConfigBase): + 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: KernelConfigBase): + 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}' 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 diff --git a/notebooks/config_builder.ipynb b/notebooks/config_builder.ipynb new file mode 100644 index 0000000..ee62027 --- /dev/null +++ b/notebooks/config_builder.ipynb @@ -0,0 +1,224 @@ +{ + "cells": [ + { + "cell_type": "code", + "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_backward import ReLULinearBackward, autotune_relu_lin_backward, generate_inputs_relu_lin_backward\n", + "from kerops.ops.assets import ASSETS_ROOT" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "6e35524d-0c68-4fd8-859c-fb581cbd96f8", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "869514dd7ec64fb2aa9df82d9adef393", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Problem sizes: 0%| | 0/3 [00:00= 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, 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", + " 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}%')\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 * '-')" + ] + }, + { + "cell_type": "markdown", + "id": "432e294a-c563-4dad-ad5c-b03093d2a61f", + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, + "source": [ + "### Conv3dWgradV1" + ] + }, + { + "cell_type": "code", + "execution_count": 69, + "id": "4bc6cd74-9594-47d6-a6f7-ed5ee36373de", + "metadata": {}, + "outputs": [], + "source": [ + "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", + " 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": "378154ee-095c-4e04-b314-db73ad4f0b7f", + "metadata": {}, + "outputs": [], + "source": [ + "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": "markdown", + "id": "2b12f8ca-7a33-471e-b6f8-0d868e355b9a", + "metadata": {}, + "source": [ + "### Conv3dWgradV2" + ] + }, + { + "cell_type": "code", + "execution_count": 150, + "id": "ec8a8fa1-6afe-4a25-9ca9-54b21816edcb", + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "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": 151, + "id": "d28d18e2-1ac4-4a11-bf12-6e64397080f9", + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "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": 3, + "id": "4e6da0ab-c024-495f-92e6-d17143eb928a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "All stages passed\n" + ] + }, + { + "data": { + "text/plain": [ + "True" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from kerops.utils import weight_grad_similarity\n", + "\n", + "S = 128\n", + "D = 128\n", + "CIN = 16\n", + "COUT = 16\n", + "BS = 1\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", + "\n", + "#grad_w = Conv3dWgradV2(grad, x, **{'D_BLOCK': 32, 'num_warps': 4, 'ACCTYPE': 'float32', 'REDUCTION_FACTOR': 32, 'CIN_BLOCK': 16, 'COUT_BLOCK': 64})\n", + "#grad_w = Conv3dWgradV2(grad, x, D_BLOCK=32, ACCTYPE='float32', num_warps=2, REDUCTION_FACTOR=32, CIN_BLOCK=32, COUT_BLOCK=32)\n", + "grad_w = Conv3dWgrad(grad, x)\n", + "\n", + "_, grad_w_reference, _ = 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", + "\n", + "grad_w_p = grad_w.permute(-1, -2, 0, 1, 2)\n", + "\n", + "weight_grad_similarity(grad_w_reference.float(), grad_w_p.float(), rtol_cos=0.0001, debug_info='print')" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "03e4903b-95f6-44ae-8a16-bbd16d022fff", + "metadata": {}, + "outputs": [], + "source": [ + "BS = 1\n", + "S = 128\n", + "D = 128\n", + "CIN = 64\n", + "COUT = 64\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()" + ] + }, + { + "cell_type": "code", + "execution_count": 51, + "id": "d2b605ae-3608-4a38-86c5-94a936062539", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CIN=16 COUT=16\n", + "Speedup - 1.748\n", + "------------------------------\n", + "CIN=16 COUT=32\n", + "Speedup - 1.207\n", + "------------------------------\n", + "CIN=32 COUT=16\n", + "Speedup - 1.690\n", + "------------------------------\n", + "CIN=32 COUT=32\n", + "Speedup - 1.112\n", + "------------------------------\n", + "CIN=32 COUT=64\n", + "Speedup - 1.074\n", + "------------------------------\n", + "CIN=64 COUT=32\n", + "Speedup - 1.136\n", + "------------------------------\n", + "CIN=64 COUT=64\n", + "Speedup - 0.952\n", + "------------------------------\n", + "CIN=64 COUT=128\n", + "Speedup - 1.003\n", + "------------------------------\n", + "CIN=128 COUT=64\n", + "Speedup - 0.951\n", + "------------------------------\n", + "CIN=128 COUT=128\n", + "Speedup - 0.960\n", + "------------------------------\n" + ] + } + ], + "source": [ + "channels = [16, 32, 64, 128]\n", + "\n", + "with_fwd = True\n", + "with_bwd = True\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", + " BS = 1\n", + " S = 128\n", + " D = 128\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", + " for _ in range(5):\n", + " if with_fwd:\n", + " nn.functional.conv3d(x, weight, None, stride=1, padding=1)\n", + " if with_bwd:\n", + " _, grad_w, _ = 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()\n", + "\n", + " times_ms_cudnn = []\n", + " \n", + " for _ in range(25):\n", + " start = perf_counter()\n", + " if with_fwd:\n", + " nn.functional.conv3d(x, weight, None, stride=1, padding=1)\n", + " if with_bwd:\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()\n", + " end = perf_counter()\n", + " times_ms_cudnn.append((end - start) * 1e3)\n", + "\n", + " for _ in range(5):\n", + " if with_fwd:\n", + " Conv3d(x, w)\n", + " if with_bwd:\n", + " Conv3dWgrad(grad, x)\n", + " torch.cuda.synchronize()\n", + "\n", + " times_ms_triton = []\n", + " \n", + " for _ in range(25):\n", + " start = perf_counter()\n", + " if with_fwd:\n", + " Conv3d(x, w)\n", + " if with_bwd:\n", + " Conv3dWgrad(grad, x)\n", + " torch.cuda.synchronize()\n", + " end = perf_counter()\n", + " times_ms_triton.append((end - start) * 1e3)\n", + "\n", + " speedup = np.mean(times_ms_cudnn) / np.mean(times_ms_triton)\n", + "\n", + " print(f'Speedup - {speedup:.3f}')\n", + " \n", + " print(30 * '-')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e87ebad3-3306-4ce2-962b-08e441b4cc57", + "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 +} diff --git a/requirements.txt b/requirements.txt index b24cb10..84ba852 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -triton==3.1.0 +triton>=3.6.0 torch diff --git a/tests/conftest.py b/tests/conftest.py index 345e0b7..f68b91f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,21 +11,31 @@ 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 + + +@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 index 002a8b0..224392a 100644 --- a/tests/ops/test_conv.py +++ b/tests/ops/test_conv.py @@ -4,26 +4,27 @@ from torch import nn from torch.nn import functional as F -from kerops.ops.conv import DWConv, DWConvWGRAD +from kerops.ops.conv import Conv3d, Conv3dWgrad from kerops.utils import allclose_two_stage, weight_grad_similarity -def test_dwconv(bsize, channels, other_1, other_2, other_3): - if channels < 8: +# 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 - + torch.manual_seed(322) - x = torch.randn(bsize, channels, other_1, other_2, other_3, device='cuda', dtype=torch.float16).to( + 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(channels, 1, 3, 3, 3, device='cuda', dtype=torch.float16) + 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_ker = weight[:, 0].permute(1, 2, 3, 0).contiguous() + 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), groups=channels) - out_check = DWConv(x, weight_ker) + 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, @@ -38,47 +39,37 @@ def test_dwconv(bsize, channels, other_1, other_2, other_3): torch.cuda.empty_cache() -# TODO: how to sample to make weight's gradient adequate? -def test_dwconv_wgrad(bsize, channels, other_1, other_2, other_3): - # DWConv from torch is EXTREAMLY slow - if channels < 8 or bsize > 3 or other_1 > 53 or other_2 > 53 or other_3 > 53: +def test_conv_dwgrad(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, channels, other_1, other_2, other_3, device='cuda', dtype=torch.float16).to( + + torch.manual_seed(1337) + 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 ) - grad = torch.randn(bsize, channels, other_1, other_2, other_3, device='cuda', dtype=torch.float16).to( + grad = torch.randn(bsize, conv_out_channels, other_1, other_2, other_3, device='cuda', dtype=torch.float16).to( memory_format=torch.channels_last_3d ) - weight = torch.empty(channels, 1, 3, 3, 3, device='cuda', dtype=torch.float32) + 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 = weight[:, 0].permute(1, 2, 3, 0).contiguous() - weight.requires_grad_(True) - - grad_w_check = DWConvWGRAD(x, grad).to(torch.float32) - with torch.amp.autocast('cuda'): - out = F.conv3d( - x, - weight.permute(3, 0, 1, 2)[:, None].contiguous(), - None, - stride=(1, 1, 1), - padding=(1, 1, 1), - groups=channels, - ) - - out.backward(grad) - - assert weight_grad_similarity( - weight.grad, - grad_w_check, - rtol_cos=1e-4, - atol_cos=1e-4, - rtol_len=1e-3 * bsize, - atol_len=1e-3, - debug_info='print', + _, wgrad_standard, _ = torch.ops.aten.convolution_backward( + grad, + x, + 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 ) + wgrad_check = Conv3dWgrad(grad, x) + wgrad_check = wgrad_check.permute(-1, -2, 0, 1, 2) + + assert weight_grad_similarity(wgrad_standard.float(), wgrad_check.float()) torch.cuda.empty_cache() diff --git a/tests/ops/test_dwconv.py b/tests/ops/test_dwconv.py new file mode 100644 index 0000000..002a8b0 --- /dev/null +++ b/tests/ops/test_dwconv.py @@ -0,0 +1,84 @@ +import math + +import torch +from torch import nn +from torch.nn import functional as F + +from kerops.ops.conv import DWConv, DWConvWGRAD +from kerops.utils import allclose_two_stage, weight_grad_similarity + + +def test_dwconv(bsize, channels, other_1, other_2, other_3): + if channels < 8: + return + + torch.manual_seed(322) + x = torch.randn(bsize, channels, other_1, other_2, other_3, device='cuda', dtype=torch.float16).to( + memory_format=torch.channels_last_3d + ) + + weight = torch.empty(channels, 1, 3, 3, 3, device='cuda', dtype=torch.float16) + nn.init.kaiming_uniform_(weight, a=math.sqrt(5)) + + weight_ker = weight[:, 0].permute(1, 2, 3, 0).contiguous() + + out_standard = F.conv3d(x, weight, None, padding=(1, 1, 1), stride=(1, 1, 1), groups=channels) + out_check = DWConv(x, weight_ker) + + 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() + + +# TODO: how to sample to make weight's gradient adequate? +def test_dwconv_wgrad(bsize, channels, other_1, other_2, other_3): + # DWConv from torch is EXTREAMLY slow + if channels < 8 or bsize > 3 or other_1 > 53 or other_2 > 53 or other_3 > 53: + return + + torch.manual_seed(322) + x = torch.randn(bsize, channels, other_1, other_2, other_3, device='cuda', dtype=torch.float16).to( + memory_format=torch.channels_last_3d + ) + grad = torch.randn(bsize, channels, other_1, other_2, other_3, device='cuda', dtype=torch.float16).to( + memory_format=torch.channels_last_3d + ) + + weight = torch.empty(channels, 1, 3, 3, 3, device='cuda', dtype=torch.float32) + nn.init.kaiming_uniform_(weight, a=math.sqrt(5)) + weight = weight[:, 0].permute(1, 2, 3, 0).contiguous() + weight.requires_grad_(True) + + grad_w_check = DWConvWGRAD(x, grad).to(torch.float32) + + with torch.amp.autocast('cuda'): + out = F.conv3d( + x, + weight.permute(3, 0, 1, 2)[:, None].contiguous(), + None, + stride=(1, 1, 1), + padding=(1, 1, 1), + groups=channels, + ) + + out.backward(grad) + + assert weight_grad_similarity( + weight.grad, + grad_w_check, + rtol_cos=1e-4, + atol_cos=1e-4, + rtol_len=1e-3 * bsize, + atol_len=1e-3, + debug_info='print', + ) + + torch.cuda.empty_cache()