Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+300
@@ -0,0 +1,300 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#define ATOMIC(NAME) \
|
||||
template <typename scalar, size_t size> \
|
||||
struct Atomic##NAME##IntegerImpl; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##IntegerImpl<scalar, 1> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
uint32_t *address_as_ui = (uint32_t *)(address - ((size_t)address & 3)); \
|
||||
uint32_t old = *address_as_ui; \
|
||||
uint32_t shift = ((size_t)address & 3) * 8; \
|
||||
uint32_t sum; \
|
||||
uint32_t assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
sum = OP(val, scalar((old >> shift) & 0xff)); \
|
||||
old = (old & ~(0x000000ff << shift)) | (sum << shift); \
|
||||
old = atomicCAS(address_as_ui, assumed, old); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##IntegerImpl<scalar, 2> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
uint32_t *address_as_ui = (uint32_t *)((char *)address - ((size_t)address & 2)); \
|
||||
uint32_t old = *address_as_ui; \
|
||||
uint32_t sum; \
|
||||
uint32_t newval; \
|
||||
uint32_t assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
sum = OP(val, (size_t)address & 2 ? scalar(old >> 16) : scalar(old & 0xffff)); \
|
||||
newval = (size_t)address & 2 ? (old & 0xffff) | (sum << 16) : (old & 0xffff0000) | sum; \
|
||||
old = atomicCAS(address_as_ui, assumed, newval); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##IntegerImpl<scalar, 4> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
uint32_t *address_as_ui = (uint32_t *)address; \
|
||||
uint32_t old = *address_as_ui; \
|
||||
uint32_t assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = atomicCAS(address_as_ui, assumed, OP(val, (scalar)old)); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##IntegerImpl<scalar, 8> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
unsigned long long *address_as_ull = (unsigned long long *)address; \
|
||||
unsigned long long old = *address_as_ull; \
|
||||
unsigned long long assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = atomicCAS(address_as_ull, assumed, OP(val, (scalar)old)); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar, size_t size> \
|
||||
struct Atomic##NAME##DecimalImpl; \
|
||||
\
|
||||
template <> \
|
||||
struct Atomic##NAME##DecimalImpl<at::Half, 2> { \
|
||||
inline __device__ void operator()(at::Half *address, at::Half 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::Half 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 <> \
|
||||
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> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
int *address_as_i = (int *)address; \
|
||||
int old = *address_as_i; \
|
||||
int assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = atomicCAS(address_as_i, assumed, __float_as_int(OP(val, __int_as_float(assumed)))); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
}; \
|
||||
\
|
||||
template <typename scalar> \
|
||||
struct Atomic##NAME##DecimalImpl<scalar, 8> { \
|
||||
inline __device__ void operator()(scalar *address, scalar val) { \
|
||||
unsigned long long int *address_as_ull = (unsigned long long int *)address; \
|
||||
unsigned long long int old = *address_as_ull; \
|
||||
unsigned long long int assumed; \
|
||||
\
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = \
|
||||
atomicCAS(address_as_ull, assumed, __double_as_longlong(OP(val, __longlong_as_double(assumed)))); \
|
||||
} while (assumed != old); \
|
||||
} \
|
||||
};
|
||||
|
||||
#define OP(X, Y) Y + X
|
||||
ATOMIC(Add)
|
||||
#undef OP
|
||||
static inline __device__ void atomAdd(uint8_t *address, uint8_t val) {
|
||||
AtomicAddIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomAdd(int8_t *address, int8_t val) {
|
||||
AtomicAddIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomAdd(int16_t *address, int16_t val) {
|
||||
AtomicAddIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomAdd(int32_t *address, int32_t val) { atomicAdd(address, val); }
|
||||
static inline __device__ void atomAdd(int64_t *address, int64_t val) {
|
||||
AtomicAddIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
#if defined(USE_ROCM) || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700 || CUDA_VERSION < 10000))
|
||||
static inline __device__ void atomAdd(at::Half *address, at::Half val) {
|
||||
AtomicAddDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
#else
|
||||
static inline __device__ void atomAdd(at::Half *address, at::Half val) {
|
||||
atomicAdd(reinterpret_cast<__half *>(address), val);
|
||||
}
|
||||
#endif
|
||||
static inline __device__ void atomAdd(float *address, float val) { atomicAdd(address, val); }
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600 || CUDA_VERSION < 8000)
|
||||
static inline __device__ void atomAdd(double *address, double val) {
|
||||
AtomicAddDecimalImpl<double, sizeof(double)>()(address, val);
|
||||
}
|
||||
#else
|
||||
static inline __device__ void atomAdd(double *address, double val) { atomicAdd(address, val); }
|
||||
#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
|
||||
ATOMIC(Mul)
|
||||
#undef OP
|
||||
static inline __device__ void atomMul(uint8_t *address, uint8_t val) {
|
||||
AtomicMulIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(int8_t *address, int8_t val) {
|
||||
AtomicMulIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(int16_t *address, int16_t val) {
|
||||
AtomicMulIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(int32_t *address, int32_t val) {
|
||||
AtomicMulIntegerImpl<int32_t, sizeof(int32_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(int64_t *address, int64_t val) {
|
||||
AtomicMulIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(float *address, float val) {
|
||||
AtomicMulDecimalImpl<float, sizeof(float)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(at::Half *address, at::Half val) {
|
||||
AtomicMulDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMul(double *address, double 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
|
||||
ATOMIC(Div)
|
||||
#undef OP
|
||||
static inline __device__ void atomDiv(uint8_t *address, uint8_t val) {
|
||||
AtomicDivIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(int8_t *address, int8_t val) {
|
||||
AtomicDivIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(int16_t *address, int16_t val) {
|
||||
AtomicDivIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(int32_t *address, int32_t val) {
|
||||
AtomicDivIntegerImpl<int32_t, sizeof(int32_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(int64_t *address, int64_t val) {
|
||||
AtomicDivIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(at::Half *address, at::Half val) {
|
||||
AtomicDivDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(float *address, float val) {
|
||||
AtomicDivDecimalImpl<float, sizeof(float)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomDiv(double *address, double 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)
|
||||
ATOMIC(Max)
|
||||
#undef OP
|
||||
static inline __device__ void atomMax(uint8_t *address, uint8_t val) {
|
||||
AtomicMaxIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(int8_t *address, int8_t val) {
|
||||
AtomicMaxIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(int16_t *address, int16_t val) {
|
||||
AtomicMaxIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(int32_t *address, int32_t val) { atomicMax(address, val); }
|
||||
static inline __device__ void atomMax(int64_t *address, int64_t val) {
|
||||
AtomicMaxIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(at::Half *address, at::Half val) {
|
||||
AtomicMaxDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(float *address, float val) {
|
||||
AtomicMaxDecimalImpl<float, sizeof(float)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMax(double *address, double 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)
|
||||
ATOMIC(Min)
|
||||
#undef OP
|
||||
static inline __device__ void atomMin(uint8_t *address, uint8_t val) {
|
||||
AtomicMinIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(int8_t *address, int8_t val) {
|
||||
AtomicMinIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(int16_t *address, int16_t val) {
|
||||
AtomicMinIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(int32_t *address, int32_t val) { atomicMin(address, val); }
|
||||
static inline __device__ void atomMin(int64_t *address, int64_t val) {
|
||||
AtomicMinIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(at::Half *address, at::Half val) {
|
||||
AtomicMinDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(float *address, float val) {
|
||||
AtomicMinDecimalImpl<float, sizeof(float)>()(address, val);
|
||||
}
|
||||
static inline __device__ void atomMin(double *address, double 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);
|
||||
}
|
||||
+115
@@ -0,0 +1,115 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#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);
|
||||
}
|
||||
};
|
||||
@@ -0,0 +1,128 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include "scatter_cuda.h"
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/Dispatch.h>
|
||||
#include <ATen/cuda/detail/IndexUtils.cuh>
|
||||
#include <ATen/cuda/detail/TensorInfo.cuh>
|
||||
|
||||
#include "reducer.cuh"
|
||||
#include "utils.cuh"
|
||||
|
||||
#define THREADS 256
|
||||
#define BLOCKS(N) (N + THREADS - 1) / THREADS
|
||||
|
||||
template <typename scalar_t, ReductionType REDUCE>
|
||||
__global__ void scatter_kernel(const scalar_t *src_data, const at::cuda::detail::TensorInfo<int64_t, int> index_info,
|
||||
scalar_t *out_data, int E, int K, int N, int numel) {
|
||||
int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
int b = thread_idx / (E * K);
|
||||
int k = thread_idx % K;
|
||||
|
||||
if (thread_idx < numel) {
|
||||
int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(thread_idx, index_info);
|
||||
int64_t idx = index_info.data[offset];
|
||||
|
||||
Reducer<scalar_t, REDUCE>::atomic_write(out_data + b * N * K + idx * K + k, src_data[thread_idx]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void scatter_arg_kernel(const scalar_t *src_data,
|
||||
const at::cuda::detail::TensorInfo<int64_t, int> index_info,
|
||||
const scalar_t *out_data, int64_t *arg_out_data, int E, int K, int N, int numel) {
|
||||
int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
int b = thread_idx / (E * K);
|
||||
int e = (thread_idx / K) % E;
|
||||
int k = thread_idx % K;
|
||||
|
||||
if (thread_idx < numel) {
|
||||
int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(thread_idx, index_info);
|
||||
int64_t idx = index_info.data[offset];
|
||||
|
||||
if (src_data[thread_idx] == out_data[b * N * K + idx * K + k]) {
|
||||
arg_out_data[b * N * K + idx * K + k] = e;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_cuda(torch::Tensor src, torch::Tensor index,
|
||||
int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size,
|
||||
std::string reduce) {
|
||||
CHECK_CUDA(src);
|
||||
CHECK_CUDA(index);
|
||||
if (optional_out.has_value()) CHECK_CUDA(optional_out.value());
|
||||
c10::cuda::MaybeSetDevice(src.get_device());
|
||||
|
||||
CHECK_INPUT(src.dim() == index.dim());
|
||||
for (auto i = 0; i < index.dim() - 1; i++) CHECK_INPUT(src.size(i) >= index.size(i));
|
||||
|
||||
src = src.contiguous();
|
||||
|
||||
torch::Tensor out;
|
||||
if (optional_out.has_value()) {
|
||||
out = optional_out.value().contiguous();
|
||||
for (auto i = 0; i < out.dim(); i++)
|
||||
if (i != dim) CHECK_INPUT(src.size(i) == out.size(i));
|
||||
} else {
|
||||
auto sizes = src.sizes().vec();
|
||||
if (dim_size.has_value())
|
||||
sizes[dim] = dim_size.value();
|
||||
else if (index.numel() == 0)
|
||||
sizes[dim] = 0;
|
||||
else {
|
||||
sizes[dim] = 1 + index.max().cpu().data_ptr<int64_t>()[0];
|
||||
}
|
||||
out = torch::empty(sizes, src.options());
|
||||
}
|
||||
|
||||
std::optional<torch::Tensor> arg_out = std::nullopt;
|
||||
int64_t *arg_out_data = nullptr;
|
||||
if (reduce2REDUCE.at(reduce) == MIN || reduce2REDUCE.at(reduce) == MAX) {
|
||||
arg_out = torch::full_like(out, src.size(dim), index.options());
|
||||
arg_out_data = arg_out.value().data_ptr<int64_t>();
|
||||
}
|
||||
|
||||
if (src.numel() == 0) {
|
||||
if (!optional_out.has_value()) out.fill_(0);
|
||||
return std::make_tuple(out, arg_out);
|
||||
}
|
||||
|
||||
auto B = 1;
|
||||
for (auto i = 0; i < dim; i++) B *= src.size(i);
|
||||
auto E = src.size(dim);
|
||||
auto K = src.numel() / (B * E);
|
||||
auto N = out.size(dim);
|
||||
|
||||
auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
|
||||
auto src_data = src.data_ptr<scalar_t>();
|
||||
auto out_data = out.data_ptr<scalar_t>();
|
||||
|
||||
AT_DISPATCH_REDUCTION_TYPES(reduce, [&] {
|
||||
if (!optional_out.has_value()) out.fill_(Reducer<scalar_t, REDUCE>::init());
|
||||
|
||||
scatter_kernel<scalar_t, REDUCE>
|
||||
<<<BLOCKS(src.numel()), THREADS, 0, stream>>>(src_data, index_info, out_data, E, K, N, src.numel());
|
||||
|
||||
if (!optional_out.has_value() && (REDUCE == MIN || REDUCE == MAX))
|
||||
out.masked_fill_(out == Reducer<scalar_t, REDUCE>::init(), (scalar_t)0);
|
||||
|
||||
if (REDUCE == MIN || REDUCE == MAX)
|
||||
scatter_arg_kernel<scalar_t><<<BLOCKS(src.numel()), THREADS, 0, stream>>>(
|
||||
src_data, index_info, out_data, arg_out_data, E, K, N, src.numel());
|
||||
});
|
||||
});
|
||||
|
||||
return std::make_tuple(out, arg_out);
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_cuda(torch::Tensor src, torch::Tensor index,
|
||||
int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size,
|
||||
std::string reduce);
|
||||
@@ -0,0 +1,29 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#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);
|
||||
}
|
||||
|
||||
__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);
|
||||
}
|
||||
|
||||
#define SHFL_UP_SYNC __shfl_up_sync
|
||||
#define SHFL_DOWN_SYNC __shfl_down_sync
|
||||
+238
@@ -0,0 +1,238 @@
|
||||
/**
|
||||
* This file includes code originally from the Pytorch Scatter repository:
|
||||
* https://github.com/rusty1s/pytorch_scatter
|
||||
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
#include "cuda/scatter_cuda.h"
|
||||
|
||||
inline std::vector<int64_t> list2vec(const c10::List<int64_t> list) {
|
||||
std::vector<int64_t> result;
|
||||
result.reserve(list.size());
|
||||
for (size_t i = 0; i < list.size(); i++) result.push_back(list[i]);
|
||||
return result;
|
||||
}
|
||||
|
||||
torch::Tensor broadcast(torch::Tensor src, torch::Tensor other, int64_t dim) {
|
||||
if (src.dim() == 1)
|
||||
for (auto i = 0; i < dim; i++) src = src.unsqueeze(0);
|
||||
for (auto i = src.dim(); i < other.dim(); i++) src = src.unsqueeze(-1);
|
||||
src = src.expand(other.sizes().vec());
|
||||
return src;
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_fw(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size,
|
||||
std::string reduce) {
|
||||
return scatter_cuda(src, index, dim, optional_out, dim_size, reduce);
|
||||
}
|
||||
|
||||
using torch::autograd::AutogradContext;
|
||||
using torch::autograd::Variable;
|
||||
using torch::autograd::variable_list;
|
||||
|
||||
class ScatterSum : public torch::autograd::Function<ScatterSum> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");
|
||||
auto out = std::get<0>(result);
|
||||
ctx->save_for_backward({index});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto index = saved[0];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
auto grad_in = torch::gather(grad_out, dim, index, false);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterMul : public torch::autograd::Function<ScatterMul> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "mul");
|
||||
auto out = std::get<0>(result);
|
||||
ctx->save_for_backward({src, index, out});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto src = saved[0];
|
||||
auto index = saved[1];
|
||||
auto out = saved[2];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
auto grad_in = torch::gather(grad_out * out, dim, index, false).div_(src);
|
||||
grad_in.masked_fill_(grad_in.isnan(), 0);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterMean : public torch::autograd::Function<ScatterMean> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
|
||||
auto old_index = index;
|
||||
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");
|
||||
auto out = std::get<0>(result);
|
||||
|
||||
auto ones = torch::ones(old_index.sizes(), src.options());
|
||||
result = scatter_fw(ones, old_index, old_index.dim() <= dim ? old_index.dim() - 1 : dim, std::nullopt,
|
||||
out.size(dim), "sum");
|
||||
auto count = std::get<0>(result);
|
||||
count.masked_fill_(count < 1, 1);
|
||||
count = broadcast(count, out, dim);
|
||||
if (out.is_floating_point())
|
||||
out.true_divide_(count);
|
||||
else
|
||||
out.div_(count, "floor");
|
||||
|
||||
ctx->save_for_backward({index, count});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto index = saved[0];
|
||||
auto count = saved[1];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
count = torch::gather(count, dim, index, false);
|
||||
auto grad_in = torch::gather(grad_out, dim, index, false);
|
||||
grad_in.true_divide_(count);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterMin : public torch::autograd::Function<ScatterMin> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "min");
|
||||
auto out = std::get<0>(result);
|
||||
auto arg_out = std::get<1>(result).value();
|
||||
ctx->save_for_backward({index, arg_out});
|
||||
ctx->mark_non_differentiable({arg_out});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out, arg_out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto index = saved[0];
|
||||
auto arg_out = saved[1];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
src_shape[dim] += 1;
|
||||
auto grad_in = torch::zeros(src_shape, grad_out.options());
|
||||
grad_in.scatter_(dim, arg_out, grad_out);
|
||||
grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
class ScatterMax : public torch::autograd::Function<ScatterMax> {
|
||||
public:
|
||||
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
||||
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
||||
dim = dim < 0 ? src.dim() + dim : dim;
|
||||
ctx->saved_data["dim"] = dim;
|
||||
ctx->saved_data["src_shape"] = src.sizes();
|
||||
|
||||
index = broadcast(index, src, dim);
|
||||
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "max");
|
||||
auto out = std::get<0>(result);
|
||||
auto arg_out = std::get<1>(result).value();
|
||||
ctx->save_for_backward({index, arg_out});
|
||||
ctx->mark_non_differentiable({arg_out});
|
||||
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
||||
return {out, arg_out};
|
||||
}
|
||||
|
||||
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
||||
auto grad_out = grad_outs[0];
|
||||
auto saved = ctx->get_saved_variables();
|
||||
auto index = saved[0];
|
||||
auto arg_out = saved[1];
|
||||
auto dim = ctx->saved_data["dim"].toInt();
|
||||
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
||||
src_shape[dim] += 1;
|
||||
auto grad_in = torch::zeros(src_shape, grad_out.options());
|
||||
grad_in.scatter_(dim, arg_out, grad_out);
|
||||
grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
|
||||
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
||||
}
|
||||
};
|
||||
|
||||
torch::Tensor scatter_sum(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
|
||||
return ScatterSum::apply(src, index, dim, optional_out, dim_size)[0];
|
||||
}
|
||||
|
||||
torch::Tensor scatter_mul(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
|
||||
return ScatterMul::apply(src, index, dim, optional_out, dim_size)[0];
|
||||
}
|
||||
|
||||
torch::Tensor scatter_mean(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
|
||||
return ScatterMean::apply(src, index, dim, optional_out, dim_size)[0];
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> scatter_min(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size) {
|
||||
auto result = ScatterMin::apply(src, index, dim, optional_out, dim_size);
|
||||
return std::make_tuple(result[0], result[1]);
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> scatter_max(torch::Tensor src, torch::Tensor index, int64_t dim,
|
||||
std::optional<torch::Tensor> optional_out,
|
||||
std::optional<int64_t> dim_size) {
|
||||
auto result = ScatterMax::apply(src, index, dim, optional_out, dim_size);
|
||||
return std::make_tuple(result[0], result[1]);
|
||||
}
|
||||
|
||||
void pybind_scatter_ext(py::module &m) {
|
||||
m.def("scatter_sum", &scatter_sum, "Scatter sum operation");
|
||||
m.def("scatter_mul", &scatter_mul, "Scatter multiplication operation");
|
||||
m.def("scatter_mean", &scatter_mean, "Scatter mean operation");
|
||||
m.def("scatter_min", &scatter_min, "Scatter min operation");
|
||||
m.def("scatter_max", &scatter_max, "Scatter max operation");
|
||||
}
|
||||
Reference in New Issue
Block a user