enable bf16 vec copy

[zhuhaozhe]

Enable bf16 vectorized copy.

BFloat16's copy get 2x performance for fp32 as our expectation.

BFloat16's vec copy dose not show performance gain compare with scalar version with op benchmark. This should caused by the memory system of operator. The system will really "read/write" a scalar at one time, although the code is written in scalar version.

benchmarks code:

import torch
import torch.utils.benchmark as benchmark

# x = torch.empty(10 * 18304 * 1024 * 16, dtype=torch.bfloat16)
x = torch.empty(10 * 18304 * 1024 * 16, dtype=torch.float)
def copy(tensors):
    for t in tensors:
        x.copy_(t)

tensors = []
for i in range(2):
    # l3 cache size 36608k = 18304 bfloat16 * 2 byte(per bfloat16)
    # tensors.append(torch.rand(10 * 18304 * 1024 * 16).bfloat16())
    tensors.append(torch.rand(10 * 18304 * 1024 * 16))

t0 = benchmark.Timer(
    stmt='copy(tensors)',
    setup='from __main__ import copy',
    globals={'tensors': tensors},
    num_threads=1)
    
print(t0.timeit(20))

Before this comit:
fp32:
3.84 s
1 measurement, 20 runs , 1 thread
bf16:
1.89 s
1 measurement, 20 runs , 1 thread

After:
fp32:
3.71 s
1 measurement, 20 runs , 1 thread
bf16:
1.85 s
1 measurement, 20 runs , 1 thread

[facebook-github-bot]

💊 CI failures summary and remediations

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

---

• 1/1 failures possibly* introduced in this PR
  • 1/1 non-scanned failure(s)

---

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.

[heitorschueroff]

LGTM!

[facebook-github-bot]

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

[codecov]

Codecov Report

Merging #54671 (df3cab1) into master (5781aec) will increase coverage by 0.00%.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master   #54671   +/-   ##
=======================================
  Coverage   77.45%   77.46%           
=======================================
  Files        1893     1893           
  Lines      185681   185681           
=======================================
+ Hits       143823   143830    +7     
+ Misses      41858    41851    -7     

[facebook-github-bot]

@heitorschueroff merged this pull request in 01b1557.