Allow at::native::offset_t to be offset using operator+= (#164570)

This will be required by CCCL 3.1.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/164570
Approved by: https://github.com/Skylion007, https://github.com/eqy
This commit is contained in:
Bernhard Manfred Gruber
2025-10-15 01:40:49 +00:00
committed by PyTorch MergeBot
parent 47524dcc48
commit a20afb6100

View File

@ -21,9 +21,15 @@ namespace {
struct offset_t {
int stride;
int begin;
__device__ int operator[](int i) {
__device__ int operator[](int i) const {
return stride * (begin + i);
}
#if CCCL_VERSION >= 3001000
__device__ offset_t& operator+=(int i) {
begin += i;
return *this;
}
#endif
};
// Segmented sort by full sort algorithm:.
// Say we are sorting a (2, 3) tensor. We have in flattened form: