Day 10 - Improving Conv2d performance

import pandas as pd
import numpy as np
from math import prod
from PIL import Image
from pathlib import Path
from tqdm.auto import tqdm

from lovely_numpy import Lo
from lovely_tensors import monkey_patch; monkey_patch()
from torch import Tensor
from torch.nn.functional import conv2d

import pycuda.driver as cuda
from pycuda.compiler import SourceModule
cuda.init()

device = cuda.Device(0)

print(f"Cuda version: {".".join([str(i) for i in cuda.get_version()])}")
print(f"Device:\t{device.name()}")

Cuda version: 12.8.0
Device: NVIDIA GeForce RTX 3080 Laptop GPU

cu_file="day_10_conv2d-experiments.cu"

day_10_conv2d-experiments.cu

#include <stdint.h>
#include <stdio.h>

/* 2D convolution, with padding to valid shape. Channel-first */
__global__ void conv2d_pad(float *in,
                           float *out,
                           float *filter,
                           int h,
                           int w,
                           int in_channels,
                           int out_channels,
                           int filter_size /* Must be an odd number */,
                           float pad) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    int filter_r = (filter_size - 1) / 2;

    // In and Out data dimensions:
    // 0 - channel
    // 1 - height
    // 2 - width

    // Filter dimensions:
    // 0 - out channels
    // 1 - in channels
    // 2 - height
    // 3 - width

    // if (x == 0 && y == 0) {
    //     printf("h: %d\n", h);
    //     printf("w: %d\n", w);
    //     printf("in_channels: %d\n", in_channels);
    //     printf("out_channels: %d\n", out_channels);
    //     printf("filter_size: %d\n", filter_size);
    //     printf("filter r: %d\n", filter_r);
    //     printf("pad: %f\n", pad);

    //     printf("Filter:\n");
    //     for (int oc = 0; oc < out_channels; oc++) {
    //         printf("Output channel %d:\n", oc);
    //         for (int ic = 0; ic < in_channels; ic++) {
    //             printf("  Input channel %d:\n", ic);
    //             float *sub_filter = filter + (filter_size * filter_size * in_channels * oc) +
    //                                 (filter_size * filter_size * ic);
    //             for (int i = 0; i < filter_size; i++) {
    //                 printf("    ");
    //                 for (int j = 0; j < filter_size; j++) {
    //                     printf("%f ", sub_filter[i * filter_size + j]);
    //                 }
    //                 printf("\n");
    //             }
    //         }
    //     }
    // }

    if (x >= w || y >= h) return;

    // Loop over the output channels
    for (int out_c = 0; out_c < out_channels; out_c++) {
        float R = 0;

        // Pointer to the 2d slice of the output
        float *sub_output = out + out_c * w * h;

        // Loop over the input channels
        for (int in_c = 0; in_c < in_channels; in_c++) {
            // Pointer to the 2d slice of the filter that corresponds to the active input and output
            // channels
            float *sub_filter = filter + (filter_size * filter_size * in_channels * out_c) +
                                (filter_size * filter_size * in_c);
            // Pinter to the current channel in the input
            float *sub_input = in + (w * h * in_c);

            // Apply the filter to the input or the pad value for outside indices.
            for (int filter_y = 0; filter_y < filter_size; filter_y++) {
                for (int filter_x = 0; filter_x < filter_size; filter_x++) {
                    float v = pad;
                    int input_x = x - filter_r + filter_x;
                    int input_y = y - filter_r + filter_y;

                    if (input_x >= 0 && input_x < w && input_y >= 0 && input_y < h) {
                        v = sub_input[input_y * w + input_x];
                    }
                    R += v * sub_filter[filter_y * filter_size + filter_x];
                }
            }
        }
        sub_output[y * w + x] = R;
    }
}

/* This version uses the z grid dimensions for out channels. */

