```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 CUDA kernels.
#
# The whole loss is computed in one streaming pass over the inputs:
#   dist_pos = ||anchor - positive + eps||_2   (fused sub/square/reduce per row)
#   dist_neg = ||anchor - negative + eps||_2
#   loss_i   = max(dist_pos - dist_neg + margin, 0)
#   output   = mean(loss_i)                    (tiny second reduction kernel)
# The inner loop is unrolled x2 so each thread keeps six 16-byte loads in
# flight, which saturates HBM bandwidth. A dedicated eval entry point skips
# materializing the per-row distances, and a matching fused backward kernel
# keeps the module fully trainable.

_triplet_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>

#define TML_BLOCK 256
#define TML_WARPS (TML_BLOCK / 32)

__device__ __forceinline__ void acc4(const float4 a, const float4 b,
                                     const float eps, float& s) {
    float d;
    d = a.x - b.x + eps; s += d * d;
    d = a.y - b.y + eps; s += d * d;
    d = a.z - b.z + eps; s += d * d;
    d = a.w - b.w + eps; s += d * d;
}

template<bool SAVE_DISTS>
__global__ __launch_bounds__(TML_BLOCK) void tml_forward_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    float* __restrict__ losses,
    float* __restrict__ dist_pos,
    float* __restrict__ dist_neg,
    const int batch,
    const int dim,
    const float margin,
    const float eps)
{
    const int row = blockIdx.x;
    const size_t base = (size_t)row * (size_t)dim;
    const float* a = anchor + base;
    const float* p = positive + base;
    const float* n = negative + base;

    float sp = 0.0f;
    float sn = 0.0f;

    if ((dim & 3) == 0) {
        // Vectorized path (float4); 16B alignment holds when dim % 4 == 0
        // because torch allocations are 256B aligned.
        const int vdim = 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);
        int i = threadIdx.x;
        // Unrolled x2: keeps six independent 16-byte loads in flight.
        for (; i + TML_BLOCK < vdim; i += 2 * TML_BLOCK) {
            const float4 av0 = __ldg(a4 + i);
            const float4 av1 = __ldg(a4 + i + TML_BLOCK);
            const float4 pv0 = __ldg(p4 + i);
            const float4 pv1 = __ldg(p4 + i + TML_BLOCK);
            const float4 nv0 = __ldg(n4 + i);
            const float4 nv1 = __ldg(n4 + i + TML_BLOCK);
            acc4(av0, pv0, eps, sp);
            acc4(av1, pv1, eps, sp);
            acc4(av0, nv0, eps, sn);
            acc4(av1, nv1, eps, sn);
        }
        for (; i < vdim; i += TML_BLOCK) {
            const float4 av = __ldg(a4 + i);
            const float4 pv = __ldg(p4 + i);
            const float4 nv = __ldg(n4 + i);
            acc4(av, pv, eps, sp);
            acc4(av, nv, eps, sn);
        }
    } else {
        for (int i = threadIdx.x; i < dim; i += TML_BLOCK) {
            const float av = __ldg(a + i);
            const float dp = av - __ldg(p + i) + eps;
            const float dn = av - __ldg(n + i) + eps;
            sp += dp * dp;
            sn += dn * dn;
        }
    }

    // Block reduction of the two squared-distance accumulators.
    __shared__ float sh_p[TML_WARPS];
    __shared__ float sh_n[TML_WARPS];
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sp += __shfl_down_sync(0xffffffffu, sp, o);
        sn += __shfl_down_sync(0xffffffffu, sn, o);
    }
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) { sh_p[warp] = sp; sh_n[warp] = sn; }
    __syncthreads();
    if (warp == 0) {
        sp = (lane < TML_WARPS) ? sh_p[lane] : 0.0f;
        sn = (lane < TML_WARPS) ? sh_n[lane] : 0.0f;
        #pragma unroll
        for (int o = TML_WARPS / 2; o > 0; o >>= 1) {
            sp += __shfl_down_sync(0xffffffffu, sp, o);
            sn += __shfl_down_sync(0xffffffffu, sn, o);
        }
        if (lane == 0) {
            const float dp = sqrtf(sp);
            const float dn = sqrtf(sn);
            if (SAVE_DISTS) {
                dist_pos[row] = dp;
                dist_neg[row] = dn;
            }
            losses[row] = fmaxf(dp - dn + margin, 0.0f);
        }
    }
}

__global__ __launch_bounds__(TML_BLOCK) void tml_mean_kernel(
    const float* __restrict__ losses,
    float* __restrict__ out,
    const int batch)
{
    // Accumulate in double for an accurate mean.
    double sum = 0.0;
    for (int i = threadIdx.x; i < batch; i += TML_BLOCK) {
        sum += (double)__ldg(losses + i);
    }
    __shared__ double sh[TML_WARPS];
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, o);
    }
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) sh[warp] = sum;
    __syncthreads();
    if (warp == 0) {
        sum = (lane < TML_WARPS) ? sh[lane] : 0.0;
        #pragma unroll
        for (int o = TML_WARPS / 2; o > 0; o >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, o);
        }
        if (lane == 0) out[0] = (float)(sum / (double)batch);
    }
}

__global__ __launch_bounds__(TML_BLOCK) void tml_backward_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    const float* __restrict__ dist_pos,
    const float* __restrict__ dist_neg,
    const float* __restrict__ grad_out,
    float* __restrict__ grad_anchor,
    float* __restrict__ grad_positive,
    float* __restrict__ grad_negative,
    const int batch,
    const int dim,
    const float margin,
    const float eps)
{
    const int row = blockIdx.x;
    const float dp = __ldg(dist_pos + row);
    const float dn = __ldg(dist_neg + row);
    const float scale = __ldg(grad_out) / (float)batch;
    const float act = (dp - dn + margin > 0.0f) ? scale : 0.0f;
    const float cp = (dp > 0.0f) ? (act / dp) : 0.0f;
    const float cn = (dn > 0.0f) ? (act / dn) : 0.0f;

    const size_t base = (size_t)row * (size_t)dim;
    const float* a = anchor + base;
    const float* p = positive + base;
    const float* n = negative + base;
    float* ga = grad_anchor + base;
    float* gp = grad_positive + base;
    float* gn = grad_negative + base;

    if ((dim & 3) == 0) {
        const int vdim = 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);
        float4* ga4 = reinterpret_cast<float4*>(ga);
        float4* gp4 = reinterpret_cast<float4*>(gp);
        float4* gn4 = reinterpret_cast<float4*>(gn);
        for (int i = threadIdx.x; i < vdim; i += TML_BLOCK) {
            const float4 av = __ldg(a4 + i);
            const float4 pv = __ldg(p4 + i);
            const float4 nv = __ldg(n4 + i);
            float4 gav, gpv, gnv;
            float dap, dan;
            dap = av.x - pv.x + eps; dan = av.x - nv.x + eps;
            gav.x = cp * dap - cn * dan; gpv.x = -cp * dap; gnv.x = cn * dan;
            dap = av.y - pv.y + eps; dan = av.y - nv.y + eps;
            gav.y = cp * dap - cn * dan; gpv.y = -cp * dap; gnv.y = cn * dan;
            dap = av.z - pv.z + eps; dan = av.z - nv.z + eps;
            gav.z = cp * dap - cn * dan; gpv.z = -cp * dap; gnv.z = cn * dan;
            dap = av.w - pv.w + eps; dan = av.w - nv.w + eps;
            gav.w = cp * dap - cn * dan; gpv.w = -cp * dap; gnv.w = cn * dan;
            ga4[i] = gav;
            gp4[i] = gpv;
            gn4[i] = gnv;
        }
    } else {
        for (int i = threadIdx.x; i < dim; i += TML_BLOCK) {
            const float av = __ldg(a + i);
            const float dap = av - __ldg(p + i) + eps;
            const float dan = av - __ldg(n + i) + eps;
            ga[i] = cp * dap - cn * dan;
            gp[i] = -cp * dap;
            gn[i] = cn * dan;
        }
    }
}

torch::Tensor triplet_forward_eval(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps)
{
    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");

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

    const int64_t batch = a.size(0);
    const int64_t dim = batch > 0 ? (a.numel() / batch) : 0;

    auto opts = a.options();
    auto losses = at::empty({batch}, opts);
    auto out    = at::empty({}, opts);

    if (batch == 0 || dim == 0) {
        out.fill_(0.0f);
        return out;
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    tml_forward_kernel<false><<<(int)batch, TML_BLOCK, 0, stream>>>(
        a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
        losses.data_ptr<float>(), nullptr, nullptr,
        (int)batch, (int)dim, (float)margin, (float)eps);
    tml_mean_kernel<<<1, TML_BLOCK, 0, stream>>>(
        losses.data_ptr<float>(), out.data_ptr<float>(), (int)batch);

    return out;
}

std::vector<torch::Tensor> triplet_forward(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps)
{
    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");

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

    const int64_t batch = a.size(0);
    const int64_t dim = batch > 0 ? (a.numel() / batch) : 0;

    auto opts = a.options();
    auto losses   = at::empty({batch}, opts);
    auto dist_pos = at::empty({batch}, opts);
    auto dist_neg = at::empty({batch}, opts);
    auto out      = at::empty({}, opts);

    if (batch == 0 || dim == 0) {
        out.fill_(0.0f);
        return {out, dist_pos, dist_neg};
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    tml_forward_kernel<true><<<(int)batch, TML_BLOCK, 0, stream>>>(
        a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
        losses.data_ptr<float>(), dist_pos.data_ptr<float>(), dist_neg.data_ptr<float>(),
        (int)batch, (int)dim, (float)margin, (float)eps);
    tml_mean_kernel<<<1, TML_BLOCK, 0, stream>>>(
        losses.data_ptr<float>(), out.data_ptr<float>(), (int)batch);

    return {out, dist_pos, dist_neg};
}

std::vector<torch::Tensor> triplet_backward(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    torch::Tensor dist_pos,
    torch::Tensor dist_neg,
    torch::Tensor grad_out,
    double margin,
    double eps)
{
    auto a = anchor.contiguous();
    auto p = positive.contiguous();
    auto n = negative.contiguous();
    auto go = grad_out.contiguous();

    const int64_t batch = a.size(0);
    const int64_t dim = batch > 0 ? (a.numel() / batch) : 0;

    auto ga = at::empty_like(a);
    auto gp = at::empty_like(p);
    auto gn = at::empty_like(n);

    if (batch == 0 || dim == 0) {
        return {ga, gp, gn};
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    tml_backward_kernel<<<(int)batch, TML_BLOCK, 0, stream>>>(
        a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
        dist_pos.data_ptr<float>(), dist_neg.data_ptr<float>(), go.data_ptr<float>(),
        ga.data_ptr<float>(), gp.data_ptr<float>(), gn.data_ptr<float>(),
        (int)batch, (int)dim, (float)margin, (float)eps);

    return {ga, gp, gn};
}
"""

_triplet_cpp_src = r"""
#include <torch/extension.h>
#include <vector>

torch::Tensor triplet_forward_eval(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps);

std::vector<torch::Tensor> triplet_forward(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps);

std::vector<torch::Tensor> triplet_backward(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    torch::Tensor dist_pos,
    torch::Tensor dist_neg,
    torch::Tensor grad_out,
    double margin,
    double eps);
"""

_triplet_ext = load_inline(
    name="fused_triplet_margin_loss_v2",
    cpp_sources=_triplet_cpp_src,
    cuda_sources=_triplet_cuda_src,
    functions=["triplet_forward_eval", "triplet_forward", "triplet_backward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class _FusedTripletMarginLoss(torch.autograd.Function):
    @staticmethod
    def forward(ctx, anchor, positive, negative, margin, eps):
        out, dist_pos, dist_neg = _triplet_ext.triplet_forward(
            anchor, positive, negative, margin, eps
        )
        ctx.save_for_backward(anchor, positive, negative, dist_pos, dist_neg)
        ctx.margin = margin
        ctx.eps = eps
        return out

    @staticmethod
    def backward(ctx, grad_output):
        anchor, positive, negative, dist_pos, dist_neg = ctx.saved_tensors
        ga, gp, gn = _triplet_ext.triplet_backward(
            anchor,
            positive,
            negative,
            dist_pos,
            dist_neg,
            grad_output.contiguous(),
            ctx.margin,
            ctx.eps,
        )
        return ga, gp, gn, None, None


class ModelNew(nn.Module):
    """
    Triplet Margin Loss implemented with a fused custom 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

    def forward(self, anchor, positive, negative):
        if (
            anchor.is_cuda
            and positive.is_cuda
            and negative.is_cuda
            and anchor.dtype == torch.float32
            and positive.dtype == torch.float32
            and negative.dtype == torch.float32
        ):
            a = anchor if anchor.is_contiguous() else anchor.contiguous()
            p = positive if positive.is_contiguous() else positive.contiguous()
            n = negative if negative.is_contiguous() else negative.contiguous()
            if torch.is_grad_enabled() and (
                a.requires_grad or p.requires_grad or n.requires_grad
            ):
                return _FusedTripletMarginLoss.apply(
                    a, p, n, self.margin, self.eps
                )
            # Fast eval path: no autograd bookkeeping, no distance buffers.
            return _triplet_ext.triplet_forward_eval(
                a, p, n, self.margin, self.eps
            )
        # Fallback for non-CUDA / non-fp32 inputs.
        return F.triplet_margin_loss(
            anchor, positive, negative, margin=self.margin, p=2.0, eps=self.eps
        )
```