adds Sort, revamps TensorMath, adds masked* operations

torch · Mar 30, 2015 · f004da6 · f004da6
1 parent 4fdbf0c
commit f004da6
Show file tree

Hide file tree

Showing 19 changed files with 2,467 additions and 909 deletions.
diff --git a/FFI.lua b/FFI.lua
@@ -7,6 +7,7 @@ typedef struct THCState
 {
   struct THCRNGState* rngState;
   struct THCBlasState* blasState;
+  struct cudaDeviceProp* deviceProperties;
 } THCState;
 
 typedef struct THCudaStorage

diff --git a/Tensor.c b/Tensor.c
@@ -5,21 +5,6 @@
 
 /* everything is as the generic Storage.c, except few things (see below) */
 
-static void THCudaTensor_maskedFill(THCState *state, THCudaTensor *tensor, THByteTensor *mask, float value)
-{
-  THError("not yet implemented for CUDA");
-}
-
-static void THCudaTensor_maskedCopy(THCState *state, THCudaTensor *tensor, THByteTensor *mask, THCudaTensor* src)
-{
-  THError("not yet implemented for CUDA");
-}
-
-void THCudaTensor_maskedSelect(THCState *state, THCudaTensor *tensor, THCudaTensor* src, THByteTensor *mask)
-{
-  THError("not yet implemented for CUDA");
-}
-
 #define real float
 #define Real Cuda
 

diff --git a/TensorMath.lua b/TensorMath.lua
@@ -169,7 +169,7 @@ end
 
 local function lastdim(argn)
    return function(arg)
-      return string.format("THCudaTensor_nDimension(%s)", arg.args[argn]:carg())
+      return string.format("THCudaTensor_nDimension(cutorch_getstate(L), %s)", arg.args[argn]:carg())
    end
 end
 
@@ -289,6 +289,33 @@ wrap("addcdiv",
         {name=Tensor},
         {name=Tensor}})
 
+wrap("maskedFill",
+     cname("maskedFill"),
+     {{name=Tensor, returned=true, method={default='nil'}},
+      {name=Tensor},
+      {name=real}})
+
+wrap("maskedCopy",
+     cname("maskedCopy"),
+     {{name=Tensor, returned=true, method={default='nil'}},
+      {name=Tensor},
+      {name=Tensor}})
+
+wrap("maskedSelect",
+     cname("maskedSelect"),
+     {{name=Tensor, returned=true, default=true},
+      {name=Tensor},
+      {name=Tensor}})
+
+wrap("sort",
+     cname("sort"),
+     {{name=Tensor, default=true, returned=true},
+        {name=Tensor, default=true, returned=true, noreadadd=true},
+        {name=Tensor},
+        {name="index", default=lastdim(3)},
+        {name="boolean", default=0}})
+
+
 do
    local Tensor = Tensor
    local real = real

diff --git a/init.c b/init.c
@@ -40,6 +40,15 @@ static int cutorch_getDeviceCount(lua_State *L)
   return 1;
 }
 
+static int cutorch_getMemoryUsage(lua_State *L) {
+  size_t freeBytes = 0;
+  size_t totalBytes = 0;
+  THCudaCheck(cudaMemGetInfo(&freeBytes, &totalBytes));
+  lua_pushnumber(L, freeBytes);
+  lua_pushnumber(L, totalBytes);
+  return 2;
+}
+
 static int cutorch_setDevice(lua_State *L)
 {
   THCState *state = cutorch_getstate(L);
@@ -158,6 +167,7 @@ static const struct luaL_Reg cutorch_stuff__ [] = {
   {"deviceReset", cutorch_deviceReset},
   {"getDeviceCount", cutorch_getDeviceCount},
   {"getDeviceProperties", cutorch_getDeviceProperties},
+  {"getMemoryUsage", cutorch_getMemoryUsage},
   {"setDevice", cutorch_setDevice},
   {"seed", cutorch_seed},
   {"seedAll", cutorch_seedAll},

diff --git a/lib/THC/CMakeLists.txt b/lib/THC/CMakeLists.txt
@@ -1,7 +1,7 @@
 SET(src
     THCGeneral.c THCStorage.c THCStorageCopy.c THCTensor.c THCTensorCopy.c)
 
-SET(src-cuda THC.cu)
+SET(src-cuda THC.cu THCReduceApplyUtils.cu THCTensorSort.cu)
 
 CUDA_ADD_LIBRARY(THC SHARED ${src} ${src-cuda})
 CUDA_ADD_CUBLAS_TO_TARGET(THC)
@@ -23,4 +23,8 @@ INSTALL(FILES
           THCTensorRandom.h
           THCTensorMath.h
           THCTensorConv.h
+          THCTensorSort.h
+          THCApply.cuh
+          THCReduce.cuh
+          THCReduceApplyUtils.cuh
           DESTINATION "${Torch_INSTALL_INCLUDE_SUBDIR}/THC")
diff --git a/lib/THC/THC.h b/lib/THC/THC.h
@@ -10,5 +10,6 @@
 #include "THCTensorRandom.h"
 #include "THCTensorMath.h"
 #include "THCTensorConv.h"
+#include "THCTensorSort.h"
 
 #endif