submission 612507

import ctypes
import hashlib
import os
import subprocess

import torch
from task import input_t, output_t

_TILE_PIXELS = 1024 * 16

_SRC = r'''
#include <cuda_runtime.h>
#include <cooperative_groups.h>
#include <cub/agent/agent_histogram.cuh>
#include <cub/device/dispatch/kernels/histogram.cuh>
#include <cub/grid/grid_queue.cuh>
#include <stdint.h>

namespace cg = cooperative_groups;

struct IdOp {
  template <cub::CacheLoadModifier MOD, typename T>
  __host__ __device__ __forceinline__ void BinSelect(T sample, int& bin, bool valid) {
    if (valid) {
      bin = (int) sample;
    }
  }
};

using Policy = cub::AgentHistogramPolicy<1024, 16, cub::BLOCK_LOAD_VECTORIZE, cub::LOAD_DEFAULT, false, cub::SMEM, false>;
using AgentT = cub::detail::histogram::AgentHistogram<Policy, 256, 1, 1, const unsigned char*, int, IdOp, IdOp, int>;

__global__ void hist_i64_kernel(const unsigned char* in, unsigned long long* out, int* dummy, int n, int tiles_per_row) {
  cg::grid_group grid = cg::this_grid();
  int gtid = int(blockIdx.x) * int(blockDim.x) + int(threadIdx.x);
  int gstride = int(gridDim.x) * int(blockDim.x);
  for (int bin = gtid; bin < 256; bin += gstride) {
    out[bin] = 0;
  }
  grid.sync();

  __shared__ typename AgentT::TempStorage temp_storage;
  int num_bins[1] = {256};
  int* d_out[1] = {dummy};
  int* d_priv[1] = {dummy};
  IdOp out_op[1] = {};
  IdOp priv_op[1] = {};
  AgentT agent(temp_storage, in, num_bins, num_bins, d_out, d_priv, out_op, priv_op);
  agent.InitBinCounters();
  agent.ConsumeTiles(n, 1, n, tiles_per_row, cub::GridQueue<int>());
  __syncthreads();
  for (int bin = threadIdx.x; bin < 256; bin += blockDim.x) {
    int count = temp_storage.Alias().histograms[0][bin];
    if (count > 0) {
      atomicAdd(out + bin, (unsigned long long) count);
    }
  }
}

extern "C" int hist_max_blocks() {
  int occ = 0;
  int dev = 0;
  int sms = 0;
  int coop = 0;
  cudaError_t err = cudaGetDevice(&dev);
  if (err != cudaSuccess) return -1;
  err = cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, dev);
  if (err != cudaSuccess || coop == 0) return -1;
  err = cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
  if (err != cudaSuccess) return -1;
  err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occ, hist_i64_kernel, 1024, 0);
  if (err != cudaSuccess) return -1;
  return occ * sms;
}

extern "C" int hist_i64_launch(uint64_t in_ptr, uint64_t out_ptr, uint64_t dummy_ptr, int blocks, int n) {
  const unsigned char* in = reinterpret_cast<const unsigned char*>(in_ptr);
  unsigned long long* out = reinterpret_cast<unsigned long long*>(out_ptr);
  int* dummy = reinterpret_cast<int*>(dummy_ptr);
  int tiles_per_row = (n + (1024 * 16) - 1) / (1024 * 16);
  void* args[] = {&in, &out, &dummy, &n, &tiles_per_row};
  cudaError_t err = cudaLaunchCooperativeKernel((void*) hist_i64_kernel, dim3(blocks, 1, 1), dim3(1024, 1, 1), args, 0);
  if (err != cudaSuccess) return (int) err;
  err = cudaGetLastError();
  return (int) err;
}
'''


