content: import time
  from abc import ABC, abstractmethod
  
  import torch
  import torch.nn.functional as F
  from torch.utils.cpp_extension import load_inline
  
  from .utils import SINGLE_IMAGE_TYPE, image_to_tensor
  
  pb_solver_cpp_source = """
  #include <torch/extension.h>
  
  #include <vector>
  
  // CUDA forward declarations
  
  void pb_solver_run_cuda(
      torch::Tensor A,
      torch::Tensor X,
      torch::Tensor B,
      torch::Tensor Xbuf,
      int num_iters
  );
  
  // C++ interface
  
  #define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
  #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
  #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
  
  void pb_solver_run(
      torch::Tensor A,
      torch::Tensor X,
      torch::Tensor B,
      torch::Tensor Xbuf,
      int num_iters
  ) {
      CHECK_INPUT(A);
      CHECK_INPUT(X);
      CHECK_INPUT(B);
      CHECK_INPUT(Xbuf);
  
      pb_solver_run_cuda(A, X, B, Xbuf, num_iters);
      return;
  }
  
  """
  
  pb_solver_cuda_source = """
  #include <torch/extension.h>
  
  #include <cuda.h>
  #include <cuda_runtime.h>
  
  #include <algorithm>
  
  
  __global__ void pb_solver_run_cuda_kernel(
      torch::PackedTensorAccessor32<int64_t,2,torch::RestrictPtrTraits> A,
      torch::PackedTensorAccessor32<float,1,torch::RestrictPtrTraits> X,
      torch::PackedTensorAccessor32<float,1,torch::RestrictPtrTraits> B,
      torch::PackedTensorAccessor32<float,1,torch::RestrictPtrTraits> Xbuf
  ) {
      const int index = blockIdx.x * blockDim.x + threadIdx.x;
      
      if (index < A.size(0)) {
          Xbuf[index] = (X[A[index][0]] + X[A[index][1]] + X[A[index][2]] + X[A[index][3]] + B[index]) / 4.;
      }
  }
  
  void pb_solver_run_cuda(
      torch::Tensor A,
      torch::Tensor X,
      torch::Tensor B,
      torch::Tensor Xbuf,
      int num_iters
  ) {
      int batch_size = A.size(0);
  
      const int threads = 1024;
      const dim3 blocks((batch_size + threads - 1) / threads);
      
      auto A_ptr = A.packed_accessor32<int64_t,2,torch::RestrictPtrTraits>();
      auto X_ptr = X.packed_accessor32<float,1,torch::RestrictPtrTraits>();
      auto B_ptr = B.packed_accessor32<float,1,torch::RestrictPtrTraits>();
      auto Xbuf_ptr = Xbuf.packed_accessor32<float,1,torch::RestrictPtrTraits>();
  
      for (int i = 0; i < num_iters; ++i) {
          pb_solver_run_cuda_kernel<<<blocks, threads>>>(
              A_ptr,
              X_ptr,
              B_ptr,
              Xbuf_ptr
          );
          cudaDeviceSynchronize();
          std::swap(X_ptr, Xbuf_ptr);
      }
      // we may waste an iteration here, but it's fine
      return;
  }
  """
  
  
  class PBBackend(ABC):
      def __init__(self, *args, **kwargs) -> None:
          pass
  
      def solve(self, num_iters, A, X, B, Xbuf) -> None:
          pass
  
  
  try:
      import triton
      import triton.language as tl
  
      @triton.jit
      def pb_triton_step_kernel(
          A_ptr,
          X_ptr,
          B_ptr,
          Xbuf_ptr,
          A_row_stride: tl.constexpr,
          n_elements: tl.constexpr,
          BLOCK_SIZE: tl.constexpr,
      ):
          pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
          block_start = pid * BLOCK_SIZE
          offsets = tl.arange(0, BLOCK_SIZE) + block_start
  
          A_start_ptr = A_ptr + block_start * A_row_stride
          A_ptrs = A_start_ptr + (
              tl.arange(0, BLOCK_SIZE)[:, None] * A_row_stride
              + tl.arange(0, A_row_stride)[None, :]
          )
          B_ptrs = B_ptr + offsets
          Xbuf_ptrs = Xbuf_ptr + offsets
  
          mask = offsets < n_elements
  
          A = tl.load(A_ptrs, mask=mask[:, None], other=0)
          X = tl.load(X_ptr + A)
          B = tl.load(B_ptrs, mask=mask, other=0.0)
  
          Xout = (tl.sum(X, axis=1) + B) / 4
          tl.store(Xbuf_ptrs, Xout, mask=mask)
  
  except:
      pb_triton_step_kernel = None
  
  
  class PBTorchCUDAKernelBackend(PBBackend):
      def __init__(self) -> None:
          self.kernel = load_inline(
              name="pb_solver",
              cpp_sources=[pb_solver_cpp_source],
              cuda_sources=[pb_solver_cuda_source],
              functions=["pb_solver_run"],
              verbose=True,
          )
  
      def solve(self, num_iters, A, X, B, Xbuf) -> None:
          self.kernel.pb_solver_run(A, X, B, Xbuf, num_iters)
  
  
  class PBTritonBackend(PBBackend):
      def step(
          self, X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, Xbuf: torch.Tensor
      ) -> None:
          n_elements = X.shape[0]
          grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
          assert pb_triton_step_kernel is not None
          pb_triton_step_kernel[grid](A, X, B, Xbuf, A.stride(0), n_elements, 1024)
  
      def solve(self, num_iters, A, X, B, Xbuf) -> None:
          for _ in range(num_iters):
              self.step(X, A, B, Xbuf)
              X, Xbuf = Xbuf, X
  
  
  class PBTorchNativeBackend(PBBackend):
      def step(
          self, X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, Xbuf: torch.Tensor
      ) -> None:
          X.copy_(X[A].sum(-1).add_(B).div_(4))
  
      def solve(self, num_iters, A, X, B, Xbuf) -> None:
          for _ in range(num_iters):
              self.step(X, A, B, Xbuf)
  
  
  class PoissonBlendingSolver:
      def __init__(self, backend: str, device: str):
          self.backend = backend
          if backend == "torch-native":
              self.pb_solver = PBTorchNativeBackend()
          elif backend == "torch-cuda":
              self.pb_solver = PBTorchCUDAKernelBackend()
          elif backend == "triton":
              self.pb_solver = PBTritonBackend()
          else:
              raise ValueError(f"Unknown backend: {backend}")
  
          self.device = device
          self.lap_kernel = torch.tensor(
              [[0, -1, 0], [-1, 4, -1], [0, -1, 0]], device=device, dtype=torch.float32
          ).view(1, 1, 3, 3)
          self.lap_kernel4 = torch.tensor(
              [
                  [[0, -1, 0], [0, 1, 0], [0, 0, 0]],
                  [[0, 0, 0], [0, 1, 0], [0, -1, 0]],
                  [[0, 0, 0], [-1, 1, 0], [0, 0, 0]],
                  [[0, 0, 0], [0, 1, -1], [0, 0, 0]],
              ],
              device=device,
              dtype=torch.float32,
          ).view(4, 1, 3, 3)
          self.neighbor_kernel = torch.tensor(
              [[0, 1, 0], [1, 0, 1], [0, 1, 0]], device=device, dtype=torch.float32
          ).view(1, 1, 3, 3)
  
      def __call__(
          self,
          src: SINGLE_IMAGE_TYPE,
          mask: SINGLE_IMAGE_TYPE,
          tgt: SINGLE_IMAGE_TYPE,
          num_iters: int,
          inplace: bool = True,
          grad_mode: str = "src",
      ):
          src = image_to_tensor(src, device=self.device)
          mask = image_to_tensor(mask, device=self.device)
          tgt = image_to_tensor(tgt, device=self.device)
  
          assert src.ndim == 3 and tgt.ndim == 3 and mask.ndim in [2, 3]
  
          if mask.ndim == 3:
              mask = mask.mean(-1) > 0.5
          else:
              mask = mask > 0.5
          mask[0, :] = 0
          mask[-1, :] = 0
          mask[:, 0] = 0
          mask[:, -1] = 0
  
          tgt_masked = torch.where(mask[..., None], torch.zeros_like(tgt), tgt)
  
          x, y = mask.nonzero(as_tuple=True)
          N = x.shape[0]
          index_map = torch.cumsum(mask.reshape(-1).long(), dim=-1).reshape(mask.shape)
          index_map[~mask] = 0
  
          if grad_mode == "src":
              src_lap = F.conv2d(
                  src.permute(2, 0, 1)[:, None],
                  weight=self.lap_kernel,
                  padding=1,
              )[:, 0].permute(1, 2, 0)
              lap = src_lap
          elif grad_mode == "max":
              src_lap = F.conv2d(
                  src.permute(2, 0, 1)[:, None],
                  weight=self.lap_kernel4,
                  padding=1,
              )
              tgt_lap = F.conv2d(
                  tgt.permute(2, 0, 1)[:, None],
                  weight=self.lap_kernel4,
                  padding=1,
              )
              lap = (
                  torch.where(src_lap.abs() > tgt_lap.abs(), src_lap, tgt_lap)
                  .sum(1)
                  .permute(1, 2, 0)
              )
          elif grad_mode == "avg":
              src_lap = F.conv2d(
                  src.permute(2, 0, 1)[:, None],
                  weight=self.lap_kernel4,
                  padding=1,
              )
              tgt_lap = F.conv2d(
                  tgt.permute(2, 0, 1)[:, None],
                  weight=self.lap_kernel4,
                  padding=1,
              )
              lap = (src_lap + tgt_lap).mul(0.5).sum(1).permute(1, 2, 0)
  
          fq_star = F.conv2d(
              tgt_masked.permute(2, 0, 1)[:, None],
              weight=self.neighbor_kernel,
              padding=1,
          )[:, 0].permute(1, 2, 0)
  
          A = torch.zeros(N + 1, 4, device=self.device, dtype=torch.long)
          X = torch.zeros(N + 1, 3, device=self.device, dtype=torch.float32)
          B = torch.zeros(N + 1, 3, device=self.device, dtype=torch.float32)
  
          A[1:] = torch.stack(
              [
                  index_map[x - 1, y],
                  index_map[x + 1, y],
                  index_map[x, y - 1],
                  index_map[x, y + 1],
              ],
              dim=-1,
          )
          X[1:] = tgt[x, y]
          B[1:] = lap[x, y] + fq_star[x, y]
  
          X_flatten = X.flatten()
          B_flatten = B.flatten()
          A_flatten = torch.stack([3 * A, 3 * A + 1, 3 * A + 2], dim=1).reshape(-1, 4)
  
          buffer = torch.zeros_like(X_flatten)
  
          self.pb_solver.solve(num_iters, A_flatten, X_flatten, B_flatten, buffer)
  
          if inplace:
              tgt[x, y] = X_flatten.view(-1, 3)[1:].clamp(0.0, 1.0)
          else:
              tgt = tgt.clone()
              tgt[x, y] = X_flatten.view(-1, 3)[1:].clamp(0.0, 1.0)
  
          return tgt
  