Enable BFloat16 for logaddexp & logaddexp2 on CUDA

[imaginary-person]

Enabled BFloat16 for logaddexp & logaddexp2 on CUDA, with a workaround suggested by @zasdfgbnm.

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 700d612 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[imaginary-person]

Compilation fails on Windows

[imaginary-person]

Hey @zasdfgbnm, while enabling BFloat16 for logaddexp & logaddexp2 on CUDA, Windows build fails with the following error -

C:/Users/circleci/project/aten/src/ATen/native/cuda/LogAddExpKernel.cu(13): error: calling a
 __host__ function("isinf< ::c10::BFloat16> ") from 
a __host__ __device__ function("at::native::logaddexp_kernel_cuda(::at::TensorIteratorBase &)::[lambda() (instance 1)]::operator 
()() const::[lambda() (instance 6)]::operator ()() const::[lambda( ::c10::BFloat16,  ::c10::BFloat16) 
(instance 1)]::operator () const") is not allowed

Can you please help fix this? Thank you!

[zasdfgbnm]

Let me try something

[zasdfgbnm]

I think ::isinf doesn't work for half as well, see: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/test/cuda_half_test.cu#L70

And you can workaround this by doing

#include <ATen/AccumulateType.h>
using accscalar_t = at::acc_type<scalar_t, /*is_cuda=*/true>;
::isinf(static_cast<accscalar_t>(a));

This workaround is already being used in https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/IGammaKernel.cu#L403

[codecov]

Codecov Report

Merging #57908 (700d612) into master (e8fb167) will increase coverage by 0.00%.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master   #57908   +/-   ##
=======================================
  Coverage   76.83%   76.83%           
=======================================
  Files        1984     1984           
  Lines      197144   197144           
=======================================
+ Hits       151471   151480    +9     
+ Misses      45673    45664    -9     

[facebook-github-bot]

@ngimel has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[ngimel]

Thanks!

[facebook-github-bot]

@ngimel merged this pull request in 067147a.