Bool tensor. Part 0: Boolean storage implementation (pytorch#16810)

Summary:
This is the first commit from a series of planned changes in order to add boolean tensors to PyTorch. The whole plan looks like this:

0. Storage Implementation (this change)
1. Tensor Creation.
2. Tensor Conversions.
3. Tensor Indexing.
4. Tensor Operations.
5. Back compatibility related changes.

This feature was requested by the community:
pytorch#4764
pytorch#4219
pytorch#4288

**Change**:
Added boolean type to the Storage class for CPU and CUDA backends.

**Tested via**:
1. unit tests
2. running this:
-> import torch
-> torch.BoolStorage
<class 'torch.BoolStorage'>
-> torch.cuda.BoolStorage
<class 'torch.cuda.BoolStorage'>
Pull Request resolved: pytorch#16810

Reviewed By: gchanan

Differential Revision: D14087246

Pulled By: izdeby

fbshipit-source-id: 042642ced1cb0fd1bb6bff05f9ca871a5c54ee5e

Author: izdeby

aten/src/ATen/DLConvertor.cpp

@@ -37,6 +37,9 @@ static DLDataType getDLDataType(const Type& type) {
     case ScalarType::Half:
       dtype.code = DLDataTypeCode::kDLFloat;
       break;
+    case ScalarType::Bool:
+      dtype.code = DLDataTypeCode::kDLUInt;
+      break;
     case ScalarType::ComplexHalf:
       throw std::logic_error("ComplexHalf is not supported by dlpack");
     case ScalarType::ComplexFloat:

aten/src/ATen/core/Type.h

@@ -44,6 +44,7 @@ struct Generator;
 static inline void noop_deleter(void*) {}
 
 enum class TypeID {
+  CPUBool,
   CPUByte,
   CPUChar,
   CPUDouble,
@@ -52,13 +53,15 @@ enum class TypeID {
   CPULong,
   CPUShort,
   CPUHalf,
+  SparseCPUBool,
   SparseCPUByte,
   SparseCPUChar,
   SparseCPUDouble,
   SparseCPUFloat,
   SparseCPUInt,
   SparseCPULong,
   SparseCPUShort,
+  CUDABool,
   CUDAByte,
   CUDAChar,
   CUDADouble,
@@ -67,6 +70,7 @@ enum class TypeID {
   CUDALong,
   CUDAShort,
   CUDAHalf,
+  SparseCUDABool,
   SparseCUDAByte,
   SparseCUDAChar,
   SparseCUDADouble,

aten/src/ATen/gen.py

@@ -178,6 +178,7 @@ def check_all_files_written(self):
 
 # scalar_name, c_type, accreal, th_scalar_type, is_floating_type
 scalar_types = [
+    ('Bool', 'uint8_t', 'BoolAccrealNotDefined', 'uint8_t', False),
     ('Byte', 'uint8_t', 'Long', 'uint8_t', False),
     ('Char', 'int8_t', 'Long', 'int8_t', False),
     ('Double', 'double', 'Double', 'double', True),

aten/src/TH/CMakeLists.txt

@@ -64,6 +64,7 @@ INSTALL(FILES
   THFilePrivate.h
   ${CMAKE_CURRENT_BINARY_DIR}/THGeneral.h
   THGenerateAllTypes.h
+  THGenerateBoolType.h
   THGenerateDoubleType.h
   THGenerateFloatType.h
   THGenerateHalfType.h

aten/src/TH/THFile.h

@@ -46,6 +46,7 @@ TH_API size_t THFile_readInt(THFile *self, THIntStorage *storage);
 TH_API size_t THFile_readLong(THFile *self, THLongStorage *storage);
 TH_API size_t THFile_readFloat(THFile *self, THFloatStorage *storage);
 TH_API size_t THFile_readDouble(THFile *self, THDoubleStorage *storage);
+TH_API size_t THFile_readBool(THFile *self, THBoolStorage *storage);
 
 TH_API size_t THFile_writeByte(THFile *self, THByteStorage *storage);
 TH_API size_t THFile_writeChar(THFile *self, THCharStorage *storage);
@@ -54,6 +55,7 @@ TH_API size_t THFile_writeInt(THFile *self, THIntStorage *storage);
 TH_API size_t THFile_writeLong(THFile *self, THLongStorage *storage);
 TH_API size_t THFile_writeFloat(THFile *self, THFloatStorage *storage);
 TH_API size_t THFile_writeDouble(THFile *self, THDoubleStorage *storage);
+TH_API size_t THFile_writeBool(THFile *self, THBoolStorage *storage);
 
 /* raw */
 TH_API size_t THFile_readByteRaw(THFile *self, uint8_t *data, size_t n);

aten/src/TH/THGenerateBoolType.h

@@ -0,0 +1,22 @@
+#ifndef TH_GENERIC_FILE
+#error "You must define TH_GENERIC_FILE before including THGenerateBoolType.h"
+#endif
+
+// TODO: define accreal type once the correct value is known.
+#define scalar_t bool
+#define ureal bool
+#define Real Bool
+#define TH_CONVERT_ACCREAL_TO_REAL(_val) (scalar_t)(_val)
+#define TH_REAL_IS_BOOL
+#line 1 TH_GENERIC_FILE
+#include TH_GENERIC_FILE
+#undef scalar_t
+#undef ureal
+#undef Real
+#undef TH_REAL_IS_BOOL
+#undef TH_CONVERT_REAL_TO_ACCREAL
+#undef TH_CONVERT_ACCREAL_TO_REAL
+
+#ifndef THGenerateManyTypes
+#undef TH_GENERIC_FILE
+#endif

aten/src/TH/THStorageFunctions.cpp

@@ -9,12 +9,18 @@
 #include <TH/generic/THStorage.cpp>
 #include <TH/THGenerateHalfType.h>
 
+#include <TH/generic/THStorage.cpp>
+#include <TH/THGenerateBoolType.h>
+
 #include <TH/generic/THStorageCopy.cpp>
 #include <TH/THGenerateAllTypes.h>
 
 #include <TH/generic/THStorageCopy.cpp>
 #include <TH/THGenerateHalfType.h>
 
+#include <TH/generic/THStorageCopy.cpp>
+#include <TH/THGenerateBoolType.h>
+
 THStorage* THStorage_new(caffe2::TypeMeta data_type) {
   THStorage* storage = c10::make_intrusive<at::StorageImpl>(
       data_type,

aten/src/TH/THStorageFunctions.h

@@ -11,11 +11,17 @@
 #include <TH/generic/THStorage.h>
 #include <TH/THGenerateHalfType.h>
 
+#include <TH/generic/THStorage.h>
+#include <TH/THGenerateBoolType.h>
+
 #include <TH/generic/THStorageCopy.h>
 #include <TH/THGenerateAllTypes.h>
 
 #include <TH/generic/THStorageCopy.h>
 #include <TH/THGenerateHalfType.h>
 
+#include <TH/generic/THStorageCopy.h>
+#include <TH/THGenerateBoolType.h>
+
 // This exists to have a data-type independent way of freeing (necessary for THPPointer).
 TH_API void THStorage_free(THStorage *storage);

aten/src/TH/THTensor.cpp

@@ -6,6 +6,9 @@
 #include <TH/generic/THTensor.cpp>
 #include <TH/THGenerateHalfType.h>
 
+#include <TH/generic/THTensor.cpp>
+#include <TH/THGenerateBoolType.h>
+
 #include <ATen/native/Resize.h>
 
 #include <numeric>

aten/src/TH/THTensor.h

@@ -13,6 +13,9 @@
 #include <TH/generic/THTensor.h>
 #include <TH/THGenerateHalfType.h>
 
+#include <TH/generic/THTensor.h>
+#include <TH/THGenerateBoolType.h>
+
 /* random numbers */
 #include <TH/THRandom.h>
 #include <TH/generic/THTensorRandom.h>

aten/src/TH/THTensor.hpp

@@ -127,3 +127,6 @@ TH_CPP_API c10::optional<std::vector<int64_t>> THTensor_compute_stride(
 
 #include <TH/generic/THTensor.hpp>
 #include <TH/THGenerateHalfType.h>
+
+#include <TH/generic/THTensor.hpp>
+#include <TH/THGenerateBoolType.h>

aten/src/TH/generic/THStorage.h

@@ -33,6 +33,7 @@
 #define THShortStorage THStorage
 #define THIntStorage THStorage
 #define THLongStorage THStorage
+#define THBoolStorage THStorage
 
 TH_API scalar_t* THStorage_(data)(const THStorage*);
 TH_API ptrdiff_t THStorage_(size)(const THStorage*);

aten/src/TH/generic/THStorageCopy.cpp

@@ -37,5 +37,6 @@ IMPLEMENT_THStorage_COPY(Long)
 IMPLEMENT_THStorage_COPY(Float)
 IMPLEMENT_THStorage_COPY(Double)
 IMPLEMENT_THStorage_COPY(Half)
+IMPLEMENT_THStorage_COPY(Bool)
 
 #endif

aten/src/TH/generic/THStorageCopy.h

@@ -14,5 +14,6 @@ TH_API void THStorage_(copyLong)(THStorage *storage, struct THLongStorage *src);
 TH_API void THStorage_(copyFloat)(THStorage *storage, struct THFloatStorage *src);
 TH_API void THStorage_(copyDouble)(THStorage *storage, struct THDoubleStorage *src);
 TH_API void THStorage_(copyHalf)(THStorage *storage, struct THHalfStorage *src);
+TH_API void THStorage_(copyBool)(THStorage *storage, struct THBoolStorage *src);
 
 #endif

aten/src/TH/generic/THTensor.h

@@ -18,6 +18,7 @@
 #define THShortTensor THTensor
 #define THIntTensor THTensor
 #define THLongTensor THTensor
+#define THBoolTensor THTensor
 
 /**** access methods ****/
 TH_API THStorage* THTensor_(storage)(const THTensor *self);

aten/src/THC/CMakeLists.txt

@@ -78,6 +78,7 @@ INSTALL(FILES
           THCDeviceTensorUtils.cuh
           THCDeviceTensorUtils-inl.cuh
           THCGenerateAllTypes.h
+          THCGenerateBoolType.h
           THCGenerateByteType.h
           THCGenerateCharType.h
           THCGenerateShortType.h

aten/src/THC/THCGenerateBoolType.h

@@ -0,0 +1,21 @@
+#ifndef THC_GENERIC_FILE
+#error "You must define THC_GENERIC_FILE before including THCGenerateBoolType.h"
+#endif
+
+// TODO: define accreal type once the correct value is known.
+#define scalar_t bool
+#define ureal bool
+#define Real Bool
+#define CReal CudaBool
+#define THC_REAL_IS_BOOL
+#line 1 THC_GENERIC_FILE
+#include THC_GENERIC_FILE
+#undef scalar_t
+#undef ureal
+#undef Real
+#undef CReal
+#undef THC_REAL_IS_BOOL
+
+#ifndef THCGenerateBoolType
+#undef THC_GENERIC_FILE
+#endif

aten/src/THC/THCStorage.cpp

@@ -8,6 +8,9 @@
 #include <THC/generic/THCStorage.cpp>
 #include <THC/THCGenerateAllTypes.h>
 
+#include <THC/generic/THCStorage.cpp>
+#include <THC/THCGenerateBoolType.h>
+
 #include <c10/util/intrusive_ptr.h>
 
 void THCStorage_resize(THCState *state, THCStorage *self, ptrdiff_t size)

aten/src/THC/THCStorage.cu

@@ -11,3 +11,6 @@
 
 #include <THC/generic/THCStorage.cu>
 #include <THC/THCGenerateAllTypes.h>
+
+#include <THC/generic/THCStorage.cu>
+#include <THC/THCGenerateBoolType.h>

aten/src/THC/THCStorage.h

@@ -9,4 +9,7 @@
 #include <THC/generic/THCStorage.h>
 #include <THC/THCGenerateAllTypes.h>
 
+#include <THC/generic/THCStorage.h>
+#include <THC/THCGenerateBoolType.h>
+
 #endif

aten/src/THC/THCStorageCopy.cpp

@@ -5,3 +5,6 @@
 
 #include <THC/generic/THCStorageCopy.cpp>
 #include <THC/THCGenerateAllTypes.h>
+
+#include <THC/generic/THCStorageCopy.cpp>
+#include <THC/THCGenerateBoolType.h>

aten/src/THC/THCStorageCopy.cu

@@ -8,3 +8,6 @@
 
 #include <THC/generic/THCStorageCopy.cu>
 #include <THC/THCGenerateAllTypes.h>
+
+#include <THC/generic/THCStorageCopy.cu>
+#include <THC/THCGenerateBoolType.h>

aten/src/THC/THCStorageCopy.h

@@ -8,4 +8,7 @@
 #include <THC/generic/THCStorageCopy.h>
 #include <THC/THCGenerateAllTypes.h>
 
+#include <THC/generic/THCStorageCopy.h>
+#include <THC/THCGenerateBoolType.h>
+
 #endif

aten/src/THC/THCTensor.cpp

@@ -7,6 +7,9 @@
 #include <THC/generic/THCTensor.cpp>
 #include <THC/THCGenerateAllTypes.h>
 
+#include <THC/generic/THCTensor.cpp>
+#include <THC/THCGenerateBoolType.h>
+
 #include <THC/THCTensorInfo.cuh>
 
 #include <ATen/native/cuda/Resize.cuh>
@@ -61,6 +64,8 @@ THCTensor *THCTensor_new(THCState *state, caffe2::TypeMeta type_meta) {
       return THCudaTensor_new(state);
     case at::ScalarType::Double:
       return THCudaDoubleTensor_new(state);
+    case at::ScalarType::Bool:
+      return THCudaBoolTensor_new(state);
     default:
       AT_ERROR("unexpected ScalarType: ", toString(scalar_type));
   }

aten/src/THC/THCTensor.cu

@@ -3,3 +3,6 @@
 
 #include <THC/generic/THCTensor.cu>
 #include <THC/THCGenerateAllTypes.h>
+
+#include <THC/generic/THCTensor.cu>
+#include <THC/THCGenerateBoolType.h>

aten/src/THC/THCTensor.h

@@ -17,4 +17,6 @@ typedef struct THC_CLASS THCDescBuff
 #include <THC/generic/THCTensor.h>
 #include <THC/THCGenerateAllTypes.h>
 
+#include <THC/generic/THCTensor.h>
+#include <THC/THCGenerateBoolType.h>
 #endif

aten/src/THC/THCTensor.hpp

@@ -56,3 +56,6 @@ THC_API bool THCTensor_maybeOverlappingIndices(THCState* state, const THCTensor*
 
 #include <THC/generic/THCTensor.hpp>
 #include <THC/THCGenerateAllTypes.h>
+
+#include <THC/generic/THCTensor.hpp>
+#include <THC/THCGenerateBoolType.h>

aten/src/THC/THCTensorCopy.cu

@@ -5,16 +5,26 @@
 #include <type_traits>
 
 // Copy operator for the pointwise apply kernel
-template <typename TypeDst, typename TypeSrc>
+template <typename T>
 struct CopyOp {
-  __device__ __forceinline__ void operator()(TypeDst* dst, TypeSrc* src) {
+  __device__ __forceinline__ void operator()(T* dst, T* src) {
 #if __CUDA_ARCH__ >= 350
-    *dst = ScalarConvert<TypeSrc, TypeDst>::to(__ldg(src));
+    *dst = ScalarConvert<T, T>::to(__ldg(src));
 #else
-    *dst = ScalarConvert<TypeSrc, TypeDst>::to(*src);
+    *dst = ScalarConvert<T, T>::to(*src);
 #endif
   }
 };
 
+template <>
+struct CopyOp <bool> {
+  __device__ __forceinline__ void operator()(bool* dst, bool* src) {
+      *dst = ScalarConvert<bool, bool>::to(*src);
+  }
+};
+
 #include <THC/generic/THCTensorCopy.cu>
 #include <THC/THCGenerateAllTypes.h>
+
+#include <THC/generic/THCTensorCopy.cu>
+#include <THC/THCGenerateBoolType.h>

aten/src/THC/THCTensorCopy.h

@@ -9,4 +9,7 @@
 #include <THC/generic/THCTensorCopy.h>
 #include <THC/THCGenerateAllTypes.h>
 
+#include <THC/generic/THCTensorCopy.h>
+#include <THC/THCGenerateBoolType.h>
+
 #endif

aten/src/THC/generic/THCStorage.h

@@ -14,6 +14,7 @@
 #define THCudaShortStorage  THCStorage
 #define THCudaIntStorage    THCStorage
 #define THCudaLongStorage   THCStorage
+#define THCudaBoolStorage   THCStorage
 
 THC_API scalar_t* THCStorage_(data)(THCState *state, const THCStorage*);
 THC_API ptrdiff_t THCStorage_(size)(THCState *state, const THCStorage*);

aten/src/THC/generic/THCStorageCopy.cpp

@@ -33,6 +33,7 @@ TH_CUDA_STORAGE_IMPLEMENT_COPY(Long)
 TH_CUDA_STORAGE_IMPLEMENT_COPY(Float)
 TH_CUDA_STORAGE_IMPLEMENT_COPY(Half)
 TH_CUDA_STORAGE_IMPLEMENT_COPY(Double)
+TH_CUDA_STORAGE_IMPLEMENT_COPY(Bool)
 
 void THStorage_(copyCuda)(THCState *state, THStorage *self, struct THCStorage *src)
 {
@@ -65,6 +66,7 @@ TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Long)
 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Float)
 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Half)
 TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Double)
+TH_CUDA_STORAGE_IMPLEMENT_COPYTO(Bool)
 
 #undef TH_CUDA_STORAGE_IMPLEMENT_COPY
 #undef TH_CUDA_STORAGE_IMPLEMENT_COPYTO

aten/src/THC/generic/THCStorageCopy.cu

@@ -28,6 +28,7 @@ THC_CUDA_STORAGE_IMPLEMENT_COPY(Long,Long)
 THC_CUDA_STORAGE_IMPLEMENT_COPY(Float,)  // i.e. float
 THC_CUDA_STORAGE_IMPLEMENT_COPY(Double,Double)
 THC_CUDA_STORAGE_IMPLEMENT_COPY(Half,Half)
+THC_CUDA_STORAGE_IMPLEMENT_COPY(Bool,Bool)
 
 #undef THC_CUDA_STORAGE_IMPLEMENT_COPY