functional tests pass

3 files changed, 114 insertions, 9 deletions

diff --git a/functional.lua b/functional.lua
index 959b945..002696d 100644
--- a/functional.lua
+++ b/functional.lua
@@ -44,6 +44,21 @@ cudnn.functional.bias2D_accGradParameters = function(handle, gradOutput, gradBia
              biasDesc[0], gradBias:data())
 end
 
+
+local function getWsPtrAndSize(workspace, maxBufSize)
+  local wsPtr, extraBufferSizeInBytes
+  if workspace then
+    if maxBufSize > workspace:nElement()*workspace:elementSize() then
+       local nElems = math.ceil(maxBufSize/workspace:elementSize())
+       workspace:resize(nElems)
+    end
+  else
+    cudnn.setSharedWorkspaceSize(maxBufSize,true)
+    wsPtr, extraBufferSizeInBytes = cudnn.getSharedWorkspace()
+  end
+  return wsPtr, extraBufferSizeInBytes
+end
+
 -- Does a 2D Convolution (updateOutput) on input, weight
 -- output is assumed to be allocated and given.
 cudnn.functional.Convolution2D_updateOutput = function(handle, input, weight, output,
@@ -118,21 +133,22 @@ cudnn.functional.Convolution2D_updateOutput = function(handle, input, weight, ou
               convDesc[0], oDesc[0],
               algType[0], bufSize:data())
    local maxBufSize = bufSize[1]
+   local wsPtr, extraBufferSizeInBytes = getWsPtrAndSize(workspace, maxBufSize)
 
-   local extraBuffer = workspace or cudnn.getSharedWorkspace()
+--[[   local extraBuffer = workspace or cudnn.getSharedWorkspace()
    local extraBufferSizeInBytes = extraBuffer:nElement() * 4 -- extraBuffer is always float
    if maxBufSize > extraBufferSizeInBytes then
       extraBuffer:resize(math.ceil(maxBufSize / 4))
       extraBufferSizeInBytes = maxBufSize
    end
-
+--]]
    -- do convolution
    errcheck('cudnnConvolutionForward', handle,
             cudnn.scalar(input, 1),
             iDesc[0], input:data(),
             weightDesc[0], weight:data(),
             convDesc[0], algType[0],
-            extraBuffer:data(), extraBufferSizeInBytes,
+            wsPtr, extraBufferSizeInBytes,
             cudnn.scalar(input, 0),
             oDesc[0], output:data());
 end
@@ -198,14 +214,15 @@ cudnn.functional.Convolution2D_updateGradInput = function(handle, input, weight,
            convDesc[0], iDesc[0],
            algType[0], bufSize:data())
    local maxBufSize = bufSize[1]
+   local wsPtr, extraBufferSizeInBytes = getWsPtrAndSize(workspace, maxBufSize)
 
-   local extraBuffer = cudnn.getSharedWorkspace()
+--[[   local extraBuffer = cudnn.getSharedWorkspace()
    local extraBufferSizeInBytes = extraBuffer:nElement() * 4 -- extraBuffer is always float
    if maxBufSize > extraBufferSizeInBytes then
       extraBuffer:resize(math.ceil(maxBufSize / 4))
       extraBufferSizeInBytes = maxBufSize
    end
-
+--]]
    -- do convolution
    errcheck('cudnnConvolutionBackwardData', handle,
                cudnn.scalar(input, 1),
@@ -213,7 +230,7 @@ cudnn.functional.Convolution2D_updateGradInput = function(handle, input, weight,
                oDesc[0], gradOutput:data(),
                convDesc[0],
                algType[0],
-               extraBuffer:data(), extraBufferSizeInBytes,
+               wsPtr, extraBufferSizeInBytes,
                cudnn.scalar(input, 0),
                iDesc[0], gradInput:data());
 end
@@ -280,14 +297,15 @@ cudnn.functional.Convolution2D_accGradParameters = function(handle, input, gradW
               convDesc[0], weightDesc[0],
               algType[0], bufSize:data())
     local maxBufSize = bufSize[1]
+    local wsPtr, extraBufferSizeInBytes = getWsPtrAndSize(workspace, maxBufSize)
 
-    local extraBuffer = cudnn.getSharedWorkspace()
+--[[    local extraBuffer = cudnn.getSharedWorkspace()
     local extraBufferSizeInBytes = extraBuffer:nElement() * 4 -- extraBuffer is always float
     if maxBufSize > extraBufferSizeInBytes then
        extraBuffer:resize(math.ceil(maxBufSize / 4))
        extraBufferSizeInBytes = maxBufSize
     end
-
+--]]
     -- do convolution
     errcheck('cudnnConvolutionBackwardFilter', handle,
              scaleT:data(),
@@ -295,7 +313,7 @@ cudnn.functional.Convolution2D_accGradParameters = function(handle, input, gradW
              oDesc[0], gradOutput:data(),
              convDesc[0],
              algType[0],
-             extraBuffer:data(), extraBufferSizeInBytes,
+             wsPtr, extraBufferSizeInBytes,
              cudnn.scalar(input, 1),
              weightDesc[0], gradWeight:data());
 end
diff --git a/init.lua b/init.lua
index 8813149..357e367 100644
--- a/init.lua
+++ b/init.lua
@@ -314,6 +314,7 @@ require('cudnn.BLSTM')
 require('cudnn.LSTM')
 require('cudnn.BGRU')
 require('cudnn.GRU')
+require('cudnn.functional')
 require('cudnn.convert')
 
 function cudnn.reset()
diff --git a/test/test.lua b/test/test.lua
index a2f9923..86c4d70 100644
--- a/test/test.lua
+++ b/test/test.lua
@@ -800,6 +800,92 @@ function cudnntest.VolumetricCrossEntropyCriterion()
 end
 
 
+function cudnntest.functional_bias2D()
+   local bs = math.random(1,32)
+   local from = math.random(1,32)
+   local to = math.random(1,64)
+   local ki = math.random(1,15)
+   local kj = math.random(1,15)
+   local si = math.random(1,ki)
+   local sj = math.random(1,kj)
+   local outi = math.random(1,64)
+   local outj = math.random(1,64)
+   local ini = (outi-1)*si+ki
+   local inj = (outj-1)*sj+kj
+   local scale = torch.uniform()
+   local input = torch.zeros(bs,from,inj,ini):cuda()
+   local mod = cudnn.SpatialConvolution(from,to,ki,kj,si,sj):cuda()
+   mod.weight:zero()
+   local groundtruth = mod:forward(input)
+   local result = groundtruth:clone():zero()
+   cudnn.functional.bias2D_updateOutput(cudnn.getHandle(), mod.bias, result)
+   local error = result:float() - groundtruth:float()
+   mytester:assertlt(error:abs():max(),
+                     testparams.precision_forward, 'error on forward ')
+
+   mod:zeroGradParameters()
+   local gradOutput = groundtruth:clone():normal()
+   mod:backward(input, gradOutput, scale)
+   local groundtruth = mod.gradBias
+   local result = groundtruth:clone():zero()
+   cudnn.functional.bias2D_accGradParameters(cudnn.getHandle(), gradOutput, result, scale)
+   error = result:float() - groundtruth:float()
+   mytester:assertlt(error:abs():max(),
+                     testparams.precision_backward, 'error on accGradParameters ')
+end
+
+function cudnntest.functional_convolution2d()
+    local a=cudnn.SpatialConvolution(3,16,5,5):cuda()
+    a.bias:zero();
+    local input = torch.randn(10,3,10,10):cuda()
+    a:zeroGradParameters()
+    a:forward(input);
+    local output = a.output:clone():normal()
+    local gradOutput = a.output:clone():normal()
+    local gradInput = a:backward(input, gradOutput):clone():normal()
+    local gradWeight = a.gradWeight:clone():zero()
+    cudnn.functional.Convolution2D_updateOutput(cudnn.getHandle(), input,
+                                                a.weight, output, a.dH,
+                                                a.dW, a.padH, a.padW)
+    mytester:assertlt((output - a.output):abs():max(),
+                     testparams.precision_forward, 'error on forward ')
+
+    cudnn.functional.Convolution2D_updateGradInput(cudnn.getHandle(), input,
+                                                   a.weight, output, gradOutput,
+                                                   gradInput,
+                                                   a.dH, a.dW, a.padH, a.padW)
+    mytester:assertlt((gradInput - a.gradInput):abs():max(),
+                     testparams.precision_forward, 'error on updateGradInput ')
+
+    cudnn.functional.Convolution2D_accGradParameters(cudnn.getHandle(), input,
+                                                   gradWeight, gradOutput,
+                                                   a.dH, a.dW, a.padH, a.padW)
+    mytester:assertlt((gradWeight - a.gradWeight):abs():max(),
+                     testparams.precision_forward, 'error on accGradParameters ')
+end
+
+function cudnntest.functional_maxpooling2d()
+    local a=cudnn.SpatialMaxPooling(2,2,2,2):cuda()
+    local input = torch.randn(10,3,10,10):cuda()
+    a:forward(input);
+    local output = a.output:clone():normal()
+    local gradOutput = a.output:clone():normal()
+    local gradInput = a:backward(input, gradOutput):clone():normal()
+    cudnn.functional.MaxPooling2D_updateOutput(cudnn.getHandle(), input,
+                                               output, a.kH, a.kW,
+                                               a.dH, a.dW, a.padH, a.padW)
+    mytester:assertlt((output - a.output):abs():max(),
+                     testparams.precision_forward, 'error on forward ')
+
+    cudnn.functional.MaxPooling2D_updateGradInput(cudnn.getHandle(), input,
+                                                   output, gradOutput, gradInput,
+                                                   a.kH, a.kW, a.dH, a.dW,
+                                                   a.padH, a.padW)
+    mytester:assertlt((gradInput - a.gradInput):abs():max(),
+                     testparams.precision_forward, 'error on updateGradInput ')
+end
+
+
 local function test_functional_activation(mode, module)
    local a = module:cuda()
    local input = torch.randn(10,12):cuda()