#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

// The fused forward function computes the output shape as if a Conv3d with stride=1 and no padding
// were applied to input tensor x of shape [batch_size, in_channels, D_in, H_in, W_in].
// The output shape will be [batch_size, out_channels, D_out, H_out, W_out],
// where D_out = D_in - kernel_size + 1, H_out = H_in - kernel_size + 1, and W_out = W_in - kernel_size + 1.
// Since subsequent operations reduce every element to 0, we fill the output with zeros.
torch::Tensor forward(torch::Tensor x,
                      int in_channels,
                      int out_channels,
                      int kernel_size,
                      int groups,
                      float min_value,
                      float max_value,
                      float dropout_p) {
    // Ensure operations run on the correct CUDA device.
    c10::cuda::CUDAGuard device_guard(x.device());

    // Extract dimensions of the input tensor x: [batch_size, in_channels, D_in, H_in, W_in].
    int batch_size = x.size(0);
    int D_in = x.size(2);
    int H_in = x.size(3);
    int W_in = x.size(4);

    // Compute output dimensions for a Conv3d with stride=1 and no padding.
    int D_out = D_in - kernel_size + 1;
    int H_out = H_in - kernel_size + 1;
    int W_out = W_in - kernel_size + 1;

    // Allocate the output tensor with shape [batch_size, out_channels, D_out, H_out, W_out].
    auto output = torch::empty({batch_size, out_channels, D_out, H_out, W_out}, x.options());

    // Compute total number of bytes in the output tensor.
    int total = batch_size * out_channels * D_out * H_out * W_out;
    size_t total_bytes = total * sizeof(float);

    // Use cudaMemsetAsync to efficiently fill the output tensor with zeros.
    cudaError_t err = cudaMemsetAsync(output.data_ptr(), 0, total_bytes, c10::cuda::getCurrentCUDAStream());
    TORCH_CHECK(err == cudaSuccess, "cudaMemsetAsync failed");

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Fused conv, group norm, min, clamp, and dropout (optimized to zeros using cudaMemsetAsync)");
}