require 'cutorch'
require 'nn'
cudnn = require 'cudnn.env'
require('cudnn.ffi')
local C = cudnn.C
local ffi = require 'ffi'

local thc = ffi.C
if ffi.os == "Windows" then thc = ffi.load("THC") end

--------------------------------------------------------------------
-- defaults, each should be overrideable via env var:
--------------------------------------------------------------------

cudnn.benchmark = false
cudnn.fastest = false

-- use new cudnn FindEx APIs
-- Warning: this option is experimental and assumes at least 2 warmup iterations!
cudnn.useFindEx = false

-- amount of memory to use on 1st iteration for FindEx
cudnn.initialWorkspaceBytes = 1024

--
cudnn.reservedGPUBytes = 1024*1024

cudnn.maxWorkspaceGPUMemPercent = 95

local maxStreamsPerDevice = 1024

--------------------------------------------------------------------
-- end defaults
--------------------------------------------------------------------

local numDevices = cutorch.getDeviceCount()
-- this tensor keeps track of whether a handle has been initialized or not
local handleStatus = torch.ByteTensor(numDevices,
                                  maxStreamsPerDevice):zero()
-- here we create an array of cudnn handle structs
cudnn.handle = ffi.new('struct cudnnContext*[?]', numDevices*maxStreamsPerDevice)
local function destroy(handle)
    local currentDevice = cutorch.getDevice()
    for i=1,numDevices do
        cutorch.setDevice(i)
        -- streams go from 0 to maxStreamsPerDevice - 1
        for j=0,maxStreamsPerDevice - 1 do
            if handleStatus[i][j + 1] == 1 then -- if handle was created
                cudnn.errcheck('cudnnDestroy', handle[(((i-1)*maxStreamsPerDevice) + j)]);
            end
        end
    end
    cutorch.setDevice(currentDevice)
end
ffi.gc(cudnn.handle, destroy)

cudnn.typemap = {
   ['torch.CudaHalfTensor']   = 'CUDNN_DATA_HALF',
   ['torch.CudaTensor']       = 'CUDNN_DATA_FLOAT',
   ['torch.CudaDoubleTensor'] = 'CUDNN_DATA_DOUBLE',
}

local sizeofmap = {
   ['torch.CudaHalfTensor']   = cutorch.hasHalf and ffi.sizeof('half') or 2,
   ['torch.CudaTensor']       = ffi.sizeof('float'),
   ['torch.CudaDoubleTensor'] = ffi.sizeof('double'),
}

function cudnn.sizeof(t)
   return sizeofmap[torch.type(t)]
end


local onemap = {
   ['torch.CudaHalfTensor']   = torch.FloatTensor({1}),
   ['torch.CudaTensor']       = torch.FloatTensor({1}),
   ['torch.CudaDoubleTensor'] = torch.DoubleTensor({1}),
}

local zeromap = {
   ['torch.CudaHalfTensor']   = torch.FloatTensor({0}),
   ['torch.CudaTensor']       = torch.FloatTensor({0}),
   ['torch.CudaDoubleTensor'] = torch.DoubleTensor({0}),
}
function cudnn.scalar(t, val)
   if val == 1 then
      return onemap[torch.type(t)]:data()
   elseif val == 0 then
      return zeromap[torch.type(t)]:data()
   else
      error('unknown scalar')
   end
end


local function fasterHalfMathTypeForCurrentDevice()
   -- get info from cutorc
   if cutorch.hasFastHalfInstructions() then
      return 'CUDNN_DATA_HALF'
   else
      return 'CUDNN_DATA_FLOAT'
   end
end

local configMaths = {}

local function configureMath(overrides)
   local currentDevice = cutorch.getDevice()
   for i=1,cutorch.getDeviceCount() do
      cutorch.setDevice(i)
      configMaths[i] = {
         ['torch.CudaHalfTensor']   = fasterHalfMathTypeForCurrentDevice(),
         ['torch.CudaTensor']       = 'CUDNN_DATA_FLOAT',
         ['torch.CudaDoubleTensor'] = 'CUDNN_DATA_DOUBLE',
      }
      -- apply overrides
      if overrides then
         for k,v in pairs(overrides) do configMaths[i][k] = v end
      end
   end
   cutorch.setDevice(currentDevice)
end
cudnn.configureMath = configureMath

-- TODO: rename to something like "configuredMathType" on next refactor
-- also, should move torch.type() inside
cudnn.configmap = function(tensortype)
   return configMaths[cutorch.getDevice()][tensortype]
end

configureMath()

