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

# Fused Triplet Margin Loss:
#   loss = mean( max( ||a - p + eps||_2 - ||a - n + eps||_2 + margin, 0 ) )
# One thread block per sample computes both squared distances with vectorized
# float4 loads, reduces them in-block, applies the hinge, and atomically
# accumulates into a double-precision buffer. A tiny finalize kernel divides
# by the batch size. This replaces ~10 memory passes of the eager PyTorch
# implementation with a single streaming pass over the input data.
triplet_margin_loss_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BLOCK_SIZE 256
#define NWARPS (BLOCK_SIZE / 32)

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

__global__ void triplet_margin_loss_vec4_kernel(
        const float4* __restrict__ anchor,
        const float4* __restrict__ positive,
        const float4* __restrict__ negative,
        double* __restrict__ accum,
        const int n_vec,
        const float margin,
        const float eps) {
    const size_t row = blockIdx.x;
    const float4* a = anchor + row * (size_t)n_vec;
    const float4* p = positive + row * (size_t)n_vec;
    const float4* n = negative + row * (size_t)n_vec;

    float sp = 0.0f;
    float sn = 0.0f;
    for (int i = threadIdx.x; i < n_vec; i += BLOCK_SIZE) {
        const float4 av = a[i];
        const float4 pv = p[i];
        const float4 nv = n[i];
        const float d0 = av.x - pv.x + eps;
        const float d1 = av.y - pv.y + eps;
        const float d2 = av.z - pv.z + eps;
        const float d3 = av.w - pv.w + eps;
        sp += (d0 * d0 + d1 * d1) + (d2 * d2 + d3 * d3);
        const float e0 = av.x - nv.x + eps;
        const float e1 = av.y - nv.y + eps;
        const float e2 = av.z - nv.z + eps;
        const float e3 = av.w - nv.w + eps;
        sn += (e0 * e0 + e1 * e1) + (e2 * e2 + e3 * e3);
    }

    __shared__ float sh_sp[NWARPS];
    __shared__ float sh_sn[NWARPS];
    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) {
        sh_sp[warp] = sp;
        sh_sn[warp] = sn;
    }
    __syncthreads();

    if (warp == 0) {
        sp = (lane < NWARPS) ? sh_sp[lane] : 0.0f;
        sn = (lane < NWARPS) ? sh_sn[lane] : 0.0f;
        sp = warp_reduce_sum(sp);
        sn = warp_reduce_sum(sn);
        if (lane == 0) {
            const float loss = fmaxf(sqrtf(sp) - sqrtf(sn) + margin, 0.0f);
            if (loss != 0.0f) {
                atomicAdd(accum, (double)loss);
            }
        }
    }
}

__global__ void triplet_margin_loss_scalar_kernel(
        const float* __restrict__ anchor,
        const float* __restrict__ positive,
        const float* __restrict__ negative,
        double* __restrict__ accum,
        const int dim_size,
        const float margin,
        const float eps) {
    const size_t row = blockIdx.x;
    const float* a = anchor + row * (size_t)dim_size;
    const float* p = positive + row * (size_t)dim_size;
    const float* n = negative + row * (size_t)dim_size;

    float sp = 0.0f;
    float sn = 0.0f;
    for (int i = threadIdx.x; i < dim_size; i += BLOCK_SIZE) {
        const float dp = a[i] - p[i] + eps;
        const float dn = a[i] - n[i] + eps;
        sp += dp * dp;
        sn += dn * dn;
    }

    __shared__ float sh_sp[NWARPS];
    __shared__ float sh_sn[NWARPS];
    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) {
        sh_sp[warp] = sp;
        sh_sn[warp] = sn;
    }
    __syncthreads();

    if (warp == 0) {
        sp = (lane < NWARPS) ? sh_sp[lane] : 0.0f;
        sn = (lane < NWARPS) ? sh_sn[lane] : 0.0f;
        sp = warp_reduce_sum(sp);
        sn = warp_reduce_sum(sn);
        if (lane == 0) {
            const float loss = fmaxf(sqrtf(sp) - sqrtf(sn) + margin, 0.0f);
            if (loss != 0.0f) {
                atomicAdd(accum, (double)loss);
            }
        }
    }
}

__global__ void mean_finalize_kernel(
        const double* __restrict__ accum,
        float* __restrict__ out,
        const double n) {
    out[0] = (float)(accum[0] / 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(),
                "All inputs must be CUDA tensors.");
    TORCH_CHECK(anchor.scalar_type() == at::kFloat &&
                positive.scalar_type() == at::kFloat &&
                negative.scalar_type() == at::kFloat,
                "All inputs must be float32.");
    TORCH_CHECK(anchor.dim() == 2, "Inputs must be 2-D.");

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

    const int64_t batch = a.size(0);
    const int64_t dim = a.size(1);
    TORCH_CHECK(p.size(0) == batch && p.size(1) == dim, "positive shape mismatch.");
    TORCH_CHECK(n.size(0) == batch && n.size(1) == dim, "negative shape mismatch.");

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

    const float eps = 1e-6f;
    const float m = static_cast<float>(margin);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (batch > 0 && dim > 0) {
        const bool aligned =
            ((reinterpret_cast<uintptr_t>(a.data_ptr<float>()) & 0xF) == 0) &&
            ((reinterpret_cast<uintptr_t>(p.data_ptr<float>()) & 0xF) == 0) &&
            ((reinterpret_cast<uintptr_t>(n.data_ptr<float>()) & 0xF) == 0);
        const dim3 grid(static_cast<unsigned int>(batch));
        if (((dim & 3) == 0) && aligned) {
            triplet_margin_loss_vec4_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
                reinterpret_cast<const float4*>(a.data_ptr<float>()),
                reinterpret_cast<const float4*>(p.data_ptr<float>()),
                reinterpret_cast<const float4*>(n.data_ptr<float>()),
                accum.data_ptr<double>(),
                static_cast<int>(dim >> 2),
                m, eps);
        } else {
            triplet_margin_loss_scalar_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
                a.data_ptr<float>(),
                p.data_ptr<float>(),
                n.data_ptr<float>(),
                accum.data_ptr<double>(),
                static_cast<int>(dim),
                m, eps);
        }
    }
    mean_finalize_kernel<<<1, 1, 0, stream>>>(
        accum.data_ptr<double>(),
        out.data_ptr<float>(),
        static_cast<double>(batch));
    return out;
}
"""

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

# Compile the inline CUDA extension
triplet_margin_loss_op = load_inline(
    name="triplet_margin_loss_fused_v1",
    cpp_sources=triplet_margin_loss_cpp_src,
    cuda_sources=triplet_margin_loss_cuda_src,
    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._op = triplet_margin_loss_op

    def forward(self, anchor, positive, negative):
        if anchor.is_cuda and anchor.dtype == torch.float32 and anchor.dim() == 2:
            return self._op.triplet_margin_loss_cuda(
                anchor, positive, negative, self.margin
            )
        # Fallback for non-CUDA / non-float32 inputs
        return F.triplet_margin_loss(
            anchor, positive, negative, margin=self.margin
        )
```