"""
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 = r'''
#include <torch/extension.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;
    
    unsigned int acc = 0;
    for (int r = 0; r < R; r++) {
        acc |= static_cast<unsigned int>(x[r * B + b]);
    }
    int cnt = __popc(acc);
    out[b] = static_cast<float>(cnt);
}

torch::Tensor or_popcount_cuda(torch::Tensor x) {
    int R = x.size(0);
    int B = x.size(1);
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());
    auto out = torch::empty({B}, options);
    
    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 = r'''
torch::Tensor or_popcount_cuda(torch::Tensor x);
'''

module = load_inline(
    name='or_popcount',
    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:
        x = x.contiguous()
        return self.module.or_popcount_cuda(x)