submission 673007

# -*- coding: utf-8 -*-
import torch
import triton
import triton.language as tl

@triton.jit
def _mcd_kernel(input_ptr, output_ptr, size, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    # 1. Carga de datos contiguos para maximizar ancho de banda
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < size
    data = tl.load(input_ptr + offsets, mask=mask, other=0)

    # 2. Atómico Global Directo
    # En Blackwell, el motor de coherencia L2 agrupa estos atómicos
    # automáticamente si el puntero base (output_ptr) es el mismo.
    tl.atomic_add(output_ptr + data, 1, mask=mask)

def custom_kernel(data):
    input_tensor, output = data[0], data[1]
    size = input_tensor.numel()
    output.zero_()

    # BLOCK_SIZE de 512 o 1024 para equilibrar hilos por SM
    BLOCK_SIZE = 1024
    grid = (triton.cdiv(size, BLOCK_SIZE),)

    _mcd_kernel[grid](
        input_tensor, output, size,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=8 # Maximiza ocupación en B200
    )
    return output