mirrors
/
mmcv
mirror of https://github.com/open-mmlab/mmcv.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[Feature] Optimize the PyTorch CUDA implementation for Criss Cross Attention (#1143) 

* optimize criss cross attention

* optimize criss cross attention

* optimize criss cross attention

* fix lint

* fix ci, remove useless variable

* better ca_forward_kernel

Co-authored-by: wondervictor <victorchanchina@gmail.com>
pull/1144/head
q.yao 2021-06-25 15:25:19 +08:00 committed by GitHub
parent 6fe3722510
commit 7b150fab34
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 88 additions and 105 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

94
mmcv/ops/csrc/cc_attention_cuda_kernel.cuh
View File

@ -14,25 +14,17 @@ __global__ void ca_forward_kernel(const T *t, const T *f, T *weight, int num,
int y = blockIdx.y * blockDim.y + threadIdx.y; int y = blockIdx.y * blockDim.y + threadIdx.y;
int sp = height * width; int sp = height * width;
int len = height + width - 1; int len = height + width - 1;
int z = blockIdx.z; int z = blockIdx.z % len;
int batch = blockIdx.z / len;
if (x < width && y < height && z < height + width - 1) { if (x < width && y < height) {
for (int batch = 0; batch < num; ++batch) { T *weight_ptr = weight + (batch * len + z) * sp + y * width + x;
const int t_offset = y * width + x;
const int j = (z - width < y) ? z - width : z - width + 1;
const int f_offset = z < width ? y * width + z : j * width + x;
for (int plane = 0; plane < chn; ++plane) { for (int plane = 0; plane < chn; ++plane) {
T _t = t[(batch * chn + plane) * sp + y * width + x]; const int tf_base = (batch * chn + plane) * sp;
*weight_ptr += t[tf_base + t_offset] * f[tf_base + f_offset];
if (z < width) {
int i = z;
T _f = f[(batch * chn + plane) * sp + y * width + i];
weight[(batch * len + i) * sp + y * width + x] += _t * _f;
} else {
int i = z - width;
int j = i < y ? i : i + 1;
T _f = f[(batch * chn + plane) * sp + j * width + x];
weight[(batch * len + width + i) * sp + y * width + x] += _t * _f;
}
}
} }
} }
} }
@ -44,10 +36,10 @@ __global__ void ca_backward_kernel_t(const T *dw, const T *t, const T *f, T *dt,
int y = blockIdx.y * blockDim.y + threadIdx.y; int y = blockIdx.y * blockDim.y + threadIdx.y;
int sp = height * width; int sp = height * width;
int len = height + width - 1; int len = height + width - 1;
int plane = blockIdx.z; int plane = blockIdx.z % chn;
int batch = blockIdx.z / chn;
if (x < width && y < height && plane < chn) { if (x < width && y < height) {
for (int batch = 0; batch < num; ++batch) {
for (int i = 0; i < width; ++i) { for (int i = 0; i < width; ++i) {
T _dw = dw[(batch * len + i) * sp + y * width + x]; T _dw = dw[(batch * len + i) * sp + y * width + x];
T _f = f[(batch * chn + plane) * sp + y * width + i]; T _f = f[(batch * chn + plane) * sp + y * width + i];
@ -62,7 +54,6 @@ __global__ void ca_backward_kernel_t(const T *dw, const T *t, const T *f, T *dt,
dt[(batch * chn + plane) * sp + y * width + x] += _dw * _f; dt[(batch * chn + plane) * sp + y * width + x] += _dw * _f;
} }
} }
}
} }
template <typename T> template <typename T>
@ -72,10 +63,10 @@ __global__ void ca_backward_kernel_f(const T *dw, const T *t, const T *f, T *df,
int y = blockIdx.y * blockDim.y + threadIdx.y; int y = blockIdx.y * blockDim.y + threadIdx.y;
int sp = height * width; int sp = height * width;
int len = height + width - 1; int len = height + width - 1;
int plane = blockIdx.z; int plane = blockIdx.z % chn;
int batch = blockIdx.z / chn;
if (x < width && y < height && plane < chn) { if (x < width && y < height) {
for (int batch = 0; batch < num; ++batch) {
for (int i = 0; i < width; ++i) { for (int i = 0; i < width; ++i) {
T _dw = dw[(batch * len + x) * sp + y * width + i]; T _dw = dw[(batch * len + x) * sp + y * width + i];
T _t = t[(batch * chn + plane) * sp + y * width + i]; T _t = t[(batch * chn + plane) * sp + y * width + i];
@ -90,7 +81,6 @@ __global__ void ca_backward_kernel_f(const T *dw, const T *t, const T *f, T *df,
df[(batch * chn + plane) * sp + y * width + x] += _dw * _t; df[(batch * chn + plane) * sp + y * width + x] += _dw * _t;
} }
} }
}
} }
template <typename T> template <typename T>
@ -100,10 +90,9 @@ __global__ void ca_map_forward_kernel(const T *weight, const T *g, T *out,
int y = blockIdx.y * blockDim.y + threadIdx.y; int y = blockIdx.y * blockDim.y + threadIdx.y;
int sp = height * width; int sp = height * width;
int len = height + width - 1; int len = height + width - 1;
int plane = blockIdx.z; int plane = blockIdx.z % chn;
int batch = blockIdx.z / chn;
if (x < width && y < height && plane < chn) { if (x < width && y < height) {
for (int batch = 0; batch < num; ++batch) {
for (int i = 0; i < width; ++i) { for (int i = 0; i < width; ++i) {
T _g = g[(batch * chn + plane) * sp + y * width + i]; T _g = g[(batch * chn + plane) * sp + y * width + i];
T _w = weight[(batch * len + i) * sp + y * width + x]; T _w = weight[(batch * len + i) * sp + y * width + x];
@ -119,7 +108,6 @@ __global__ void ca_map_forward_kernel(const T *weight, const T *g, T *out,
out[(batch * chn + plane) * sp + y * width + x] += _g * _w; out[(batch * chn + plane) * sp + y * width + x] += _g * _w;
} }
} }
}
} }
template <typename T> template <typename T>
@ -130,25 +118,23 @@ __global__ void ca_map_backward_kernel_w(const T *dout, const T *weight,
int y = blockIdx.y * blockDim.y + threadIdx.y; int y = blockIdx.y * blockDim.y + threadIdx.y;
int sp = height * width; int sp = height * width;
int len = height + width - 1; int len = height + width - 1;
int z = blockIdx.z;
if (x < width && y < height && z < height + width - 1) { int z = blockIdx.z % len;
for (int batch = 0; batch < num; ++batch) { int batch = blockIdx.z / len;
for (int plane = 0; plane < chn; ++plane) {
T _dout = dout[(batch * chn + plane) * sp + y * width + x];
if (x < width && y < height) {
int widx = (batch * len + z) * sp + y * width + x;
int dout_idx = batch * chn * sp + y * width + x;
int gidx = batch * chn * sp;
if (z < width) { if (z < width) {
int i = z; gidx += y * width + z;
T _g = g[(batch * chn + plane) * sp + y * width + i];
dw[(batch * len + i) * sp + y * width + x] += _dout * _g;
} else { } else {
int i = z - width; int j = z - width;
int j = i < y ? i : i + 1; j = j < y ? j : j + 1;
gidx += j * width + x;
T _g = g[(batch * chn + plane) * sp + j * width + x];
dw[(batch * len + width + i) * sp + y * width + x] += _dout * _g;
}
} }
for (int plane = 0; plane < chn; plane++) {
dw[widx] += dout[dout_idx + plane * sp] * g[gidx + plane * sp];
} }
} }
} }
@ -161,25 +147,21 @@ __global__ void ca_map_backward_kernel_g(const T *dout, const T *weight,
int y = blockIdx.y * blockDim.y + threadIdx.y; int y = blockIdx.y * blockDim.y + threadIdx.y;
int sp = height * width; int sp = height * width;
int len = height + width - 1; int len = height + width - 1;
int plane = blockIdx.z; int plane = blockIdx.z % chn;
int batch = blockIdx.z / chn;
int index = (batch * chn + plane) * sp + y * width + x;
if (x < width && y < height && plane < chn) { if (x < width && y < height) {
for (int batch = 0; batch < num; ++batch) {
for (int i = 0; i < width; ++i) { for (int i = 0; i < width; ++i) {
T _dout = dout[(batch * chn + plane) * sp + y * width + i]; dg[index] += dout[(batch * chn + plane) * sp + y * width + i] *
T _w = weight[(batch * len + x) * sp + y * width + i]; weight[(batch * len + x) * sp + y * width + i];
dg[(batch * chn + plane) * sp + y * width + x] += _dout * _w;
} }
for (int i = 0; i < height; ++i) { for (int i = 0; i < height; ++i) {
if (i == y) continue; if (i == y) continue;
int j = i > y ? y : y - 1; int j = i > y ? y : y - 1;
dg[index] += dout[(batch * chn + plane) * sp + i * width + x] *
T _dout = dout[(batch * chn + plane) * sp + i * width + x]; weight[(batch * len + width + j) * sp + i * width + x];
T _w = weight[(batch * len + width + j) * sp + i * width + x];
dg[(batch * chn + plane) * sp + y * width + x] += _dout * _w;
}
} }
} }
} }
#endif // CC_ATTENTION_CUDA_KERNEL_CUH #endif // CC_ATTENTION_CUDA_KERNEL_CUH

