```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 that computes the entire TripletMarginLoss (default p=2,
# eps=1e-6, swap=False, reduction='mean') in a single pass over the data:
#   dist_pos = ||a - p + eps||_2, dist_neg = ||a - n + eps||_2
#   loss = mean(max(dist_pos - dist_neg + margin, 0))
# Each block handles one sample (row), performs a block-level reduction of the
# two squared-distance sums, and atomically accumulates the clamped per-sample
# loss into a scalar accumulator (double precision for an accurate mean).
triplet_margin_loss_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void triplet_margin_loss_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    double* __restrict__ out,
    float margin,
    float eps,
    long batch_size,
    long dim)
{
    const long row = blockIdx.x;
    const float* a = anchor + row * dim;
    const float* p = positive + row * dim;
    const float* n = negative + row * dim;

    float sum_pos = 0.0f;
    float sum_neg = 0.0f;

    if ((dim & 3) == 0) {
        // Vectorized path: 16-byte aligned rows, float4 loads.
        const long dim_vec = dim >> 2;
        const float4* a4 = reinterpret_cast<const float4*>(a);
        const float4* p4 = reinterpret_cast<const float4*>(p);
        const float4* n4 = reinterpret_cast<const float4*>(n);
        for (long i = threadIdx.x; i < dim_vec; i += blockDim.x) {
            const float4 av = a4[i];
            const float4 pv = p4[i];
            const float4 nv = n4[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;
            const float dn0 = av.x - nv.x + eps;
            const float dn1 = av.y - nv.y + eps;
            const float dn2 = av.z - nv.z + eps;
            const float dn3 = av.w - nv.w + eps;
            sum_pos += dp0 * dp0 + dp1 * dp1 + dp2 * dp2 + dp3 * dp3;
            sum_neg += dn0 * dn0 + dn1 * dn1 + dn2 * dn2 + dn3 * dn3;
        }
    } else {
        // Scalar fallback path for arbitrary dims.
        for (long i = threadIdx.x; i < dim; i += blockDim.x) {
            const float dp = a[i] - p[i] + eps;
            const float dn = a[i] - n[i] + eps;
            sum_pos += dp * dp;
            sum_neg += dn * dn;
        }
    }

    // Block reduction (in double precision for accuracy).
    double d_pos = (double)sum_pos;
    double d_neg = (double)sum_neg;
    for (int offset = 16; offset > 0; offset >>= 1) {
        d_pos += __shfl_down_sync(0xffffffffu, d_pos, offset);
        d_neg += __shfl_down_sync(0xffffffffu, d_neg, offset);
    }

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

    if (warp == 0) {
        const int nwarps = (blockDim.x + 31) >> 5;
        double v_pos = (lane < nwarps) ? s_pos[lane] : 0.0;
        double v_neg = (lane < nwarps) ? s_neg[lane] : 0.0;
        for (int offset = 16; offset > 0; offset >>= 1) {
            v_pos += __shfl_down_sync(0xffffffffu, v_pos, offset);
            v_neg += __shfl_down_sync(0xffffffffu, v_neg, offset);
        }
        if (lane == 0) {
            const double dist_pos = sqrt(v_pos);
            const double dist_neg = sqrt(v_neg);
            const double loss = dist_pos - dist_neg + (double)margin;
            if (loss > 0.0) {
                atomicAdd(out, loss / (double)batch_size);
            }
        }
    }
}

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");
    auto a = anchor.contiguous();
    auto p = positive.contiguous();
    auto n = negative.contiguous();

    const long batch_size = a.size(0);
    const long dim = a.numel() / batch_size;

    // Accumulate in double, return in the input dtype (float32 scalar).
    auto out_d = torch::zeros({}, a.options().dtype(torch::kDouble));

    const int block_size = 256;
    const long num_blocks = batch_size;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    triplet_margin_loss_kernel<<<num_blocks, block_size, 0, stream>>>(
        a.data_ptr<float>(),
        p.data_ptr<float>(),
        n.data_ptr<float>(),
        out_d.data_ptr<double>(),
        (float)margin,
        (float)eps,
        batch_size,
        dim);

    return out_d.to(anchor.scalar_type());
}
"""

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

# Compile the fused TripletMarginLoss CUDA kernel.
triplet_margin_loss = load_inline(
    name="triplet_margin_loss_fused",
    cpp_sources=triplet_margin_loss_cpp_src,
    cuda_sources=triplet_margin_loss_cuda_src,
    functions=["triplet_margin_loss_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized Triplet Margin Loss model using a single fused CUDA kernel
    (pairwise distances + hinge + mean reduction in one kernel launch).

    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  # default eps of torch.nn.TripletMarginLoss (p=2)
        self.triplet_margin_loss = triplet_margin_loss

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