```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


triplet_margin_loss_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cstdint>

#define TML_THREADS 256
#define TML_ROWS_PER_BLOCK (TML_THREADS / 32)

__global__ void triplet_margin_loss_fused_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    float* __restrict__ out,
    int batch,
    int n,
    float margin,
    float eps) {
  const int lane = threadIdx.x & 31;
  const int warp = threadIdx.x >> 5;
  const int row = blockIdx.x * TML_ROWS_PER_BLOCK + warp;
  __shared__ float sloss[TML_ROWS_PER_BLOCK];

  float sp0 = 0.0f, sp1 = 0.0f, sp2 = 0.0f, sp3 = 0.0f;
  float sn0 = 0.0f, sn1 = 0.0f, sn2 = 0.0f, sn3 = 0.0f;

  if (row < batch) {
    const float* a = anchor + static_cast<long long>(row) * n;
    const float* p = positive + static_cast<long long>(row) * n;
    const float* g = negative + static_cast<long long>(row) * n;

    const bool vec4 = ((n & 3) == 0) &&
        (((reinterpret_cast<uintptr_t>(a) | reinterpret_cast<uintptr_t>(p) |
           reinterpret_cast<uintptr_t>(g)) & 15u) == 0);

    if (vec4) {
      const int n4 = n >> 2;
      const float4* a4 = reinterpret_cast<const float4*>(a);
      const float4* p4 = reinterpret_cast<const float4*>(p);
      const float4* g4 = reinterpret_cast<const float4*>(g);

      for (int i = lane; i < n4; i += 32) {
        const float4 av = a4[i];
        const float4 pv = p4[i];
        const float4 gv = g4[i];

        const float dp0 = av.x - pv.x + eps;
        const float dp1 = av.y - pv.y + eps;
        const float dp2 = av.z - pv.z + eps;
        const float dp3 = av.w - pv.w + eps;
        sp0 += dp0 * dp0;
        sp1 += dp1 * dp1;
        sp2 += dp2 * dp2;
        sp3 += dp3 * dp3;

        const float dn0 = av.x - gv.x + eps;
        const float dn1 = av.y - gv.y + eps;
        const float dn2 = av.z - gv.z + eps;
        const float dn3 = av.w - gv.w + eps;
        sn0 += dn0 * dn0;
        sn1 += dn1 * dn1;
        sn2 += dn2 * dn2;
        sn3 += dn3 * dn3;
      }
    } else {
      for (int i = lane; i < n; i += 32) {
        const float av = a[i];
        const float dp = av - p[i] + eps;
        const float dn = av - g[i] + eps;
        sp0 += dp * dp;
        sn0 += dn * dn;
      }
    }
  }

  float sp = (sp0 + sp1) + (sp2 + sp3);
  float sn = (sn0 + sn1) + (sn2 + sn3);

#pragma unroll
  for (int offset = 16; offset > 0; offset >>= 1) {
    sp += __shfl_down_sync(0xffffffffu, sp, offset);
    sn += __shfl_down_sync(0xffffffffu, sn, offset);
  }

  if (lane == 0) {
    float loss = 0.0f;
    if (row < batch) {
      const float dpos = sqrtf(sp);
      const float dneg = sqrtf(sn);
      loss = dpos - dneg + margin;
      loss = loss > 0.0f ? loss : 0.0f;
    }
    sloss[warp] = loss;
  }
  __syncthreads();

  if (threadIdx.x == 0) {
    float total = 0.0f;
#pragma unroll
    for (int i = 0; i < TML_ROWS_PER_BLOCK; ++i) {
      total += sloss[i];
    }
    atomicAdd(out, total);
  }
}

__global__ void triplet_margin_loss_scale_kernel(float* out, float scale) {
  out[0] *= scale;
}

torch::Tensor triplet_margin_loss_cuda(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps) {
  TORCH_CHECK(anchor.is_cuda(), "anchor must be a CUDA tensor");
  TORCH_CHECK(positive.is_cuda() && negative.is_cuda(), "positive and negative must be CUDA tensors");
  TORCH_CHECK(anchor.scalar_type() == at::kFloat && positive.scalar_type() == at::kFloat &&
                  negative.scalar_type() == at::kFloat,
              "anchor, positive and negative must be float32");
  TORCH_CHECK(anchor.device() == positive.device() && anchor.device() == negative.device(),
              "anchor, positive and negative must be on the same device");
  TORCH_CHECK(anchor.dim() == 2, "expected 2D inputs [batch, features]");
  TORCH_CHECK(anchor.sizes().vec() == positive.sizes().vec() &&
                  anchor.sizes().vec() == negative.sizes().vec(),
              "anchor, positive and negative must have the same shape");

  const c10::cuda::CUDAGuard device_guard(anchor.device());
  auto a = anchor.contiguous();
  auto p = positive.contiguous();
  auto g = negative.contiguous();

  const int64_t batch = a.size(0);
  const int64_t n = a.size(1);
  auto out = torch::zeros({}, a.options());
  if (batch == 0) {
    return out;
  }

  const int blocks = static_cast<int>((batch + TML_ROWS_PER_BLOCK - 1) / TML_ROWS_PER_BLOCK);
  auto stream = at::cuda::getCurrentCUDAStream();
  triplet_margin_loss_fused_kernel<<<blocks, TML_THREADS, 0, stream.stream()>>>(
      a.data_ptr<float>(),
      p.data_ptr<float>(),
      g.data_ptr<float>(),
      out.data_ptr<float>(),
      static_cast<int>(batch),
      static_cast<int>(n),
      static_cast<float>(margin),
      static_cast<float>(eps));

  const float scale = 1.0f / static_cast<float>(batch);
  triplet_margin_loss_scale_kernel<<<1, 1, 0, stream.stream()>>>(out.data_ptr<float>(), scale);
  return out;
}
"""

triplet_margin_loss_cpp_source = (
    "torch::Tensor triplet_margin_loss_cuda(torch::Tensor anchor, torch::Tensor positive, "
    "torch::Tensor negative, double margin, double eps);"
)

_triplet_margin_loss = load_inline(
    name="triplet_margin_loss_fused",
    cpp_sources=triplet_margin_loss_cpp_source,
    cuda_sources=triplet_margin_loss_cuda_source,
    functions=["triplet_margin_loss_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    def __init__(self, margin=1.0):
        super().__init__()
        self.margin = float(margin)
        self.eps = 1e-6

    def forward(self, anchor, positive, negative):
        if (
            anchor.is_cuda
            and positive.is_cuda
            and negative.is_cuda
            and anchor.dtype == torch.float32
            and positive.dtype == torch.float32
            and negative.dtype == torch.float32
            and anchor.dim() == 2
            and anchor.shape == positive.shape == negative.shape
            and anchor.size(0) > 0
        ):
            return _triplet_margin_loss.triplet_margin_loss_cuda(
                anchor, positive, negative, self.margin, self.eps
            )

        return F.triplet_margin_loss(
            anchor,
            positive,
            negative,
            margin=self.margin,
            p=2.0,
            eps=self.eps,
            swap=False,
            reduction="mean",
        )
```