"""
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 bitpop_forward_out(torch::Tensor x, torch::Tensor out);
'''

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

__global__ void bitpop_i32_24_kernel(const int32_t* __restrict__ x, float* __restrict__ out) {
    int b = threadIdx.x;
    if (b < 24) {
        uint32_t v = (uint32_t)x[b] | (uint32_t)x[24+b] | (uint32_t)x[48+b] | (uint32_t)x[72+b];
        out[b] = (float)__popc(v);
    }
}

__global__ void bitpop_i64_24_kernel(const int64_t* __restrict__ x, float* __restrict__ out) {
    int b = threadIdx.x;
    if (b < 24) {
        uint32_t v = (uint32_t)x[b] | (uint32_t)x[24+b] | (uint32_t)x[48+b] | (uint32_t)x[72+b];
        out[b] = (float)__popc(v);
    }
}

__global__ void bitpop_f32_24_kernel(const float* __restrict__ x, float* __restrict__ out) {
    int b = threadIdx.x;
    if (b < 24) {
        uint32_t v = (uint32_t)((int64_t)x[b]) | (uint32_t)((int64_t)x[24+b]) |
                     (uint32_t)((int64_t)x[48+b]) | (uint32_t)((int64_t)x[72+b]);
        out[b] = (float)__popc(v);
    }
}

torch::Tensor bitpop_forward_out(torch::Tensor x, torch::Tensor out) {
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (x.scalar_type() == at::kInt) {
        bitpop_i32_24_kernel<<<1, 32, 0, stream>>>(x.data_ptr<int32_t>(), out.data_ptr<float>());
    } else if (x.scalar_type() == at::kLong) {
        bitpop_i64_24_kernel<<<1, 32, 0, stream>>>(x.data_ptr<int64_t>(), out.data_ptr<float>());
    } else if (x.scalar_type() == at::kFloat) {
        bitpop_f32_24_kernel<<<1, 32, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<float>());
    }
    return out;
}
'''

_ext = load_inline(
    name='bitpop_ext_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['bitpop_forward_out'],
    verbose=False,
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3'],
)

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._out is None or self._out.device != x.device:
            self._out = torch.empty((24,), device=x.device, dtype=torch.float32)
        return _ext.bitpop_forward_out(x, self._out)