```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 Triplet Margin Loss CUDA kernel.
#
# TripletMarginLoss with default settings (p=2, eps=1e-6, swap=False,
# reduction='mean') computes:
#   dist_pos = sqrt(sum((anchor - positive + eps)^2))
#   dist_neg = sqrt(sum((anchor - negative + eps)^2))
#   loss     = mean(max(0, margin + dist_pos - dist_neg))
#
# This kernel fuses the subtraction, eps-add, squaring, row reduction, sqrt,
# hinge, and clamp into a single pass over the data (one thread block per
# triplet/row), using float4 vectorized loads to saturate memory bandwidth.
triplet_margin_loss_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void triplet_margin_loss_rows_kernel(
    const float4* __restrict__ anchor,
    const float4* __restrict__ positive,
    const float4* __restrict__ negative,
    float* __restrict__ row_loss,
    int vec_len,
    float margin,
    float eps)
{
    const long long row = blockIdx.x;
    const float4* a_row = anchor   + row * (long long)vec_len;
    const float4* p_row = positive + row * (long long)vec_len;
    const float4* n_row = negative + row * (long long)vec_len;

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

    for (int i = threadIdx.x; i < vec_len; i += blockDim.x) {
        float4 a = __ldg(a_row + i);
        float4 p = __ldg(p_row + i);
        float4 n = __ldg(n_row + i);

        float d0 = a.x - p.x + eps;
        float d1 = a.y - p.y + eps;
        float d2 = a.z - p.z + eps;
        float d3 = a.w - p.w + eps;
        sum_pos += d0 * d0 + d1 * d1 + d2 * d2 + d3 * d3;

        float e0 = a.x - n.x + eps;
        float e1 = a.y - n.y + eps;
        float e2 = a.z - n.z + eps;
        float e3 = a.w - n.w + eps;
        sum_neg += e0 * e0 + e1 * e1 + e2 * e2 + e3 * e3;
    }

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

    __shared__ float sh_pos[32];
    __shared__ float sh_neg[32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int num_warps = (blockDim.x + 31) >> 5;

    if (lane == 0) {
        sh_pos[warp] = sum_pos;
        sh_neg[warp] = sum_neg;
    }
    __syncthreads();

    if (warp == 0) {
        sum_pos = (lane < num_warps) ? sh_pos[lane] : 0.0f;
        sum_neg = (lane < num_warps) ? sh_neg[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum_pos += __shfl_down_sync(0xffffffffu, sum_pos, offset);
            sum_neg += __shfl_down_sync(0xffffffffu, sum_neg, offset);
        }
        if (lane == 0) {
            float loss = margin + sqrtf(sum_pos) - sqrtf(sum_neg);
            row_loss[row] = loss > 0.0f ? loss : 0.0f;
        }
    }
}

torch::Tensor triplet_margin_loss_forward(
    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");
    TORCH_CHECK(anchor.is_contiguous() && positive.is_contiguous() &&
                negative.is_contiguous(), "inputs must be contiguous");

    const int64_t batch = anchor.size(0);
    const int64_t dim   = anchor.size(1);
    TORCH_CHECK(dim % 4 == 0, "feature dim must be divisible by 4");

    auto row_loss = torch::empty({batch}, anchor.options());

    if (batch == 0) {
        return row_loss;
    }

    const int vec_len = static_cast<int>(dim / 4);
    const int threads = 256;
    const float eps = 1e-6f;

    triplet_margin_loss_rows_kernel<<<static_cast<unsigned int>(batch), threads>>>(
        reinterpret_cast<const float4*>(anchor.data_ptr<float>()),
        reinterpret_cast<const float4*>(positive.data_ptr<float>()),
        reinterpret_cast<const float4*>(negative.data_ptr<float>()),
        row_loss.data_ptr<float>(),
        vec_len,
        static_cast<float>(margin),
        eps);

    return row_loss;
}
"""

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

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_forward"],
    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)
        # Kept as a safe fallback for non-CUDA / non-contiguous / odd shapes.
        self.loss_fn = nn.TripletMarginLoss(margin=margin)

    def forward(self, anchor, positive, negative):
        if (anchor.is_cuda and anchor.dtype == torch.float32
                and positive.dtype == torch.float32
                and negative.dtype == torch.float32
                and anchor.dim() == 2 and anchor.size(1) % 4 == 0
                and anchor.is_contiguous() and positive.is_contiguous()
                and negative.is_contiguous()):
            row_loss = triplet_margin_loss_ext.triplet_margin_loss_forward(
                anchor, positive, negative, self.margin)
            return row_loss.mean()
        return self.loss_fn(anchor, positive, negative)
```