".github/vscode:/vscode.git/clone" did not exist on "db75402ade445ae32e7d4de1884149bbb8784f3a"
utils.cuh 669 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
6
7
#pragma once

#include <torch/extension.h>

#define CHECK_CUDA(x)                                                          \
  AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")
#define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch")
rusty1s's avatar
rusty1s committed
8

rusty1s's avatar
typo  
rusty1s committed
9
10
__device__ __inline__ at::Half
__shfl_sync(const unsigned mask, const at::Half var, const int srcLane) {
rusty1s's avatar
rusty1s committed
11
12
13
14
15
16
17
18
  return __shfl_sync(mask, (__half)var, srcLane);
}

__device__ __inline__ at::Half __shfl_down_sync(const unsigned mask,
                                                const at::Half var,
                                                const unsigned int delta) {
  return __shfl_down_sync(mask, (__half)var, delta);
}