"""
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 []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void popcount_or_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;
    int64_t acc = 0;
    for (int r = 0; r < R; ++r) {
        int32_t val = x[r * B + b];
        acc |= static_cast<int64_t>(val) & 0xFFFFFFFFLL;
    }
    int pop = __popcll(acc);
    out[b] = static_cast<float>(pop);
}

torch::Tensor popcount_or_cuda(torch::Tensor x) {
    auto x_int = x.to(torch::kInt32).contiguous();
    int R = x_int.size(0);
    int B = x_int.size(1);
    auto out = torch::empty({B}, x.options().dtype(torch::kFloat32).device(x.device()));
    dim3 block(32);
    dim3 grid((B + 31) / 32);
    popcount_or_kernel<<<grid, block>>>(x_int.data_ptr<int32_t>(), out.data_ptr<float>(), R, B);
    return out;
}
'''

cpp_source = '''
torch::Tensor popcount_or_cuda(torch::Tensor x);
'''

module = load_inline(
    name='popcount_or',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['popcount_or_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.popcount_or_cuda(x)
```