Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "1702c85beca5e69fe068db94e955319d668ae417"
Commit 0799bc08 authored by limm's avatar limm
Browse files

suport v2.1.0

parent 50e05e1e
Hide whitespace changes
Inline Side-by-side
Showing with 347 additions and 78 deletions
csrc/cuda/index_info.cuh 0 → 100644
View file @ 0799bc08
#pragma once
#include <ATen/cuda/detail/TensorInfo.cuh>
// We need our own `IndexToOffset` implementation since we do not want to
// access the last element of the `indexptr`.
template <typename scalar_t> struct IndexPtrToOffset {
static inline __host__ __device__ int
get(int idx, const at::cuda::detail::TensorInfo<scalar_t, int> &info) {
int offset = idx % (info.sizes[info.dims - 1] - 1);
offset *= info.strides[info.dims - 1];
idx /= info.sizes[info.dims - 1] - 1;
for (int i = info.dims - 2; i >= 0; --i) {
offset += (idx % info.sizes[i]) * info.strides[i];
idx /= info.sizes[i];
}
return offset;
}
};
csrc/cuda/reducer.cuh 0 → 100644
View file @ 0799bc08
#pragma once
#include <limits>
#include <map>
#include "atomics.cuh"
enum ReductionType { SUM, MEAN, MUL, DIV, MIN, MAX };
const std::map<std::string, ReductionType> reduce2REDUCE = {
{"sum", SUM}, {"mean", MEAN}, {"mul", MUL},
{"div", DIV}, {"min", MIN}, {"max", MAX},
};
#define AT_DISPATCH_REDUCTION_TYPES(reduce, ...) \
[&] { \
switch (reduce2REDUCE.at(reduce)) { \
case SUM: { \
static constexpr ReductionType REDUCE = SUM; \
return __VA_ARGS__(); \
} \
case MEAN: { \
static constexpr ReductionType REDUCE = MEAN; \
return __VA_ARGS__(); \
} \
case MUL: { \
static constexpr ReductionType REDUCE = MUL; \
return __VA_ARGS__(); \
} \
case DIV: { \
static constexpr ReductionType REDUCE = DIV; \
return __VA_ARGS__(); \
} \
case MIN: { \
static constexpr ReductionType REDUCE = MIN; \
return __VA_ARGS__(); \
} \
case MAX: { \
static constexpr ReductionType REDUCE = MAX; \
return __VA_ARGS__(); \
} \
} \
}()
template <typename scalar_t, ReductionType REDUCE> struct Reducer {
static inline __host__ __device__ scalar_t init() {
if (REDUCE == MUL || REDUCE == DIV)
return (scalar_t)1;
else if (REDUCE == MIN)
return std::numeric_limits<scalar_t>::max();
else if (REDUCE == MAX)
return std::numeric_limits<scalar_t>::lowest();
else
return (scalar_t)0;
}
static inline __host__ __device__ void update(scalar_t *val,
scalar_t new_val) {
if (REDUCE == SUM || REDUCE == MEAN)
*val = *val + new_val;
else if (REDUCE == MUL)
*val = *val * new_val;
else if (REDUCE == DIV)
*val = *val / new_val;
else if ((REDUCE == MIN && new_val < *val) ||
(REDUCE == MAX && new_val > *val)) {
*val = new_val;
}
}
static inline __host__ __device__ void update(scalar_t *val, scalar_t new_val,
int64_t *arg, int64_t new_arg) {
if (REDUCE == SUM || REDUCE == MEAN)
*val = *val + new_val;
else if (REDUCE == MUL)
*val = *val * new_val;
else if (REDUCE == DIV)
*val = *val / new_val;
else if ((REDUCE == MIN && new_val < *val) ||
(REDUCE == MAX && new_val > *val)) {
*val = new_val;
*arg = new_arg;
}
}
static inline __host__ __device__ void write(scalar_t *address, scalar_t val,
int64_t *arg_address,
int64_t arg, int count) {
if (REDUCE == SUM || REDUCE == MUL || REDUCE == DIV)
*address = val;
else if (REDUCE == MEAN)
*address = val / (scalar_t)(count > 0 ? count : 1);
else if (REDUCE == MIN || REDUCE == MAX) {
if (count > 0) {
*address = val;
*arg_address = arg;
} else
*address = (scalar_t)0;
}
}
static inline __device__ void atomic_write(scalar_t *address, scalar_t val) {
if (REDUCE == SUM || REDUCE == MEAN)
atomAdd(address, val);
else if (REDUCE == MUL)
atomMul(address, val);
else if (REDUCE == DIV)
atomDiv(address, val);
else if (REDUCE == MIN)
atomMin(address, val);
else if (REDUCE == MAX)
atomMax(address, val);
}
};
csrc/hip/scatter_hip.hip → csrc/cuda/scatter_cuda.cu
View file @ 0799bc08
#include "hip/hip_runtime.h" #include "scatter_cuda.h"
#include "scatter_hip.h"
#include <ATen/hip/HIPContext.h> #include <ATen/cuda/CUDAContext.h>
#include <ATen/hip/detail/IndexUtils.cuh> #include <ATen/cuda/detail/IndexUtils.cuh>
#include <ATen/hip/detail/TensorInfo.cuh> #include <ATen/cuda/detail/TensorInfo.cuh>
#include "reducer.cuh" #include "reducer.cuh"
#include "utils.cuh" #include "utils.cuh"
#define THREADS 1024 #define THREADS 256
#define BLOCKS(N) (N + THREADS - 1) / THREADS #define BLOCKS(N) (N + THREADS - 1) / THREADS
template <typename scalar_t, ReductionType REDUCE> template <typename scalar_t, ReductionType REDUCE>
...@@ -64,7 +63,7 @@ scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim, ...@@ -64,7 +63,7 @@ scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim,
CHECK_CUDA(index); CHECK_CUDA(index);
if (optional_out.has_value()) if (optional_out.has_value())
CHECK_CUDA(optional_out.value()); CHECK_CUDA(optional_out.value());
hipSetDevice(src.get_device()); cudaSetDevice(src.get_device());
CHECK_INPUT(src.dim() == index.dim()); CHECK_INPUT(src.dim() == index.dim());
for (auto i = 0; i < index.dim() - 1; i++) for (auto i = 0; i < index.dim() - 1; i++)
...@@ -112,7 +111,7 @@ scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim, ...@@ -112,7 +111,7 @@ scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim,
auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index); auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
auto stream = at::cuda::getCurrentCUDAStream(); auto stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
......
csrc/hip/scatter_hip.h → csrc/cuda/scatter_cuda.h
View file @ 0799bc08
#pragma once #pragma once
#include <torch/extension.h> #include "../extensions.h"
std::tuple<torch::Tensor, torch::optional<torch::Tensor>> std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim, scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim,
......
csrc/hip/segment_coo_hip.hip → csrc/cuda/segment_coo_cuda.cu
View file @ 0799bc08
#include "hip/hip_runtime.h" #include "segment_coo_cuda.h"
#include "segment_coo_hip.h"
#include <ATen/hip/HIPContext.h> #include <ATen/cuda/CUDAContext.h>
#include <ATen/hip/detail/IndexUtils.cuh> #include <ATen/cuda/detail/IndexUtils.cuh>
#include <ATen/hip/detail/TensorInfo.cuh> #include <ATen/cuda/detail/TensorInfo.cuh>
#include "reducer.cuh" #include "reducer.cuh"
#include "utils.cuh" #include "utils.cuh"
...@@ -37,8 +36,8 @@ segment_coo_kernel(const scalar_t *src_data, ...@@ -37,8 +36,8 @@ segment_coo_kernel(const scalar_t *src_data,
#pragma unroll #pragma unroll
for (int i = 1; i < 32; i *= 2) { for (int i = 1; i < 32; i *= 2) {
// Parallel reduction inside a single warp. // Parallel reduction inside a single warp.
tmp = __shfl_up_sync(FULL_MASK, val, i); tmp = SHFL_UP_SYNC(FULL_MASK, val, i);
next_idx = __shfl_up_sync(FULL_MASK, idx, i); next_idx = SHFL_UP_SYNC(FULL_MASK, idx, i);
if (lane_idx >= i && row_idx / D == (row_idx - i) / D) { if (lane_idx >= i && row_idx / D == (row_idx - i) / D) {
assert(idx >= next_idx); assert(idx >= next_idx);
if (idx == next_idx) if (idx == next_idx)
...@@ -46,7 +45,7 @@ segment_coo_kernel(const scalar_t *src_data, ...@@ -46,7 +45,7 @@ segment_coo_kernel(const scalar_t *src_data,
} }
} }
next_idx = __shfl_down_sync(FULL_MASK, idx, 1); next_idx = SHFL_DOWN_SYNC(FULL_MASK, idx, 1);
if (lane_idx == 32 - 1 || row_idx / D != (row_idx + 1) / D || if (lane_idx == 32 - 1 || row_idx / D != (row_idx + 1) / D ||
idx != next_idx) idx != next_idx)
Reducer<scalar_t, REDUCE>::atomic_write(out_data + out_idx, val); Reducer<scalar_t, REDUCE>::atomic_write(out_data + out_idx, val);
...@@ -158,7 +157,7 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index, ...@@ -158,7 +157,7 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index,
CHECK_CUDA(index); CHECK_CUDA(index);
if (optional_out.has_value()) if (optional_out.has_value())
CHECK_CUDA(optional_out.value()); CHECK_CUDA(optional_out.value());
hipSetDevice(src.get_device()); cudaSetDevice(src.get_device());
CHECK_INPUT(src.dim() >= index.dim()); CHECK_INPUT(src.dim() >= index.dim());
...@@ -215,7 +214,7 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index, ...@@ -215,7 +214,7 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index,
auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index); auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
auto stream = at::cuda::getCurrentCUDAStream(); auto stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
...@@ -331,7 +330,7 @@ torch::Tensor gather_coo_cuda(torch::Tensor src, torch::Tensor index, ...@@ -331,7 +330,7 @@ torch::Tensor gather_coo_cuda(torch::Tensor src, torch::Tensor index,
CHECK_CUDA(index); CHECK_CUDA(index);
if (optional_out.has_value()) if (optional_out.has_value())
CHECK_CUDA(optional_out.value()); CHECK_CUDA(optional_out.value());
hipSetDevice(src.get_device()); cudaSetDevice(src.get_device());
CHECK_INPUT(src.dim() >= index.dim()); CHECK_INPUT(src.dim() >= index.dim());
...@@ -366,7 +365,7 @@ torch::Tensor gather_coo_cuda(torch::Tensor src, torch::Tensor index, ...@@ -366,7 +365,7 @@ torch::Tensor gather_coo_cuda(torch::Tensor src, torch::Tensor index,
auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index); auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
auto stream = at::cuda::getCurrentCUDAStream(); auto stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
......
csrc/hip/segment_coo_hip.h → csrc/cuda/segment_coo_cuda.h
View file @ 0799bc08
#pragma once #pragma once
#include <torch/extension.h> #include "../extensions.h"
std::tuple<torch::Tensor, torch::optional<torch::Tensor>> std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
segment_coo_cuda(torch::Tensor src, torch::Tensor index, segment_coo_cuda(torch::Tensor src, torch::Tensor index,
...@@ -9,7 +9,3 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index, ...@@ -9,7 +9,3 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index,
torch::Tensor gather_coo_cuda(torch::Tensor src, torch::Tensor index, torch::Tensor gather_coo_cuda(torch::Tensor src, torch::Tensor index,
torch::optional<torch::Tensor> optional_out); torch::optional<torch::Tensor> optional_out);
template<typename T>
__device__ T __ldg(const T* ptr) {
return *ptr;
}
csrc/hip/segment_csr_hip.hip → csrc/cuda/segment_csr_cuda.cu
View file @ 0799bc08
#include "hip/hip_runtime.h" #include "segment_csr_cuda.h"
#include "segment_csr_hip.h"
#include <ATen/hip/HIPContext.h> #include <ATen/cuda/CUDAContext.h>
#include <ATen/hip/detail/IndexUtils.cuh> #include <ATen/cuda/detail/IndexUtils.cuh>
#include <ATen/hip/detail/TensorInfo.cuh> #include <ATen/cuda/detail/TensorInfo.cuh>
#include "index_info.cuh" #include "index_info.cuh"
#include "reducer.cuh" #include "reducer.cuh"
...@@ -47,9 +46,9 @@ segment_csr_kernel(const scalar_t *src_data, ...@@ -47,9 +46,9 @@ segment_csr_kernel(const scalar_t *src_data,
for (int i = TB / 2; i > 0; i /= 2) { for (int i = TB / 2; i > 0; i /= 2) {
// Parallel reduction inside a single warp. // Parallel reduction inside a single warp.
if (REDUCE == MIN || REDUCE == MAX) if (REDUCE == MIN || REDUCE == MAX)
arg_tmp = __shfl_down_sync(FULL_MASK, arg, i); arg_tmp = SHFL_DOWN_SYNC(FULL_MASK, arg, i);
Reducer<scalar_t, REDUCE>::update( Reducer<scalar_t, REDUCE>::update(
&val, __shfl_down_sync(FULL_MASK, val, i), &arg, arg_tmp); &val, SHFL_DOWN_SYNC(FULL_MASK, val, i), &arg, arg_tmp);
} }
if (lane_idx == 0) { if (lane_idx == 0) {
...@@ -103,7 +102,7 @@ segment_csr_cuda(torch::Tensor src, torch::Tensor indptr, ...@@ -103,7 +102,7 @@ segment_csr_cuda(torch::Tensor src, torch::Tensor indptr,
CHECK_CUDA(indptr); CHECK_CUDA(indptr);
if (optional_out.has_value()) if (optional_out.has_value())
CHECK_CUDA(optional_out.value()); CHECK_CUDA(optional_out.value());
hipSetDevice(src.get_device()); cudaSetDevice(src.get_device());
CHECK_INPUT(src.dim() >= indptr.dim()); CHECK_INPUT(src.dim() >= indptr.dim());
...@@ -148,7 +147,7 @@ segment_csr_cuda(torch::Tensor src, torch::Tensor indptr, ...@@ -148,7 +147,7 @@ segment_csr_cuda(torch::Tensor src, torch::Tensor indptr,
auto indptr_info = at::cuda::detail::getTensorInfo<int64_t, int>(indptr); auto indptr_info = at::cuda::detail::getTensorInfo<int64_t, int>(indptr);
auto stream = at::cuda::getCurrentCUDAStream(); auto stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
...@@ -223,7 +222,7 @@ torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr, ...@@ -223,7 +222,7 @@ torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr,
CHECK_CUDA(indptr); CHECK_CUDA(indptr);
if (optional_out.has_value()) if (optional_out.has_value())
CHECK_CUDA(optional_out.value()); CHECK_CUDA(optional_out.value());
hipSetDevice(src.get_device()); cudaSetDevice(src.get_device());
CHECK_INPUT(src.dim() >= indptr.dim()); CHECK_INPUT(src.dim() >= indptr.dim());
...@@ -265,7 +264,7 @@ torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr, ...@@ -265,7 +264,7 @@ torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr,
auto indptr_info = at::cuda::detail::getTensorInfo<int64_t, int>(indptr); auto indptr_info = at::cuda::detail::getTensorInfo<int64_t, int>(indptr);
auto stream = at::cuda::getCurrentCUDAStream(); auto stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
......
csrc/hip/segment_csr_hip.h → csrc/cuda/segment_csr_cuda.h
View file @ 0799bc08
#pragma once #pragma once
#include <torch/extension.h> #include "../extensions.h"
std::tuple<torch::Tensor, torch::optional<torch::Tensor>> std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
segment_csr_cuda(torch::Tensor src, torch::Tensor indptr, segment_csr_cuda(torch::Tensor src, torch::Tensor indptr,
...@@ -9,7 +9,3 @@ segment_csr_cuda(torch::Tensor src, torch::Tensor indptr, ...@@ -9,7 +9,3 @@ segment_csr_cuda(torch::Tensor src, torch::Tensor indptr,
torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr, torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr,
torch::optional<torch::Tensor> optional_out); torch::optional<torch::Tensor> optional_out);
template<typename T>
__device__ T __ldg(const T* ptr) {
return *ptr;
}
csrc/cuda/utils.cuh 0 → 100644
View file @ 0799bc08
#pragma once
#include "../extensions.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")
__device__ __inline__ at::Half __shfl_up_sync(const unsigned mask,
const at::Half var,
const unsigned int delta) {
return __shfl_up_sync(mask, var.operator __half(), delta);
}
__device__ __inline__ at::Half __shfl_down_sync(const unsigned mask,
const at::Half var,
const unsigned int delta) {
return __shfl_down_sync(mask, var.operator __half(), delta);
}
#ifdef USE_ROCM
__device__ __inline__ at::Half __shfl_up(const at::Half var, const unsigned int delta) {
return __shfl_up(var.operator __half(), delta);
}
__device__ __inline__ at::Half __shfl_down(const at::Half var, const unsigned int delta) {
return __shfl_down(var.operator __half(), delta);
}
#endif
#ifdef USE_ROCM
__device__ __inline__ at::Half __ldg(const at::Half* ptr) {
return __ldg(reinterpret_cast<const __half*>(ptr));
}
#define SHFL_UP_SYNC(mask, var, delta) __shfl_up(var, delta)
#define SHFL_DOWN_SYNC(mask, var, delta) __shfl_down(var, delta)
#else
#define SHFL_UP_SYNC __shfl_up_sync
#define SHFL_DOWN_SYNC __shfl_down_sync
#endif
csrc/extensions.h 0 → 100644
View file @ 0799bc08
#include "macros.h"
#include <torch/torch.h>
csrc/hip/atomics.cuh
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#pragma once #pragma once
#define ATOMIC(NAME) \ #define ATOMIC(NAME) \
...@@ -68,8 +70,8 @@ ...@@ -68,8 +70,8 @@
\ \
template <typename scalar, size_t size> struct Atomic##NAME##DecimalImpl; \ template <typename scalar, size_t size> struct Atomic##NAME##DecimalImpl; \
\ \
template <typename scalar> struct Atomic##NAME##DecimalImpl<scalar, 2> { \ template <> struct Atomic##NAME##DecimalImpl<at::Half, 2> { \
inline __device__ void operator()(scalar *address, scalar val) { \ inline __device__ void operator()(at::Half *address, at::Half val) { \
unsigned int *address_as_ui = \ unsigned int *address_as_ui = \
(unsigned int *)((char *)address - ((size_t)address & 2)); \ (unsigned int *)((char *)address - ((size_t)address & 2)); \
unsigned int old = *address_as_ui; \ unsigned int old = *address_as_ui; \
...@@ -87,6 +89,25 @@ ...@@ -87,6 +89,25 @@
} \ } \
}; \ }; \
\ \
template <> struct Atomic##NAME##DecimalImpl<at::BFloat16, 2> { \
inline __device__ void operator()(at::BFloat16 *address, at::BFloat16 val){\
unsigned int *address_as_ui = \
(unsigned int *)((char *)address - ((size_t)address & 2)); \
unsigned int old = *address_as_ui; \
unsigned int assumed; \
\
do { \
assumed = old; \
at::BFloat16 hsum; \
hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff); \
hsum = OP(hsum, val); \
old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16) \
: (old & 0xffff0000) | hsum.x; \
old = atomicCAS(address_as_ui, assumed, old); \
} while (assumed != old); \
} \
}; \
\
template <typename scalar> struct Atomic##NAME##DecimalImpl<scalar, 4> { \ template <typename scalar> struct Atomic##NAME##DecimalImpl<scalar, 4> { \
inline __device__ void operator()(scalar *address, scalar val) { \ inline __device__ void operator()(scalar *address, scalar val) { \
int *address_as_i = (int *)address; \ int *address_as_i = (int *)address; \
...@@ -135,19 +156,19 @@ static inline __device__ void atomAdd(int32_t *address, int32_t val) { ...@@ -135,19 +156,19 @@ static inline __device__ void atomAdd(int32_t *address, int32_t val) {
static inline __device__ void atomAdd(int64_t *address, int64_t val) { static inline __device__ void atomAdd(int64_t *address, int64_t val) {
AtomicAddIntegerImpl<int64_t, sizeof(int64_t)>()(address, val); AtomicAddIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
} }
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700 || TORCH_HIP_VERSION < 10000) #if defined(USE_ROCM) || (defined(__DTK_ARCH__) && (__DTK_ARCH__ < 700 || DTK_VERSION < 10000))
static inline __device__ void atomAdd(at::Half *address, at::Half val) { static inline __device__ void atomAdd(at::Half *address, at::Half val) {
AtomicAddDecimalImpl<at::Half, sizeof(at::Half)>()(address, val); AtomicAddDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
} }
#else #else
static inline __device__ void atomAdd(at::Half *address, at::Half val) { static inline __device__ void atomAdd(at::Half *address, at::Half val) {
AtomicAddDecimalImpl<at::Half, sizeof(at::Half)>()(address, val); atomicAdd(reinterpret_cast<__half *>(address), val);
} }
#endif #endif
static inline __device__ void atomAdd(float *address, float val) { static inline __device__ void atomAdd(float *address, float val) {
atomicAdd(address, val); atomicAdd(address, val);
} }
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600 || TORCH_HIP_VERSION < 8000) #if defined(__DTK_ARCH__) && (__DTK_ARCH__ < 600 || DTK_VERSION < 8000)
static inline __device__ void atomAdd(double *address, double val) { static inline __device__ void atomAdd(double *address, double val) {
AtomicAddDecimalImpl<double, sizeof(double)>()(address, val); AtomicAddDecimalImpl<double, sizeof(double)>()(address, val);
} }
...@@ -156,6 +177,9 @@ static inline __device__ void atomAdd(double *address, double val) { ...@@ -156,6 +177,9 @@ static inline __device__ void atomAdd(double *address, double val) {
atomicAdd(address, val); atomicAdd(address, val);
} }
#endif #endif
static inline __device__ void atomAdd(at::BFloat16 *address, at::BFloat16 val) {
AtomicAddDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
}
#define OP(X, Y) Y *X #define OP(X, Y) Y *X
ATOMIC(Mul) ATOMIC(Mul)
...@@ -184,6 +208,9 @@ static inline __device__ void atomMul(at::Half *address, at::Half val) { ...@@ -184,6 +208,9 @@ static inline __device__ void atomMul(at::Half *address, at::Half val) {
static inline __device__ void atomMul(double *address, double val) { static inline __device__ void atomMul(double *address, double val) {
AtomicMulDecimalImpl<double, sizeof(double)>()(address, val); AtomicMulDecimalImpl<double, sizeof(double)>()(address, val);
} }
static inline __device__ void atomMul(at::BFloat16 *address, at::BFloat16 val) {
AtomicMulDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
}
#define OP(X, Y) Y / X #define OP(X, Y) Y / X
ATOMIC(Div) ATOMIC(Div)
...@@ -212,6 +239,9 @@ static inline __device__ void atomDiv(float *address, float val) { ...@@ -212,6 +239,9 @@ static inline __device__ void atomDiv(float *address, float val) {
static inline __device__ void atomDiv(double *address, double val) { static inline __device__ void atomDiv(double *address, double val) {
AtomicDivDecimalImpl<double, sizeof(double)>()(address, val); AtomicDivDecimalImpl<double, sizeof(double)>()(address, val);
} }
static inline __device__ void atomDiv(at::BFloat16 *address, at::BFloat16 val) {
AtomicDivDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
}
#define OP(X, Y) max(Y, X) #define OP(X, Y) max(Y, X)
ATOMIC(Max) ATOMIC(Max)
...@@ -240,6 +270,9 @@ static inline __device__ void atomMax(float *address, float val) { ...@@ -240,6 +270,9 @@ static inline __device__ void atomMax(float *address, float val) {
static inline __device__ void atomMax(double *address, double val) { static inline __device__ void atomMax(double *address, double val) {
AtomicMaxDecimalImpl<double, sizeof(double)>()(address, val); AtomicMaxDecimalImpl<double, sizeof(double)>()(address, val);
} }
static inline __device__ void atomMax(at::BFloat16 *address, at::BFloat16 val) {
AtomicMaxDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
}
#define OP(X, Y) min(Y, X) #define OP(X, Y) min(Y, X)
ATOMIC(Min) ATOMIC(Min)
...@@ -268,3 +301,6 @@ static inline __device__ void atomMin(float *address, float val) { ...@@ -268,3 +301,6 @@ static inline __device__ void atomMin(float *address, float val) {
static inline __device__ void atomMin(double *address, double val) { static inline __device__ void atomMin(double *address, double val) {
AtomicMinDecimalImpl<double, sizeof(double)>()(address, val); AtomicMinDecimalImpl<double, sizeof(double)>()(address, val);
} }
static inline __device__ void atomMin(at::BFloat16 *address, at::BFloat16 val) {
AtomicMinDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
}
csrc/hip/index_info.cuh
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#pragma once #pragma once
#include <ATen/hip/detail/TensorInfo.cuh> #include <ATen/hip/detail/TensorInfo.cuh>
......
csrc/hip/reducer.cuh
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#pragma once #pragma once
#include <limits> #include <limits>
#include <map> #include <map>
#include "atomics.cuh" #include "../hip/atomics.cuh"
enum ReductionType { SUM, MEAN, MUL, DIV, MIN, MAX }; enum ReductionType { SUM, MEAN, MUL, DIV, MIN, MAX };
...@@ -16,27 +18,27 @@ const std::map<std::string, ReductionType> reduce2REDUCE = { ...@@ -16,27 +18,27 @@ const std::map<std::string, ReductionType> reduce2REDUCE = {
[&] { \ [&] { \
switch (reduce2REDUCE.at(reduce)) { \ switch (reduce2REDUCE.at(reduce)) { \
case SUM: { \ case SUM: { \
const ReductionType REDUCE = SUM; \ static constexpr ReductionType REDUCE = SUM; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
case MEAN: { \ case MEAN: { \
const ReductionType REDUCE = MEAN; \ static constexpr ReductionType REDUCE = MEAN; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
case MUL: { \ case MUL: { \
const ReductionType REDUCE = MUL; \ static constexpr ReductionType REDUCE = MUL; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
case DIV: { \ case DIV: { \
const ReductionType REDUCE = DIV; \ static constexpr ReductionType REDUCE = DIV; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
case MIN: { \ case MIN: { \
const ReductionType REDUCE = MIN; \ static constexpr ReductionType REDUCE = MIN; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
case MAX: { \ case MAX: { \
const ReductionType REDUCE = MAX; \ static constexpr ReductionType REDUCE = MAX; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
} \ } \
......
csrc/hip/scatter_cuda.h 0 → 100644
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#pragma once
#include "../extensions.h"
std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim,
torch::optional<torch::Tensor> optional_out,
torch::optional<int64_t> dim_size, std::string reduce);
csrc/hip/scatter_hip_hip.hip → csrc/hip/scatter_cuda.hip
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#include "hip/hip_runtime.h" #include "hip/hip_runtime.h"
#include "scatter_hip.h" #include "../hip/scatter_cuda.h"
#include <ATen/hip/HIPContext.h> #include <ATen/hip/HIPContext.h>
#include <ATen/hip/detail/IndexUtils.cuh> #include <ATen/hip/detail/IndexUtils.cuh>
#include <ATen/hip/detail/TensorInfo.cuh> #include <ATen/hip/detail/TensorInfo.cuh>
#include "reducer.cuh" #include "../hip/reducer.cuh"
#include "utils.cuh" #include "../hip/utils.cuh"
#define THREADS 1024 #define THREADS 256
#define BLOCKS(N) (N + THREADS - 1) / THREADS #define BLOCKS(N) (N + THREADS - 1) / THREADS
template <typename scalar_t, ReductionType REDUCE> template <typename scalar_t, ReductionType REDUCE>
...@@ -112,7 +114,7 @@ scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim, ...@@ -112,7 +114,7 @@ scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim,
auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index); auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA(); auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
......
csrc/hip/segment_coo_cuda.h 0 → 100644
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#pragma once
#include "../extensions.h"
std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
segment_coo_cuda(torch::Tensor src, torch::Tensor index,
torch::optional<torch::Tensor> optional_out,
torch::optional<int64_t> dim_size, std::string reduce);
torch::Tensor gather_coo_cuda(torch::Tensor src, torch::Tensor index,
torch::optional<torch::Tensor> optional_out);
csrc/hip/segment_coo_hip_hip.hip → csrc/hip/segment_coo_cuda.hip
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#include "hip/hip_runtime.h" #include "hip/hip_runtime.h"
#include "segment_coo_hip.h" #include "../hip/segment_coo_cuda.h"
#include <ATen/hip/HIPContext.h> #include <ATen/hip/HIPContext.h>
#include <ATen/hip/detail/IndexUtils.cuh> #include <ATen/hip/detail/IndexUtils.cuh>
#include <ATen/hip/detail/TensorInfo.cuh> #include <ATen/hip/detail/TensorInfo.cuh>
#include "reducer.cuh" #include "../hip/reducer.cuh"
#include "utils.cuh" #include "../hip/utils.cuh"
#define THREADS 256 #define THREADS 256
#define BLOCKS(TB, N) (TB * N + THREADS - 1) / THREADS #define BLOCKS(TB, N) (TB * N + THREADS - 1) / THREADS
...@@ -37,8 +39,8 @@ segment_coo_kernel(const scalar_t *src_data, ...@@ -37,8 +39,8 @@ segment_coo_kernel(const scalar_t *src_data,
#pragma unroll #pragma unroll
for (int i = 1; i < 32; i *= 2) { for (int i = 1; i < 32; i *= 2) {
// Parallel reduction inside a single warp. // Parallel reduction inside a single warp.
tmp = __shfl_up_sync(FULL_MASK, val, i); tmp = SHFL_UP_SYNC(FULL_MASK, val, i);
next_idx = __shfl_up_sync(FULL_MASK, idx, i); next_idx = SHFL_UP_SYNC(FULL_MASK, idx, i);
if (lane_idx >= i && row_idx / D == (row_idx - i) / D) { if (lane_idx >= i && row_idx / D == (row_idx - i) / D) {
assert(idx >= next_idx); assert(idx >= next_idx);
if (idx == next_idx) if (idx == next_idx)
...@@ -46,7 +48,7 @@ segment_coo_kernel(const scalar_t *src_data, ...@@ -46,7 +48,7 @@ segment_coo_kernel(const scalar_t *src_data,
} }
} }
next_idx = __shfl_down_sync(FULL_MASK, idx, 1); next_idx = SHFL_DOWN_SYNC(FULL_MASK, idx, 1);
if (lane_idx == 32 - 1 || row_idx / D != (row_idx + 1) / D || if (lane_idx == 32 - 1 || row_idx / D != (row_idx + 1) / D ||
idx != next_idx) idx != next_idx)
Reducer<scalar_t, REDUCE>::atomic_write(out_data + out_idx, val); Reducer<scalar_t, REDUCE>::atomic_write(out_data + out_idx, val);
...@@ -215,7 +217,7 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index, ...@@ -215,7 +217,7 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index,
auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index); auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA(); auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
...@@ -366,7 +368,7 @@ torch::Tensor gather_coo_cuda(torch::Tensor src, torch::Tensor index, ...@@ -366,7 +368,7 @@ torch::Tensor gather_coo_cuda(torch::Tensor src, torch::Tensor index,
auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index); auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA(); auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
......
csrc/hip/segment_csr_cuda.h 0 → 100644
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#pragma once
#include "../extensions.h"
std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
segment_csr_cuda(torch::Tensor src, torch::Tensor indptr,
torch::optional<torch::Tensor> optional_out,
std::string reduce);
torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr,
torch::optional<torch::Tensor> optional_out);
csrc/hip/segment_csr_hip_hip.hip → csrc/hip/segment_csr_cuda.hip
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#include "hip/hip_runtime.h" #include "hip/hip_runtime.h"
#include "segment_csr_hip.h" #include "../hip/segment_csr_cuda.h"
#include <ATen/hip/HIPContext.h> #include <ATen/hip/HIPContext.h>
#include <ATen/hip/detail/IndexUtils.cuh> #include <ATen/hip/detail/IndexUtils.cuh>
#include <ATen/hip/detail/TensorInfo.cuh> #include <ATen/hip/detail/TensorInfo.cuh>
#include "index_info.cuh" #include "../hip/index_info.cuh"
#include "reducer.cuh" #include "../hip/reducer.cuh"
#include "utils.cuh" #include "../hip/utils.cuh"
#define THREADS 256 #define THREADS 256
#define BLOCKS(TB, N) (TB * N + THREADS - 1) / THREADS #define BLOCKS(TB, N) (TB * N + THREADS - 1) / THREADS
...@@ -47,9 +49,9 @@ segment_csr_kernel(const scalar_t *src_data, ...@@ -47,9 +49,9 @@ segment_csr_kernel(const scalar_t *src_data,
for (int i = TB / 2; i > 0; i /= 2) { for (int i = TB / 2; i > 0; i /= 2) {
// Parallel reduction inside a single warp. // Parallel reduction inside a single warp.
if (REDUCE == MIN || REDUCE == MAX) if (REDUCE == MIN || REDUCE == MAX)
arg_tmp = __shfl_down_sync(FULL_MASK, arg, i); arg_tmp = SHFL_DOWN_SYNC(FULL_MASK, arg, i);
Reducer<scalar_t, REDUCE>::update( Reducer<scalar_t, REDUCE>::update(
&val, __shfl_down_sync(FULL_MASK, val, i), &arg, arg_tmp); &val, SHFL_DOWN_SYNC(FULL_MASK, val, i), &arg, arg_tmp);
} }
if (lane_idx == 0) { if (lane_idx == 0) {
...@@ -148,7 +150,7 @@ segment_csr_cuda(torch::Tensor src, torch::Tensor indptr, ...@@ -148,7 +150,7 @@ segment_csr_cuda(torch::Tensor src, torch::Tensor indptr,
auto indptr_info = at::cuda::detail::getTensorInfo<int64_t, int>(indptr); auto indptr_info = at::cuda::detail::getTensorInfo<int64_t, int>(indptr);
auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA(); auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
...@@ -265,7 +267,7 @@ torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr, ...@@ -265,7 +267,7 @@ torch::Tensor gather_csr_cuda(torch::Tensor src, torch::Tensor indptr,
auto indptr_info = at::cuda::detail::getTensorInfo<int64_t, int>(indptr); auto indptr_info = at::cuda::detail::getTensorInfo<int64_t, int>(indptr);
auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA(); auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, src.scalar_type(), "_", [&] { AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
auto src_data = src.data_ptr<scalar_t>(); auto src_data = src.data_ptr<scalar_t>();
auto out_data = out.data_ptr<scalar_t>(); auto out_data = out.data_ptr<scalar_t>();
......
csrc/hip/utils.cuh
View file @ 0799bc08
// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#pragma once #pragma once
#include <torch/extension.h> #include "../extensions.h"
#define CHECK_CUDA(x) \ #define CHECK_CUDA(x) \
AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor") AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")
...@@ -9,11 +11,32 @@ ...@@ -9,11 +11,32 @@
__device__ __inline__ at::Half __shfl_up_sync(const unsigned mask, __device__ __inline__ at::Half __shfl_up_sync(const unsigned mask,
const at::Half var, const at::Half var,
const unsigned int delta) { const unsigned int delta) {
return __shfl_up_sync(mask, (__half)var, delta); return __shfl_up_sync(mask, var.operator __half(), delta);
} }
__device__ __inline__ at::Half __shfl_down_sync(const unsigned mask, __device__ __inline__ at::Half __shfl_down_sync(const unsigned mask,
const at::Half var, const at::Half var,
const unsigned int delta) { const unsigned int delta) {
return __shfl_down_sync(mask, (__half)var, delta); return __shfl_down_sync(mask, var.operator __half(), delta);
} }
__device__ __inline__ at::Half __shfl_up(const at::Half var,
const unsigned int delta) {
return __shfl_up(var.operator __half(), delta);
}
__device__ __inline__ at::Half __shfl_down(const at::Half var,
const unsigned int delta) {
return __shfl_down(var.operator __half(), delta);
}
#ifdef USE_ROCM
__device__ __inline__ at::Half __ldg(const at::Half* ptr) {
return __ldg(reinterpret_cast<const __half*>(ptr));
}
#define SHFL_UP_SYNC(mask, var, delta) __shfl_up(var, delta)
#define SHFL_DOWN_SYNC(mask, var, delta) __shfl_down(var, delta)
#else
#define SHFL_UP_SYNC __shfl_up_sync
#define SHFL_DOWN_SYNC __shfl_down_sync
#endif
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment