Im2Col.h

Go to the documentation of this file.

#pragma once
 
#include <cstddef>
#include <cstring>
#include <immintrin.h>

 
namespace Deep
{
    inline int ConvOutDim(int inDim, int kernel, int stride, int pad) noexcept
    {
        return (inDim + 2 * pad - kernel) / stride + 1;
    }

    inline void Im2Col(const float *input,
                       int channels, int height, int width,
                       int kH, int kW,
                       int strideH, int strideW,
                       int padH, int padW,
                       float *columns) noexcept
    {
        const int outH = ConvOutDim(height, kH, strideH, padH);
        const int outW = ConvOutDim(width, kW, strideW, padW);
        const int colWidth = outH * outW;
 
        for (int c = 0; c < channels; ++c)
        {
            for (int kh = 0; kh < kH; ++kh)
            {
                for (int kw = 0; kw < kW; ++kw)
                {
                    const int rowIdx = (c * kH + kh) * kW + kw;
                    float *destRow = columns + (size_t)rowIdx * colWidth;
 
                    for (int oh = 0; oh < outH; ++oh)
                    {
                        const int inRow = oh * strideH - padH + kh;
                        float *dst = destRow + (size_t)oh * outW;
 
                        if (inRow < 0 || inRow >= height)
                        {
                            std::memset(dst, 0, (size_t)outW * sizeof(float));
                            continue;
                        }
 
                        const float *src =
                            input + ((size_t)c * height + inRow) * width;
 
                        if (strideW == 1)
                        {
                            const int first = padW > kw ? padW - kw : 0;
                            const int last = width + padW - kw < outW
                                                 ? width + padW - kw
                                                 : outW;
 
                            int ow = 0;
 
                            for (; ow < first; ++ow)
                                dst[ow] = 0.0f;
 
                            const float *s = src + first - padW + kw;
                            float *d = dst + first;
                            const int count = last - first;
                            int i = 0;
 
#if defined(__AVX512F__)
                            for (; i + 16 <= count; i += 16)
                                _mm512_storeu_ps(
                                    d + i,
                                    _mm512_loadu_ps(s + i));
#elif defined(__AVX2__) || defined(__AVX__)
                            for (; i + 8 <= count; i += 8)
                                _mm256_storeu_ps(
                                    d + i,
                                    _mm256_loadu_ps(s + i));
#elif defined(__SSE__)
                            for (; i + 4 <= count; i += 4)
                                _mm_storeu_ps(
                                    d + i,
                                    _mm_loadu_ps(s + i));
#endif
 
                            for (; i < count; ++i)
                                d[i] = s[i];
 
                            for (ow = last; ow < outW; ++ow)
                                dst[ow] = 0.0f;
                        }
                        else
                        {
                            for (int ow = 0; ow < outW; ++ow)
                            {
                                const int inCol =
                                    ow * strideW - padW + kw;
 
                                dst[ow] =
                                    (inCol < 0 || inCol >= width)
                                        ? 0.0f
                                        : src[inCol];
                            }
                        }
                    }
                }
            }
        }
    }

    inline void Col2Im(const float *columns,
                       int channels, int height, int width,
                       int kH, int kW,
                       int strideH, int strideW,
                       int padH, int padW,
                       float *outputImage) noexcept
    {
        const int outH = ConvOutDim(height, kH, strideH, padH);
        const int outW = ConvOutDim(width, kW, strideW, padW);
        const int colWidth = outH * outW;
 
        for (int c = 0; c < channels; ++c)
        {
            for (int kh = 0; kh < kH; ++kh)
            {
                for (int kw = 0; kw < kW; ++kw)
                {
                    const int rowIdx = (c * kH + kh) * kW + kw;
                    const float *srcRow =
                        columns + (size_t)rowIdx * colWidth;
 
                    for (int oh = 0; oh < outH; ++oh)
                    {
                        const int inRow = oh * strideH - padH + kh;
 
                        if (inRow < 0 || inRow >= height)
                            continue;
 
                        float *dst =
                            outputImage +
                            ((size_t)c * height + inRow) * width;
 
                        const float *src =
                            srcRow + (size_t)oh * outW;
 
                        if (strideW == 1)
                        {
                            const int first = padW > kw ? padW - kw : 0;
                            const int last = width + padW - kw < outW
                                                 ? width + padW - kw
                                                 : outW;
 
                            const float *s = src + first;
                            float *d = dst + first - padW + kw;
                            const int count = last - first;
                            int i = 0;
 
#if defined(__AVX512F__)
                            for (; i + 16 <= count; i += 16)
                            {
                                const __m512 a = _mm512_loadu_ps(s + i);
                                const __m512 b = _mm512_loadu_ps(d + i);
                                _mm512_storeu_ps(
                                    d + i,
                                    _mm512_add_ps(a, b));
                            }
#elif defined(__AVX2__) || defined(__AVX__)
                            for (; i + 8 <= count; i += 8)
                            {
                                const __m256 a = _mm256_loadu_ps(s + i);
                                const __m256 b = _mm256_loadu_ps(d + i);
                                _mm256_storeu_ps(
                                    d + i,
                                    _mm256_add_ps(a, b));
                            }
#elif defined(__SSE__)
                            for (; i + 4 <= count; i += 4)
                            {
                                const __m128 a = _mm_loadu_ps(s + i);
                                const __m128 b = _mm_loadu_ps(d + i);
                                _mm_storeu_ps(
                                    d + i,
                                    _mm_add_ps(a, b));
                            }
#endif
 
                            for (; i < count; ++i)
                                d[i] += s[i];
                        }
                        else
                        {
                            for (int ow = 0; ow < outW; ++ow)
                            {
                                const int inCol =
                                    ow * strideW - padW + kw;
 
                                if (inCol >= 0 && inCol < width)
                                    dst[inCol] += src[ow];
                            }
                        }
                    }
                }
            }
        }
    }
}