import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a masked cumulative sum, only summing elements that satisfy a condition.

    Parameters:
        dim (int): The dimension along which to perform the masked cumulative sum.
    """

    def __init__(self, dim):
        super(Model, self).__init__()
        self.dim = dim

    def forward(self, x, mask):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, *input_shape).
            mask (torch.Tensor): Boolean mask of the same shape as x.

        Returns:
            torch.Tensor: Cumulative sum of elements where mask is True.
        """
        return torch.cumsum(x * mask, dim=self.dim)

batch_size = 32768
input_shape = (32768,)
dim = 1

def get_inputs():
    x = torch.rand(batch_size, *input_shape)
    mask = torch.randint(0, 2, x.shape).bool()  # Random boolean mask
    return [x, mask]

def get_init_inputs():
    return [dim]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor masked_cumsum_forward(torch::Tensor x, torch::Tensor mask);
'''

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

template<typename MaskT>
__device__ __forceinline__ float mval(MaskT m) { return m ? 1.0f : 0.0f; }
template<>
__device__ __forceinline__ float mval<float>(float m) { return m; }

template<int BT, int ITEMS, typename MaskT>
__global__ void masked_cumsum_kernel2(const float* __restrict__ x, const MaskT* __restrict__ mask,
                                      float* __restrict__ out, int rows, int cols) {
    int row = blockIdx.x;
    int base = row * cols;
    float vals[ITEMS];
    MaskT masks[ITEMS];
    typedef cub::BlockLoad<float, BT, ITEMS, cub::BLOCK_LOAD_WARP_TRANSPOSE> BlockLoadF;
    typedef cub::BlockLoad<MaskT, BT, ITEMS, cub::BLOCK_LOAD_WARP_TRANSPOSE> BlockLoadM;
    typedef cub::BlockScan<float, BT> BlockScan;
    typedef cub::BlockStore<float, BT, ITEMS, cub::BLOCK_STORE_WARP_TRANSPOSE> BlockStoreF;
    union TempStorage {
        typename BlockLoadF::TempStorage loadf;
        typename BlockLoadM::TempStorage loadm;
        typename BlockScan::TempStorage scan;
        typename BlockStoreF::TempStorage storef;
    };
    __shared__ TempStorage temp;
    float carry = 0.0f;
    const int TILE = BT * ITEMS;
    for (int off = 0; off < cols; off += TILE) {
        int valid = cols - off;
        if (valid > TILE) valid = TILE;
        BlockLoadF(temp.loadf).Load(x + base + off, vals, valid, 0.0f);
        __syncthreads();
        BlockLoadM(temp.loadm).Load(mask + base + off, masks, valid, (MaskT)0);
        __syncthreads();
        #pragma unroll
        for (int i = 0; i < ITEMS; ++i) vals[i] *= mval<MaskT>(masks[i]);
        float aggregate;
        BlockScan(temp.scan).InclusiveSum(vals, vals, aggregate);
        #pragma unroll
        for (int i = 0; i < ITEMS; ++i) vals[i] += carry;
        carry += aggregate;
        __syncthreads();
        BlockStoreF(temp.storef).Store(out + base + off, vals, valid);
        __syncthreads();
    }
}

torch::Tensor masked_cumsum_forward(torch::Tensor x, torch::Tensor mask) {
    auto out = torch::empty_like(x);
    int rows = (int)x.size(0), cols = (int)x.size(1);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (mask.scalar_type() == at::kBool) {
        masked_cumsum_kernel2<256, 32, bool><<<rows, 256, 0, stream>>>(x.data_ptr<float>(), mask.data_ptr<bool>(), out.data_ptr<float>(), rows, cols);
    } else {
        masked_cumsum_kernel2<256, 32, float><<<rows, 256, 0, stream>>>(x.data_ptr<float>(), mask.data_ptr<float>(), out.data_ptr<float>(), rows, cols);
    }
    return out;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, dim):
        super(ModelNew, self).__init__()
        self.dim = dim
    def forward(self, x, mask):
        return _ext.masked_cumsum_forward(x.contiguous(), mask.contiguous())