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

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

# Fused Triplet Margin Loss kernel:
#   loss_i = max(||a_i - p_i + eps||_2 - ||a_i - n_i + eps||_2 + margin, 0)
#   output = mean(loss_i)
# Each block handles one sample (row), reads each input exactly once with
# vectorized float4 loads, block-reduces the two sums of squared differences,
# and accumulates the per-sample loss into a double-precision global total.
_triplet_loss_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BLOCK_THREADS 256
#define NUM_WARPS (BLOCK_THREADS / 32)

__device__ __forceinline__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffffu, v, offset);
    }
    return v;
}

__device__ __forceinline__ void block_reduce_and_accumulate(
    float sp, float sn, double* __restrict__ total, float margin)
{
    __shared__ float s_sp[NUM_WARPS];
    __shared__ float s_sn[NUM_WARPS];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    sp = warp_reduce_sum(sp);
    sn = warp_reduce_sum(sn);
    if (lane == 0) {
        s_sp[warp] = sp;
        s_sn[warp] = sn;
    }
    __syncthreads();

    if (warp == 0) {
        sp = (lane < NUM_WARPS) ? s_sp[lane] : 0.0f;
        sn = (lane < NUM_WARPS) ? s_sn[lane] : 0.0f;
        sp = warp_reduce_sum(sp);
        sn = warp_reduce_sum(sn);
        if (lane == 0) {
            const float dpos = sqrtf(sp);
            const float dneg = sqrtf(sn);
            const float loss = fmaxf(dpos - dneg + margin, 0.0f);
            atomicAdd(total, (double)loss);
        }
    }
}

__global__ void triplet_margin_loss_vec4_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    double* __restrict__ total,
    const int64_t dim,
    const float margin,
    const float eps)
{
    const int64_t base = (int64_t)blockIdx.x * dim;
    const float4* a4 = reinterpret_cast<const float4*>(anchor + base);
    const float4* p4 = reinterpret_cast<const float4*>(positive + base);
    const float4* n4 = reinterpret_cast<const float4*>(negative + base);
    const int64_t vdim = dim >> 2;

    float sp = 0.0f, sn = 0.0f;
    #pragma unroll 4
    for (int64_t i = threadIdx.x; i < vdim; i += BLOCK_THREADS) {
        const float4 a = __ldg(a4 + i);
        const float4 p = __ldg(p4 + i);
        const float4 n = __ldg(n4 + i);
        float d;
        d = (a.x - p.x) + eps; sp += d * d;
        d = (a.y - p.y) + eps; sp += d * d;
        d = (a.z - p.z) + eps; sp += d * d;
        d = (a.w - p.w) + eps; sp += d * d;
        d = (a.x - n.x) + eps; sn += d * d;
        d = (a.y - n.y) + eps; sn += d * d;
        d = (a.z - n.z) + eps; sn += d * d;
        d = (a.w - n.w) + eps; sn += d * d;
    }
    block_reduce_and_accumulate(sp, sn, total, margin);
}

__global__ void triplet_margin_loss_scalar_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    double* __restrict__ total,
    const int64_t dim,
    const float margin,
    const float eps)
{
    const int64_t base = (int64_t)blockIdx.x * dim;
    const float* a = anchor + base;
    const float* p = positive + base;
    const float* n = negative + base;

    float sp = 0.0f, sn = 0.0f;
    for (int64_t i = threadIdx.x; i < dim; i += BLOCK_THREADS) {
        const float av = __ldg(a + i);
        const float pv = __ldg(p + i);
        const float nv = __ldg(n + i);
        const float dp = (av - pv) + eps;
        const float dn = (av - nv) + eps;
        sp += dp * dp;
        sn += dn * dn;
    }
    block_reduce_and_accumulate(sp, sn, total, margin);
}

__global__ void finalize_mean_kernel(const double* __restrict__ total,
                                     float* __restrict__ out,
                                     const double inv_n)
{
    *out = (float)((*total) * inv_n);
}

torch::Tensor triplet_margin_loss_cuda(torch::Tensor anchor,
                                       torch::Tensor positive,
                                       torch::Tensor negative,
                                       double margin)
{
    TORCH_CHECK(anchor.is_cuda() && positive.is_cuda() && negative.is_cuda(),
                "inputs must be CUDA tensors");
    TORCH_CHECK(anchor.scalar_type() == at::kFloat &&
                positive.scalar_type() == at::kFloat &&
                negative.scalar_type() == at::kFloat,
                "inputs must be float32");
    TORCH_CHECK(anchor.dim() == 2, "inputs must be 2D");

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

    const int64_t num_rows = a.size(0);
    const int64_t dim = a.size(1);

    auto total = torch::zeros({}, a.options().dtype(torch::kDouble));
    auto out = torch::empty({}, a.options());

    const dim3 grid((unsigned int)num_rows);
    const dim3 block(BLOCK_THREADS);
    const float eps = 1e-6f;
    const float m = (float)margin;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if ((dim & 3) == 0) {
        triplet_margin_loss_vec4_kernel<<<grid, block, 0, stream>>>(
            a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
            total.data_ptr<double>(), dim, m, eps);
    } else {
        triplet_margin_loss_scalar_kernel<<<grid, block, 0, stream>>>(
            a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
            total.data_ptr<double>(), dim, m, eps);
    }
    finalize_mean_kernel<<<1, 1, 0, stream>>>(
        total.data_ptr<double>(), out.data_ptr<float>(), 1.0 / (double)num_rows);

    return out;
}
"""

_fused_triplet_loss = load_inline(
    name="fused_triplet_margin_loss",
    cpp_sources=_triplet_loss_cpp_src,
    cuda_sources=_triplet_loss_cuda_src,
    functions=["triplet_margin_loss_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Triplet Margin Loss with a fused custom CUDA kernel.

    Computes, in a single pass over the inputs:
        loss_i = max(||anchor_i - positive_i + eps||_2
                     - ||anchor_i - negative_i + eps||_2 + margin, 0)
        output = mean(loss_i)
    matching torch.nn.TripletMarginLoss defaults
    (p=2, eps=1e-6, swap=False, reduction='mean').
    """
    def __init__(self, margin=1.0):
        super(ModelNew, self).__init__()
        self.margin = float(margin)

    def forward(self, anchor, positive, negative):
        if (
            anchor.is_cuda
            and anchor.dim() == 2
            and anchor.dtype == torch.float32
            and positive.dtype == torch.float32
            and negative.dtype == torch.float32
            and not (anchor.requires_grad or positive.requires_grad
                     or negative.requires_grad)
        ):
            return _fused_triplet_loss.triplet_margin_loss_cuda(
                anchor, positive, negative, self.margin
            )
        # Safe fallback for anything the custom kernel does not cover.
        return torch.nn.functional.triplet_margin_loss(
            anchor, positive, negative, margin=self.margin
        )
```