class _HistHelper:
    def __init__(self, lib):
        self.lib = lib
        self.dev = None
        self.dummy = None
        self.dummy_cap = 0
        self.max_blocks = int(self.lib.hist_max_blocks())
        if self.max_blocks <= 0:
            raise RuntimeError('hist_max_blocks failed')

    def _ensure(self, x: torch.Tensor) -> int:
        n = int(x.numel())
        tiles = (n + _TILE_PIXELS - 1) // _TILE_PIXELS
        blocks = tiles if tiles < self.max_blocks else self.max_blocks
        need = blocks * 256
        if self.dev != x.device:
            self.dev = x.device
            self.dummy = None
            self.dummy_cap = 0
        if need > self.dummy_cap:
            self.dummy = torch.empty(need, device=x.device, dtype=torch.int32)
            self.dummy_cap = need
        return blocks

    def run(self, x: torch.Tensor, out: torch.Tensor) -> bool:
        blocks = self._ensure(x)
        err = self.lib.hist_i64_launch(
            x.data_ptr(),
            out.data_ptr(),
            self.dummy.data_ptr(),
            blocks,
            int(x.numel()),
        )
        return err == 0


_EXACT_SIZE = 10485760
_EXACT_BLOCKS = 128

_SRC_EXACT = r'''
#include <cuda_runtime.h>
#include <cooperative_groups.h>
#include <cub/block/block_histogram.cuh>
#include <stdint.h>

namespace cg = cooperative_groups;

using BlockHistT = cub::BlockHistogram<unsigned char, 1024, 16, 256, cub::BLOCK_HISTO_ATOMIC>;

union Bytes16 {
  uint4 v;
  unsigned char bytes[16];
};

__global__ void hist_exact_kernel(const unsigned char* in, unsigned long long* out) {
  cg::grid_group grid = cg::this_grid();
  int gtid = int(blockIdx.x) * int(blockDim.x) + int(threadIdx.x);
  int gstride = int(gridDim.x) * int(blockDim.x);
  for (int bin = gtid; bin < 256; bin += gstride) {
    out[bin] = 0;
  }
  grid.sync();

  __shared__ typename BlockHistT::TempStorage temp_storage;
  __shared__ unsigned int hist[256];
  BlockHistT block_hist(temp_storage);
  block_hist.InitHistogram(hist);
  __syncthreads();

  int tile = int(blockIdx.x);
  const int tile_stride = int(gridDim.x) << 10;
  const uint4* ptr = reinterpret_cast<const uint4*>(in) + (static_cast<unsigned long long>(tile) << 10) + threadIdx.x;
  Bytes16 curr;
  curr.v = __ldcs(ptr);
  tile += int(gridDim.x);
  ptr += tile_stride;
  for (; tile < 640; tile += int(gridDim.x), ptr += tile_stride) {
    Bytes16 next;
    next.v = __ldcs(ptr);
    block_hist.Composite(curr.bytes, hist);
    curr.v = next.v;
  }
  block_hist.Composite(curr.bytes, hist);
  __syncthreads();

  for (int bin = threadIdx.x; bin < 256; bin += blockDim.x) {
    unsigned int count = hist[bin];
    if (count > 0) {
      atomicAdd(out + bin, static_cast<unsigned long long>(count));
    }
  }
}

extern "C" int hist_exact_max_blocks() {
  int occ = 0;
  int dev = 0;
  int sms = 0;
  int coop = 0;
  cudaError_t err = cudaGetDevice(&dev);
  if (err != cudaSuccess) return -1;
  err = cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, dev);
  if (err != cudaSuccess || coop == 0) return -1;
  err = cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
  if (err != cudaSuccess) return -1;
  err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occ, hist_exact_kernel, 1024, 0);
  if (err != cudaSuccess) return -1;
  return occ * sms;
}

extern "C" int hist_exact_launch(uint64_t in_ptr, uint64_t out_ptr, int blocks) {
  const unsigned char* in = reinterpret_cast<const unsigned char*>(in_ptr);
  unsigned long long* out = reinterpret_cast<unsigned long long*>(out_ptr);
  void* args[] = {&in, &out};
  cudaError_t err = cudaLaunchCooperativeKernel((void*) hist_exact_kernel, dim3(blocks, 1, 1), dim3(1024, 1, 1), args, 0);
  if (err != cudaSuccess) return (int) err;
  err = cudaGetLastError();
  return (int) err;
}
'''


class _ExactHelper:
    def __init__(self, lib):
        self.lib = lib
        self.max_blocks = int(self.lib.hist_exact_max_blocks())
        if self.max_blocks <= 0:
            raise RuntimeError('hist_exact_max_blocks failed')

    def run(self, x: torch.Tensor, out: torch.Tensor) -> bool:
        blocks = _EXACT_BLOCKS if self.max_blocks >= _EXACT_BLOCKS else self.max_blocks
        err = self.lib.hist_exact_launch(x.data_ptr(), out.data_ptr(), blocks)
        return err == 0


