Sync Guardrails Within Benchmarking Loop

[DhruvSrikanth]

In triton's do_bench() (found here), do we need need have a cross-device synchronize before each fn() call in the speed trial loop? Specifically here?

I found something like this to have a far lower time than the computation actually takes. One way around this is by adding a synchronize step to fn() itself instead of directly passing forward(**get_inputs()), however, this seems like something that should be covered within the benchmarking function for robustness.

import torch
import torch.nn as nn
import torch.utils.cpp_extension
import os

# JIT compile custom CUDA kernel for fused conv+relu+bias with shared memory and warp optimizations
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>

template <typename scalar_t>
__global__ void fused_conv_relu_bias_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ output,
    int batch_size, int in_channels, int out_channels,
    int height, int width, int kernel_size, int output_h, int output_w) {
    
    extern __shared__ char shared_memory[];
    scalar_t* shared_input = reinterpret_cast<scalar_t*>(shared_memory);
    
    const int TILE_SIZE = 16;
    const int PADDING = kernel_size / 2;
    const int SHARED_WIDTH = TILE_SIZE + 2 * PADDING;
    const int SHARED_HEIGHT = TILE_SIZE + 2 * PADDING;
    
    int tx = threadIdx.x % TILE_SIZE;
    int ty = threadIdx.y % TILE_SIZE;
    int out_x = blockIdx.x * TILE_SIZE + tx;
    int out_y = blockIdx.y * TILE_SIZE + ty;
    int out_c = blockIdx.z % out_channels;
    int out_b = blockIdx.z / out_channels;
    
    if (out_b >= batch_size) return;
    
    // Thread coarsening - each thread processes multiple outputs
    scalar_t weight_cache[9]; // Assuming kernel_size=3
    
    // Prefetch weights to registers for better reuse
    #pragma unroll
    for (int ic = 0; ic < in_channels; ++ic) {
        #pragma unroll
        for (int kh = 0; kh < kernel_size; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < kernel_size; ++kw) {
                int weight_idx = ((out_c * in_channels + ic) * kernel_size + kh) * kernel_size + kw;
                weight_cache[(ic * kernel_size + kh) * kernel_size + kw] = weight[weight_idx];
            }
        }
    }
    
    // Collaborative loading of input tile into shared memory
    for (int ic = 0; ic < in_channels; ++ic) {
        for (int i = ty; i < SHARED_HEIGHT; i += TILE_SIZE) {
            for (int j = tx; j < SHARED_WIDTH; j += TILE_SIZE) {
                int h_pos = blockIdx.y * TILE_SIZE + i - PADDING;
                int w_pos = blockIdx.x * TILE_SIZE + j - PADDING;
                
                scalar_t val = 0.0f;
                if (h_pos >= 0 && h_pos < height && w_pos >= 0 && w_pos < width && out_b < batch_size) {
                    val = input[((out_b * in_channels + ic) * height + h_pos) * width + w_pos];
                }
                
                shared_input[(ic * SHARED_HEIGHT + i) * SHARED_WIDTH + j] = val;
            }
        }
    }
    
    __syncthreads();
    
    if (out_x >= output_w || out_y >= output_h) return;
    
    scalar_t sum = 0.0f;
    
    // Compute convolution using shared memory with loop unrolling for inner loops
    for (int ic = 0; ic < in_channels; ++ic) {
        #pragma unroll
        for (int kh = 0; kh < kernel_size; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < kernel_size; ++kw) {
                int sh = ty + kh;
                int sw = tx + kw;
                
                scalar_t input_val = shared_input[(ic * SHARED_HEIGHT + sh) * SHARED_WIDTH + sw];
                scalar_t weight_val = weight_cache[(ic * kernel_size + kh) * kernel_size + kw];
                
                sum = __fmaf_rn(input_val, weight_val, sum);
            }
        }
    }
    
    // Add bias and apply ReLU in one step
    sum += bias[out_c];
    
    // Write output with ReLU
    if (out_x < output_w && out_y < output_h) {
        int out_idx = ((out_b * out_channels + out_c) * output_h + out_y) * output_w + out_x;
        output[out_idx] = sum > 0.0f ? sum : 0.0f;
    }
}

