Extra GPU-CPU memory transfer when broadcasting operations between integer tensors about tf-keras

[kumariko]

@sachinprasadhs Was able to reproduce the issue on CPU and GPU please find the gist attached.Thanks!

[tilakrayal]

@unrealwill,
I tried to execute the code by provinding the alternative x_train and y-train data, and found the difference in the time interval.
Case1: tf.random.uniform((600000,100) -- 35s
Case2: tf.random.uniform((60000,100) -- 3s
Case3 : tf.random.uniform((6000,100) -- 1s

Kindly find the gist of it here and I suspect the computation time was taken longer for the x_train and y_train data when providing the tf.random.uniform((600000,100). Thank you!

[unrealwill]

@tilakrayal
This is not about computation time. It's about the presence of extra GPU to CPU memory transfer : "The green block "MemCpyH2D" " in the picture in the other thread shouldn't be.

It's quite a binary problem, either it's no longer there in new tensorflow versions, in which case you can close the bug, or it's still there and in which case it still needs to be solved by defining the proper missing ops for the int types : tensorflow/tensorflow#54197 (comment)

I no longer use tensorflow so can't test anymore, and don't care anymore. I encountered this problem when I was implementing some InstantNerf neural network with tensorflow : the hashing operations involved used some operations on int32 tensors therefore encountered the problem, which manifested through slow performance due to this extra gpu-cpu transfer. Some elementary ops (like broadcasted addition and mul between integer tensor) that shouldn't necessitate gpu-cpu transfer suddenly do because the kernel is not defined for int tensors. But because it's done transparently, the ops get done but on the cpu so you may miss the probem, everything work, but it really kills performance for no other reasons than a kernel not being defined (But if defining such a basic kernel increase the tf binary size so much it means you have other problems).