incorrect
31_bitmask_or_popcount.py
started
2026-05-11T08:53:45Z
correctness issue:
timingkernelbench.timing
kernel mean / median / std
—
ref mean / median / std
0.26 / — / 7.93e-04 μs
roofline / occupancynsight_hardware_counters
reference PyTorch source (the prompt the agent sees)
"""
R int32 bitmasks per batch element: reduce with bitwise OR across virtual ranks, then popcount (32-bit).
Output shape (B,) float32: number of set bits in the unsigned lower 32 bits of the OR result.
"""
import torch
import torch.nn as nn
from kernelbench.distributed_collectives import default_device
R = 4
B = 24
def _popcount_u32(v: torch.Tensor) -> torch.Tensor:
"""v: int32 or int64 tensor (B,), return float32 popcount in 0..32."""
b = v.shape[0]
u = v.to(torch.int64) & 0xFFFFFFFF
c = torch.zeros(b, dtype=torch.int64, device=v.device)
for k in range(32):
c = c + ((u >> k) & 1)
return c.to(torch.float32)
class Model(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: (R, B) int32
acc = x[0].to(torch.int64)
for r in range(1, R):
acc = acc | x[r].to(torch.int64)
acc = acc & 0xFFFFFFFF
return _popcount_u32(acc.to(torch.int32))
def get_inputs():
dev = default_device()
p = popcorn_pri
ints = torch.randint(
-(2**31), 2**31 - 1, (p.jitter_int(R), p.jitter_int(B)), device=dev
)
return [ints.to(torch.int32)]
def get_init_inputs():
return []
turn 0531.5s turn · 38.8s LLM · 492.8s tools · tokens in=1,233 · out=407 · total=1,856 · FINAL
reasoning
The task is to replace the forward computation with a custom CUDA kernel for this specific operation: reducing R=4 int32 bitmasks per batch element with bitwise OR across virtual ranks, then popcount ...
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = '''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void or_popcount_kernel(const int32_t* __restrict__ x, float* __restrict__ out, int R, int B) {\n int b = blockIdx.x * blockDim.x + threadIdx.x;\n if (b >= B) return;\n \n int32_t acc = x[b]; // first rank\n for (int r = 1; r < R; ++r) {\n acc |= x[r * B + b];\n }\n acc &= 0xFFFFFFFF;\n \n // popcount\n int count = __popc((unsigned int)acc);\n out[b] = (float)count;\n}\n\ntorch::Tensor or_popcount_cuda(torch::Tensor x) {\n int R = x.size(0);\n int B = x.size(1);\n auto out = torch::empty({B}, x.options().dtype(torch::kFloat32));\n \n dim3 block(256);\n dim3 grid((B + 255) / 256);\n \n or_popcount_kernel<<<grid, block>>>(\n x.data_ptr<int32_t>(),\n out.data_ptr<float>(),\n R, B\n );\n \n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor or_popcount_cuda(torch::Tensor x);\n'''\n\nmodule = load_inline(\n name='or_popcount_ext',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['or_popcount_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.module.or_popcount_cuda(x)\n"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: correctness check did not pass (? trials).
runtime_error: expected scalar type Int but found Float
final submitted kernel (62 lines, 1,421 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void or_popcount_kernel(const int32_t* __restrict__ x, float* __restrict__ out, int R, int B) {
int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= B) return;
int32_t acc = x[b]; // first rank
for (int r = 1; r < R; ++r) {
acc |= x[r * B + b];
}
acc &= 0xFFFFFFFF;
// popcount
int count = __popc((unsigned int)acc);
out[b] = (float)count;
}
torch::Tensor or_popcount_cuda(torch::Tensor x) {
int R = x.size(0);
int B = x.size(1);
auto out = torch::empty({B}, x.options().dtype(torch::kFloat32));
dim3 block(256);
dim3 grid((B + 255) / 256);
or_popcount_kernel<<<grid, block>>>(
x.data_ptr<int32_t>(),
out.data_ptr<float>(),
R, B
);
return out;
}
'''
cpp_source = '''
torch::Tensor or_popcount_cuda(torch::Tensor x);
'''
module = load_inline(
name='or_popcount_ext',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['or_popcount_cuda'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.module = module
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.module.or_popcount_cuda(x)