torch::Tensor fused_conv_relu_bias_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias) {
    
    auto batch_size = input.size(0);
    auto in_channels = input.size(1);
    auto height = input.size(2);
    auto width = input.size(3);
    auto out_channels = weight.size(0);
    auto kernel_size = weight.size(2);
    
    auto output_h = height - kernel_size + 1;
    auto output_w = width - kernel_size + 1;
    
    auto output = torch::zeros({batch_size, out_channels, output_h, output_w}, 
                              input.options());
    
    const int TILE_SIZE = 16;
    const int PADDING = kernel_size / 2;
    const int SHARED_WIDTH = TILE_SIZE + 2 * PADDING;
    const int SHARED_HEIGHT = TILE_SIZE + 2 * PADDING;
    
    dim3 threads(TILE_SIZE, TILE_SIZE);
    dim3 blocks(
        (output_w + TILE_SIZE - 1) / TILE_SIZE,
        (output_h + TILE_SIZE - 1) / TILE_SIZE,
        batch_size * out_channels
    );
    
    // Calculate shared memory size
    int shared_mem_size = in_channels * SHARED_HEIGHT * SHARED_WIDTH * sizeof(float);
    
    // Use higher priority stream for better scheduling
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cudaStreamSetAttribute(stream, cudaStreamAttributePriority, 
                          (void*)(size_t)cudaStreamPriorityHigh);
    
    AT_DISPATCH_FLOATING_TYPES(input.type(), "fused_conv_relu_bias_cuda", ([&] {
        fused_conv_relu_bias_kernel<scalar_t><<<blocks, threads, shared_mem_size, stream>>>(
            input.data_ptr<scalar_t>(),
            weight.data_ptr<scalar_t>(),
            bias.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            batch_size, in_channels, out_channels,
            height, width, kernel_size, output_h, output_w);
    }));
    
    return output;
}
"""

cuda_module = None


class Model(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=0)
        self.bias = nn.Parameter(torch.randn(bias_shape))

        global cuda_module
        if cuda_module is None and torch.cuda.is_available():
            try:
                cuda_module = torch.utils.cpp_extension.load_inline(
                    name="fused_ops",
                    cpp_sources="",
                    cuda_sources=cuda_source,
                    functions=["fused_conv_relu_bias_cuda"],
                    with_cuda=True,
                    extra_cuda_cflags=["--use_fast_math", "-O3", "--ptxas-options=-v"],
                    verbose=False,
                )
            except:
                pass

        # Fallback to compiled forward if CUDA extension fails
        self.optimized_forward = torch.compile(self._forward, mode="reduce-overhead")
        self.stream = (
            torch.cuda.Stream(priority=-1) if torch.cuda.is_available() else None
        )

    def _forward(self, x):
        x = self.conv(x)
        x = torch.relu(x)
        x = x + self.bias
        return x

    def forward(self, x):
        if x.device.type == "cpu" and torch.cuda.is_available() and not x.is_pinned():
            x = x.pin_memory().to("cuda", non_blocking=True)
        elif x.device.type == "cpu" and torch.cuda.is_available():
            x = x.to("cuda", non_blocking=True)

        if cuda_module is not None and x.device.type == "cuda":
            with torch.cuda.stream(self.stream):
                try:
                    return cuda_module.fused_conv_relu_bias_cuda(
                        x, self.conv.weight, self.bias.squeeze()
                    )
                except:
                    return self.optimized_forward(x)
        else:
            if self.stream is not None:
                with torch.cuda.stream(self.stream):
                    return self.optimized_forward(x)
            else:
                return self.optimized_forward(x)

batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
bias_shape = (out_channels, 1, 1)

def get_inputs():
    inputs = torch.randn(batch_size, in_channels, height, width)
    return [inputs.pin_memory() if torch.cuda.is_available() else inputs]


def get_init_inputs():
    return [in_channels, out_channels, kernel_size, bias_shape]

Curious to get the communities take on something like this