Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

CUDA BF16 norm #48806

Closed
wants to merge 2 commits into from
Closed

CUDA BF16 norm #48806

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dc37d67
CUDA BF16 norm
zasdfgbnm Dec 3, 2020
219512e
Update test_torch.py
zasdfgbnm Dec 4, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 0 additions & 2 deletions aten/src/ATen/native/cuda/ReduceNormKernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,14 +40,12 @@ static void norm_kernel_cuda(TensorIterator& iter, Scalar p) {
// type promotion that does cast and reduction in a single kernel
return norm_kernel_cuda_impl<at::Half, float, float>(iter, p);
}
#ifdef __HIP_PLATFORM_HCC__
else if(iter.dtype() == kBFloat16) {
return norm_kernel_cuda_impl<at::BFloat16, float>(iter, p);
} else if (iter.dtype(1) == kBFloat16 && iter.dtype() == kFloat) {
// type promotion that does cast and reduction in a single kernel
return norm_kernel_cuda_impl<at::BFloat16, float, float>(iter, p);
}
#endif
AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "norm_cuda", [&]() {
norm_kernel_cuda_impl<scalar_t>(iter, p);
});
Expand Down
14 changes: 6 additions & 8 deletions test/test_torch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6328,10 +6328,6 @@ def test_copy_broadcast(self, device) -> None:

_float_types_no_half = [torch.float, torch.double]

# _float_types2 adds bfloat16 type to _float_types only on ROCm. Should eventually be unified
# with _float_types when bfloat16 bringup is complete on all platforms
_float_types2 = _float_types + [torch.bfloat16] if TEST_WITH_ROCM else _float_types

_signed_types = [
torch.half, torch.bfloat16, torch.float, torch.double,
torch.int8, torch.short, torch.int, torch.long
Expand Down Expand Up @@ -6689,10 +6685,12 @@ def inner(self, device, dtype):
('narrow', '', _small_3d, lambda t, d: [1, 3, 2], 1e-5, 1e-5, 1e-5, _types, _cpu_types, False),
('narrow', 'neg_dim', _small_3d, lambda t, d: [-1, 3, 2], 1e-5, 1e-5, 1e-5, _types, _cpu_types, False),
('nonzero', '', _small_3d, lambda t, d: [], 1e-5, 1e-5, 1e-5, _types, _cpu_types, False),
('norm', '', _small_3d, lambda t, d: [], 1e-1, 1e-1, 1e-5, _float_types2, _cpu_types, False),
('norm', '3_norm', _small_3d, lambda t, d: [3], 1e-1, 1e-1, 1e-5, _float_types2, _cpu_types, False),
('norm', '3_norm_dim', _small_3d, lambda t, d: [3, 0], 1e-1, 1e-1, 1e-5, _float_types2, _cpu_types, False),
('norm', '3_norm_neg_dim', _small_3d, lambda t, d: [3, -2], 1e-1, 1e-1, 1e-5, _float_types2, _cpu_types, False),
('norm', '', _small_3d, lambda t, d: [], 1e-1, 1e-1, 1e-5, torch.testing.get_all_fp_dtypes(), _cpu_types, False),
('norm', '3_norm', _small_3d, lambda t, d: [3], 1e-1, 1e-1, 1e-5, torch.testing.get_all_fp_dtypes(), _cpu_types, False),
('norm', '3_norm_dim', _small_3d, lambda t, d: [3, 0], 1e-1, 1e-1, 1e-5,
torch.testing.get_all_fp_dtypes(), _cpu_types, False),
('norm', '3_norm_neg_dim', _small_3d, lambda t, d: [3, -2], 1e-1, 1e-1, 1e-5,
torch.testing.get_all_fp_dtypes(), _cpu_types, False),
('new_ones', '', _small_3d, lambda t, d: [1, 2, 3, 4, 5], 1e-5, 1e-5, 1e-5, _types, _cpu_types, False),
('permute', '', _new_t((1, 2, 3, 4)), lambda t, d: [2, 1, 3, 0], 1e-5, 1e-5, 1e-5, _types, _cpu_types, False),
('put_', '', _new_t((2, 5, 3)),
Expand Down