```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 (v3).
#
# 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)
#
# v3 changes vs v2 (which regressed slightly vs v1):
#  - Reverted to the fastest measured main-loop config: 256 threads/block,
#    __ldg float4 loads, no manual unrolling (the kernel is purely DRAM-bound,
#    and this config reached the highest achieved bandwidth).
#  - Eliminated the separate mean-reduction kernel (11.6 us): the last block
#    to finish (tracked with a global ticket counter) reduces the per-row
#    losses in double precision and writes the final mean. This removes a
#    full kernel launch + drain from the critical path.
# 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>
#include <limits>

#define TML_BLOCK 256
#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,
    float* __restrict__ out,
    int* __restrict__ counter,
    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), requires 16B alignment which 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);
        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);
            const float dp0 = av.x - pv.x + eps;
            const float dp1 = av.y - pv.y + eps;
            const float dp2 = av.z - pv.z + eps;
            const float dp3 = av.w - pv.w + eps;
            const float dn0 = av.x - nv.x + eps;
            const float dn1 = av.y - nv.y + eps;
            const float dn2 = av.z - nv.z + eps;
            const float dn3 = av.w - nv.w + eps;
            sp += dp0 * dp0 + dp1 * dp1 + dp2 * dp2 + dp3 * dp3;
            sn += dn0 * dn0 + dn1 * dn1 + dn2 * dn2 + dn3 * dn3;
        }
    } 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;
        }
    }

    __shared__ float sh_p[TML_WARPS];
    __shared__ float sh_n[TML_WARPS];
    __shared__ double sh_d[TML_WARPS];
    __shared__ unsigned int sh_ticket;

    // Block reduction of the two squared-distance accumulators.
    #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);
            // Make the losses[] write visible, then take a ticket.
            __threadfence();
            sh_ticket = (unsigned int)atomicAdd(counter, 1);
        }
    }
    __syncthreads();

    // The last block to finish reduces all per-row losses into the mean.
    // This fuses away the separate (latency-dominated) mean kernel.
    if (sh_ticket == (unsigned int)(gridDim.x - 1)) {
        double sum = 0.0;
        if ((batch & 3) == 0) {
            const float4* l4 = reinterpret_cast<const float4*>(losses);
            const int vb = batch >> 2;
            for (int i = threadIdx.x; i < vb; i += TML_BLOCK) {
                const float4 v = __ldcg(l4 + i);
                sum += ((double)v.x + (double)v.y) + ((double)v.z + (double)v.w);
            }
        } else {
            for (int i = threadIdx.x; i < batch; i += TML_BLOCK) {
                sum += (double)__ldcg(losses + i);
            }
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, o);
        }
        if (lane == 0) sh_d[warp] = sum;
        __syncthreads();
        if (warp == 0) {
            sum = (lane < TML_WARPS) ? sh_d[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 __launch_bounds__(TML_BLOCK) 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);
    auto counter  = at::empty({1}, opts.dtype(at::kInt));

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (batch > 0) {
        cudaMemsetAsync(counter.data_ptr(), 0, sizeof(int), stream);
        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>(), out.data_ptr<float>(),
            counter.data_ptr<int>(),
            (int)batch, (int)dim, (float)margin, (float)eps);
    } else {
        out.fill_(std::numeric_limits<double>::quiet_NaN());
    }

    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) {
        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_v3",
    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
        )
```