```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
from torch.utils.cpp_extension import load_inline

# Fused CUDA kernel for TripletMarginLoss (p=2, swap=False, reduction='mean'):
#   loss = mean_i( clamp(margin + ||a_i - p_i + eps||_2 - ||a_i - n_i + eps||_2, 0) )
# Each block handles one sample: threads compute partial squared L2 sums with
# vectorized float4 loads, a shared-memory block reduction produces both norms,
# and the per-sample hinge loss is accumulated with a double atomicAdd. The mean
# is finalized on the host side. This collapses PyTorch's multi-pass pipeline
# (sub, add-eps, square, sum, sqrt, sub, clamp, mean -> many temporary tensors)
# into a single memory-bound kernel.
triplet_margin_loss_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

template <bool VEC>
__global__ void triplet_margin_loss_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    double* __restrict__ out,
    const long D,
    const float margin,
    const float eps)
{
    const long sample = blockIdx.x;
    const float* __restrict__ a = anchor + sample * D;
    const float* __restrict__ p = positive + sample * D;
    const float* __restrict__ n = negative + sample * D;

    float local_pos = 0.0f;
    float local_neg = 0.0f;

    if (VEC) {
        const long vec_len = D >> 2;
        const float4* __restrict__ a4 = reinterpret_cast<const float4*>(a);
        const float4* __restrict__ p4 = reinterpret_cast<const float4*>(p);
        const float4* __restrict__ n4 = reinterpret_cast<const float4*>(n);
        #pragma unroll 4
        for (long i = threadIdx.x; i < vec_len; i += blockDim.x) {
            const float4 av = a4[i];
            const float4 pv = p4[i];
            const float4 nv = n4[i];
            float d0 = av.x - pv.x + eps;
            float d1 = av.y - pv.y + eps;
            float d2 = av.z - pv.z + eps;
            float d3 = av.w - pv.w + eps;
            local_pos += d0 * d0 + d1 * d1 + d2 * d2 + d3 * d3;
            d0 = av.x - nv.x + eps;
            d1 = av.y - nv.y + eps;
            d2 = av.z - nv.z + eps;
            d3 = av.w - nv.w + eps;
            local_neg += d0 * d0 + d1 * d1 + d2 * d2 + d3 * d3;
        }
    } else {
        for (long i = threadIdx.x; i < D; i += blockDim.x) {
            const float dp = a[i] - p[i] + eps;
            const float dn = a[i] - n[i] + eps;
            local_pos += dp * dp;
            local_neg += dn * dn;
        }
    }

    // Warp-level reduction
    for (int offset = 16; offset > 0; offset >>= 1) {
        local_pos += __shfl_down_sync(0xffffffffu, local_pos, offset);
        local_neg += __shfl_down_sync(0xffffffffu, local_neg, offset);
    }

    __shared__ float s_pos[32];
    __shared__ float s_neg[32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) {
        s_pos[warp] = local_pos;
        s_neg[warp] = local_neg;
    }
    __syncthreads();

    if (warp == 0) {
        const int num_warps = (blockDim.x + 31) >> 5;
        float sp = (lane < num_warps) ? s_pos[lane] : 0.0f;
        float sn = (lane < num_warps) ? s_neg[lane] : 0.0f;
        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) {
            const float loss = margin + sqrtf(sp) - sqrtf(sn);
            if (loss > 0.0f) {
                atomicAdd(out, (double)loss);
            }
        }
    }
}

torch::Tensor triplet_margin_loss_cuda(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps)
{
    TORCH_CHECK(anchor.is_cuda() && positive.is_cuda() && negative.is_cuda(),
                "All inputs must be CUDA tensors");
    TORCH_CHECK(anchor.scalar_type() == torch::kFloat32 &&
                positive.scalar_type() == torch::kFloat32 &&
                negative.scalar_type() == torch::kFloat32,
                "All inputs must be float32");
    TORCH_CHECK(anchor.numel() == positive.numel() && anchor.numel() == negative.numel(),
                "Input shapes must match");

    const c10::cuda::CUDAGuard device_guard(anchor.device());

    auto a = anchor.contiguous();
    auto p = positive.contiguous();
    auto n = negative.contiguous();

    const long batch = a.size(0);
    const long D = a.numel() / batch;

    auto out = torch::zeros({}, torch::TensorOptions()
                                    .dtype(torch::kDouble)
                                    .device(a.device()));

    if (batch > 0 && D > 0) {
        const int threads = 256;
        const dim3 blocks((unsigned int)batch);
        cudaStream_t stream = at::cuda::getCurrentCUDAStream();

        if ((D & 3) == 0) {
            triplet_margin_loss_kernel<true><<<blocks, threads, 0, stream>>>(
                a.data_ptr<float>(),
                p.data_ptr<float>(),
                n.data_ptr<float>(),
                out.data_ptr<double>(),
                D,
                (float)margin,
                (float)eps);
        } else {
            triplet_margin_loss_kernel<false><<<blocks, threads, 0, stream>>>(
                a.data_ptr<float>(),
                p.data_ptr<float>(),
                n.data_ptr<float>(),
                out.data_ptr<double>(),
                D,
                (float)margin,
                (float)eps);
        }
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    // 'mean' reduction over the batch, cast back to the input dtype
    return out.div((double)batch).to(a.scalar_type());
}
"""

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

# Compile the inline CUDA extension
triplet_margin_loss_ext = load_inline(
    name="triplet_margin_loss_ext",
    cpp_sources=triplet_margin_loss_cpp_source,
    cuda_sources=triplet_margin_loss_cuda_source,
    functions=["triplet_margin_loss_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized Triplet Margin Loss using a single fused CUDA kernel.

    Parameters:
        margin (float): The margin between the positive and negative samples.
    """
    def __init__(self, margin=1.0):
        super(ModelNew, self).__init__()
        self.margin = float(margin)
        self.eps = 1e-6  # matches torch.nn.TripletMarginLoss default
        self.ext = triplet_margin_loss_ext

    def forward(self, anchor, positive, negative):
        return self.ext.triplet_margin_loss_cuda(
            anchor, positive, negative, self.margin, self.eps
        )
```