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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor bitmask_popcount_forward(torch::Tensor x);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdint.h>

__global__ void bitmask_popcount_kernel(const int32_t* __restrict__ x,
                                        float* __restrict__ out,
                                        int B) {
    int b = threadIdx.x;
    if (b < B) {
        uint32_t v = (uint32_t)x[b]
                   | (uint32_t)x[B + b]
                   | (uint32_t)x[(B << 1) + b]
                   | (uint32_t)x[3 * B + b];
        out[b] = (float)__popc(v);
    }
}

torch::Tensor bitmask_popcount_forward(torch::Tensor x) {
    const int B = (int)x.size(1);
    auto out = torch::empty({B}, x.options().dtype(torch::kFloat32));
    bitmask_popcount_kernel<<<1, 32, 0, at::cuda::getCurrentCUDAStream()>>>(
        reinterpret_cast<const int32_t*>(x.data_ptr<int>()),
        out.data_ptr<float>(),
        B);
    return out;
}
"""

_ext = load_inline(
    name="bitmask_popcount_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["bitmask_popcount_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.bitmask_popcount_forward(x)