experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Register AddN T=uint32 kernel on GPU

Browse Source 
PiperOrigin-RevId: 358242898
Change-Id: I9518bf1ccb53c65b0e434efb5d791d32d51de12e
This commit is contained in:
Eugene Zhulenev 2021-02-18 12:49:56 -08:00 committed by TensorFlower Gardener
parent dfc6473f5d
commit 8d1b9a9f82
2 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/core/kernels/aggregate_ops.cc
View File

@ -263,6 +263,7 @@ REGISTER_ADDN_CPU(Variant);
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define REGISTER_ADDN_GPU(type) REGISTER_ADDN(type, GPU) #define REGISTER_ADDN_GPU(type) REGISTER_ADDN(type, GPU)
TF_CALL_int64(REGISTER_ADDN_GPU); TF_CALL_int64(REGISTER_ADDN_GPU);
TF_CALL_uint32(REGISTER_ADDN_GPU);
TF_CALL_variant(REGISTER_ADDN_GPU); TF_CALL_variant(REGISTER_ADDN_GPU);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_ADDN_GPU); TF_CALL_GPU_NUMBER_TYPES(REGISTER_ADDN_GPU);
TF_CALL_COMPLEX_TYPES(REGISTER_ADDN_GPU); TF_CALL_COMPLEX_TYPES(REGISTER_ADDN_GPU);

1
tensorflow/core/kernels/aggregate_ops_gpu.cu.cc
View File

@ -155,6 +155,7 @@ struct Add9Functor<GPUDevice, T> {
template struct functor::Add9Functor<GPUDevice, type>; template struct functor::Add9Functor<GPUDevice, type>;
TF_CALL_int64(REGISTER_FUNCTORS); TF_CALL_int64(REGISTER_FUNCTORS);
TF_CALL_uint32(REGISTER_FUNCTORS);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_FUNCTORS); TF_CALL_GPU_NUMBER_TYPES(REGISTER_FUNCTORS);
TF_CALL_COMPLEX_TYPES(REGISTER_FUNCTORS); TF_CALL_COMPLEX_TYPES(REGISTER_FUNCTORS);