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 BFloat16 TopK #44755

Closed
wants to merge 5 commits into from
Closed

CUDA BFloat16 TopK #44755

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fbd83d1
CUDA BFloat16 TopK
zasdfgbnm Sep 15, 2020
f1d0fbe
Update THCDeviceUtils.cuh
zasdfgbnm Sep 16, 2020
56cf998
Merge branch 'master' of github.com:pytorch/pytorch into bfloat-topk
zasdfgbnm Oct 3, 2020
cb8ea5d
fix ldg
zasdfgbnm Oct 3, 2020
ce32659
Merge branch 'bfloat-topk' of github.com:pytorch/pytorch into bfloat-…
zasdfgbnm Oct 3, 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
11 changes: 11 additions & 0 deletions aten/src/THC/THCDeviceUtils.cuh
View file
Open in desktop
Expand Up @@ -7,6 +7,8 @@
#include <c10/util/Half.h>
#endif

#include <c10/util/BFloat16.h>

/* The largest consecutive integer representable in float32 (2^24) */
#define FLOAT32_MAX_CONSECUTIVE_INT 16777216.0f

Expand Down Expand Up @@ -39,6 +41,15 @@ __device__ __forceinline__ T doLdg(const T* p) {
#endif
}

template <>
__device__ __forceinline__ c10::BFloat16 doLdg<c10::BFloat16>(const c10::BFloat16* p) {
#if __CUDA_ARCH__ >= 350
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do you need #if here? torch is only supported on CUDA_ARCH >= 350

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is actually equivalent to #ifndef __HIP_PLATFORM_HCC__?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Then it should say so? __ldg doesn't provide performance benefit these days, but I guess you still need to load short and construct bfloat16 from bits on cuda, and hip is able to handle it natively?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I will change this to #ifndef __HIP_PLATFORM_HCC__, and maybe remove it later (needs benchmark). On HIP, it is just *p, so it's OK.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to let you know, changing the #if does break the build for NVIDIA GRID K520 GPU. I understand that is not a supported CUDA architecture though.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry about that, but as you note it is not a supported architecture.

return c10::BFloat16(__ldg(reinterpret_cast<const short *>(p)), c10::BFloat16::from_bits());
#else
return *p;
#endif
}

#include <ATen/cuda/DeviceUtils.cuh>

#endif // THC_DEVICE_UTILS_INC
4 changes: 0 additions & 4 deletions aten/src/THC/generic/THCTensorTopK.cu
View file
Open in desktop
Expand Up @@ -9,9 +9,6 @@ void THCTensor_(topk)(THCState* state,
THCudaLongTensor *indices,
THCTensor *input_,
int64_t k, int dim, int dir, int sorted) {
#if defined(THC_REAL_IS_BFLOAT16) && !defined(__HIP_PLATFORM_HCC__)
TORCH_CHECK(false, "topk not suppported with BFloat16");
#else
THAssert(topK != NULL && indices != NULL && input_ != NULL);
THCAssertSameGPU(THCTensor_(checkGPU)(state, 3, topK, indices, input_));
dim = at::maybe_wrap_dim(dim, input_);
Expand Down Expand Up @@ -186,7 +183,6 @@ void THCTensor_(topk)(THCState* state,
THCudaLongTensor_free(state, input);

THCudaCheck(cudaGetLastError());
#endif // THC_REAL_IS_BFLOAT16 && !__HIP_PLATFORM_HCC__
}

#endif // THC_GENERIC_FILE
12 changes: 4 additions & 8 deletions test/test_torch.py
View file
Open in desktop
Expand Up @@ -14394,8 +14394,7 @@ def test_topk_integral(self, device, dtype):
self.assertEqual(sort_topk, topk[0]) # check values
self.assertEqual(sort_topk, a[topk[1]]) # check indices

@dtypesIfCUDA(*([torch.half, torch.float, torch.double]
+ ([torch.bfloat16] if TEST_WITH_ROCM else [])))
@dtypesIfCUDA(*torch.testing.get_all_fp_dtypes())
@dtypes(torch.float, torch.double)
def test_topk_nonfinite(self, device, dtype):
x = torch.tensor([float('nan'), float('inf'), 1e4, 0, -1e4, -float('inf')], device=device, dtype=dtype)
Expand All @@ -14422,9 +14421,6 @@ def test_topk_4d(self, device):
self.assertEqual(val, expected_val, atol=0, rtol=0)
self.assertEqual(ind, expected_ind, atol=0, rtol=0)




def test_is_signed(self, device):
self.assertEqual(torch.IntTensor(5).to(device).is_signed(), True)
self.assertEqual(torch.ByteTensor(5).to(device).is_signed(), False)
Expand Down Expand Up @@ -20040,11 +20036,11 @@ def inner(self, device, dtype):
('transpose', 'neg_dim', _new_t((1, 2, 3, 4)), lambda t, d: [-1, -2], ),
('tolist', '', _small_3d, lambda t, d: [], 1e-5, 1e-5, 1e-5, _types, _cpu_types, False),
('topk', 'dim_sort', _small_3d_unique, lambda t, d: [2, 1, False, True],
1e-5, 1e-5, 1e-5, _types2, _cpu_types, False),
1e-5, 1e-5, 1e-5, torch.testing.get_all_dtypes(include_complex=False, include_bool=False), _cpu_types, False),
('topk', 'neg_dim_sort', _small_3d_unique, lambda t, d: [2, -1, False, True],
1e-5, 1e-5, 1e-5, _types2, _cpu_types, False),
1e-5, 1e-5, 1e-5, torch.testing.get_all_dtypes(include_complex=False, include_bool=False), _cpu_types, False),
('topk', 'dim_desc_sort', _small_3d_unique, lambda t, d: [2, 1, True, True],
1e-5, 1e-5, 1e-5, _types2, _cpu_types, False),
1e-5, 1e-5, 1e-5, torch.testing.get_all_dtypes(include_complex=False, include_bool=False), _cpu_types, False),
('trace', '', _medium_2d, lambda t, d: [], 1e-3, 1e-5, 1e-5, _types, _cpu_types, False),
('tril', '', _medium_2d, lambda t, d: [],),
('tril', 'zero_stride', _medium_2d, lambda t, d: [], 1e-5, 1e-5, 1e-5, _types, _cpu_types, False),
Expand Down