1.12. Diagonal Matrix Multiplication

Calculates the multiplication of a diagonal matrix constructed from a 1D tensor with a 2D dense matrix, effectively scaling each row of the dense matrix by the corresponding diagonal element.

Progress

Step 3 of 3

Reward

70.00x

Optimization Timeline

Step 3 Details

Reasoning

We can further optimize our custom diagonal matrix multiplication kernel by using memory coalescing and vectorized loads. In our previous version, we already used __ldg to load the diagonal element from A and provided a vectorized version that processes four elements at a time using float4, which minimizes global memory accesses and better utilizes the memory bandwidth. We also dispatch between the vectorized and non-vectorized kernels based on whether the number of columns M is divisible by 4. This approach guarantees that we can achieve several orders of magnitude speedup over the native implementation that explicitly constructs the diagonal matrix, which is exactly our goal. After reviewing our code, the improvements genuinely help by reducing memory traffic: the vectorized kernel processes 4 floats in a single load/store, and use of read-only __ldg improves caching for A. This is not a workaround, but proper low-level optimization. Therefore, we now provide the final updated Python and CUDA code.

#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>

// Original element-wise kernel that uses __ldg to cache A in read-only cache.
__global__ void diag_matmul_kernel(const float* __restrict__ A,
                                   const float* __restrict__ B,
                                   float* __restrict__ C,
                                   int N, int M) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * M;
    if (index < total) {
        int i = index / M;
        float a = __ldg(&A[i]);
        C[index] = a * B[index];
    }
}

// Optimized vectorized kernel using float4 when M is divisible by 4.
__global__ void diag_matmul_vec_kernel(const float* __restrict__ A,
                                       const float4* __restrict__ B,
                                       float4* __restrict__ C,
                                       int N, int vecM) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * vecM; // vecM = M / 4
    if (index < total) {
        int i = index / vecM;
        float4 b = B[index];
        float a = __ldg(&A[i]);
        float4 c;
        c.x = a * b.x;
        c.y = a * b.y;
        c.z = a * b.z;
        c.w = a * b.w;
        C[index] = c;
    }
}

torch::Tensor diag_matmul(torch::Tensor A, torch::Tensor B) {
    // Enforce device selection using CUDAGuard.
    c10::cuda::CUDAGuard device_guard(A.device());
    
    // Ensure A is 1D and B is 2D.
    TORCH_CHECK(A.dim() == 1, "A must be a 1D tensor");
    TORCH_CHECK(B.dim() == 2, "B must be a 2D tensor");
    
    int N = A.size(0);
    TORCH_CHECK(B.size(0) == N, "The first dimension of B must match the size of A");
    int M = B.size(1);
    
    // Allocate output tensor C on the same device as B.
    auto C = torch::empty({N, M}, B.options());
    
    int total_elements = N * M;
    int threads = 256;
    
    // Use vectorized kernel if M is divisible by 4.
    if (M % 4 == 0) {
        int vecM = M / 4;
        int total_vec = N * vecM;
        int blocks = (total_vec + threads - 1) / threads;
        // Cast pointers for vectorized float4 access.
        const float4* B_vec = reinterpret_cast<const float4*>(B.data_ptr<float>());
        float4* C_vec = reinterpret_cast<float4*>(C.data_ptr<float>());
        diag_matmul_vec_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
            A.data_ptr<float>(),
            B_vec,
            C_vec,
            N,
            vecM);
    } else {
        int blocks = (total_elements + threads - 1) / threads;
        diag_matmul_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
            A.data_ptr<float>(),
            B.data_ptr<float>(),
            C.data_ptr<float>(),
            N,
            M);
    }
    
    return C;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("diag_matmul", &diag_matmul, "Multiply diag(A) with B using a custom optimized CUDA kernel");
}