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

template <typename T>
__device__ __forceinline__ uint32_t elem_to_u32(T a) {
    return (uint32_t)((int64_t)a);
}

template <>
__device__ __forceinline__ uint32_t elem_to_u32<int32_t>(int32_t a) {
    return (uint32_t)a;
}

template <typename T>
__global__ void or_pop_kernel_t(const T* __restrict__ x, float* __restrict__ y) {
    int b = threadIdx.x;
    if (b < 24) {
        uint32_t v = elem_to_u32<T>(x[b]) | elem_to_u32<T>(x[24 + b]) |
                     elem_to_u32<T>(x[48 + b]) | elem_to_u32<T>(x[72 + b]);
        y[b] = (float)__popc(v);
    }
}

__global__ void or_pop_kernel_i32(const int32_t* __restrict__ x, float* __restrict__ y) {
    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];
        y[b] = (float)__popc(v);
    }
}

torch::Tensor or_pop_forward(torch::Tensor x) {
    auto y = torch::empty({24}, x.options().dtype(at::kFloat));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (x.scalar_type() == at::kInt) {
        or_pop_kernel_i32<<<1, 32, 0, stream>>>(x.data_ptr<int32_t>(), y.data_ptr<float>());
    } else if (x.scalar_type() == at::kLong) {
        or_pop_kernel_t<int64_t><<<1, 32, 0, stream>>>(x.data_ptr<int64_t>(), y.data_ptr<float>());
    } else if (x.scalar_type() == at::kFloat) {
        or_pop_kernel_t<float><<<1, 32, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>());
    } else if (x.scalar_type() == at::kDouble) {
        or_pop_kernel_t<double><<<1, 32, 0, stream>>>(x.data_ptr<double>(), y.data_ptr<float>());
    }
    return y;
}
"""

_ext = load_inline(
    name="or_popcount_r4_b24_final_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["or_pop_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.or_pop_forward(x)