"""
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.h>
#include <cuda_runtime.h>

__global__ void popcount_kernel(const int32_t* __restrict__ x, float* __restrict__ out, int R, int B) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < B) {
        // Load R elements and compute bitwise OR
        uint32_t acc = 0;
        for (int r = 0; r < R; ++r) {
            // Access x[r, idx] - row-major layout: x[r * B + idx]
            uint32_t val = static_cast<uint32_t>(x[r * B + idx]);
            acc |= val;
        }
        // Popcount using CUDA intrinsic
        int count = __popc(acc);
        out[idx] = static_cast<float>(count);
    }
}

torch::Tensor popcount_forward(torch::Tensor x) {
    // x: (R, B) int32
    TORCH_CHECK(x.dim() == 2, "x must be 2D");
    
    int R = x.size(0);
    int B = x.size(1);
    
    // Ensure contiguous int32 tensor
    torch::Tensor x_contiguous;
    if (x.dtype() == torch::kInt32 && x.is_contiguous()) {
        x_contiguous = x;
    } else {
        x_contiguous = x.toType(torch::kInt32).contiguous();
    }
    
    auto out = torch::empty({B}, torch::TensorOptions().dtype(torch::kFloat32).device(x.device()));
    
    int threads = 256;
    int blocks = (B + threads - 1) / threads;
    
    popcount_kernel<<<blocks, threads>>>(
        x_contiguous.data_ptr<int32_t>(),
        out.data_ptr<float>(),
        R, B
    );
    
    return out;
}
"""

ext = load_inline(
    name='popcount_ext_final',
    cpp_sources=['torch::Tensor popcount_forward(torch::Tensor x);'],
    cuda_sources=[cuda_source],
    functions=['popcount_forward'],
    extra_cuda_cflags=['-O3'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.ext = ext
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.ext.popcount_forward(x)