Add CUDA support for arange

Also enables CUDA for range
zhouyonglong · Jul 19, 2017 · bd6263c · bd6263c
1 parent cfcf2af
commit bd6263c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/generic/THCTensorMath.cu b/generic/THCTensorMath.cu
@@ -449,4 +449,15 @@ void THCTensor_(range)(THCState *state, THCTensor *r_, accreal xmin, accreal xma
   THCudaCheck(cudaGetLastError());
 }
 
+void THCTensor_(arange)(THCState* state, THCTensor *r_, accreal xmin, accreal xmax, accreal step) {
+#if defined(THC_REAL_IS_FLOAT) || defined(THC_REAL_IS_DOUBLE) || defined(THC_REAL_IS_HALF)
+  int m = fmod(xmax - xmin, step) == 0;
+#else
+  int m = (xmax - xmin) % step == 0;
+#endif
+  if (m)
+    xmax -= step;
+  THCTensor_(range)(state, r_, xmin, xmax, step);
+}
+
 #endif
diff --git a/generic/THCTensorMath.h b/generic/THCTensorMath.h
@@ -26,5 +26,6 @@ THC_API void THCTensor_(logspace)(THCState *state, THCTensor *r_, real a, real b
 #endif
 
 THC_API void THCTensor_(range)(THCState *state, THCTensor *r_, accreal xmin, accreal xmax, accreal step);
+THC_API void THCTensor_(arange)(THCState *state, THCTensor *r_, accreal xmin, accreal xmax, accreal step);
 
 #endif
-Original file line number
+Diff line change
@@ Expand Up @@
     #endif
     THC_API void THCTensor_(range)(THCState *state, THCTensor *r_, accreal xmin, accreal xmax, accreal step);
+    THC_API void THCTensor_(arange)(THCState *state, THCTensor *r_, accreal xmin, accreal xmax, accreal step);
     #endif