__global__ void conv2d_pad_z_out(float *in,
                                 float *out,
                                 float *filter,
                                 int h,
                                 int w,
                                 int in_channels,
                                 int out_channels,
                                 int filter_size /* Must be an odd number */,
                                 float pad) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    int out_ch = blockIdx.z;

    int filter_r = (filter_size - 1) / 2;

    // In and Out data dimensions:
    // 0 - channel
    // 1 - height
    // 2 - width

    // Filter dimensions:
    // 0 - out channels
    // 1 - in channels
    // 2 - height
    // 3 - width

    if (x >= w || y >= h) return;

    // Loop over the output channels

    float R = 0;

    // // Pointer to the 2d slice of the output
    float *sub_output = out + out_ch * w * h;

    // Loop over the input channels
    for (int in_c = 0; in_c < in_channels; in_c++) {
        // Pointer to the 2d slice of the filter that corresponds to the active input and output
        // channels
        float *sub_filter = filter + (filter_size * filter_size * in_channels * out_ch) +
                            (filter_size * filter_size * in_c);
        // Pinter to the current channel in the input
        float *sub_input = in + (w * h * in_c);

        // Apply the filter to the input or the pad value for outside indices.
        for (int filter_y = 0; filter_y < filter_size; filter_y++) {
            for (int filter_x = 0; filter_x < filter_size; filter_x++) {
                float v = pad;
                int input_x = x - filter_r + filter_x;
                int input_y = y - filter_r + filter_y;

                if (input_x >= 0 && input_x < w && input_y >= 0 && input_y < h) {
                    v = sub_input[input_y * w + input_x];
                }
                R += v * sub_filter[filter_y * filter_size + filter_x];
            }
        }
    }
    sub_output[y * w + x] = R;
}

## Compiler options for more compile-time warnings.
warn_options=[
    '-Xcompiler', '-Wall',
    '-Xcompiler', '-Wextra',
    '-Xcompiler', '-Wsign-conversion',
    '-Xcompiler', '-Wcast-qual',
    '-Xcompiler', '-Wunused-parameter',
    '-Xcompiler', '-Wdouble-promotion',
    '-Xcompiler', '-Wformat=2',
    '-Xcompiler', '-Wfloat-equal',
    '-Xcompiler', '-Wshadow'
]

def benchmark_conv2d_pad(ctx, kernel, input, filter, pad, block_size, grid_size, repeat=10, warmup=True):
    # input, channel-first
    # - Channel
    # - Height
    # - Width
    assert len(input.shape) == 3

    # Filter shape should be
    # - Out channels
    # - In  channels
    # - Height
    # - Width
    assert len(filter.shape) == 4

    in_ch, h, w = input.shape
    out_ch, in_ch2, fh, fw = filter.shape

    assert fh == fw, f"Only square filters supported, got shape={filter.shape}"

    assert in_ch == in_ch2

    out_shape = (out_ch, h, w)

    gpu_input = cuda.mem_alloc_like(input)
    gpu_filter = cuda.mem_alloc_like(filter)

    out = np.empty(out_shape, dtype=np.float32)

    cuda.memcpy_htod(gpu_input, input)
    cuda.memcpy_htod(gpu_filter, filter)
    ctx.synchronize()

    timing=0
    for _ in range(repeat):
        start = cuda.Event()
        end = cuda.Event()

        gpu_out = cuda.mem_alloc_like(out)

        if warmup:
            kernel(gpu_input, gpu_out, gpu_filter,
                   np.int32(h),
                   np.int32(w),
                   np.int32(in_ch),
                   np.int32(out_ch),
                   np.int32(fh),
                   np.float32(pad),
                   grid=grid_size,
                   block=block_size)
            ctx.synchronize()

        start.record()
        kernel(gpu_input, gpu_out, gpu_filter,
               np.int32(h),
               np.int32(w),
               np.int32(in_ch),
               np.int32(out_ch),
               np.int32(fh),
               np.float32(pad),
               grid=grid_size,
               block=block_size)
        end.record()
        end.synchronize()

        timing += end.time_since(start)
    timing /= repeat

    cuda.memcpy_dtoh(out, gpu_out)
    return out, timing;

in_chan_range = [1, 3, 8, 32, 128, 512]
out_chan_range = [1, 4, 8, 32, 128, 512]

filter_size = [1, 3, 5, 9]

img_size_range = [64, 128, 256, 512, 1024]

kernel_range = ["conv2d_pad", "conv2d_pad_z_out"]

# Let's sample from the available options.
n_samples = 50


# Generate all possible combinations
combinations = []
for in_ch in in_chan_range:
    for out_ch in out_chan_range:
        for fs in filter_size:
            for img_size in img_size_range:
                    n = in_ch * out_ch * img_size * img_size

                    # Skip combinatoins that are too large
                    if n < 1024*1024*32*32:
                        combinations.append((in_ch, out_ch, fs, img_size))

