BFloat16 on cuda: add triu/tril support

[Maykeye]

🚀 The feature, motivation and pitch

Right now if you try to use torch.triu on bfloat16 tensor (I hit it when was training simple network with AMP) you'll get a error that triu is not support

In [8]: torch.__version__
Out[8]: '2.0.1+cu117'

In [9]: torch.arange(4).reshape(2,2).bfloat16().cuda().triu()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[9], line 1
----> 1 torch.arange(4).reshape(2,2).bfloat16().cuda().triu()

RuntimeError: "triu_tril_cuda_template" not implemented for 'BFloat16'

It would be nice to have it.

Alternatives

Can be replaced with multiplication against torch.ones.triu()

   self.register_buffer("triu0", torch.ones(n, n).triu()) # __init__
   ...
   y = x * self.triu0 # forward()
   

Additional context

No response

[Maykeye]

Oh.. Fixed by merged #101414 : if I try my example in freshly build docker container, it will work

>>> torch.__version__
'2.1.0a0+git22ca1a1'
>>> torch.arange(4).reshape(2,2).bfloat16().cuda().triu()
tensor([[0., 1.],
        [0., 3.]], device='cuda:0', dtype=torch.bfloat16)

[18335100284]

🚀 The feature, motivation and pitch

Right now if you try to use torch.triu on bfloat16 tensor (I hit it when was training simple network with AMP) you'll get a error that triu is not support

In [8]: torch.__version__
Out[8]: '2.0.1+cu117'

In [9]: torch.arange(4).reshape(2,2).bfloat16().cuda().triu()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[9], line 1
----> 1 torch.arange(4).reshape(2,2).bfloat16().cuda().triu()

RuntimeError: "triu_tril_cuda_template" not implemented for 'BFloat16'

It would be nice to have it.

Alternatives

Can be replaced with multiplication against torch.ones.triu()

   self.register_buffer("triu0", torch.ones(n, n).triu()) # __init__
   ...
   y = x * self.triu0 # forward()
   

Additional context

No response

Very nice idea, but are there any potential problems with this implementation? @Maykeye

[guotong1988]

So torch==2.1.0 have fixed this?

[guotong1988]

How to install 2.1.0a0+git22ca1a1? Thank you @Maykeye