function cudnn.getHandle()
    local device = cutorch.getDevice()
    local stream = cutorch.getStream() -- starts from 0
    assert(stream < maxStreamsPerDevice, 'cudnn bindings only support max of : '
               .. maxStreamsPerDevice .. ' streams per device')
    -- lazy initialization of handles
    if handleStatus[device][stream + 1] == 0 then
        local status = C['cudnnCreate'](cudnn.handle
                                        + (((device-1) * maxStreamsPerDevice)
                                                + stream))
        if status ~= ffi.C.CUDNN_STATUS_SUCCESS then
            local str = ffi.string(C.cudnnGetErrorString(status))
            error('Error in CuDNN: ' .. str)
        end
        handleStatus[device][stream + 1] = 1 -- mark handle as initialized
    end
    return cudnn.handle[(((device-1)*maxStreamsPerDevice) + stream)]
end

function cudnn.call(f, ...)
    C.cudnnSetStream(cudnn.getHandle(),
                     thc.THCState_getCurrentStream(cutorch.getState()))
    return C[f](...)
end

local errcheck = function(f, ...)
   local status = cudnn.call(f, ...)
   if status ~= ffi.C.CUDNN_STATUS_SUCCESS then
      local str = ffi.string(C.cudnnGetErrorString(status))
      error('Error in CuDNN: ' .. str .. ' ('..f..')')
      return false
   end
   return true
end
cudnn.errcheck = errcheck

function cudnn.toDescriptor(t)
   local typename = torch.typename(t)
   assert(cudnn.typemap[typename])
   local descriptor = ffi.new('struct cudnnTensorStruct*[1]')
   -- create descriptor
   errcheck('cudnnCreateTensorDescriptor', descriptor)
   -- set gc hook
   local function destroy(d)
      errcheck('cudnnDestroyTensorDescriptor', d[0]);
   end
   ffi.gc(descriptor, destroy)
   -- view 2D and 3D as 4D
   if t:dim() == 2 then
      t = t:view(t:size(1), t:size(2), 1, 1)
   elseif t:dim() == 3 then
      t = t:view(t:size(1), t:size(2), t:size(3), 1)
   end
   -- set descriptor
   local size = torch.LongTensor(t:size()):int()
   local stride = torch.LongTensor(t:stride()):int()

   errcheck('cudnnSetTensorNdDescriptor', descriptor[0], cudnn.typemap[typename],
            t:dim(), size:data(), stride:data())
   return descriptor
end

function cudnn.createDescriptors(count, descs_type, create_func, destroy_func)
   local ds = ffi.new(descs_type, count)
   for i = 0, count - 1 do
      errcheck(create_func, ds + i)
   end
   local function destroyDescriptors(ds)
      for i = 0, count - 1 do
         errcheck(destroy_func, ds[i])
      end
   end
   ffi.gc(ds, destroyDescriptors)
   return ds
end



function cudnn.setConvolutionDescriptor(data, desc)
   if not data.arrayLength then data.arrayLength = #data.padA end
   if not data.upscaleA then data.upscaleA =  torch.IntStorage(data.arrayLength):fill(1) end
   if not data.mode then data.mode = 'CUDNN_CROSS_CORRELATION' end

   local myDesc = desc or cudnn.createDescriptors(
      1, 'struct cudnnConvolutionStruct*[?]',
      'cudnnCreateConvolutionDescriptor', 'cudnnDestroyConvolutionDescriptor')
   -- make sure we have references to these tensors so gc doesn't clean them up
   local padATensor = torch.IntTensor(data.padA)
   local filterStrideATensor = torch.IntTensor(data.filterStrideA)
   local upscaleATensor = torch.IntTensor(data.upscaleA)
   errcheck('cudnnSetConvolutionNdDescriptor', myDesc[0],
            data.arrayLength,
            padATensor:data(),
            filterStrideATensor:data(),
            upscaleATensor:data(),
            data.mode,
            data.dataType)
   return myDesc
end

function cudnn.setFilterDescriptor(data, filterDesc)
   local myDesc = filterDesc or cudnn.createDescriptors(
      1, 'struct cudnnFilterStruct*[?]',
      'cudnnCreateFilterDescriptor', 'cudnnDestroyFilterDescriptor')
   local dims = data.nbDims or #data.filterDimA
   -- make sure we have references to these tensors so gc doesn't clean them up
   local filterDimATensor = torch.IntTensor(data.filterDimA)
   errcheck('cudnnSetFilterNdDescriptor', myDesc[0],
            data.dataType, data.format or 'CUDNN_TENSOR_NCHW',
            dims, filterDimATensor:data());
   return myDesc
end

local sharedBuffer = {}
local nextBufferSize = {}

