```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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


_conv3d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

constexpr int TILE_M = 128;
constexpr int TILE_N = 64;
constexpr int THREADS = 256;
constexpr int K_TOTAL = 81;
constexpr int A_PAD = 83;
constexpr int B_PAD = 68;
constexpr int SHARED_FLOATS =
    TILE_M * A_PAD + K_TOTAL * B_PAD + 2 * TILE_M + TILE_N + K_TOTAL;
constexpr size_t SHARED_BYTES = SHARED_FLOATS * sizeof(float);

__global__ void __launch_bounds__(THREADS, 2) conv3d_3x3x3_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int batch,
    const int out_channels,
    const int in_d,
    const int in_h,
    const int in_w,
    const int out_d,
    const int out_h,
    const int out_w,
    const int total_m,
    const bool has_bias) {
  extern __shared__ float shared_mem[];

  float* As = shared_mem;
  float* Bs = As + TILE_M * A_PAD;
  int* input_base = reinterpret_cast<int*>(Bs + K_TOTAL * B_PAD);
  int* output_base = input_base + TILE_M;
  float* bias_values = reinterpret_cast<float*>(output_base + TILE_M);
  int* kernel_offsets = reinterpret_cast<int*>(bias_values + TILE_N);

  const int tid = threadIdx.x;
  const int m0 = blockIdx.x * TILE_M;
  const int channel0 = blockIdx.y * TILE_N;

  const int out_plane = out_h * out_w;
  const int out_spatial = out_d * out_plane;
  const int in_plane = in_h * in_w;
  const int in_volume = in_d * in_plane;

  if (tid < TILE_M) {
    const int global_m = m0 + tid;
    int in_base = -1;
    int out_base = 0;

    if (global_m < total_m) {
      const int n = global_m / out_spatial;
      const int spatial_pos = global_m - n * out_spatial;
      const int d = spatial_pos / out_plane;
      const int plane_pos = spatial_pos - d * out_plane;
      const int h = plane_pos / out_w;
      const int w = plane_pos - h * out_w;

      in_base = (((n * 3) * in_d + d) * in_h + h) * in_w + w;
      out_base = n * out_channels * out_spatial + spatial_pos;
    }

    input_base[tid] = in_base;
    output_base[tid] = out_base;
  }

  if (tid < K_TOTAL) {
    const int c = tid / 27;
    const int remainder_c = tid - c * 27;
    const int kd = remainder_c / 9;
    const int remainder_d = remainder_c - kd * 9;
    const int kh = remainder_d / 3;
    const int kw = remainder_d - kh * 3;
    kernel_offsets[tid] =
        c * in_volume + kd * in_plane + kh * in_w + kw;
  }

  for (int idx = tid; idx < TILE_N * K_TOTAL; idx += THREADS) {
    const int channel = idx / K_TOTAL;
    const int k = idx - channel * K_TOTAL;
    const int global_channel = channel0 + channel;
    float value = 0.0f;

    if (global_channel < out_channels) {
      value = __ldg(weight + global_channel * K_TOTAL + k);
    }

    Bs[k * B_PAD + channel] = value;
  }

  if (tid < TILE_N) {
    const int global_channel = channel0 + tid;
    float value = 0.0f;
    if (has_bias && global_channel < out_channels) {
      value = __ldg(bias + global_channel);
    }
    bias_values[tid] = value;
  }

  for (int idx = tid; idx < TILE_M * K_TOTAL; idx += THREADS) {
    const int k = idx >> 7;
    const int m = idx & (TILE_M - 1);
    const int base = input_base[m];
    float value = 0.0f;

    if (base >= 0) {
      value = __ldg(input + base + kernel_offsets[k]);
    }

    As[m * A_PAD + k] = value;
  }

  __syncthreads();

  const int row0 = (tid >> 3) * 4;
  const int col0 = (tid & 7) * 8;

  float accum[4][8] = {};

#pragma unroll 9
  for (int k = 0; k < K_TOTAL; ++k) {
    const float a0 = As[(row0 + 0) * A_PAD + k];
    const float a1 = As[(row0 + 1) * A_PAD + k];
    const float a2 = As[(row0 + 2) * A_PAD + k];
    const float a3 = As[(row0 + 3) * A_PAD + k];

    const float4 b0 =
        *reinterpret_cast<const float4*>(Bs + k * B_PAD + col0);
    const float4 b1 =
        *reinterpret_cast<const float4*>(Bs + k * B_PAD + col0 + 4);

    accum[0][0] = fmaf(a0, b0.x, accum[0][0]);
    accum[0][1] = fmaf(a0, b0.y, accum[0][1]);
    accum[0][2] = fmaf(a0, b0.z, accum[0][2]);
    accum[0][3] = fmaf(a0, b0.w, accum[0][3]);
    accum[0][4] = fmaf(a0, b1.x, accum[0][4]);
    accum[0][5] = fmaf(a0, b1.y, accum[0][5]);
    accum[0][6] = fmaf(a0, b1.z, accum[0][6]);
    accum[0][7] = fmaf(a0, b1.w, accum[0][7]);

    accum[1][0] = fmaf(a1, b0.x, accum[1][0]);
    accum[1][1] = fmaf(a1, b0.y, accum[1][1]);
    accum[1][2] = fmaf(a1, b0.z, accum[1][2]);
    accum[1][3] = fmaf(a1, b0.w, accum[1][3]);
    accum[1][4] = fmaf(a1, b1.x, accum[1][4]);
    accum[1][5] = fmaf(a1, b1.y, accum[1][5]);
    accum[1][6] = fmaf(a1, b1.z, accum[1][6]);
    accum[1][7] = fmaf(a1, b1.w, accum[1][7]);

    accum[2][0] = fmaf(a2, b0.x, accum[2][0]);
    accum[2][1] = fmaf(a2, b0.y, accum[2][1]);
    accum[2][2] = fmaf(a2, b0.z, accum[2][2]);
    accum[2][3] = fmaf(a2, b0.w, accum[2][3]);
    accum[2][4] = fmaf(a2, b1.x, accum[2][4]);
    accum[2][5] = fmaf(a2, b1.y, accum[2][5]);
    accum[2][6] = fmaf(a2, b1.z, accum[2][6]);
    accum[2][7] = fmaf(a2, b1.w, accum[2][7]);

    accum[3][0] = fmaf(a3, b0.x, accum[3][0]);
    accum[3][1] = fmaf(a3, b0.y, accum[3][1]);
    accum[3][2] = fmaf(a3, b0.z, accum[3][2]);
    accum[3][3] = fmaf(a3, b0.w, accum[3][3]);
    accum[3][4] = fmaf(a3, b1.x, accum[3][4]);
    accum[3][5] = fmaf(a3, b1.y, accum[3][5]);
    accum[3][6] = fmaf(a3, b1.z, accum[3][6]);
    accum[3][7] = fmaf(a3, b1.w, accum[3][7]);
  }

  __syncthreads();

#pragma unroll
  for (int i = 0; i < 4; ++i) {
#pragma unroll
    for (int j = 0; j < 8; ++j) {
      As[(row0 + i) * A_PAD + col0 + j] = accum[i][j];
    }
  }

  __syncthreads();

  for (int idx = tid; idx < TILE_M * TILE_N; idx += THREADS) {
    const int channel = idx >> 7;
    const int m = idx & (TILE_M - 1);
    const int global_m = m0 + m;
    const int global_channel = channel0 + channel;

    if (global_m < total_m && global_channel < out_channels) {
      float value = As[m * A_PAD + channel];
      if (has_bias) {
        value += bias_values[channel];
      }
      output[output_base[m] + global_channel * out_spatial] = value;
    }
  }
}

}  // namespace

torch::Tensor conv3d_3x3x3_forward(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias) {
  TORCH_CHECK(input.dim() == 5, "input must be a 5D NCDHW tensor");
  TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
  TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
  TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
  TORCH_CHECK(weight.scalar_type() == at::kFloat, "weight must be float32");
  TORCH_CHECK(weight.dim() == 5, "weight must be a 5D tensor");
  TORCH_CHECK(weight.size(1) == 3, "custom kernel requires 3 input channels");
  TORCH_CHECK(weight.size(2) == 3 && weight.size(3) == 3 &&
                  weight.size(4) == 3,
              "custom kernel requires a 3x3x3 kernel");

  auto x = input.contiguous();
  auto w = weight.contiguous();

  const int batch = static_cast<int>(x.size(0));
  const int in_d = static_cast<int>(x.size(2));
  const int in_h = static_cast<int>(x.size(3));
  const int in_w = static_cast<int>(x.size(4));
  const int out_channels = static_cast<int>(w.size(0));

  const int out_d = in_d - 2;
  const int out_h = in_h - 2;
  const int out_w = in_w - 2;

  TORCH_CHECK(out_d > 0 && out_h > 0 && out_w > 0,
              "input spatial dimensions must be at least 3");

  const bool use_bias = bias.defined() && bias.numel() > 0;
  torch::Tensor b;
  const float* bias_ptr = nullptr;

  if (use_bias) {
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must be float32");
    TORCH_CHECK(bias.numel() == out_channels,
                "bias must have one element per output channel");
    b = bias.contiguous();
    bias_ptr = b.data_ptr<float>();
  }

  const int64_t out_spatial64 =
      static_cast<int64_t>(out_d) * out_h * out_w;
  const int64_t total_m64 = static_cast<int64_t>(batch) * out_spatial64;
  const int64_t output_numel64 =
      total_m64 * static_cast<int64_t>(out_channels);
  const int64_t input_numel64 = x.numel();

  TORCH_CHECK(total_m64 <= INT32_MAX,
              "custom kernel requires fewer than INT32_MAX output positions");
  TORCH_CHECK(output_numel64 <= INT32_MAX,
              "custom kernel requires an output smaller than INT32_MAX elements");
  TORCH_CHECK(input_numel64 <= INT32_MAX,
              "custom kernel requires an input smaller than INT32_MAX elements");

  auto output = torch::empty(
      {batch, out_channels, out_d, out_h, out_w}, x.options());

  const int total_m = static_cast<int>(total_m64);
  const dim3 grid(
      (total_m + TILE_M - 1) / TILE_M,
      (out_channels + TILE_N - 1) / TILE_N);

  const cudaError_t attribute_error = cudaFuncSetAttribute(
      conv3d_3x3x3_kernel,
      cudaFuncAttributeMaxDynamicSharedMemorySize,
      static_cast<int>(SHARED_BYTES));
  TORCH_CHECK(attribute_error == cudaSuccess,
              "failed to configure shared memory: ",
              cudaGetErrorString(attribute_error));

  cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
  conv3d_3x3x3_kernel<<<grid, THREADS, SHARED_BYTES, stream>>>(
      x.data_ptr<float>(),
      w.data_ptr<float>(),
      bias_ptr,
      output.data_ptr<float>(),
      batch,
      out_channels,
      in_d,
      in_h,
      in_w,
      out_d,
      out_h,
      out_w,
      total_m,
      use_bias);

  const cudaError_t launch_error = cudaGetLastError();
  TORCH_CHECK(launch_error == cudaSuccess,
              "custom 3D convolution kernel launch failed: ",
              cudaGetErrorString(launch_error));

  return output;
}
"""

_conv3d_cpp_source = (
    "torch::Tensor conv3d_3x3x3_forward("
    "torch::Tensor input, torch::Tensor weight, torch::Tensor bias);"
)

_custom_conv3d = load_inline(
    name="custom_implicit_gemm_conv3d_3x3x3",
    cpp_sources=_conv3d_cpp_source,
    cuda_sources=_conv3d_cuda_source,
    functions=["conv3d_3x3x3_forward"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _CustomConv3d3x3x3(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, bias):
        has_bias = bias is not None
        ctx.has_bias = has_bias

        if has_bias:
            ctx.save_for_backward(input, weight, bias)
            bias_argument = bias
        else:
            ctx.save_for_backward(input, weight)
            bias_argument = torch.empty(
                0, device=input.device, dtype=input.dtype
            )

        return _custom_conv3d.conv3d_3x3x3_forward(
            input, weight, bias_argument
        )

    @staticmethod
    def backward(ctx, grad_output):
        if ctx.has_bias:
            input, weight, bias = ctx.saved_tensors
        else:
            input, weight = ctx.saved_tensors
            bias = None

        needs_input_grad, needs_weight_grad, needs_bias_grad = (
            ctx.needs_input_grad
        )

        detached_input = input.detach()
        detached_weight = weight.detach()
        detached_bias = bias.detach() if bias is not None else None

        grad_targets = []
        target_names = []

        if needs_input_grad:
            detached_input.requires_grad_(True)
            grad_targets.append(detached_input)
            target_names.append("input")

        if needs_weight_grad:
            detached_weight.requires_grad_(True)
            grad_targets.append(detached_weight)
            target_names.append("weight")

        if needs_bias_grad and detached_bias is not None:
            detached_bias.requires_grad_(True)
            grad_targets.append(detached_bias)
            target_names.append("bias")

        if not grad_targets:
            return None, None, None

        with torch.enable_grad():
            reference_output = F.conv3d(
                detached_input,
                detached_weight,
                detached_bias,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
            )
            gradients = torch.autograd.grad(
                reference_output,
                grad_targets,
                grad_output,
                allow_unused=True,
            )

        gradient_map = dict(zip(target_names, gradients))
        return (
            gradient_map.get("input"),
            gradient_map.get("weight"),
            gradient_map.get("bias"),
        )


class ModelNew(nn.Module):
    """
    3D convolution accelerated by a custom implicit-GEMM CUDA kernel for the
    common float32, 3-input-channel, 3x3x3, unit-stride, zero-padding case.
    Other configurations use the standard PyTorch implementation.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        groups: int = 1,
        bias: bool = False,
    ):
        super(ModelNew, self).__init__()
        self.conv3d = nn.Conv3d(
            in_channels,
            out_channels,
            (kernel_size, kernel_size, kernel_size),
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv3d

        use_custom_kernel = (
            x.is_cuda
            and conv.weight.is_cuda
            and x.dtype == torch.float32
            and conv.in_channels == 3
            and conv.groups == 1
            and conv.kernel_size == (3, 3, 3)
            and conv.stride == (1, 1, 1)
            and conv.padding == (0, 0, 0)
            and conv.dilation == (1, 1, 1)
            and not torch.is_autocast_enabled()
        )

        if use_custom_kernel:
            return _CustomConv3d3x3x3.apply(
                x, conv.weight, conv.bias
            )

        return conv(x)
```