The AI CUDA Engineer 👷

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch as th


def module_fn(
    x: torch.Tensor,
    clusters: torch.Tensor,
    clusters2: torch.Tensor,
    bn_weight: torch.Tensor,
    bn_bias: torch.Tensor,
    bn_running_mean: torch.Tensor,
    bn_running_var: torch.Tensor,
    feature_size: int,
    cluster_size: int,
    is_training: bool,
) -> torch.Tensor:
    """
    Functional version of the NetVLAD without ghost clusters

    Args:
        x: Input tensor of shape (batch_size, num_features, feature_size)
        clusters: Weight tensor for cluster assignments
        clusters2: Weight tensor for visual words
        bn_weight: BatchNorm weight
        bn_bias: BatchNorm bias
        bn_running_mean: BatchNorm running mean
        bn_running_var: BatchNorm running var
        feature_size: Size of each feature
        cluster_size: Number of clusters (excluding ghost clusters)
        is_training: Whether in training mode

    Returns:
        Output tensor of shape (batch_size, cluster_size * feature_size)
    """
    max_sample = x.size()[1]
    x = x.view(-1, feature_size)  # B x N x D -> BN x D

    if x.device != clusters.device:
        msg = f"x.device {x.device} != cluster.device {clusters.device}"
        raise ValueError(msg)

    assignment = th.matmul(x, clusters)  # (BN x D) x (D x (K+G)) -> BN x (K+G)
    assignment = F.batch_norm(
        assignment,
        bn_running_mean,
        bn_running_var,
        bn_weight,
        bn_bias,
        training=is_training,
    )

    assignment = F.softmax(assignment, dim=1)  # BN x (K+G) -> BN x (K+G)
    # remove ghost assigments
    assignment = assignment[:, :cluster_size]
    assignment = assignment.view(-1, max_sample, cluster_size)  # -> B x N x K
    a_sum = th.sum(assignment, dim=1, keepdim=True)  # B x N x K -> B x 1 x K
    a = a_sum * clusters2

    assignment = assignment.transpose(1, 2)  # B x N x K -> B x K x N

    x = x.view(-1, max_sample, feature_size)  # BN x D -> B x N x D
    vlad = th.matmul(assignment, x)  # (B x K x N) x (B x N x D) -> B x K x D
    vlad = vlad.transpose(1, 2)  # -> B x D x K
    vlad = vlad - a

    # L2 intra norm
    vlad = F.normalize(vlad)

    # flattening + L2 norm
    vlad = vlad.reshape(-1, cluster_size * feature_size)  # -> B x DK
    vlad = F.normalize(vlad)
    return vlad  # B x DK


class Model(nn.Module):
    def __init__(self, cluster_size, feature_size, ghost_clusters):
        super(Model, self).__init__()

        self.feature_size = feature_size
        self.cluster_size = cluster_size
        self.ghost_clusters = ghost_clusters

        init_sc = 1 / math.sqrt(feature_size)
        clusters = cluster_size + ghost_clusters

        # The `clusters` weights are the `(w,b)` in the paper
        self.clusters = nn.Parameter(init_sc * th.randn(feature_size, clusters))

        # Extract batchnorm parameters
        bn = nn.BatchNorm1d(clusters)
        self.bn_weight = nn.Parameter(bn.weight.data.clone())
        self.bn_bias = nn.Parameter(bn.bias.data.clone())
        self.bn_running_mean = nn.Parameter(bn.running_mean.data.clone())
        self.bn_running_var = nn.Parameter(bn.running_var.data.clone())

        # The `clusters2` weights are the visual words `c_k` in the paper
        self.clusters2 = nn.Parameter(init_sc * th.randn(1, feature_size, cluster_size))
        self.out_dim = self.cluster_size * feature_size

    def forward(self, x, fn=module_fn):
        return fn(
            x,
            self.clusters,
            self.clusters2,
            self.bn_weight,
            self.bn_bias,
            self.bn_running_mean,
            self.bn_running_var,
            self.feature_size,
            self.cluster_size,
            self.training,
        )


batch_size = 32
num_features = 100
num_clusters = 32
feature_size = 512
ghost_clusters = 0


def get_inputs():
    return [torch.randn(batch_size, num_features, feature_size)]


def get_init_inputs():
    return [num_clusters, feature_size, ghost_clusters]

#include <torch/extension.h>
#include <ATen/ATen.h>
#include <vector>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

// Optimized block size for shared memory tiling
#define TILE_DIM 32

