def _s2_ifft(x, for_grad, b_in, b_out):
'''
:param x: [l * m, batch, complex] (b_in**2, nbatch, 2)
:return: [batch, beta, alpha, complex] (nbatch, 2 b_out, 2 * b_out, 2)
'''
device = x.get_device()
nbatch = x.size(1)
plan = _setup_fft_plan(b_out, nbatch)
wigner = _setup_wigner(
b_out, nl=b_in, weighted=for_grad,
device=device) # [beta, l * m] (2 * b_out - 1, nspec)
cuda_kernel = _setup_s2ifft_cuda_kernel(b=b_out, nl=b_in, nbatch=nbatch)
stream = cuda_utils.Stream(ptr=torch.cuda.current_stream().cuda_stream)
output = torch.cuda.FloatTensor(nbatch, 2 * b_out, 2 * b_out, 2)
cuda_kernel(block=(1024, 1, 1),
grid=(cuda_utils.get_blocks(nbatch * (2 * b_out)**2,
1024), 1, 1),
args=[x.data_ptr(),
wigner.data_ptr(),
output.data_ptr()],
stream=stream)
# [batch, beta, m, complex] (nbatch, 2 * b_out, 2 * b_out, 2)
plan(output, output, 1) # [batch, beta, alpha, complex]
return output
def _s2_fft(x, for_grad, b_in, b_out):
'''
this function performs in-place operations on x
:param x: [batch, beta, alpha, complex] (nbatch, 2 * b_in, 2 * b_in, 2)
:return: [l * m, batch, complex] (b_out**2, nbatch, 2)
'''
device = x.get_device()
nspec = b_out**2
nbatch = x.size(0)
plan = _setup_fft_plan(b_in, nbatch)
wigner = _setup_wigner(b_in,
nl=b_out,
weighted=not for_grad,
device=device)
cuda_kernel = _setup_s2fft_cuda_kernel(b=b_in, nspec=nspec, nbatch=nbatch)
plan(x, x, -1) # [batch, beta, m, complex]
stream = cuda_utils.Stream(ptr=torch.cuda.current_stream().cuda_stream)
output = torch.cuda.FloatTensor(nspec, nbatch, 2)
cuda_kernel(block=(1024, 1, 1),
grid=(cuda_utils.get_blocks(nspec * nbatch, 1024), 1, 1),
args=[x.data_ptr(),
wigner.data_ptr(),
output.data_ptr()],
stream=stream)
# [l * m, batch, complex]
return output
def s2_mm(x, y):
'''
:param x: [l * m, batch, feature_in, complex]
:param y: [l * m, feature_in, feature_out, complex]
:return: [l * m * n, batch, feature_out, complex]
'''
assert y.size(3) == 2
assert x.size(3) == 2
nbatch = x.size(1)
nfeature_in = x.size(2)
nfeature_out = y.size(2)
assert y.size(1) == nfeature_in
assert y.size(0) == x.size(0)
nl = round(x.size(0)**0.5)
nspec = (4 * nl**2 - 1) * nl // 3
assert x.size(0) == nl**2
assert y.size(0) == nl**2
cuda_kernel = _setup_s2mm_cuda_kernel(nbatch=nbatch,
nspec=nspec,
nfeature_in=nfeature_in,
nfeature_out=nfeature_out)
stream = cuda_utils.Stream(ptr=torch.cuda.current_stream().cuda_stream)
output = torch.cuda.FloatTensor(nspec, nbatch, nfeature_out, 2)
cuda_kernel(block=(cuda_utils.CUDA_NUM_THREADS, 1, 1),
grid=(cuda_utils.get_blocks(nspec * nbatch * nfeature_out,
1024), 1, 1),
args=[x.data_ptr(),
y.data_ptr(),
output.data_ptr()],
stream=stream)
# [l * m * n, batch, feature_out, complex]
return output
def backward(self, gradz): #pylint: disable=W
x, y = self.saved_tensors
nl = round(x.size(0)**0.5)
nbatch = x.size(1)
nfeature_in = x.size(2)
nfeature_out = y.size(2)
nspec = (4 * nl**2 - 1) * nl // 3
gradx_cuda_kernel = _setup_s2mm_gradx_cuda_kernel(
nbatch=nbatch,
nspec=nspec,
nl=nl,
nfeature_in=nfeature_in,
nfeature_out=nfeature_out)
grady_cuda_kernel = _setup_s2mm_grady_cuda_kernel(
nbatch=nbatch,
nspec=nspec,
nl=nl,
nfeature_in=nfeature_in,
nfeature_out=nfeature_out)
stream = cuda_utils.Stream(ptr=torch.cuda.current_stream().cuda_stream)
gradx = grady = None
if self.needs_input_grad[0]:
gradx = torch.cuda.FloatTensor(nl**2, nbatch, nfeature_in, 2)
gradx_cuda_kernel(
block=(cuda_utils.CUDA_NUM_THREADS, 1, 1),
grid=(cuda_utils.get_blocks(nl**2 * nbatch * nfeature_in,
1024), 1, 1),
args=[gradz.data_ptr(),
y.data_ptr(),
gradx.data_ptr()],
stream=stream)
if self.needs_input_grad[1]:
grady = torch.cuda.FloatTensor(nl**2, nfeature_in, nfeature_out, 2)
grady_cuda_kernel(
block=(cuda_utils.CUDA_NUM_THREADS, 1, 1),
grid=(cuda_utils.get_blocks(nl**2 * nfeature_in * nfeature_out,
1024), 1, 1),
args=[gradz.data_ptr(),
x.data_ptr(),
grady.data_ptr()],
stream=stream)
return gradx, grady
def _setup_so3ifft_cuda_kernel(b_in, b_out, nbatch, real_output):
kernel = '''
#define B_IN {}
#define B_OUT {}
#define NSPEC {}
#define NBATCH {}
'''.format(b_in, b_out,
b_in * (4 * b_in**2 - 1) // 3, nbatch)
if real_output:
kernel += '''
#define REAL_OUT
'''
kernel += '''
#define MOD(i, n) (((i) + (n)) % (n))
#define MAX(x, y) ((x) < (y) ? (y) : (x))
#define CEIL_DIV(x, y) (((x) + (y) - 1) / (y))
extern "C"
__global__ void main_(const float* in, const float* wig, float* out)
{
int m = (blockIdx.z / (2 * B_OUT - 1)) - (B_OUT - 1);
int n = (blockIdx.z % (2 * B_OUT - 1)) - (B_OUT - 1);
#ifdef REAL_OUT
if (n < 0 || (n == 0 && m < 0)) {
return; // note: this return does not depend on threadIdx
}
#endif
int l_min = MAX(abs(m), abs(n));
int batch = blockIdx.y * 32 + threadIdx.y;
float sum_re = 0.0;
float sum_im = 0.0;
for (int tile = 0; tile < CEIL_DIV(B_IN - l_min, 32); ++tile) {
__shared__ float tileA[2][32][32];
__shared__ float tileB[32][32+1];
int l = l_min + tile * 32 + threadIdx.x;
int lmn = (4 * l*l - 1) * l / 3 + (l+m) * (2 * l + 1) + (l+n);
int i = (lmn * NBATCH + batch) * 2;
tileA[0][threadIdx.y][threadIdx.x] = l < B_IN && batch < NBATCH ? in[i + 0] : 0.0;
tileA[1][threadIdx.y][threadIdx.x] = l < B_IN && batch < NBATCH ? in[i + 1] : 0.0;
int beta = blockIdx.x * 32 + threadIdx.y;
tileB[threadIdx.x][threadIdx.y] = l < B_IN && beta < 2*B_OUT ? wig[beta * NSPEC + lmn] : 0.0;
__syncthreads();
for (int l = 0; l < 32; ++l) {
sum_re += tileA[0][threadIdx.y][l] * tileB[l][threadIdx.x];
sum_im += tileA[1][threadIdx.y][l] * tileB[l][threadIdx.x];
}
__syncthreads();
}
int beta = blockIdx.x * 32 + threadIdx.x;
if (beta < 2*B_OUT && batch < NBATCH) {
int i = (((batch * 2*B_OUT + beta) * 2*B_OUT + MOD(m, 2*B_OUT)) * 2*B_OUT + MOD(n, 2*B_OUT)) * 2;
out[i + 0] = sum_re;
out[i + 1] = sum_im;
#ifdef REAL_OUT
i = (((batch * 2*B_OUT + beta) * 2*B_OUT + MOD(-m, 2*B_OUT)) * 2*B_OUT + MOD(-n, 2*B_OUT)) * 2;
out[i + 0] = sum_re;
out[i + 1] = -sum_im;
#endif
}
}
'''
kernel = cuda_utils.compile_kernel(kernel, b'so3ifft.cu', 'main_')
stream = cuda_utils.Stream(ptr=torch.cuda.current_stream().cuda_stream)
def fun(x, wigner, output):
output[:] = 0
kernel(block=(32, 32, 1),
grid=(math.ceil(2 * b_out / 32), math.ceil(nbatch / 32),
(2 * b_out - 1)**2),
args=[x.data_ptr(),
wigner.data_ptr(),
output.data_ptr()],
stream=stream)
return fun
def _setup_so3fft_cuda_kernel(b_in, b_out, nbatch, real_input):
kernel = '''
#define B_IN {}
#define B_OUT {}
#define NSPEC {}
#define NBATCH {}
'''.format(b_in, b_out,
b_out * (4 * b_out**2 - 1) // 3, nbatch)
if real_input:
kernel += '''
#define REAL_IN
'''
kernel += '''
#define MOD(i, n) (((i) + (n)) % (n))
#define MAX(x, y) ((x) < (y) ? (y) : (x))
#define CEIL_DIV(x, y) (((x) + (y) - 1) / (y))
extern "C"
__global__ void main_(const float* in, const float* wig, float* out)
{
// blockIdx = (l, batch, mn)
// blockDim = (32, 32, 1)
// threadIdx = (sub l, sub batch, 0)
// gridDim = (b / 32, nbatch / 32, (2b-1)**2)
int m = (blockIdx.z / (2 * B_OUT - 1)) - (B_OUT - 1);
int n = (blockIdx.z % (2 * B_OUT - 1)) - (B_OUT - 1);
int l_min = MAX(abs(m), abs(n));
if (blockIdx.x * 32 + 31 < l_min) {
// for blocks fully out of l-range
return; // note: this return does not depend on threadIdx
}
#ifdef REAL_IN
if (n < 0 || (n == 0 && m < 0)) {
return; // note: this return does not depend on threadIdx
}
#endif
int batch = blockIdx.y * 32 + threadIdx.y;
int l = blockIdx.x * 32 + threadIdx.x;
int lmn = (4 * l*l - 1) * l / 3 + (l+m) * (2 * l + 1) + (l+n);
float sum_re = 0.0;
float sum_im = 0.0;
for (int tile = 0; tile < CEIL_DIV(2 * B_IN, 32); ++tile) {
__shared__ float tileA[32][32][2];
__shared__ float tileB[32][32];
int beta = tile * 32 + threadIdx.x;
#ifdef REAL_IN
// `in` shape is (NBATCH, 2 * B_IN, 2 * B_IN, B_IN + 1, 2)
// http://www.fftw.org/fftw3_doc/Multi_002dDimensional-DFTs-of-Real-Data.html
int i = (((batch * 2*B_IN + beta) * 2*B_IN + MOD(m, 2*B_IN)) * (B_IN + 1) + n) * 2;
#else
int i = (((batch * 2*B_IN + beta) * 2*B_IN + MOD(m, 2*B_IN)) * 2*B_IN + MOD(n, 2*B_IN)) * 2;
#endif
tileA[threadIdx.y][threadIdx.x][0] = beta < 2*B_IN && batch < NBATCH ? in[i + 0] : 0.0;
tileA[threadIdx.y][threadIdx.x][1] = beta < 2*B_IN && batch < NBATCH ? in[i + 1] : 0.0;
beta = tile * 32 + threadIdx.y;
tileB[threadIdx.y][threadIdx.x] = beta < 2*B_IN && l_min <= l && l < B_OUT ? wig[beta * NSPEC + lmn] : 0.0;
__syncthreads();
for (int beta = 0; beta < 32; ++beta) {
sum_re += tileA[threadIdx.y][beta][0] * tileB[beta][threadIdx.x];
sum_im += tileA[threadIdx.y][beta][1] * tileB[beta][threadIdx.x];
}
__syncthreads();
}
// About this if: some blocks are used to compute but not to save the results
if (l_min <= l && l < B_OUT && batch < NBATCH) {
out[(lmn * NBATCH + batch) * 2 + 0] = sum_re;
out[(lmn * NBATCH + batch) * 2 + 1] = sum_im;
#ifdef REAL_IN
lmn = (4 * l*l - 1) * l / 3 + (l-m) * (2 * l + 1) + (l-n);
float fudge = (m - n) % 2 == 0 ? 1.0 : -1.0;
out[(lmn * NBATCH + batch) * 2 + 0] = fudge * sum_re;
out[(lmn * NBATCH + batch) * 2 + 1] = -fudge * sum_im;
#endif
}
}
'''
kernel = cuda_utils.compile_kernel(kernel, b'so3fft.cu', 'main_')
stream = cuda_utils.Stream(ptr=torch.cuda.current_stream().cuda_stream)
def fun(x, wigner, output):
kernel(block=(32, 32, 1),
grid=(math.ceil(b_out / 32), math.ceil(nbatch / 32),
(2 * b_out - 1)**2),
args=[x.data_ptr(),
wigner.data_ptr(),
output.data_ptr()],
stream=stream)
return fun