← Back to Leaderboard

The AI CUDA Engineer 👷

16_Matmul_with_transposed_Atiled_double_output_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

Kernel Information

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 and block dimensions
#define BLOCK_M 16        // Number of C-rows computed per block
#define BLOCK_N 32        // Number of C-columns computed per block (each thread computes 2 outputs)
#define TILE 16           // Tile width for the K dimension

// Kernel computes C = A.T * B, where A is (K, M), B is (K, N) and C is (M, N).
// Each thread computes two adjacent elements in C to improve reuse of loaded tiles.
__global__ void tiledDoubleOutputKernel(const float* __restrict__ A,
                                          const float* __restrict__ B,
                                          float* __restrict__ C,
                                          int K, int M, int N) {
    // Map each block to a tile of C of size BLOCK_M x BLOCK_N.
    // Each thread computes two adjacent outputs in the horizontal (column) direction.
    int row = blockIdx.y * BLOCK_M + threadIdx.y;
    int col = blockIdx.x * BLOCK_N + threadIdx.x * 2;  // two outputs per thread

    float out0 = 0.0f, out1 = 0.0f;

    // Declare shared memory tiles for A and B
    // A_tile: holds a tile of A.T (which is logically A transposed). Each element is loaded as A[k, row] = A[k * M + row].
    __shared__ float A_tile[BLOCK_M][TILE];   // Dimensions: (BLOCK_M x TILE)
    // B_tile: holds a tile of B, dimensions: (TILE x BLOCK_N)
    __shared__ float B_tile[TILE][BLOCK_N];

    int numTiles = (K + TILE - 1) / TILE;
    for (int t = 0; t < numTiles; t++) {
        int tileStart = t * TILE;

        // Each thread loads one element of the A tile.
        int a_k = tileStart + threadIdx.x;  // threadIdx.x in [0, TILE-1]
        if (a_k < K && row < M)
            A_tile[threadIdx.y][threadIdx.x] = A[a_k * M + row];
        else
            A_tile[threadIdx.y][threadIdx.x] = 0.0f;

        // Each thread loads two elements of the B tile.
        int b_k = tileStart + threadIdx.y;  // threadIdx.y in [0, TILE-1]
        int global_col0 = col;
        int global_col1 = col + 1;
        if (b_k < K) {
            if (global_col0 < N)
                B_tile[threadIdx.y][threadIdx.x * 2] = B[b_k * N + global_col0];
            else
                B_tile[threadIdx.y][threadIdx.x * 2] = 0.0f;
            if (global_col1 < N)
                B_tile[threadIdx.y][threadIdx.x * 2 + 1] = B[b_k * N + global_col1];
            else
                B_tile[threadIdx.y][threadIdx.x * 2 + 1] = 0.0f;
        } else {
            B_tile[threadIdx.y][threadIdx.x * 2] = 0.0f;
            B_tile[threadIdx.y][threadIdx.x * 2 + 1] = 0.0f;
        }

        __syncthreads();

        // Compute partial dot-products for the two outputs
        #pragma unroll
        for (int s = 0; s < TILE; s++) {
            float a_val = A_tile[threadIdx.y][s];
            out0 += a_val * B_tile[s][threadIdx.x * 2];
            out1 += a_val * B_tile[s][threadIdx.x * 2 + 1];
        }

        __syncthreads();
    }

    // Write the computed outputs to global memory
    if (row < M) {
        if (col < N)
            C[row * N + col] = out0;
        if (col + 1 < N)
            C[row * N + col + 1] = out1;
    }
}

// The forward function exposed via PyBind11
// Inputs:
//   A: Tensor of shape (K, M) [CUDA, float32]
//   B: Tensor of shape (K, N) [CUDA, float32]
// Returns:
//   C: Tensor of shape (M, N) computed as A.T * B.

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: A and B must have the same first dimension (K)");
    int N = B.size(1);

    auto C = torch::zeros({M, N}, torch::device(A.device()).dtype(A.dtype()));

    // Define block dimensions:
    // We use 16 threads for the row dimension and 16 threads for the column dimension,
    // with each thread computing two adjacent output elements (total BLOCK_N = 32 columns per block).
    dim3 blockDim(TILE, BLOCK_M); // blockDim.x = 16, blockDim.y = 16
    dim3 gridDim((N + BLOCK_N - 1) / BLOCK_N, (M + BLOCK_M - 1) / BLOCK_M);

    const float* A_ptr = A.data_ptr<float>();
    const float* B_ptr = B.data_ptr<float>();
    float* C_ptr = C.data_ptr<float>();

    tiledDoubleOutputKernel<<<gridDim, blockDim>>>(A_ptr, B_ptr, C_ptr, K, M, N);
    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, "Compute C = A.T * B using tiled kernel with double output per thread (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 1.210 inst/cycle 0.000 5