_EXACT_HELPER = None
_HELPER = None


def _build_exact_helper() -> _ExactHelper | None:
    cap = torch.cuda.get_device_capability()
    arch = f"{cap[0]}{cap[1]}"
    key = hashlib.sha256(_SRC_EXACT.encode()).hexdigest()[:16]
    base = f"/tmp/glia_hist_opt028a_{arch}_{key}"
    cu_path = base + ".cu"
    so_path = base + ".so"

    if not os.path.exists(so_path):
        with open(cu_path, 'w') as f:
            f.write(_SRC_EXACT)
        try:
            subprocess.run(
                [
                    '/usr/local/cuda/bin/nvcc',
                    '-O3',
                    '--shared',
                    '-Xcompiler=-fPIC',
                    f'-gencode=arch=compute_{arch},code=sm_{arch}',
                    '-std=c++17',
                    cu_path,
                    '-o',
                    so_path,
                ],
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        except Exception:
            return None

    try:
        lib = ctypes.CDLL(so_path)
    except OSError:
        return None

    lib.hist_exact_max_blocks.argtypes = []
    lib.hist_exact_max_blocks.restype = ctypes.c_int
    lib.hist_exact_launch.argtypes = [ctypes.c_uint64, ctypes.c_uint64, ctypes.c_int]
    lib.hist_exact_launch.restype = ctypes.c_int
    try:
        return _ExactHelper(lib)
    except Exception:
        return None


def _get_exact_helper() -> _ExactHelper | None:
    global _EXACT_HELPER
    if _EXACT_HELPER is False:
        return None
    if _EXACT_HELPER is None:
        helper = _build_exact_helper()
        if helper is None:
            _EXACT_HELPER = False
            return None
        _EXACT_HELPER = helper
    return _EXACT_HELPER


def _build_helper() -> _HistHelper | None:
    cap = torch.cuda.get_device_capability()
    arch = f"{cap[0]}{cap[1]}"
    key = hashlib.sha256(_SRC.encode()).hexdigest()[:16]
    base = f"/tmp/glia_hist_opt011_{arch}_{key}"
    cu_path = base + ".cu"
    so_path = base + ".so"

    if not os.path.exists(so_path):
        with open(cu_path, 'w') as f:
            f.write(_SRC)
        try:
            subprocess.run(
                [
                    '/usr/local/cuda/bin/nvcc',
                    '-O3',
                    '--shared',
                    '-Xcompiler=-fPIC',
                    f'-gencode=arch=compute_{arch},code=sm_{arch}',
                    '-std=c++17',
                    cu_path,
                    '-o',
                    so_path,
                ],
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        except Exception:
            return None

    try:
        lib = ctypes.CDLL(so_path)
    except OSError:
        return None

    lib.hist_max_blocks.argtypes = []
    lib.hist_max_blocks.restype = ctypes.c_int
    lib.hist_i64_launch.argtypes = [
        ctypes.c_uint64,
        ctypes.c_uint64,
        ctypes.c_uint64,
        ctypes.c_int,
        ctypes.c_int,
    ]
    lib.hist_i64_launch.restype = ctypes.c_int
    try:
        return _HistHelper(lib)
    except Exception:
        return None


def _get_helper() -> _HistHelper | None:
    global _HELPER
    if _HELPER is False:
        return None
    if _HELPER is None:
        helper = _build_helper()
        if helper is None:
            _HELPER = False
            return None
        _HELPER = helper
    return _HELPER


def custom_kernel(data: input_t) -> output_t:
    x, out = data
    if x.numel() == _EXACT_SIZE and (x.data_ptr() & 15) == 0:
        helper = _get_exact_helper()
        if helper is not None and helper.run(x, out):
            return out

        global _EXACT_HELPER
        _EXACT_HELPER = False

    helper = _get_helper()
    if helper is not None and helper.run(x, out):
        return out

    global _HELPER
    _HELPER = False
    out.copy_(torch.histc(x.float(), bins=256, min=0, max=256))
    return out