← Back to Leaderboard

The AI CUDA Engineer 👷

94_MSELoss • mse_unrolled_optimized_edit_1

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


def module_fn(predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """
    Computes the Mean Squared Error loss for regression tasks.

    Args:
        predictions (torch.Tensor): Predicted values.
        targets (torch.Tensor): Target values.

    Returns:
        torch.Tensor: Mean Squared Error loss.
    """
    return F.mse_loss(predictions, targets, reduction="mean")


class Model(nn.Module):
    """
    A model that computes the Mean Squared Error loss for regression tasks.

    Parameters:
        None
    """

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

    def forward(self, predictions, targets, fn=module_fn):
        return fn(predictions, targets)


batch_size = 128
input_shape = (4096,)
dim = 1


def get_inputs():
    return [
        torch.randn(batch_size, *input_shape),
        torch.randn(batch_size, *input_shape),
    ]


def get_init_inputs():
    return []

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that computes the Mean Squared Error loss for regression tasks.

    Parameters:
        None
    """
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, predictions, targets):
        return torch.mean((predictions - targets) ** 2)

batch_size = 128
input_shape = (4096, )
dim = 1

def get_inputs():
    return [torch.randn(batch_size, *input_shape), torch.randn(batch_size, *input_shape)]

def get_init_inputs():
    return []
Download Evaluation Download PyTorch Download CUDA Download Profiles

Kernel Information

Operation Name 94_MSELoss
Level ID 1
Task ID 94
Kernel Name mse_unrolled_optimized_edit_1
CUDA Speedup (Native) 1.034x
CUDA Speedup (Compile) 2.036x
CUDA Runtime 0.016 ms
PyTorch Runtime (Native) 0.017 ms
PyTorch Runtime (Compile) 0.033 ms
Correct True
Max Diff (vs. Reference) 0.000000
Model azure-gpt-4o-2024-08-06
Temperature 0.00
View Experiment Progress Details

Related Kernels (Level 1, Task 94 • 94_MSELoss)

Rank Kernel Name Runtime (ms) Speedup Native Speedup Compile
🥇 optimized_thread_indexing_base 0.02 1.03 2.04
🥇 coalesced_shfl_mse_base 0.02 1.03 2.04
🥇 efficient_mse_base 0.02 1.03 2.04
🥇 mse_unrolled_optimized_base 0.02 1.03 2.04
🥇 mse_min_sync_edit_1 0.02 1.03 2.04
🥇 vectorized_ldg_mse_base 0.02 1.03 2.04
🥇 optimized_grid_stride_warp_reduce_base 0.02 1.03 2.04
🥇 mse_1d_optimized_indexing_base 0.02 1.03 2.04
🥇 mse_unrolled_optimized_edit_1 0.02 1.03 2.04
🥇 mse_warp_reduction_base 0.02 1.03 2.04
🥇 mse_unroll_pragma_base_base 0.02 1.03 2.04
🥇 mse_blocksize_experiment_base 0.02 1.03 2.04
🥇 mse_ldg_vectorized_edit_edit_1 0.02 1.03 2.04
🥇 mse_ldg_vectorized_edit_base 0.02 1.03 2.04
15 optimized_block_size_mse_base 0.02 0.97 1.92
15 stride_mse_loss_base 0.02 0.97 1.92
15 warp_uniform_mse_base 0.02 0.97 1.92
15 block_size_experimentation_base_base 0.02 0.97 1.92
15 optimized_mse_forward_base 0.02 0.97 1.92
15 warp_aligned_mse_base_base 0.02 0.97 1.92 
#include <pybind11/pybind11.h>
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

static const int BLOCK_SIZE = 512;
static const int VECTOR_SIZE = 4;  // Process 4 elements per thread

template <typename scalar_t>
__global__ void mse_forward_kernel_unrolled(
    const scalar_t* __restrict__ preds,
    const scalar_t* __restrict__ tgts,
    double* __restrict__ sum_out,
    const int64_t num_elements
) {
    __shared__ double shm[BLOCK_SIZE + 32];
    
    const int tid = threadIdx.x;
    const int gid = blockIdx.x * blockDim.x + tid;
    const int grid_stride = gridDim.x * blockDim.x;
    
    // Initialize local sum
    double local_sum = 0.0;
    
    // Vector loading - process 4 elements per iteration
    #pragma unroll
    for (int base_idx = gid * VECTOR_SIZE; base_idx < num_elements; base_idx += grid_stride * VECTOR_SIZE) {
        double diff[VECTOR_SIZE];
        
        // Manually unrolled vector load and computation
        #pragma unroll
        for (int i = 0; i < VECTOR_SIZE; i++) {
            int idx = base_idx + i;
            if (idx < num_elements) {
                diff[i] = static_cast<double>(preds[idx]) - static_cast<double>(tgts[idx]);
                local_sum += diff[i] * diff[i];
            }
        }
    }
    
    // Store in shared memory
    shm[tid] = local_sum;
    __syncthreads();
    
    // Reduction with manual unrolling for different power-of-2 sizes
    #pragma unroll
    if (tid < 256) { shm[tid] += shm[tid + 256]; } __syncthreads();
    #pragma unroll
    if (tid < 128) { shm[tid] += shm[tid + 128]; } __syncthreads();
    #pragma unroll
    if (tid < 64) { shm[tid] += shm[tid + 64]; } __syncthreads();
    
    // Warp-level reduction (no sync needed)
    if (tid < 32) {
        volatile double* vmem = shm;
        #pragma unroll
        if (tid < 32) vmem[tid] += vmem[tid + 32];
        #pragma unroll
        if (tid < 16) vmem[tid] += vmem[tid + 16];
        #pragma unroll
        if (tid < 8) vmem[tid] += vmem[tid + 8];
        #pragma unroll
        if (tid < 4) vmem[tid] += vmem[tid + 4];
        #pragma unroll
        if (tid < 2) vmem[tid] += vmem[tid + 2];
        #pragma unroll
        if (tid < 1) vmem[tid] += vmem[tid + 1];
    }
    
    // Single atomic add per block
    if (tid == 0) {
        atomicAdd(sum_out, shm[0]);
    }
}

torch::Tensor forward(torch::Tensor predictions, torch::Tensor targets) {
    TORCH_CHECK(predictions.is_cuda(), "predictions must be a CUDA tensor");
    TORCH_CHECK(targets.is_cuda(), "targets must be a CUDA tensor");
    TORCH_CHECK(predictions.numel() == targets.numel(),
                "predictions and targets must have the same number of elements");

    const int64_t num_elements = predictions.numel();
    
    // Calculate optimal grid size
    const int sm_count = 108; // H100 SM count
    const int blocks_per_sm = 4;
    const int num_blocks = sm_count * blocks_per_sm;
    
    auto accumulator = torch::zeros({1}, predictions.options().dtype(at::kDouble));

    AT_DISPATCH_FLOATING_TYPES(predictions.scalar_type(), "mse_forward_cuda", [&] {
        mse_forward_kernel_unrolled<scalar_t><<<num_blocks, BLOCK_SIZE>>>(
            predictions.data_ptr<scalar_t>(),
            targets.data_ptr<scalar_t>(),
            accumulator.data_ptr<double>(),
            num_elements
        );
    });

    auto result = accumulator.div_(static_cast<double>(num_elements));
    return result.to(predictions.dtype());
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Mean Squared Error (MSE) forward (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 0.950 inst/cycle 0.000 5
Executed Ipc Elapsed 0.512 inst/cycle 0.000 5
Issue Slots Busy 25.448 % 0.023 5
Issued Ipc Active 1.018 inst/cycle 0.000 5
SM Busy 25.448 % 0.023 5
Memory Throughput 744068824382.830 byte/second 144810831848252030976.000 5
Mem Busy 20.926 % 0.095 5
Max Bandwidth 22.346 % 0.106 5
L1/TEX Hit Rate 74.940 % 0.000 5
L2 Hit Rate 18.558 % 0.001 5
Mem Pipes Busy 12.862 % 0.037 5
Warp Cycles Per Issued Instruction 36.464 cycle 0.349 5
Warp Cycles Per Executed Instruction 39.036 cycle 0.399 5
Avg. Active Threads Per Warp 31.910 0.000 5
Avg. Not Predicated Off Threads Per Warp 27.960 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 4.000 block 0.000 5
Block Limit Shared Mem 12.000 block 0.000 5
Block Limit Warps 4.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 57.684 % 0.161 5
Achieved Active Warps Per SM 36.918 warp 0.066 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 CPIStall Check the Warp Stall Sampling (All Cycles) table for the top stall locations in your source based on sampling data. The Kernel Profiling Guide (https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#metrics-reference) provides more details on each stall reason.
WRN Occupancy This kernel's theoretical occupancy is not impacted by any block limit. The difference between calculated theoretical (100.0%) and measured achieved occupancy (57.5%) can be the result of warp scheduling overheads or workload imbalances during the kernel execution. Load imbalances can occur between warps within a block as well as across blocks of the same kernel. See the CUDA Best Practices Guide (https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#occupancy) for more details on optimizing occupancy.
Operation / Metric Value Unit
aten::to
CPU Time 1234229.79 μs
Device Time 229253.39 μs
Self CPU Time 47952.08 μ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 1186277.71 μs
Device Time 229253.39 μs
Self CPU Time 195255.34 μ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 4437893.10 μs
Device Time 6928855.62 μs
Self CPU Time 243705.66 μ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 4194189.20 μs
Device Time 6928855.62 μs
Self CPU Time 349565.12 μs
Self Device Time 6928777.50 μ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 4778176.09 μs
Device Time 531835.74 μs
Self CPU Time 4778176.09 μs
Self Device Time 531835.74 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
void mse_forward_kernel_unrolled<float>(float const*, float const*, double*, long)
CPU Time 0.00 μs
Device Time 438371.20 μs
Self CPU Time 0.00 μs
Self Device Time 438371.20 μ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 6726859.30 μs
Self CPU Time 0.00 μs
Self Device Time 6726859.30 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
Status: Failed 
45243 warnings and 9 errors generated when compiling for host.
Error while processing /home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu.
Suppressed 45284 warnings (45237 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.
Found compiler error(s).
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:18:21 bugprone-narrowing-conversions
18 | const int tid = threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:19:21: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
19 | const int gid = blockIdx.x * blockDim.x + tid;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:20:29: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
20 | const int grid_stride = gridDim.x * blockDim.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:47:5: error: expected a for, while, or do-while loop to follow '#pragma unroll' [clang-diagnostic-error]
47 | if (tid < 256) { shm[tid] += shm[tid + 256]; } __syncthreads();
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:49:5: error: expected a for, while, or do-while loop to follow '#pragma unroll' [clang-diagnostic-error]
49 | if (tid < 128) { shm[tid] += shm[tid + 128]; } __syncthreads();
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:51:5: error: expected a for, while, or do-while loop to follow '#pragma unroll' [clang-diagnostic-error]
51 | if (tid < 64) { shm[tid] += shm[tid + 64]; } __syncthreads();
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:57:9: error: expected a for, while, or do-while loop to follow '#pragma unroll' [clang-diagnostic-error]
57 | if (tid < 32) vmem[tid] += vmem[tid + 32];
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:59:9: error: expected a for, while, or do-while loop to follow '#pragma unroll' [clang-diagnostic-error]
59 | if (tid < 16) vmem[tid] += vmem[tid + 16];
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:61:9: error: expected a for, while, or do-while loop to follow '#pragma unroll' [clang-diagnostic-error]
61 | if (tid < 8) vmem[tid] += vmem[tid + 8];
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:63:9: error: expected a for, while, or do-while loop to follow '#pragma unroll' [clang-diagnostic-error]
63 | if (tid < 4) vmem[tid] += vmem[tid + 4];
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:65:9: error: expected a for, while, or do-while loop to follow '#pragma unroll' [clang-diagnostic-error]
65 | if (tid < 2) vmem[tid] += vmem[tid + 2];
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:67:9: error: expected a for, while, or do-while loop to follow '#pragma unroll' [clang-diagnostic-error]
67 | if (tid < 1) vmem[tid] += vmem[tid + 1];
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_1/task_94/b4_s2_mse_unrolled_optimized/edit_1/edit_1.cu:91:5: warning: inside a lambda, '__func__' expands to the name of the function call operator; consider capturing the name of the enclosing function explicitly [bugprone-lambda-function-name]
91 | AT_DISPATCH_FLOATING_TYPES(predictions.scalar_type(), "mse_forward_cuda", [&] {
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:237:34: note: expanded from macro 'AT_DISPATCH_FLOATING_TYPES'
237 | AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:233:3: note: expanded from macro 'AT_DISPATCH_CASE_FLOATING_TYPES'
233 | AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__) \
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:74:3: note: expanded from macro 'AT_DISPATCH_CASE'
74 | AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, scalar_t, __VA_ARGS__)
| ^
note: (skipping 1 expansions in backtrace; use -fmacro-backtrace-limit=0 to see all)
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:58:7: note: expanded from macro 'AT_PRIVATE_CHECK_SELECTIVE_BUILD'
58 | AT_ERROR( \
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/c10/util/Exception.h:711:32: note: expanded from macro 'AT_ERROR'
711 | C10_EXPAND_MSVC_WORKAROUND(TORCH_CHECK(false, ::c10::str(__VA_ARGS__))); \
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/c10/util/Exception.h:536:9: note: expanded from macro 'TORCH_CHECK'
536 | __func__, \
| ^