// Tiled matrix multiplication kernel using shared memory with loop unrolling and optimized block size
// Computes: C = A (MxD) * B (DxN), where A is x (reshaped) and B is clusters
__global__ void tiledMatMulOptimizedBlock(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C,
                              int M, int D, int N) {
  // Each thread computes one element in the output matrix C
  // using tiling to reduce global memory accesses
  int row = blockIdx.y * TILE_DIM + threadIdx.y;
  int col = blockIdx.x * TILE_DIM + threadIdx.x;

  __shared__ float tile_A[TILE_DIM][TILE_DIM];
  __shared__ float tile_B[TILE_DIM][TILE_DIM];

  float sum = 0.0f;

  // Loop over tiles in the inner dimension
  for (int t = 0; t < (D + TILE_DIM - 1) / TILE_DIM; t++) {
    // Load element of A into shared memory if within bounds
    int a_col = t * TILE_DIM + threadIdx.x;
    if (row < M && a_col < D) {
      tile_A[threadIdx.y][threadIdx.x] = A[row * D + a_col];
    } else {
      tile_A[threadIdx.y][threadIdx.x] = 0.0f;
    }

    // Load element of B into shared memory if within bounds
    int b_row = t * TILE_DIM + threadIdx.y;
    if (b_row < D && col < N) {
      tile_B[threadIdx.y][threadIdx.x] = B[b_row * N + col];
    } else {
      tile_B[threadIdx.y][threadIdx.x] = 0.0f;
    }

    __syncthreads();

    // Multiply the two tiles together with loop unrolling
    #pragma unroll
    for (int i = 0; i < TILE_DIM; i++) {
      sum += tile_A[threadIdx.y][i] * tile_B[i][threadIdx.x];
    }
    __syncthreads();
  }

  // Write the result to global memory if within bounds
  if (row < M && col < N) {
    C[row * N + col] = sum;
  }
}

// Forward function for NetVLAD (no ghost clusters) using the shared memory optimized kernel
torch::Tensor forward(
    torch::Tensor x,                // B x N x D
    torch::Tensor clusters,         // D x (K+G)
    torch::Tensor clusters2,        // 1 x D x K
    torch::Tensor bn_weight,        // (K+G)
    torch::Tensor bn_bias,          // (K+G)
    torch::Tensor bn_running_mean,  // (K+G)
    torch::Tensor bn_running_var,   // (K+G)
    int64_t feature_size,           // D
    int64_t cluster_size,           // K
    bool is_training
) {
    CHECK_INPUT(x);
    CHECK_INPUT(clusters);
    CHECK_INPUT(clusters2);
    CHECK_INPUT(bn_weight);
    CHECK_INPUT(bn_bias);
    CHECK_INPUT(bn_running_mean);
    CHECK_INPUT(bn_running_var);

    int64_t B = x.size(0);
    int64_t N_points = x.size(1); // Number of descriptors/features per sample
    int64_t D = feature_size;
    int64_t K = cluster_size;
    int64_t KplusG = clusters.size(1); // Total clusters including ghost clusters

    // Reshape x: B x N x D -> (B*N) x D
    auto x_reshaped = x.reshape({B * N_points, D});

    // Ensure device consistency
    if (x_reshaped.device() != clusters.device()) {
        TORCH_CHECK(false, "x.device() != clusters.device()");
    }

    // Allocate memory for assignment: (B*N) x (K+G)
    auto assignment = torch::empty({B * N_points, KplusG}, x.options());

    // Launch tiled matrix multiplication kernel
    // x_reshaped: (B*N) x D, clusters: D x (K+G), output: (B*N) x (K+G)
    int M = B * N_points;     // Number of rows
    int inner = D;            // Shared inner dimension
    int N_dim = KplusG;       // Number of columns
    
    dim3 block(TILE_DIM, TILE_DIM);
    dim3 grid((N_dim + TILE_DIM - 1) / TILE_DIM, (M + TILE_DIM - 1) / TILE_DIM);

    tiledMatMulOptimizedBlock<<<grid, block>>>(x_reshaped.data_ptr<float>(), clusters.data_ptr<float>(), assignment.data_ptr<float>(), M, inner, N_dim);

    // Apply BatchNorm
    assignment = torch::batch_norm(
        assignment,
        bn_weight,
        bn_bias,
        bn_running_mean,
        bn_running_var,
        is_training,
        0.1,
        1e-5,
        true
    );

    // Apply Softmax along dim=1
    assignment = torch::softmax(assignment, 1); // (B*N) x (K+G)

    // Remove ghost assignments: keep first K columns
    assignment = assignment.narrow(1, 0, K); // (B*N) x K

    // Reshape assignment: (B*N) x K -> B x N x K
    assignment = assignment.reshape({B, N_points, K});

    // Compute a_sum = assignment.sum(dim=1, keepdim=true) // B x 1 x K
    auto a_sum = assignment.sum(1, /*keepdim=*/true);

    // Expand clusters2 to match batch size: clusters2 is originally 1 x D x K
    auto clusters2_exp = clusters2.expand({B, D, K}); // B x D x K

    // Compute a = a_sum * clusters2
    auto a = clusters2_exp * a_sum; // B x D x K

    // Transpose assignment: B x N x K -> B x K x N
    assignment = assignment.transpose(1, 2);

    // Reshape x back to B x N x D
    auto x_orig = x.reshape({B, N_points, D});

    // Compute vlad = assignment @ x_orig; assignment: B x K x N, x_orig: B x N x D
    auto vlad = torch::bmm(assignment, x_orig); // B x K x D

    // Transpose vlad to B x D x K
    vlad = vlad.transpose(1, 2);

    // Compute vlad = vlad - a
    vlad = vlad - a;

    // L2 intra-normalization along feature dimension D
    vlad = torch::nn::functional::normalize(
        vlad, torch::nn::functional::NormalizeFuncOptions().p(2).dim(1));

    // Flatten vlad: B x D x K -> B x (D*K)
    vlad = vlad.reshape({B, D * K});

    // L2 normalization along the flattened dimension
    vlad = torch::nn::functional::normalize(
        vlad, torch::nn::functional::NormalizeFuncOptions().p(2).dim(1));

    return vlad;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "NetVLAD forward with optimized block size and loop unrolling");
}