```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 in a single kernel launch:
#   loss = mean( max( ||a - p + eps||_2 - ||a - n + eps||_2 + margin, 0 ) )
#
# v5 refinements over the validated v4 (2.14 ms, ~96% of DRAM peak):
#  * 512-thread blocks with x4 unroll (12 independent 16-byte streaming
#    loads in flight per thread) and __launch_bounds__(512, 2) so at least
#    two blocks are resident per SM: block load/compute/finalize phases
#    overlap, keeping the memory pipeline busy during reductions/exits.
#    For the benchmark shape (dim=8192 -> n_vec=2048 float4) the main loop
#    is exactly ONE iteration, so there is no loop overhead.
#  * Keeps everything proven: __ldcs evict-first streaming loads, one block
#    per row, t = a + eps formed once per element and reused for both
#    distances, block-level warp-shuffle reduction, double-precision global
#    accumulation, "last block" mean finalize, and the persistent
#    self-resetting 16-byte workspace (no per-call memset).
triplet_margin_loss_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <mutex>

#define BLOCK_SIZE 512
#define NWARPS (BLOCK_SIZE / 32)
#define UNROLL 4

__device__ __forceinline__ 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;
}

// Accumulate squared distances for one float4 group.
// t = a + eps is formed once and used for both the positive and negative side.
__device__ __forceinline__ void acc_pair4(const float4 a, const float4 p,
                                          const float4 n, const float eps,
                                          float& sp, float& sn) {
    const float t0 = a.x + eps, t1 = a.y + eps, t2 = a.z + eps, t3 = a.w + eps;
    const float dp0 = t0 - p.x, dp1 = t1 - p.y, dp2 = t2 - p.z, dp3 = t3 - p.w;
    const float dn0 = t0 - n.x, dn1 = t1 - n.y, dn2 = t2 - n.z, dn3 = t3 - n.w;
    sp += (dp0 * dp0 + dp1 * dp1) + (dp2 * dp2 + dp3 * dp3);
    sn += (dn0 * dn0 + dn1 * dn1) + (dn2 * dn2 + dn3 * dn3);
}

// Block-wide reduction of the two squared-distance sums, hinge application,
// global accumulation, "last block" mean finalize, and workspace reset.
__device__ __forceinline__ void block_hinge_finalize(
        float sp, float sn, const float margin,
        double* __restrict__ accum, unsigned int* __restrict__ counter,
        float* __restrict__ out, const double inv_n) {
    __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);
            }
            __threadfence();
            if (atomicAdd(counter, 1u) == gridDim.x - 1u) {
                // Last block: every other block's contribution is visible.
                const double total = atomicAdd(accum, 0.0);
                out[0] = (float)(total * inv_n);
                // Re-zero the persistent workspace so the next launch needs
                // no explicit memset. We are the only remaining writer.
                *accum = 0.0;
                *counter = 0u;
                __threadfence();
            }
        }
    }
}

__global__ void __launch_bounds__(BLOCK_SIZE, 2)
triplet_margin_loss_vec4_kernel(
        const float4* __restrict__ anchor,
        const float4* __restrict__ positive,
        const float4* __restrict__ negative,
        double* __restrict__ accum,
        unsigned int* __restrict__ counter,
        float* __restrict__ out,
        const double inv_n,
        const int n_vec,
        const float margin,
        const float eps) {
    const float4* a = anchor + (size_t)blockIdx.x * n_vec;
    const float4* p = positive + (size_t)blockIdx.x * n_vec;
    const float4* n = negative + (size_t)blockIdx.x * n_vec;

    float sp[UNROLL], sn[UNROLL];
#pragma unroll
    for (int u = 0; u < UNROLL; ++u) { sp[u] = 0.0f; sn[u] = 0.0f; }

    int i = threadIdx.x;
    // Unrolled x4 main loop: 12 independent 16-byte streaming loads in flight.
    // For n_vec == BLOCK_SIZE * UNROLL (the benchmark shape) this is exactly
    // one iteration.
    for (; i + (UNROLL - 1) * BLOCK_SIZE < n_vec; i += UNROLL * BLOCK_SIZE) {
        float4 av[UNROLL], pv[UNROLL], nv[UNROLL];
#pragma unroll
        for (int u = 0; u < UNROLL; ++u) av[u] = __ldcs(a + i + u * BLOCK_SIZE);
#pragma unroll
        for (int u = 0; u < UNROLL; ++u) pv[u] = __ldcs(p + i + u * BLOCK_SIZE);
#pragma unroll
        for (int u = 0; u < UNROLL; ++u) nv[u] = __ldcs(n + i + u * BLOCK_SIZE);
#pragma unroll
        for (int u = 0; u < UNROLL; ++u) acc_pair4(av[u], pv[u], nv[u], eps, sp[u], sn[u]);
    }
    // Scalar-vec tail (at most UNROLL-1 iterations).
    for (; i < n_vec; i += BLOCK_SIZE) {
        const float4 av = __ldcs(a + i);
        const float4 pv = __ldcs(p + i);
        const float4 nv = __ldcs(n + i);
        acc_pair4(av, pv, nv, eps, sp[0], sn[0]);
    }

    const float spT = (sp[0] + sp[1]) + (sp[2] + sp[3]);
    const float snT = (sn[0] + sn[1]) + (sn[2] + sn[3]);
    block_hinge_finalize(spT, snT, margin, accum, counter, out, inv_n);
}

__global__ void __launch_bounds__(BLOCK_SIZE)
triplet_margin_loss_scalar_kernel(
        const float* __restrict__ anchor,
        const float* __restrict__ positive,
        const float* __restrict__ negative,
        double* __restrict__ accum,
        unsigned int* __restrict__ counter,
        float* __restrict__ out,
        const double inv_n,
        const int dim_size,
        const float margin,
        const float eps) {
    const size_t base = (size_t)blockIdx.x * dim_size;
    const float* a = anchor + base;
    const float* p = positive + base;
    const float* n = negative + base;

    float sp0 = 0.f, sp1 = 0.f, sp2 = 0.f, sp3 = 0.f;
    float sn0 = 0.f, sn1 = 0.f, sn2 = 0.f, sn3 = 0.f;

    int i = threadIdx.x;
    for (; i + 3 * BLOCK_SIZE < dim_size; i += 4 * BLOCK_SIZE) {
        const float a0 = __ldcs(a + i), a1 = __ldcs(a + i + BLOCK_SIZE),
                    a2 = __ldcs(a + i + 2 * BLOCK_SIZE), a3 = __ldcs(a + i + 3 * BLOCK_SIZE);
        const float p0 = __ldcs(p + i), p1 = __ldcs(p + i + BLOCK_SIZE),
                    p2 = __ldcs(p + i + 2 * BLOCK_SIZE), p3 = __ldcs(p + i + 3 * BLOCK_SIZE);
        const float n0 = __ldcs(n + i), n1 = __ldcs(n + i + BLOCK_SIZE),
                    n2 = __ldcs(n + i + 2 * BLOCK_SIZE), n3 = __ldcs(n + i + 3 * BLOCK_SIZE);
        float t, d;
        t = a0 + eps; d = t - p0; sp0 += d * d; d = t - n0; sn0 += d * d;
        t = a1 + eps; d = t - p1; sp1 += d * d; d = t - n1; sn1 += d * d;
        t = a2 + eps; d = t - p2; sp2 += d * d; d = t - n2; sn2 += d * d;
        t = a3 + eps; d = t - p3; sp3 += d * d; d = t - n3; sn3 += d * d;
    }
    for (; i < dim_size; i += BLOCK_SIZE) {
        const float t = __ldcs(a + i) + eps;
        const float dp = t - __ldcs(p + i);
        const float dn = t - __ldcs(n + i);
        sp0 += dp * dp;
        sn0 += dn * dn;
    }

    const float sp = (sp0 + sp1) + (sp2 + sp3);
    const float sn = (sn0 + sn1) + (sn2 + sn3);
    block_hinge_finalize(sp, sn, margin, accum, counter, out, inv_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.");

    const c10::cuda::CUDAGuard device_guard(anchor.device());

    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.");
    TORCH_CHECK(batch > 0 && dim > 0, "Empty inputs are not supported by the fused kernel.");

    auto out = torch::empty({}, a.options());

    // Persistent 16-byte scratch per device: bytes [0,8) hold the double sum
    // accumulator, bytes [8,12) hold the block-completion counter. The kernel
    // leaves both zeroed when it finishes, so zero-init happens only once.
    // (Deliberately leaked to avoid teardown-order issues at process exit.)
    static torch::Tensor* g_ws = new torch::Tensor[64];
    static std::mutex g_ws_mutex;
    const int dev = static_cast<int>(a.get_device());
    TORCH_CHECK(dev >= 0 && dev < 64, "Unsupported device index.");
    {
        std::lock_guard<std::mutex> lock(g_ws_mutex);
        if (!g_ws[dev].defined()) {
            g_ws[dev] = torch::zeros({16}, a.options().dtype(at::kByte));
        }
    }
    void* ws_ptr = g_ws[dev].data_ptr();
    double* accum = reinterpret_cast<double*>(ws_ptr);
    unsigned int* counter = reinterpret_cast<unsigned int*>(
        static_cast<unsigned char*>(ws_ptr) + 8);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const float eps = 1e-6f;
    const float m = static_cast<float>(margin);
    const double inv_n = 1.0 / static_cast<double>(batch);
    const dim3 grid(static_cast<unsigned int>(batch));

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

    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, counter, out.data_ptr<float>(), inv_n,
            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, counter, out.data_ptr<float>(), inv_n,
            static_cast<int>(dim), m, eps);
    }
    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_v5",
    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._fn = triplet_margin_loss_op.triplet_margin_loss_cuda

    def forward(self, anchor, positive, negative):
        if (anchor.is_cuda and anchor.dtype == torch.float32 and anchor.dim() == 2
                and anchor.shape[0] > 0 and anchor.shape[1] > 0):
            return self._fn(anchor, positive, negative, self.margin)
        # Fallback for CPU tensors / exotic inputs
        return F.triplet_margin_loss(
            anchor, positive, negative, margin=self.margin
        )
```