OpenMP: spatial linear

4 files changed, 35 insertions, 8 deletions

diff --git a/README.txt b/README.txt
index 50cea03..6f183e9 100644
--- a/README.txt
+++ b/README.txt
@@ -5,3 +5,9 @@ $ luarocks --from=http://data.neuflow.org/lua/rocks install nnx
 USE:
 > require 'nnx'
 > n1 = nn.SpatialLinear(16,4)
+
+-- run tests:
+> nnx.test_all()
+...
+> nnx.test_omp()
+...
diff --git a/SpatialLinear.lua b/SpatialLinear.lua
index 1cf850e..56231d1 100644
--- a/SpatialLinear.lua
+++ b/SpatialLinear.lua
@@ -1,4 +1,4 @@
-local SpatialLinear, parent = torch.class('nn.SpatialLinear', 'nn.Module')
+local SpatialLinear, parent = torch.class('nn.SpatialLinear', 'nn.OmpModule')
 
 function SpatialLinear:__init(fanin, fanout)
    parent.__init(self)
diff --git a/generic/SpatialLinear.c b/generic/SpatialLinear.c
index 487d2f8..d97721a 100644
--- a/generic/SpatialLinear.c
+++ b/generic/SpatialLinear.c
@@ -9,6 +9,7 @@ static int nn_(SpatialLinear_forward)(lua_State *L)
   THTensor *bias = luaT_getfieldcheckudata(L, 1, "bias", torch_(Tensor_id));
   THTensor *weight = luaT_getfieldcheckudata(L, 1, "weight", torch_(Tensor_id));
   THTensor *output = luaT_getfieldcheckudata(L, 1, "output", torch_(Tensor_id));
+  int threads = luaT_getfieldcheckint(L, 1, "threads");
 
   // dims
   int iwidth = input->size[2];
@@ -18,26 +19,37 @@ static int nn_(SpatialLinear_forward)(lua_State *L)
   int oheight = iheight;
   int ochannels = output->size[0];
 
-  // select planes
-  THTensor *outputPlane = THTensor_(new)();
-  THTensor *inputPlane = THTensor_(new)();
-
   // process each plane
   int ok,ik;
+  omp_set_num_threads(threads);
+  omp_lock_t lock; omp_init_lock(&lock);
+  #pragma omp parallel for private(ok,ik)
   for (ok=0; ok<ochannels; ok++) {
-    // get output plane
+
+    // select planes
+    omp_set_lock(&lock);
+    THTensor *outputPlane = THTensor_(new)();
+    THTensor *inputPlane = THTensor_(new)();
     THTensor_(select)(outputPlane, output, 0, ok);
     THTensor_(fill)(outputPlane, THTensor_(get1d)(bias,ok));
+    omp_unset_lock(&lock);
+
     for (ik=0; ik<ichannels; ik++) {
       // get input plane
       THTensor_(select)(inputPlane, input, 0, ik);
       THTensor_(cadd)(outputPlane, THTensor_(get2d)(weight,ok,ik), inputPlane);
     }
+
+    // cleanup
+    omp_set_lock(&lock);
+    THTensor_(free)(inputPlane);
+    THTensor_(free)(outputPlane);
+    omp_unset_lock(&lock);
   }
 
   // cleanup
-  THTensor_(free)(inputPlane);
-  THTensor_(free)(outputPlane);
+  omp_destroy_lock(&lock);
+
   return 1;
 }
 
@@ -52,6 +64,7 @@ static int nn_(SpatialLinear_backward)(lua_State *L)
   THTensor *gradWeight = luaT_getfieldcheckudata(L, 1, "gradWeight", torch_(Tensor_id));
   THTensor *gradBias = luaT_getfieldcheckudata(L, 1, "gradBias", torch_(Tensor_id));
   int weightDecay = luaT_getfieldcheckint(L, 1, "weightDecay");
+  int threads = luaT_getfieldcheckint(L, 1, "threads");
 
   // dims
   int iwidth = input->size[2];
diff --git a/test/test-omp.lua b/test/test-omp.lua
index 05a941a..fc36316 100644
--- a/test/test-omp.lua
+++ b/test/test-omp.lua
@@ -87,6 +87,14 @@ function nnx.test_omp(threads)
 --       backward('SpatialMaxPooling_backward')
    end
 
+   function test_SpatialLinear()
+      ts = {}
+      times['SpatialLinear_forward'] = ts
+      n = nn.SpatialLinear(maps,maps2)
+      vec = lab.randn(maps,height,width)
+      forward('SpatialLinear_forward')
+   end
+
    -- run all tests
    lunit.main()