← Back to Leaderboard

The AI CUDA Engineer 👷

16_Matmul_with_transposed_A • pipelined_tiled_matmul_base_base

Level 1 • Task 16
import torch
import torch.nn as nn
import torch.nn.functional as F


def module_fn(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    """
    Performs a single matrix multiplication with transposed A (C = A.T * B).

    Args:
        A: Input tensor of shape (K, M).
        B: Input tensor of shape (K, N).

    Returns:
        Output tensor of shape (M, N).
    """
    return torch.matmul(A.T, B)


class Model(nn.Module):
    """
    Simple model that performs a single matrix multiplication (C = A * B)
    """

    def __init__(self):
        super(Model, self).__init__()

    def forward(self, A: torch.Tensor, B: torch.Tensor, fn=module_fn) -> torch.Tensor:
        return fn(A, B)


M = 1024
K = 4096
N = 2048


def get_inputs():
    A = torch.randn(K, M)
    B = torch.randn(K, N)
    return [A, B]


def get_init_inputs():
    return []  # No special initialization inputs needed

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a single matrix multiplication (C = A * B)
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        """
        Performs matrix multiplication.

        Args:
            A: Input tensor of shape (M, K).
            B: Input tensor of shape (K, N).

        Returns:
            Output tensor of shape (M, N).
        """
        return torch.matmul(A.T, B)

M = 1024
K = 4096
N = 2048

def get_inputs():
    A = torch.randn(K, M)
    B = torch.randn(K, N)
    return [A, B]

def get_init_inputs():
    return []  # No special initialization inputs needed
Download Evaluation Download PyTorch Download CUDA Download Profiles

Kernel Information

Operation Name 16_Matmul_with_transposed_A
Level ID 1
Task ID 16
Kernel Name pipelined_tiled_matmul_base_base
CUDA Speedup (Native) 0.130x
CUDA Speedup (Compile) 0.148x
CUDA Runtime 2.702 ms
PyTorch Runtime (Native) 0.351 ms
PyTorch Runtime (Compile) 0.400 ms
Correct True
Max Diff (vs. Reference) 0.000000
Model bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0
Temperature 0.50
View Experiment Progress Details

Related Kernels (Level 1, Task 16 • 16_Matmul_with_transposed_A)

Rank Kernel Name Runtime (ms) Speedup Native Speedup Compile
🥇 tiled_double_output_base 2.29 0.15 0.17
🥈 pipelined_tiled_matmul_base_base 2.70 0.13 0.15
🥉 hybrid_tiled_linear_matmul_base 2.76 0.13 0.14
4 modular_tiled_matmul_base_base 2.77 0.13 0.14
5 unrolled_tiled_matmul_base_base 2.81 0.13 0.14
6 optimized_tiled_matmul_base 2.81 0.13 0.14
6 tiled_shared_ldg_aligned_base 2.81 0.13 0.14
6 optimized_tiled_matmul_base 2.81 0.13 0.14
6 hybrid_tiling_grid_stride_base 2.81 0.13 0.14
10 syncthreads_optimized_tiling_edit_1 3.00 0.12 0.13
10 atomic_operations_optimized_tiling_base 3.00 0.12 0.13
12 streams_partitioned_matmul_edit_1 3.02 0.12 0.13
13 tiled_shared_unroll_base_base 3.02 0.12 0.13
14 streams_partitioned_matmul_base 3.03 0.12 0.13
15 modular_device_functions_tiling_2_base 3.04 0.12 0.13
15 modular_tiled_kernel_edit_1 3.04 0.12 0.13
15 modular_tiled_kernel_base 3.04 0.12 0.13
18 optimized_matmul_combined_kernel_edit_1 3.04 0.12 0.13
18 tiled_shared_const_memory_base 3.04 0.12 0.13
18 optimized_matmul_combined_kernel_base 3.04 0.12 0.13 
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdexcept>

#define TILE_M 16
#define TILE_N 16
#define BLOCK_K 32
#define NUM_STREAMS 2

__global__ void pipelinedMatMulKernel(const float* __restrict__ A,
                                     const float* __restrict__ B,
                                     float* __restrict__ C,
                                     int K, int M, int N,
                                     int chunk_start, int chunk_size) {
    __shared__ float As[2][BLOCK_K][TILE_M];
    __shared__ float Bs[2][BLOCK_K][TILE_N];
    
    int row = blockIdx.x * TILE_M + threadIdx.x;
    int col = blockIdx.y * TILE_N + threadIdx.y;
    
    float sum = 0.0f;
    
    int tid = threadIdx.y * blockDim.x + threadIdx.x;
    int totalThreads = blockDim.x * blockDim.y;
    
    int current_buffer = 0;
    
    if (chunk_start < K) {
        for (int index = tid; index < BLOCK_K * TILE_M; index += totalThreads) {
            int t = index / TILE_M;
            int m = index % TILE_M;
            int global_k = chunk_start + t;
            int global_m = blockIdx.x * TILE_M + m;
            As[0][t][m] = (global_k < K && global_m < M) ? __ldg(&A[global_k * M + global_m]) : 0.0f;
        }
        
        for (int index = tid; index < BLOCK_K * TILE_N; index += totalThreads) {
            int t = index / TILE_N;
            int n = index % TILE_N;
            int global_k = chunk_start + t;
            int global_n = blockIdx.y * TILE_N + n;
            Bs[0][t][n] = (global_k < K && global_n < N) ? __ldg(&B[global_k * N + global_n]) : 0.0f;
        }
    }
    
    __syncthreads();
    
    for (int k = chunk_start; k < min(chunk_start + chunk_size, K); k += BLOCK_K) {
        if (k + BLOCK_K < min(chunk_start + chunk_size, K)) {
            #pragma unroll
            for (int index = tid; index < BLOCK_K * TILE_M; index += totalThreads) {
                int t = index / TILE_M;
                int m = index % TILE_M;
                int global_k = k + BLOCK_K + t;
                int global_m = blockIdx.x * TILE_M + m;
                As[1-current_buffer][t][m] = (global_k < K && global_m < M) ? 
                    __ldg(&A[global_k * M + global_m]) : 0.0f;
            }
            
            #pragma unroll
            for (int index = tid; index < BLOCK_K * TILE_N; index += totalThreads) {
                int t = index / TILE_N;
                int n = index % TILE_N;
                int global_k = k + BLOCK_K + t;
                int global_n = blockIdx.y * TILE_N + n;
                Bs[1-current_buffer][t][n] = (global_k < K && global_n < N) ? 
                    __ldg(&B[global_k * N + global_n]) : 0.0f;
            }
        }
        
        #pragma unroll 8
        for (int t = 0; t < BLOCK_K; t++) {
            sum += As[current_buffer][t][threadIdx.x] * Bs[current_buffer][t][threadIdx.y];
        }
        
        current_buffer = 1 - current_buffer;
        __syncthreads();
    }
    
    if (row < M && col < N) {
        atomicAdd(&C[row * N + col], sum);
    }
}

torch::Tensor forward(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda(), "Input A must be a CUDA tensor");
    TORCH_CHECK(B.is_cuda(), "Input B must be a CUDA tensor");
    TORCH_CHECK(A.dtype() == torch::kFloat32, "Input A must be float32");
    TORCH_CHECK(B.dtype() == torch::kFloat32, "Input B must be float32");

    int K = A.size(0);
    int M = A.size(1);
    TORCH_CHECK(B.size(0) == K, "Dimension mismatch");
    int N = B.size(1);

    auto C = torch::zeros({M, N}, torch::device(A.device()).dtype(A.dtype()));
    
    const float* A_ptr = A.data_ptr<float>();
    const float* B_ptr = B.data_ptr<float>();
    float* C_ptr = C.data_ptr<float>();

    dim3 block(TILE_M, TILE_N);
    dim3 grid((M + TILE_M - 1) / TILE_M, (N + TILE_N - 1) / TILE_N);

    cudaStream_t streams[NUM_STREAMS];
    for (int i = 0; i < NUM_STREAMS; i++) {
        cudaStreamCreate(&streams[i]);
    }

    int chunk_size = (K + NUM_STREAMS - 1) / NUM_STREAMS;
    
    for (int i = 0; i < NUM_STREAMS; i++) {
        int chunk_start = i * chunk_size;
        pipelinedMatMulKernel<<<grid, block, 0, streams[i]>>>(
            A_ptr, B_ptr, C_ptr, K, M, N, chunk_start, chunk_size);
    }

    for (int i = 0; i < NUM_STREAMS; i++) {
        cudaStreamSynchronize(streams[i]);
        cudaStreamDestroy(streams[i]);
    }

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }

    return C;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Pipelined matrix multiplication (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 2.660 inst/cycle 0.000 5
Executed Ipc Elapsed 2.632 inst/cycle 0.000 5
Issue Slots Busy 66.602 % 0.000 5
Issued Ipc Active 2.660 inst/cycle 0.000 5
SM Busy 66.602 % 0.000 5
Memory Throughput 22625973777.932 byte/second 7962316041119429.000 5
Mem Busy 90.928 % 0.003 5
Max Bandwidth 85.492 % 0.003 5
L1/TEX Hit Rate 0.072 % 0.000 5
L2 Hit Rate 96.862 % 0.034 5
Mem Pipes Busy 85.492 % 0.003 5
Warp Cycles Per Issued Instruction 22.906 cycle 0.000 5
Warp Cycles Per Executed Instruction 22.908 cycle 0.000 5
Avg. Active Threads Per Warp 32.000 0.000 5
Avg. Not Predicated Off Threads Per Warp 31.420 0.000 5
Max Active Clusters 0.000 cluster 0.000 5
Max Cluster Size 8.000 block 0.000 5
Overall GPU Occupancy 0.000 % 0.000 5
Cluster Occupancy 0.000 % 0.000 5
Block Limit SM 32.000 block 0.000 5
Block Limit Registers 8.000 block 0.000 5
Block Limit Shared Mem 14.000 block 0.000 5
Block Limit Warps 8.000 block 0.000 5
Theoretical Active Warps per SM 64.000 warp 0.000 5
Theoretical Occupancy 100.000 % 0.000 5
Achieved Occupancy 95.340 % 0.000 5
Achieved Active Warps Per SM 61.014 warp 0.000 5
Analysis Rules
Rule Description
INF HighPipeUtilization ALU is the highest-utilized pipeline (28.4%) based on active cycles, taking into account the rates of its different instructions. It executes integer and logic operations. It is well-utilized, but should not be a bottleneck.
INF Occupancy This kernel's theoretical occupancy is not impacted by any block limit.
Operation / Metric Value Unit
aten::to
CPU Time 410509.49 μs
Device Time 5212.86 μs
Self CPU Time 46.69 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
aten::_to_copy
CPU Time 410462.80 μs
Device Time 5212.86 μs
Self CPU Time 124.03 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
aten::empty_strided
CPU Time 404748.82 μs
Device Time 0.00 μs
Self CPU Time 100.74 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
cudaDeviceGetStreamPriorityRange
CPU Time 404282.25 μs
Device Time 0.00 μs
Self CPU Time 404282.25 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
cudaStreamSynchronize
CPU Time 9777767.92 μs
Device Time 25833.09 μs
Self CPU Time 9777767.92 μs
Self Device Time 25833.09 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
aten::zero_
CPU Time 70481.24 μs
Device Time 295252.83 μs
Self CPU Time 14023.37 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
aten::fill_
CPU Time 56458.78 μs
Device Time 295252.83 μs
Self CPU Time 19383.71 μs
Self Device Time 295252.83 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
pipelinedMatMulKernel(float const*, float const*, float*, int, int, int, int, int)
CPU Time 0.00 μs
Device Time 10199109.18 μs
Self CPU Time 0.00 μs
Self Device Time 10199109.18 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<int>, at::detail::Array<char*, 1> >(int, at::native::FillFunctor<int>, at::detail::Array<char*, 1>)
CPU Time 0.00 μs
Device Time 280939.04 μs
Self CPU Time 0.00 μs
Self Device Time 280939.04 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
Status: Completed 
45290 warnings generated when compiling for host.
Suppressed 45322 warnings (45275 in non-user code, 47 NOLINT).
Use -header-filter=.* to display errors from all non-system headers. Use -system-headers to display errors from system headers as well.
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:11:39 bugprone-easily-swappable-parameters
11 | __global__ void pipelinedMatMulKernel(const float* __restrict__ A,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~
12 | const float* __restrict__ B,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:11:65: note: the first parameter in the range is 'A'
11 | __global__ void pipelinedMatMulKernel(const float* __restrict__ A,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:12:64: note: the last parameter in the range is 'B'
12 | const float* __restrict__ B,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:14:52: warning: 2 adjacent parameters of 'pipelinedMatMulKernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
14 | int K, int M, int N,
| ^~~~~~
15 | int chunk_start, int chunk_size) {
| ~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:14:56: note: the first parameter in the range is 'N'
14 | int K, int M, int N,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:15:42: note: the last parameter in the range is 'chunk_start'
15 | int chunk_start, int chunk_size) {
| ^~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:19:15: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
19 | int row = blockIdx.x * TILE_M + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:20:15: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
20 | int col = blockIdx.y * TILE_N + threadIdx.y;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:24:15: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
24 | int tid = threadIdx.y * blockDim.x + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:25:24: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
25 | int totalThreads = blockDim.x * blockDim.y;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:34:28: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
34 | int global_m = blockIdx.x * TILE_M + m;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:42:28: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
42 | int global_n = blockIdx.y * TILE_N + n;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:56:32: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
56 | int global_m = blockIdx.x * TILE_M + m;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:66:32: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
66 | int global_n = blockIdx.y * TILE_N + n;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:86:37: warning: the parameter 'A' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
86 | torch::Tensor forward(torch::Tensor A, torch::Tensor B) {
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:86:54: warning: the parameter 'B' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
86 | torch::Tensor forward(torch::Tensor A, torch::Tensor B) {
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:92:13: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
92 | int K = A.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:93:13: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
93 | int M = A.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250202_optimize_b10_s4_e0_sweep/level_1/task_16/b9_s3_pipelined_tiled_matmul_base/base/base.cu:95:13: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
95 | int N = B.size(1);
| ^