212 lines
4.9 KiB
C++
212 lines
4.9 KiB
C++
|
#include "convolution.h"
|
||
|
#include <mpi.h>
|
||
|
#include <stdio.h>
|
||
|
#include <immintrin.h>
|
||
|
#include <omp.h>
|
||
|
|
||
|
#include "util.h"
|
||
|
|
||
|
#define MAX_NUM_THREAD 40 // 0-39
|
||
|
|
||
|
#define CEIL_DIV(x,y) ( ((x) + (y) - 1) / (y) )
|
||
|
#define CEIL(x,y) ( CEIL_DIV((x),(y)) * (y) )
|
||
|
#define MIN(a,b) ( ((a) < (b)) ? (a) : (b) )
|
||
|
|
||
|
#define likely(x) __builtin_expect((x),1)
|
||
|
#define unlikely(x) __builtin_expect((x),0)
|
||
|
|
||
|
static float *input, *output, *filter;
|
||
|
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 void conv_core(float *in, float *flt, float *out)
|
||
|
{
|
||
|
for (int c = 0; c < C; ++c)
|
||
|
{
|
||
|
float *inForCh = &in[c * H * W];
|
||
|
float *fltForCh = &flt[c * R * S];
|
||
|
|
||
|
#pragma omp parallel num_threads(MAX_NUM_THREAD)
|
||
|
#pragma omp for collapse(2) schedule(static) nowait
|
||
|
for (int oh = 0; oh < OH; oh++)
|
||
|
{
|
||
|
for (int ow = 0; ow < OW; ow++)
|
||
|
{
|
||
|
float o = 0.0f;
|
||
|
|
||
|
for (int r = 0; r < R; r++)
|
||
|
{
|
||
|
for (int s = 0; s < S; s++)
|
||
|
{
|
||
|
int h = oh * stride - pad + r * dilation;
|
||
|
if (unlikely(h < 0 || h >= H))
|
||
|
continue;
|
||
|
|
||
|
int w = ow * stride - pad + s * dilation;
|
||
|
if (unlikely(w < 0 || w >= W))
|
||
|
continue;
|
||
|
|
||
|
float i = inForCh[h * W + w];
|
||
|
float f = fltForCh[r * S + s];
|
||
|
|
||
|
o += i * f;
|
||
|
}
|
||
|
}
|
||
|
|
||
|
out[oh * OW + ow] += o;
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
#define MPI_ASYNC 1
|
||
|
|
||
|
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)
|
||
|
{
|
||
|
#if MPI_ASYNC == 1
|
||
|
MPI_Request reqs[3] = { MPI_REQUEST_NULL, MPI_REQUEST_NULL, MPI_REQUEST_NULL };
|
||
|
#endif
|
||
|
|
||
|
////////////////////////////////////////////////////////////////////////////////
|
||
|
// scatter
|
||
|
if (mpi_rank == 0)
|
||
|
{
|
||
|
input = _input;
|
||
|
output = _output;
|
||
|
filter = _filter;
|
||
|
}
|
||
|
|
||
|
#if MPI_ASYNC == 1
|
||
|
MPI_Ibcast(filter, K * C * R * S, MPI_FLOAT, 0, MPI_COMM_WORLD, &reqs[0]);
|
||
|
#else
|
||
|
MPI_Bcast(filter, K * C * R * S, MPI_FLOAT, 0, MPI_COMM_WORLD);
|
||
|
#endif
|
||
|
|
||
|
if (mpi_rank == 0)
|
||
|
{
|
||
|
const int slicedN = CEIL_DIV(_N, mpi_world_size);
|
||
|
const int sizeOfN = slicedN * C * H * W;
|
||
|
float *inputForNodes = input + (N * C * H * W);
|
||
|
|
||
|
for (int i = 1; i < mpi_world_size; i++)
|
||
|
{
|
||
|
#if MPI_ASYNC == 1
|
||
|
MPI_Isend(inputForNodes + sizeOfN * (i - 1), sizeOfN, MPI_FLOAT, i, 1, MPI_COMM_WORLD, &reqs[1]);
|
||
|
#else
|
||
|
MPI_Send(inputForNodes + sizeOfN * (i - 1), sizeOfN, MPI_FLOAT, i, 1, MPI_COMM_WORLD);
|
||
|
#endif
|
||
|
}
|
||
|
}
|
||
|
else
|
||
|
{
|
||
|
#if MPI_ASYNC == 1
|
||
|
MPI_Irecv(input, N * C * H * W, MPI_FLOAT, 0, 1, MPI_COMM_WORLD, &reqs[1]);
|
||
|
#else
|
||
|
MPI_Recv(input, N * C * H * W, MPI_FLOAT, 0, 1, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
||
|
#endif
|
||
|
|
||
|
zero_tensor(output, N, K, OH, OW);
|
||
|
}
|
||
|
|
||
|
#if MPI_ASYNC == 1
|
||
|
MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE);
|
||
|
#endif
|
||
|
|
||
|
////////////////////////////////////////////////////////////////////////////////
|
||
|
// calc
|
||
|
for (int n = 0; n < N; ++n)
|
||
|
{
|
||
|
for (int k = 0; k < K; ++k)
|
||
|
{
|
||
|
conv_core(
|
||
|
&input[n * C * H * W],
|
||
|
&filter[k * C * R * S],
|
||
|
&output[n * K * OH * OW + k * OH * OW]);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
////////////////////////////////////////////////////////////////////////////////
|
||
|
// gather
|
||
|
if (mpi_rank == 0)
|
||
|
{
|
||
|
const int slicedN = CEIL_DIV(_N, mpi_world_size);
|
||
|
const int sizeOfN = slicedN * K * OH * OW;
|
||
|
float *outputForNodes = output + (N * K * OH * OW);
|
||
|
|
||
|
for (int i = 1; i < mpi_world_size; i++)
|
||
|
{
|
||
|
#if MPI_ASYNC == 1
|
||
|
MPI_Irecv(outputForNodes + sizeOfN * (i - 1), sizeOfN, MPI_FLOAT, i, 2, MPI_COMM_WORLD, &reqs[2]);
|
||
|
#else
|
||
|
MPI_Recv(outputForNodes + sizeOfN * (i - 1), sizeOfN, MPI_FLOAT, i, 2, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
||
|
#endif
|
||
|
}
|
||
|
}
|
||
|
else
|
||
|
{
|
||
|
#if MPI_ASYNC == 1
|
||
|
MPI_Isend(output, N * K * OH * OW, MPI_FLOAT, 0, 2, MPI_COMM_WORLD, &reqs[2]);
|
||
|
#else
|
||
|
MPI_Send(output, N * K * OH * OW, MPI_FLOAT, 0, 2, MPI_COMM_WORLD);
|
||
|
#endif
|
||
|
}
|
||
|
|
||
|
#if MPI_ASYNC == 1
|
||
|
MPI_Wait(&reqs[2], MPI_STATUS_IGNORE);
|
||
|
#endif
|
||
|
}
|
||
|
|
||
|
void convolution_init(
|
||
|
int _N, int _C, int _H, int _W,
|
||
|
int _K, int _R, int _S,
|
||
|
int _pad, int _dilation, int _stride)
|
||
|
{
|
||
|
C = _C; H = _H; W = _W;
|
||
|
K = _K; R = _R; S = _S;
|
||
|
pad = _pad;
|
||
|
dilation = _dilation;
|
||
|
stride = _stride;
|
||
|
|
||
|
OH = (H + 2 * pad - dilation * (R - 1) - 1) / stride + 1;
|
||
|
OW = (W + 2 * pad - dilation * (S - 1) - 1) / stride + 1;
|
||
|
|
||
|
MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
|
||
|
MPI_Comm_size(MPI_COMM_WORLD, &mpi_world_size);
|
||
|
|
||
|
N = CEIL_DIV(_N, mpi_world_size);
|
||
|
|
||
|
if (mpi_rank == 0)
|
||
|
{
|
||
|
N = _N - N * (mpi_world_size - 1);
|
||
|
}
|
||
|
else
|
||
|
{
|
||
|
alloc_tensor(&input, N, C, H, W);
|
||
|
alloc_tensor(&filter, K, C, R, S);
|
||
|
alloc_tensor(&output, N, K, OH, OW);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
void convolution_final(
|
||
|
int _N, int _C, int _H, int _W,
|
||
|
int _K, int _R, int _S,
|
||
|
int _pad, int _dilation, int _stride)
|
||
|
{
|
||
|
|
||
|
if (mpi_rank != 0)
|
||
|
{
|
||
|
free(input);
|
||
|
free(filter);
|
||
|
free(output);
|
||
|
}
|
||
|
}
|