-- may reassign currentSize
local function allocateStorage(buf, ifGreater)

   if buf.nextSize < 0 then
      buf.nextSize = buf.currentSize
   end

   local elSize = 8
   -- get number of elements in the buf, rounded up
   local newelem = math.floor((buf.nextSize+elSize-1)/elSize)

   if buf.storage then
      if (newelem == buf.storage:size()) or (ifGreater and newelem < buf.storage:size()) then
      else
         -- resize to just to make sure we return memory
         buf.storage:resize(0)
         buf.storage:resize(newelem)
      end
   else
      -- this is to be replaced with new cutorch tempbuf stuff
      -- may reassign currentSize again
      buf.storage = torch.CudaDoubleStorage(newelem)
   end

   buf.currentSize = buf.storage:size()*elSize
   buf.data = buf.storage:data()
   buf.nextSize = -1
end

local function sharedBufForStream(device, stream)
   device = device or cutorch.getDevice()
   stream = stream or cutorch.getStream() -- starts from 0
   if not sharedBuffer[device] then sharedBuffer[device] = {} end
   local buf = sharedBuffer[device][stream]
   if not buf then
      buf = {
         currentSize = cudnn.initialWorkspaceBytes,
         nextSize = -1
      }
      allocateStorage(buf)
      sharedBuffer[device][stream] = buf
   end
   return buf
end

function cudnn.getSharedWorkspace(device, stream)
   device = device or cutorch.getDevice()
   stream = stream or cutorch.getStream()
   local buf = sharedBufForStream(device, stream)
   return buf.data, buf.currentSize
end

-- Creates a clone of luaStr that can be used to prevent side
-- effects when passing char* to C functions.
function cudnn.externalizeString(luaStr)
    local cStr = ffi.new("char[?]", #luaStr+1)
    ffi.copy(cStr, luaStr)
    return cStr
end

function cudnn.adjustSharedWorkspaceSize(bytesDelta, device, stream)
   local buf = sharedBufForStream(device, stream)
   buf.nextSize = buf.currentSize + bytesDelta
   allocateStorage(buf)
end

function cudnn.setNextWorkspaceSize(bytes, device, stream)
   local buf = sharedBufForStream(device, stream)
   buf.nextSize = bytes
   return buf
end

function cudnn.setSharedWorkspaceSize(bytes, ifGreater, device, stream)
   bytes = bytes or cudnn.initialWorkspaceBytes
   local buf = cudnn.setNextWorkspaceSize(bytes, device, stream)
   allocateStorage(buf, ifGreater)
end

cudnn.find = require('cudnn.find')

require('cudnn.SpatialConvolution')
require('cudnn.VolumetricConvolution')
require('cudnn.SpatialFullConvolution')
require('cudnn.VolumetricFullConvolution')
require('cudnn.Pooling')
require('cudnn.SpatialMaxPooling')
require('cudnn.SpatialAveragePooling')
require('cudnn.Pooling3D')
require('cudnn.VolumetricMaxPooling')
require('cudnn.VolumetricAveragePooling')
require('cudnn.Pointwise')
require('cudnn.ReLU')
require('cudnn.ClippedReLU')
require('cudnn.Tanh')
require('cudnn.Sigmoid')
require('cudnn.SpatialSoftMax')
require('cudnn.SpatialLogSoftMax')
require('cudnn.VolumetricSoftMax')
require('cudnn.VolumetricLogSoftMax')
require('cudnn.SoftMax')
require('cudnn.LogSoftMax')
require('cudnn.SpatialCrossMapLRN')
require('cudnn.BatchNormalization')
require('cudnn.SpatialBatchNormalization')
require('cudnn.VolumetricBatchNormalization')
require('cudnn.SpatialCrossEntropyCriterion')
require('cudnn.VolumetricCrossEntropyCriterion')
require('cudnn.TemporalConvolution')
require('cudnn.RNN')
require('cudnn.RNNTanh')
require('cudnn.RNNReLU')
require('cudnn.BLSTM')
require('cudnn.LSTM')
require('cudnn.BGRU')
require('cudnn.GRU')
require('cudnn.functional')
require('cudnn.convert')

function cudnn.reset()
-- this resets everything
   if cudnn.verbose then
      print("cudnn::reset for device #", cutorch.getDevice())
   end
   cutorch.synchronize()
   -- make sure shared buffers that may have been cached, have 0 size
   for i=1,numDevices do
      sharedBuffer[i] = {}
   end
   collectgarbage()
   -- this resets internal algorithm finder state machine and cache
   cudnn.find.reset()
end

return cudnn