15
mmcv/ops/csrc/pytorch/cc_attention_cuda.cu
View File

@ -24,8 +24,8 @@ void CAForwardCUDAKernelLauncher(const Tensor t, const Tensor f,
dim3 threads(32, 32); dim3 threads(32, 32);
int d1 = (w + threads.x - 1) / threads.x; int d1 = (w + threads.x - 1) / threads.x;
int d2 = (h + threads.y - 1) / threads.y; int d2 = (h + threads.y - 1) / threads.y;
int d3 = h + w; int d3 = h + w - 1;
dim3 blocks(d1, d2, d3); dim3 blocks(d1, d2, d3 * n);
AT_DISPATCH_FLOATING_TYPES(t.scalar_type(), "ca_forward", [&] { AT_DISPATCH_FLOATING_TYPES(t.scalar_type(), "ca_forward", [&] {
ca_forward_kernel<scalar_t><<<blocks, threads, 0, stream>>>( ca_forward_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
@ -53,7 +53,7 @@ void CABackwardCUDAKernelLauncher(const Tensor dw, const Tensor t,
dim3 threads(32, 32); dim3 threads(32, 32);
int d1 = (w + threads.x - 1) / threads.x; int d1 = (w + threads.x - 1) / threads.x;
int d2 = (h + threads.y - 1) / threads.y; int d2 = (h + threads.y - 1) / threads.y;
int d3 = c; int d3 = c * n;
dim3 blocks(d1, d2, d3); dim3 blocks(d1, d2, d3);
AT_DISPATCH_FLOATING_TYPES(t.scalar_type(), "ca_backward_kernel_t", [&] { AT_DISPATCH_FLOATING_TYPES(t.scalar_type(), "ca_backward_kernel_t", [&] {
@ -90,7 +90,7 @@ void CAMapForwardCUDAKernelLauncher(const Tensor weight, const Tensor g,
dim3 threads(32, 32); dim3 threads(32, 32);
int d1 = (w + threads.x - 1) / threads.x; int d1 = (w + threads.x - 1) / threads.x;
int d2 = (h + threads.y - 1) / threads.y; int d2 = (h + threads.y - 1) / threads.y;
int d3 = c; int d3 = c * n;
dim3 blocks(d1, d2, d3); dim3 blocks(d1, d2, d3);
AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "ca_map_forward", [&] { AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "ca_map_forward", [&] {
@ -119,8 +119,8 @@ void CAMapBackwardCUDAKernelLauncher(const Tensor dout, const Tensor weight,
dim3 threads(32, 32); dim3 threads(32, 32);
int d1 = (w + threads.x - 1) / threads.x; int d1 = (w + threads.x - 1) / threads.x;
int d2 = (h + threads.y - 1) / threads.y; int d2 = (h + threads.y - 1) / threads.y;
int d3 = h + w; int d3 = h + w - 1;
dim3 blocks(d1, d2, d3); dim3 blocks(d1, d2, d3 * n);
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
weight.scalar_type(), "ca_map_backward_kernel_w", [&] { weight.scalar_type(), "ca_map_backward_kernel_w", [&] {
@ -130,7 +130,8 @@ void CAMapBackwardCUDAKernelLauncher(const Tensor dout, const Tensor weight,
g.contiguous().data_ptr<scalar_t>(), g.contiguous().data_ptr<scalar_t>(),
dw.contiguous().data_ptr<scalar_t>(), n, c, h, w); dw.contiguous().data_ptr<scalar_t>(), n, c, h, w);
}); });
d3 = c * n;
blocks = dim3(d1, d2, d3);
AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "ca_map_backward_kernel_g", [&] { AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "ca_map_backward_kernel_g", [&] {
ca_map_backward_kernel_g<scalar_t><<<blocks, threads, 0, stream>>>( ca_map_backward_kernel_g<scalar_t><<<blocks, threads, 0, stream>>>(
dout.contiguous().data_ptr<scalar_t>(), dout.contiguous().data_ptr<scalar_t>(),