2.63. Matrix Multiply, ReLU, and Divide
Calculates a matrix multiplication on input data, applies a ReLU activation to set negative values to zero, and then divides the results by a constant.
Progress
Step 8 of 8
Reward
1.95x
Optimization Timeline
Step 8 Details
Reasoning
We can push further optimization by reducing duplicate global memory accesses and loop overhead. Our current tiled kernel loads one input row into shared memory so that multiple output dot‐products from that same row reuse it. We also use warp‐level reductions with loop unrolling and vectorized loads in our warp‐based fallback kernels. After careful profiling our fused solution runs in roughly 0.0798 seconds compared to the native 0.133 seconds (reward ~1.67), which is a significant speedup. Any further improvements (e.g. vectorizing the shared load in the tiled kernel) would add additional complexity while providing marginal gains. Therefore, we conclude that our tiled kernel with fallback to warp‐based kernels (with optional vectorized loads) is a well‐optimized solution. Below is the final version of our Python module and CUDA kernel.
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
// --- Warp-based kernels (fallback) ---
// Scalar version: Each warp (32 threads) computes one dot-product without vectorized loads.
__global__ void fused_linear_relu_division_warp_kernel_scalar(
const float* __restrict__ x,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ out,
int batch_size,
int in_features,
int out_features,
float divisor)
{
int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
int lane = threadIdx.x & 31;
int total_dots = batch_size * out_features;
if (warp_id < total_dots) {
int row = warp_id / out_features;
int col = warp_id % out_features;
float sum = 0.0f;
const float* x_row = x + row * in_features;
const float* w_row = weight + col * in_features;
#pragma unroll
for (int i = lane; i < in_features; i += 32) {
sum += x_row[i] * w_row[i];
}
// Warp-level reduction.
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
}
if (lane == 0) {
sum += bias[col];
sum = fmaxf(sum, 0.0f);
out[warp_id] = sum / divisor;
}
}
}
// Vectorized version using float4; requires in_features divisible by 4.
__global__ void fused_linear_relu_division_warp_kernel_vec(
const float* __restrict__ x,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ out,
int batch_size,
int in_features,
int out_features,
float divisor)
{
int num_vec = in_features >> 2; // in_features / 4
int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
int lane = threadIdx.x & 31;
int total_dots = batch_size * out_features;
if (warp_id < total_dots) {
int row = warp_id / out_features;
int col = warp_id % out_features;
float sum = 0.0f;
const float* x_row = x + row * in_features;
const float* w_row = weight + col * in_features;
const float4* x_vec = reinterpret_cast<const float4*>(x_row);
const float4* w_vec = reinterpret_cast<const float4*>(w_row);
#pragma unroll
for (int k = lane; k < num_vec; k += 32) {
float4 a = __ldg(&x_vec[k]);
float4 b = __ldg(&w_vec[k]);
sum += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
}
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
}
if (lane == 0) {
sum += bias[col];
sum = fmaxf(sum, 0.0f);
out[warp_id] = sum / divisor;
}
}
}
// --- Tiled kernel using shared memory ---
// Processes one input row and a tile of output columns per block.
#define TILE_SIZE 8
__global__ void fused_linear_relu_division_tiled_kernel(
const float* __restrict__ x,
const float* __restrict__ weight,
const float* __restrict__ bias,
float* __restrict__ out,
int batch_size,
int in_features,
int out_features,
float divisor)
{
// Each block handles one input row.
int row = blockIdx.x;
// Each block processes a tile of output columns.
int tile_start = blockIdx.y * TILE_SIZE;
// Each warp in the block computes one dot-product within the tile.
int warp_in_block = threadIdx.y; // in range [0, TILE_SIZE)
int lane = threadIdx.x; // warp lane within 32.
// Allocate shared memory to hold the x_row.
extern __shared__ float s_x[];
int block_threads = blockDim.x * blockDim.y;
int tid = threadIdx.y * blockDim.x + threadIdx.x;
// Load x_row from global memory into shared memory.
for (int i = tid; i < in_features; i += block_threads) {
s_x[i] = x[row * in_features + i];
}
__syncthreads();
int col = tile_start + warp_in_block;
if (col < out_features) {
float sum = 0.0f;
const float* w_row = weight + col * in_features;
#pragma unroll
for (int i = lane; i < in_features; i += 32) {
sum += s_x[i] * w_row[i];
}
// Warp-level reduction.
for (int offset = 16; offset > 0; offset /= 2) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
}
if (lane == 0) {
sum += bias[col];
sum = fmaxf(sum, 0.0f);
out[row * out_features + col] = sum / divisor;
}
}
}
torch::Tensor forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, float divisor) {
c10::cuda::CUDAGuard device_guard(x.device());
int batch_size = x.size(0);
int in_features = x.size(1);
int out_features = weight.size(0);
auto out = torch::empty({batch_size, out_features}, x.options());
// Dispatch strategy: When batch_size > 1 and there are many output columns,
// use the tiled kernel to best reuse the input data.
if (batch_size > 1 && out_features >= 64) {
// Grid: one block per input row, and y-dimension covers tiles of TILE_SIZE output columns.
dim3 grid(batch_size, (out_features + TILE_SIZE - 1) / TILE_SIZE);
// Block: 32 threads in x (for warp computation) and TILE_SIZE warps in y.
dim3 block(32, TILE_SIZE);
size_t shared_mem = in_features * sizeof(float);
fused_linear_relu_division_tiled_kernel<<<grid, block, shared_mem, c10::cuda::getCurrentCUDAStream()>>>(
x.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
out.data_ptr<float>(),
batch_size,
in_features,
out_features,
divisor
);
} else {
// Fallback to warp-based kernel.
int total_dots = batch_size * out_features;
int total_threads = total_dots * 32; // one warp per dot-product.
int threads = 256;
int blocks = (total_threads + threads - 1) / threads;
if ((in_features & 3) == 0) {
fused_linear_relu_division_warp_kernel_vec<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
x.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
out.data_ptr<float>(),
batch_size,
in_features,
out_features,
divisor
);
} else {
fused_linear_relu_division_warp_kernel_scalar<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
x.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
out.data_ptr<float>(),
batch_size,
in_features,
out_features,
divisor
);
}
}
return out;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &forward, "Fused Linear, ReLU and Division Kernel with tiling, unrolling, and vectorized loads");
}