184 lines
5.3 KiB
C++
184 lines
5.3 KiB
C++
/*
|
|
* N: # of input sets
|
|
* C: # of channels
|
|
* H: input height
|
|
* W: input width
|
|
* K: # of filters -> # of output sets
|
|
* R: filter height
|
|
* S: filter width
|
|
*
|
|
* OH: output height
|
|
* OW: output width
|
|
*/
|
|
#include "convolution.h"
|
|
#include "omp.h"
|
|
#include "util.h"
|
|
#include <mpi.h>
|
|
#include <stdio.h>
|
|
#include <fstream>
|
|
|
|
#define I_BSIZE 128
|
|
#define J_BSIZE 1024
|
|
#define K_BSIZE 1024
|
|
|
|
#define MIN(x, y) (((x) < (y)) ? (x) : (y))
|
|
|
|
static float *input, *output, *filter, *col_buf;
|
|
static int N, C, H, W;
|
|
static int K, R, S;
|
|
static int OH, OW;
|
|
static int pad;
|
|
static int dilation;
|
|
static int stride;
|
|
static int mpi_rank, mpi_world_size;
|
|
static MPI_Status status;
|
|
|
|
void im2col_cpu(float *data_im, int batches, int channels, int height, int width,
|
|
int kernel_h, int kernel_w, int pad, int stride,
|
|
int dilation, float *data_col)
|
|
{
|
|
register const long dil_kernel_h = (kernel_h - 1) * dilation + 1;
|
|
register const long dil_kernel_w = (kernel_w - 1) * dilation + 1;
|
|
register const long height_col = (height + 2 * pad - dil_kernel_h) / stride + 1;
|
|
register const long width_col = (width + 2 * pad - dil_kernel_w) / stride + 1;
|
|
register const long channels_col = channels * kernel_h * kernel_w;
|
|
register const long batch_stride = channels_col * height_col * width_col;
|
|
register const long sbatch_stride = C * H * W;
|
|
|
|
#pragma omp parallel for num_threads(40)
|
|
for (int b = 0; b < batches; ++b) {
|
|
for (long c = 0; c < channels_col; ++c) {
|
|
long w_offset = c % kernel_w;
|
|
long h_offset = (c / kernel_w) % kernel_h;
|
|
long c_im = c / kernel_h / kernel_w;
|
|
|
|
const long hc0 = h_offset * dilation - pad;
|
|
const long wc0 = w_offset * dilation - pad;
|
|
for (long h = 0; h < height_col; ++h) {
|
|
long h_pad = h * stride + hc0;
|
|
|
|
const long row_offset = (c * height_col + h) * width_col;
|
|
const long srow_offset = (c_im * height + h_pad) * width;
|
|
for (long w = 0; w < width_col; ++w) {
|
|
long w_pad = w * stride + wc0;
|
|
if ((((unsigned long)h_pad) < ((unsigned long)height)) && (((unsigned long)w_pad) < ((unsigned long)width)))
|
|
data_col[b * batch_stride + row_offset + w] = data_im[b * sbatch_stride + srow_offset + w_pad];
|
|
else {
|
|
data_col[b * batch_stride + row_offset + w] = 0.;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/*
|
|
* _A: filter
|
|
* _B: input
|
|
* _C: output
|
|
* _BA: batch size
|
|
* _M: K
|
|
* _N: OH * OW
|
|
* _K: C * R * S
|
|
*/
|
|
static void sgemm_cpu (float *_A, float *_B, float *_C, int _BA, int _M, int _N, int _K)
|
|
{
|
|
register const long batch_stride = _M * _N;
|
|
register const long sbatch_stride = _K * _N;
|
|
|
|
#pragma omp parallel for num_threads(40) shared(_A, _B, _C)
|
|
for (long b = 0; b < _BA; b ++)
|
|
for (long ii = 0; ii < _M; ii += I_BSIZE)
|
|
for (long jj = 0; jj < _N; jj += J_BSIZE)
|
|
for (long kk = 0; kk < _K; kk += K_BSIZE)
|
|
for (long k = kk; k < MIN(_K, kk + K_BSIZE); k++)
|
|
for (long i = ii; i < MIN(_M, ii + I_BSIZE); i++) {
|
|
float tmp = _A[i * _K + k];
|
|
for (long j = jj; j < MIN(_N, jj + J_BSIZE); j++) {
|
|
_C[b * batch_stride + i * _N + j] += tmp * _B[b * sbatch_stride + k * _N + j];
|
|
}
|
|
}
|
|
}
|
|
|
|
void convolution(
|
|
float *_input, float *_output, float *_filter,
|
|
int _N, int _C, int _H, int _W,
|
|
int _K, int _R, int _S,
|
|
int _pad, int _dilation, int _stride)
|
|
{
|
|
input = _input;
|
|
output = _output;
|
|
filter = _filter;
|
|
|
|
int size[mpi_world_size];
|
|
int is[mpi_world_size], ie[mpi_world_size];
|
|
|
|
for (int i=0; i<mpi_world_size; i++) {
|
|
is[i] = N / mpi_world_size * i;
|
|
ie[i] = N / mpi_world_size * (i+1);
|
|
}
|
|
|
|
ie[mpi_world_size-1] = N;
|
|
|
|
for (int i=0; i<mpi_world_size; i++) {
|
|
size[i] = ie[i] - is[i];
|
|
}
|
|
|
|
OH = (H + 2 * pad - dilation * (R - 1) - 1) / stride + 1;
|
|
OW = (W + 2 * pad - dilation * (S - 1) - 1) / stride + 1;
|
|
|
|
alloc_tensor(&col_buf, 1, size[mpi_rank], R * S * C, OH * OW);
|
|
|
|
// Scatter
|
|
if (mpi_rank == 0) {
|
|
for (int i=1; i<mpi_world_size; i++) {
|
|
MPI_Send(input + (is[i]*C*H*W), size[i]*C*H*W, MPI_FLOAT, i, 0, MPI_COMM_WORLD);
|
|
}
|
|
} else {
|
|
alloc_tensor(&input, size[mpi_rank], C, H, W);
|
|
alloc_tensor(&output, size[mpi_rank], K, OH, OW);
|
|
alloc_tensor(&filter, K, C, R, S);
|
|
MPI_Recv(input, size[mpi_rank]*C*H*W, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &status);
|
|
}
|
|
|
|
// Broadcast
|
|
MPI_Bcast(filter, K*C*R*S, MPI_FLOAT, 0, MPI_COMM_WORLD);
|
|
|
|
// Calculation
|
|
im2col_cpu(input, size[mpi_rank], C, H, W, R, S, pad, stride, dilation, col_buf);
|
|
|
|
sgemm_cpu(filter, col_buf, output, size[mpi_rank], K, OH * OW, C * R * S);
|
|
// convolution_omp(size[mpi_rank]);
|
|
|
|
// Gather
|
|
if (mpi_rank == 0) {
|
|
for (int i=1; i<mpi_world_size; i++) {
|
|
MPI_Recv(output + (is[i]*K*OH*OW), size[i]*K*OH*OW, MPI_FLOAT, i, 1, MPI_COMM_WORLD, &status);
|
|
}
|
|
} else {
|
|
MPI_Send(output, size[mpi_rank]*K*OH*OW, MPI_FLOAT, 0, 1, MPI_COMM_WORLD);
|
|
}
|
|
}
|
|
|
|
void convolution_init(
|
|
int _N, int _C, int _H, int _W,
|
|
int _K, int _R, int _S,
|
|
int _pad, int _dilation, int _stride)
|
|
{
|
|
N = _N; C = _C; H = _H; W = _W;
|
|
K = _K; R = _R; S = _S;
|
|
pad = _pad;
|
|
dilation = _dilation;
|
|
stride = _stride;
|
|
|
|
MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
|
|
MPI_Comm_size(MPI_COMM_WORLD, &mpi_world_size);
|
|
}
|
|
|
|
void convolution_final(
|
|
int _N, int _C, int _H, int _W,
|
|
int _K, int _R, int _S,
|
|
int _pad, int _dilation, int _stride)
|
|
{
|
|
}
|