Executed Ipc Elapsed 1.172 inst/cycle 0.000 5
Issue Slots Busy 30.218 % 0.000 5
Issued Ipc Active 1.210 inst/cycle 0.000 5
SM Busy 30.218 % 0.000 5
Memory Throughput 164672136936.690 byte/second 3210715795722371072.000 5
Mem Busy 94.260 % 0.002 5
Max Bandwidth 64.868 % 0.001 5
L1/TEX Hit Rate 62.490 % 0.000 5
L2 Hit Rate 75.282 % 0.151 5
Mem Pipes Busy 38.146 % 0.000 5
Warp Cycles Per Issued Instruction 49.334 cycle 0.000 5
Warp Cycles Per Executed Instruction 49.334 cycle 0.000 5
Avg. Active Threads Per Warp 32.000 0.000 5
Avg. Not Predicated Off Threads Per Warp 31.220 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 16.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 93.186 % 0.000 5
Achieved Active Warps Per SM 59.640 warp 0.000 5
Analysis Rules
Rule Description
WRN HighPipeUtilization All compute pipelines are under-utilized. Either this kernel is very small or it doesn't issue enough warps per scheduler. Check the Launch Statistics and Scheduler Statistics sections for further details.
INF Occupancy This kernel's theoretical occupancy is not impacted by any block limit.
Operation / Metric Value Unit
aten::to
CPU Time 566679.47 μs
Device Time 5358.45 μs
Self CPU Time 38.75 μ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 566640.72 μs
Device Time 5358.45 μs
Self CPU Time 114.67 μ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::zero_
CPU Time 9344278.90 μs
Device Time 342579.32 μs
Self CPU Time 16584.76 μ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 9327696.15 μs
Device Time 342579.32 μs
Self CPU Time 21853.07 μs
Self Device Time 342579.32 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
cudaLaunchKernel
CPU Time 9324122.82 μs
Device Time 9261.65 μs
Self CPU Time 9324122.82 μs
Self Device Time 9261.65 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
tiledDoubleOutputKernel(float const*, float const*, float*, int, int, int)
CPU Time 0.00 μs
Device Time 9616609.16 μs
Self CPU Time 0.00 μs
Self Device Time 9616609.16 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
cudaEventRecord
CPU Time 17908.58 μs
Device Time 17909.89 μs
Self CPU Time 17908.58 μs
Self Device Time 17909.89 μ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 325897.52 μs
Self CPU Time 0.00 μs
Self Device Time 325897.52 μ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/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:13:41 bugprone-easily-swappable-parameters
13 | __global__ void tiledDoubleOutputKernel(const float* __restrict__ A,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~
14 | const float* __restrict__ B,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:13:67: note: the first parameter in the range is 'A'
13 | __global__ void tiledDoubleOutputKernel(const float* __restrict__ A,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:14:69: note: the last parameter in the range is 'B'
14 | const float* __restrict__ B,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:16:50: warning: 2 adjacent parameters of 'tiledDoubleOutputKernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
16 | int K, int M, int N) {
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:16:54: note: the first parameter in the range is 'M'
16 | int K, int M, int N) {
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:16:61: note: the last parameter in the range is 'N'
16 | int K, int M, int N) {
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/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.y * BLOCK_M + threadIdx.y;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/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.x * BLOCK_N + threadIdx.x * 2; // two outputs per thread
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:35:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
35 | int a_k = tileStart + threadIdx.x; // threadIdx.x in [0, TILE-1]
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:42:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
42 | int b_k = tileStart + threadIdx.y; // threadIdx.y in [0, TILE-1]
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:47:17: warning: result of multiplication in type 'unsigned int' is used as a pointer offset after an implicit widening conversion to type 'size_t' [bugprone-implicit-widening-of-multiplication-result]
47 | B_tile[threadIdx.y][threadIdx.x * 2] = B[b_k * N + global_col0];
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:47:37: note: make conversion explicit to silence this warning
4 | B_tile[threadIdx.y][threadIdx.x * 2] = B[b_k * N + global_col0];
| ^~~~~~~~~~~~~~~
| static_cast<size_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:47:37: note: perform multiplication in a wider type
47 | B_tile[threadIdx.y][threadIdx.x * 2] = B[b_k * N + global_col0];
| ^~~~~~~~~~~
| static_cast<size_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:49:17: warning: result of multiplication in type 'unsigned int' is used as a pointer offset after an implicit widening conversion to type 'size_t' [bugprone-implicit-widening-of-multiplication-result]
49 | B_tile[threadIdx.y][threadIdx.x * 2] = 0.0f;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:49:37: note: make conversion explicit to silence this warning
49 | B_tile[threadIdx.y][threadIdx.x * 2] = 0.0f;
| ^~~~~~~~~~~~~~~
| static_cast<size_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:49:37: note: perform multiplication in a wider type
49 | B_tile[threadIdx.y][threadIdx.x * 2] = 0.0f;
| ^~~~~~~~~~~
| static_cast<size_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:55:13: warning: result of multiplication in type 'unsigned int' is used as a pointer offset after an implicit widening conversion to type 'size_t' [bugprone-implicit-widening-of-multiplication-result]
55 | B_tile[threadIdx.y][threadIdx.x * 2] = 0.0f;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:55:33: note: make conversion explicit to silence this warning
55 | B_tile[threadIdx.y][threadIdx.x * 2] = 0.0f;
| ^~~~~~~~~~~~~~~
| static_cast<size_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:55:33: note: perform multiplication in a wider type
55 | B_tile[threadIdx.y][threadIdx.x * 2] = 0.0f;
| ^~~~~~~~~~~
| static_cast<size_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:65:29: warning: result of multiplication in type 'unsigned int' is used as a pointer offset after an implicit widening conversion to type 'size_t' [bugprone-implicit-widening-of-multiplication-result]
65 | out0 += a_val * B_tile[s][threadIdx.x * 2];
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:65:39: note: make conversion explicit to silence this warning
65 | out0 += a_val * B_tile[s][threadIdx.x * 2];
| ^~~~~~~~~~~~~~~
| static_cast<size_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:65:39: note: perform multiplication in a wider type
65 | out0 += a_val * B_tile[s][threadIdx.x * 2];
| ^~~~~~~~~~~
| static_cast<size_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:88: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]
88 | torch::Tensor forward(torch::Tensor A, torch::Tensor B) {
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:88: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]
88 | torch::Tensor forward(torch::Tensor A, torch::Tensor B) {
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:94:13: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
94 | int K = A.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/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 M = A.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250208_optimize_b5_s4_e1_sweep/level_1/task_16/b5_s2_tiled_double_output/base/base.cu:97:13: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
97 | int N = B.size(1);
| ^