```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 (v2).
#
# The whole loss is computed in one streaming pass over the inputs:
#   dist_pos = ||anchor - positive + eps||_2   (fused subtract/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)
#
# v2 optimizations over v1:
#  - The kernel is issue-rate bound (not DRAM bound), so instructions per byte
#    were cut: t = a + eps is computed once per element and reused for both
#    distances (5 flops/element instead of 6).
#  - 4x unrolled float4 loads with all loads batched up front for max MLP.
#  - Streaming loads (__ldcs, evict-first) since data has zero reuse.
#  - Smaller blocks (128 threads) so the block-reduction epilogue is amortized
#    over twice as many bytes per thread.
# 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 128
#define TML_WARPS (TML_BLOCK / 32)

__global__ void __launch_bounds__(TML_BLOCK) 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 sp0 = 0.0f, sp1 = 0.0f, sn0 = 0.0f, sn1 = 0.0f;

    if ((dim & 3) == 0) {
        // Vectorized path (float4); 16B alignment holds when dim % 4 == 0
        // because torch allocations are 256B aligned and rows are dim*4 bytes.
        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;
        // 4x unrolled main loop: batch all 12 loads before the math.
        for (; i + 3 * TML_BLOCK < vdim; i += 4 * TML_BLOCK) {
            float4 av[4], pv[4], nv[4];
            #pragma unroll
            for (int k = 0; k < 4; ++k) {
                av[k] = __ldcs(a4 + i + k * TML_BLOCK);
                pv[k] = __ldcs(p4 + i + k * TML_BLOCK);
                nv[k] = __ldcs(n4 + i + k * TML_BLOCK);
            }
            #pragma unroll
            for (int k = 0; k < 4; ++k) {
                const float tx = av[k].x + eps;
                const float ty = av[k].y + eps;
                const float tz = av[k].z + eps;
                const float tw = av[k].w + eps;
                const float dpx = tx - pv[k].x, dnx = tx - nv[k].x;
                const float dpy = ty - pv[k].y, dny = ty - nv[k].y;
                const float dpz = tz - pv[k].z, dnz = tz - nv[k].z;
                const float dpw = tw - pv[k].w, dnw = tw - nv[k].w;
                sp0 = fmaf(dpx, dpx, sp0); sn0 = fmaf(dnx, dnx, sn0);
                sp1 = fmaf(dpy, dpy, sp1); sn1 = fmaf(dny, dny, sn1);
                sp0 = fmaf(dpz, dpz, sp0); sn0 = fmaf(dnz, dnz, sn0);
                sp1 = fmaf(dpw, dpw, sp1); sn1 = fmaf(dnw, dnw, sn1);
            }
        }
        for (; i < vdim; i += TML_BLOCK) {
            const float4 av = __ldcs(a4 + i);
            const float4 pv = __ldcs(p4 + i);
            const float4 nv = __ldcs(n4 + i);
            const float tx = av.x + eps, ty = av.y + eps;
            const float tz = av.z + eps, tw = av.w + eps;
            const float dpx = tx - pv.x, dnx = tx - nv.x;
            const float dpy = ty - pv.y, dny = ty - nv.y;
            const float dpz = tz - pv.z, dnz = tz - nv.z;
            const float dpw = tw - pv.w, dnw = tw - nv.w;
            sp0 = fmaf(dpx, dpx, sp0); sn0 = fmaf(dnx, dnx, sn0);
            sp1 = fmaf(dpy, dpy, sp1); sn1 = fmaf(dny, dny, sn1);
            sp0 = fmaf(dpz, dpz, sp0); sn0 = fmaf(dnz, dnz, sn0);
            sp1 = fmaf(dpw, dpw, sp1); sn1 = fmaf(dnw, dnw, sn1);
        }
    } else {
        for (int i = threadIdx.x; i < dim; i += TML_BLOCK) {
            const float t = __ldcs(a + i) + eps;
            const float dp = t - __ldcs(p + i);
            const float dn = t - __ldcs(n + i);
            sp0 = fmaf(dp, dp, sp0);
            sn0 = fmaf(dn, dn, sn0);
        }
    }

    float sp = sp0 + sp1;
    float sn = sn0 + sn1;

    // 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);
            dist_pos[row] = dp;
            dist_neg[row] = dn;
            losses[row] = fmaxf(dp - dn + margin, 0.0f);
        }
    }
}

__global__ 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__ 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;
        }
    }
}

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);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    tml_forward_kernel<<<(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);

    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>

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", "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
        ):
            return _FusedTripletMarginLoss.apply(
                anchor.contiguous(),
                positive.contiguous(),
                negative.contiguous(),
                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
        )
```