n_samples = min(n_samples, len(combinations))
sampled_combinations = np.random.choice(len(combinations), size=n_samples, replace=False)
test_cases = [combinations[i] for i in sampled_combinations]

tile_width = 32

data = []

# test_cases = [(3, 4, 32,32)]

ctx = device.make_context()
try:
    mod = SourceModule(
        Path(cu_file).read_text(),
        options=warn_options)


    kernels = {
        k: mod.get_function(k) for k in ["conv2d_pad", "conv2d_pad_z_out"]
    }

    for tc in tqdm(test_cases):
        ch_in, ch_out, fs, pixels = tc

        array_in = np.random.randn(ch_in, pixels, pixels).astype(np.float32)
        filter = np.random.randn(ch_out, ch_in, fs, fs).astype(np.float32)

        torch_out = conv2d(Tensor(array_in), Tensor(filter), padding="same")

        timings = {}

        for kernel_name, kernel in kernels.items():

            block_size = (tile_width,tile_width,1)
            grid_size = (
                ((pixels + tile_width - 1) // tile_width),
                ((pixels + tile_width - 1) // tile_width),
                ch_out if kernel_name == "conv2d_pad_z_out" else 1
            )


            out, timing = benchmark_conv2d_pad(ctx, kernel, array_in, filter, 0, block_size, grid_size, repeat=3, warmup=True)

            if np.isclose(out, torch_out).mean() < 0.9:
                print("### Result mismatch ###")
                print(f"Kernel: {kernel_name}")
                print(f"Input shape: {array_in.shape}")
                print(f"Filter shape: {filter.shape}")
                print(f"Result shape: {(filter.shape[0], array_in.shape[1], array_in.shape[2])}")
                print(f"Grid size: {grid_size}")
                print(f"Block size: {block_size}")
                print(f"Total threads: {prod((*grid_size, *block_size))}")

            timings[kernel_name] = timing

        data.append({
            'in_ch': ch_in,
            'out_ch': ch_out,
            'filter_size': fs,
            'img_size': pixels,
            # 'kernel': kernel_name,
        } | timings)


finally:
    ctx.pop()
    ctx.detach()

results = pd.DataFrame(data)

# Sort by conv2d_pad timing
results_sorted = results.sort_values(by='conv2d_pad')

# Create a plot comparing the two kernels
import matplotlib.pyplot as plt
import seaborn as sns

plt.figure(figsize=(12, 8))
sns.set_style("whitegrid")

# Create labels for x-axis that include dimensions
results_sorted['dimensions'] = results_sorted.apply(
    lambda row: f"{int(row['img_size'])}×{int(row['img_size'])}×{int(row['in_ch'])} -> {int(row['out_ch'])}, f:{int(row['filter_size'])}×{int(row['filter_size'])}",
    axis=1
)

# Melt the dataframe to get it in the right format for seaborn
melted_results = pd.melt(
    results_sorted,
    id_vars=['in_ch', 'out_ch', 'filter_size', 'img_size', 'dimensions'],
    value_vars=['conv2d_pad', 'conv2d_pad_z_out'],
    var_name='kernel',
    value_name='time'
)

# Create a barplot with dimensions as x-ticks
ax = sns.barplot(x='dimensions', y='time', hue='kernel', data=melted_results)

# Add labels
plt.xlabel('Input and Filter Dimensions')
plt.ylabel('Time (seconds)')
plt.title('Performance Comparison of conv2d_pad vs conv2d_pad_z_out')
plt.xticks(rotation=90)

# Add a legend
plt.legend(title='Kernel')

# Show the plot
plt.tight_layout()
plt.show()

# Also display the sorted results table
results_sorted

	in_ch	out_ch	filter_size	img_size	conv2d_pad	conv2d_pad_z_out	dimensions
37	1	1	3	256	0.014720	0.015701	256×256×1 -> 1, f:3×3
39	1	4	3	64	0.017749	0.015360	64×64×1 -> 4, f:3×3
45	8	1	1	256	0.018432	0.019456	256×256×8 -> 1, f:1×1
24	1	1	1	512	0.022528	0.016043	512×512×1 -> 1, f:1×1
2	3	1	3	256	0.024587	0.025227	256×256×3 -> 1, f:3×3
9	1	1	9	256	0.034133	0.035211	256×256×1 -> 1, f:9×9
22	1	8	1	512	0.036864	0.050528	512×512×1 -> 8, f:1×1
36	1	1	5	512	0.037888	0.033451	512×512×1 -> 1, f:5×5
14	3	4	1	512	0.046763	0.057003	512×512×3 -> 4, f:1×1
28	1	32	1	256	0.049152	0.049835	256×256×1 -> 32, f:1×1
0	1	32	3	64	0.056320	0.021163	64×64×1 -> 32, f:3×3
10	32	1	3	64	0.058709	0.056320	64×64×32 -> 1, f:3×3
4	8	4	3	128	0.060075	0.038571	128×128×8 -> 4, f:3×3
43	32	1	5	64	0.107861	0.103424	64×64×32 -> 1, f:5×5
20	32	4	1	256	0.112299	0.121173	256×256×32 -> 4, f:1×1
3	1	128	1	256	0.138240	0.177835	256×256×1 -> 128, f:1×1
31	3	8	5	256	0.140288	0.112619	256×256×3 -> 8, f:5×5
13	3	128	1	128	0.140288	0.080555	128×128×3 -> 128, f:1×1
30	3	8	9	64	0.182613	0.033792	64×64×3 -> 8, f:9×9
44	8	8	5	64	0.184320	0.038229	64×64×8 -> 8, f:5×5
11	128	1	3	128	0.199339	0.201387	128×128×128 -> 1, f:3×3
1	8	8	5	128	0.212992	0.091136	128×128×8 -> 8, f:5×5
40	8	4	9	64	0.241664	0.069632	64×64×8 -> 4, f:9×9
47	512	1	1	128	0.244053	0.250539	128×128×512 -> 1, f:1×1
23	32	1	9	64	0.255317	0.242347	64×64×32 -> 1, f:9×9
21	1	4	5	1024	0.268971	0.349867	1024×1024×1 -> 4, f:5×5
32	1	32	1	1024	0.386688	0.860160	1024×1024×1 -> 32, f:1×1
26	8	8	9	128	0.479573	0.186368	128×128×8 -> 8, f:9×9
38	1	32	5	512	0.537600	0.596309	512×512×1 -> 32, f:5×5
6	3	8	3	1024	0.596651	1.027040	1024×1024×3 -> 8, f:3×3
12	3	32	9	128	0.694272	0.252245	128×128×3 -> 32, f:9×9
35	8	4	3	1024	0.829440	1.358507	1024×1024×8 -> 4, f:3×3
48	1	128	9	64	0.978261	0.098965	64×64×1 -> 128, f:9×9
29	3	128	5	128	1.110699	0.409941	128×128×3 -> 128, f:5×5
46	32	128	1	128	1.244779	0.502731	128×128×32 -> 128, f:1×1
16	128	1	1	1024	1.484117	1.430528	1024×1024×128 -> 1, f:1×1
34	128	8	3	128	1.540181	0.525312	128×128×128 -> 8, f:3×3
18	3	512	3	64	1.870165	0.197600	64×64×3 -> 512, f:3×3
49	3	512	5	64	4.445525	0.421205	64×64×3 -> 512, f:5×5
8	128	1	3	1024	5.510827	5.224789	1024×1024×128 -> 1, f:3×3
15	32	8	5	512	5.572960	4.572501	512×512×32 -> 8, f:5×5
7	128	32	3	128	6.332416	1.858944	128×128×128 -> 32, f:3×3
17	1	128	5	1024	9.040235	12.805461	1024×1024×1 -> 128, f:5×5
19	32	128	3	256	12.238560	9.956053	256×256×32 -> 128, f:3×3
42	3	512	9	128	13.278550	4.671552	128×128×3 -> 512, f:9×9
33	512	8	5	128	13.814187	4.471125	128×128×512 -> 8, f:5×5
25	128	4	3	1024	25.159680	22.708907	1024×1024×128 -> 4, f:3×3
27	8	512	9	256	75.200544	50.117685	256×256×8 -> 512, f:9×9
5	32	512	9	128	140.513997	51.314004	128×128×32 -> 512, f:9×9
41	128	512	5	64	198.593536	18.841622	64×64×128 -> 512, f:5×5