megenginelite-sys 1.8.2

/**
 * \file dnn/src/fallback/convolution/run_conv.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */
#include "src/fallback/convolution/run_conv.h"

#include "midout.h"
#include "src/common/utils.h"

MIDOUT_DECL(megdnn_fallback_conv)

namespace {

bool can_run_xcorr_single_channel_templated(
        size_t /* IH */, size_t /* IW */, size_t FH, size_t FW, size_t /* OH */,
        size_t /* OW */, size_t /* PH */, size_t /* PW */, size_t /* SH */,
        size_t /* SW */) {
    return FH == FW && FH >= 1 && FH <= 7;
}

template <int ker_size>
void run_xcorr_single_channel_templated_impl(
        const float* __restrict src, const float* __restrict filter,
        float* __restrict dst, size_t IH, size_t IW, size_t OH, size_t OW, size_t PH,
        size_t PW, size_t SH, size_t SW, bool add_mode) {
#define divup(x, y) (((x) + (y)-1) / (y))
#define clear(oh, ow)            \
    if (!add_mode) {             \
        dst[(oh)*OW + (ow)] = 0; \
    }
#define update(oh, ow, fh, fw)                            \
    dst[(oh)*OW + (ow)] += filter[(fh)*ker_size + (fw)] * \
                           src[((oh)*SH + (fh)-PH) * IW + ((ow)*SW + (fw)-PW)]
    // OH = (IH-ker_size)/stride+1
    // OW = (IW-ker_size)/stride+1
    // good region:
    //  oh*stride-anchor >= 0
    //  oh*stride-anchor+ker_size <= IH
    //  oh >= anchor/stride
    //  oh <= (IH+anchor-ker_size)/stride
    size_t oh_start = divup(PH, SH);
    size_t oh_end = IH + PH >= ker_size ? (IH + PH - ker_size) / SH + 1 : 0;
    size_t ow_start = divup(PW, SW);
    size_t ow_end = IW + PW >= ker_size ? (IW + PW - ker_size) / SW + 1 : 0;
    if (oh_start > oh_end)
        oh_start = oh_end = 0;
    if (ow_start > ow_end)
        ow_start = ow_end = 0;

    for (size_t oh = 0; oh < oh_start; ++oh)
        for (size_t ow = 0; ow < OW; ++ow) {
            clear(oh, ow);
            int ih = oh * SH - PH;
            int iw = ow * SW - PW;
            for (int fh = 0; fh < ker_size; ++fh)
                if (ih + fh >= 0 && ih + fh < (int)IH)
                    for (int fw = 0; fw < ker_size; ++fw)
                        if (iw + fw >= 0 && iw + fw < (int)IW) {
                            update(oh, ow, fh, fw);
                        }
        }
    for (size_t oh = oh_start; oh < oh_end; ++oh) {
        for (size_t ow = 0; ow < ow_start; ++ow) {
            clear(oh, ow);
            int iw = ow * SW - PW;
            for (int fh = 0; fh < ker_size; ++fh)
                for (int fw = 0; fw < ker_size; ++fw) {
                    if (iw + fw >= 0 && iw + fw < (int)IW)
                        update(oh, ow, fh, fw);
                }
        }
        for (size_t ow = ow_start; ow < ow_end; ++ow) {
            clear(oh, ow);
            for (int fh = 0; fh < ker_size; ++fh)
                for (int fw = 0; fw < ker_size; ++fw) {
                    update(oh, ow, fh, fw);
                }
        }
        for (size_t ow = ow_end; ow < OW; ++ow) {
            clear(oh, ow);
            int iw = ow * SW - PW;
            for (int fh = 0; fh < ker_size; ++fh)
                for (int fw = 0; fw < ker_size; ++fw) {
                    if (iw + fw >= 0 && iw + fw < (int)IW)
                        update(oh, ow, fh, fw);
                }
        }
    }
    for (size_t oh = oh_end; oh < OH; ++oh) {
        for (size_t ow = 0; ow < OW; ++ow) {
            clear(oh, ow);
            int ih = oh * SH - PH;
            int iw = ow * SW - PW;
            for (int fh = 0; fh < ker_size; ++fh)
                if (ih + fh >= 0 && ih + fh < (int)IH)
                    for (int fw = 0; fw < ker_size; ++fw)
                        if (iw + fw >= 0 && iw + fw < (int)IW) {
                            update(oh, ow, fh, fw);
                        }
        }
    }
#undef divup
#undef clear
#undef update
}

void run_xcorr_single_channel_templated(
        const float* src, const float* filter, float* dst, size_t IH, size_t IW,
        size_t FH, size_t FW, size_t OH, size_t OW, size_t PH, size_t PW, size_t SH,
        size_t SW, bool add_mode) {
    (void)FW;
#define DISPATCH(ker_size)                                                       \
    if (FH == ker_size) {                                                        \
        MIDOUT_BEGIN(megdnn_fallback_conv, ker_size) {                           \
            run_xcorr_single_channel_templated_impl<ker_size>(                   \
                    src, filter, dst, IH, IW, OH, OW, PH, PW, SH, SW, add_mode); \
        }                                                                        \
        MIDOUT_END();                                                            \
        return;                                                                  \
    }
    DISPATCH(1)
    DISPATCH(2)
    DISPATCH(3)
    DISPATCH(4)
    DISPATCH(5)
    DISPATCH(6)
    DISPATCH(7)
#undef DISPATCH
    megdnn_throw("internal error in conv template dispatching: impossible");
}

void run_xcorr_single_channel_nontemplated(
        const float* src, const float* filter, float* dst, size_t IH, size_t IW,
        size_t FH_, size_t FW_, size_t OH, size_t OW, size_t PH, size_t PW, size_t SH,
        size_t SW, bool add_mode) {
#define divup(x, y) (((x) + (y)-1) / (y))
#define clear(oh, ow)            \
    if (!add_mode) {             \
        dst[(oh)*OW + (ow)] = 0; \
    }
#define update(oh, ow, fh, fw)                      \
    dst[(oh)*OW + (ow)] += filter[(fh)*FW + (fw)] * \
                           src[((oh)*SH + (fh)-PH) * IW + ((ow)*SW + (fw)-PW)]
    // OH = (IH-ker_size)/stride+1
    // OW = (IW-ker_size)/stride+1
    // good region:
    //  oh*stride-anchor >= 0
    //  oh*stride-anchor+ker_size <= IH
    //  oh >= anchor/stride
    //  oh <= (IH+anchor-ker_size)/stride
    int FH = FH_, FW = FW_;
    size_t oh_start = divup(PH, SH);
    size_t oh_end = IH + PH >= FH_ ? (IH + PH - FH) / SH + 1 : 0;
    size_t ow_start = divup(PW, SW);
    size_t ow_end = IW + PW >= FW_ ? (IW + PW - FW) / SW + 1 : 0;
    if (oh_start > oh_end)
        oh_start = oh_end = 0;
    if (ow_start > ow_end)
        ow_start = ow_end = 0;
    for (size_t oh = 0; oh < oh_start; ++oh)
        for (size_t ow = 0; ow < OW; ++ow) {
            clear(oh, ow);
            int ih = oh * SH - PH;
            int iw = ow * SW - PW;
            for (int fh = 0; fh < FH; ++fh)
                if (ih + fh >= 0 && ih + fh < (int)IH)
                    for (int fw = 0; fw < FW; ++fw)
                        if (iw + fw >= 0 && iw + fw < (int)IW) {
                            update(oh, ow, fh, fw);
                        }
        }
    for (size_t oh = oh_start; oh < oh_end; ++oh) {
        for (size_t ow = 0; ow < ow_start; ++ow) {
            clear(oh, ow);
            int iw = ow * SW - PW;
            for (int fh = 0; fh < FH; ++fh)
                for (int fw = 0; fw < FW; ++fw) {
                    if (iw + fw >= 0 && iw + fw < (int)IW)
                        update(oh, ow, fh, fw);
                }
        }
        for (size_t ow = ow_start; ow < ow_end; ++ow) {
            clear(oh, ow);
            for (int fh = 0; fh < FH; ++fh)
                for (int fw = 0; fw < FW; ++fw) {
                    update(oh, ow, fh, fw);
                }
        }
        for (size_t ow = ow_end; ow < OW; ++ow) {
            clear(oh, ow);
            int iw = ow * SW - PW;
            for (int fh = 0; fh < FH; ++fh)
                for (int fw = 0; fw < FW; ++fw) {
                    if (iw + fw >= 0 && iw + fw < (int)IW)
                        update(oh, ow, fh, fw);
                }
        }
    }
    for (size_t oh = oh_end; oh < OH; ++oh) {
        for (size_t ow = 0; ow < OW; ++ow) {
            clear(oh, ow);
            int ih = oh * SH - PH;
            int iw = ow * SW - PW;
            for (int fh = 0; fh < FH; ++fh)
                if (ih + fh >= 0 && ih + fh < (int)IH)
                    for (int fw = 0; fw < FW; ++fw)
                        if (iw + fw >= 0 && iw + fw < (int)IW) {
                            update(oh, ow, fh, fw);
                        }
        }
    }
#undef divup
#undef clear
#undef update
}

void run_xcorr_single_channel(
        const float* src, const float* filter, float* dst, size_t IH, size_t IW,
        size_t FH, size_t FW, size_t OH, size_t OW, size_t PH, size_t PW, size_t SH,
        size_t SW, bool add_mode) {
    if (can_run_xcorr_single_channel_templated(
                IH, IW, FH, FW, OH, OW, PH, PW, SH, SW)) {
        run_xcorr_single_channel_templated(
                src, filter, dst, IH, IW, FH, FW, OH, OW, PH, PW, SH, SW, add_mode);
    } else {
        MIDOUT_BEGIN(megdnn_fallback_conv, void) {
            run_xcorr_single_channel_nontemplated(
                    src, filter, dst, IH, IW, FH, FW, OH, OW, PH, PW, SH, SW, add_mode);
        }
        MIDOUT_END();
    }
}

/*================ ConvolutionBackwardData =============*/

template <int ker_size>
void conv_backdata_single_channel_templated_impl(
        const float* __restrict diff, const float* __restrict filter,
        float* __restrict grad, size_t IH, size_t IW, size_t OH, size_t OW, size_t PH,
        size_t PW, size_t SH, size_t SW) {
#define divup(x, y) (((x) + (y)-1) / (y))
#define update(oh, ow, fh, fw, val) \
    grad[(oh + fh) * OW + (ow + fw)] += filter[(fh)*ker_size + (fw)] * val
    size_t ih_start = divup(PH, SH);
    size_t ih_end = OH + PH >= ker_size ? (OH + PH - ker_size) / SH + 1 : 0;
    size_t iw_start = divup(PW, SW);
    size_t iw_end = OW + PW >= ker_size ? (OW + PW - ker_size) / SW + 1 : 0;
    if (ih_start > ih_end)
        ih_start = ih_end = 0;
    if (iw_start > iw_end)
        iw_start = iw_end = 0;
    for (size_t ih = 0; ih < ih_start; ++ih)
        for (size_t iw = 0; iw < IW; ++iw) {
            int oh = ih * SH - PH;
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < ker_size; ++fh)
                if (oh + fh >= 0 && oh + fh < (int)OH)
                    for (int fw = 0; fw < ker_size; ++fw)
                        if (ow + fw >= 0 && ow + fw < (int)OW) {
                            update(oh, ow, fh, fw, val);
                        }
        }
    for (size_t ih = ih_start; ih < ih_end; ++ih) {
        int oh = ih * SH - PH;
        for (size_t iw = 0; iw < iw_start; ++iw) {
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < ker_size; ++fh)
                for (int fw = 0; fw < ker_size; ++fw) {
                    if (ow + fw >= 0 && ow + fw < (int)OW)
                        update(oh, ow, fh, fw, val);
                }
        }
        for (size_t iw = iw_start; iw < iw_end; ++iw) {
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < ker_size; ++fh)
                for (int fw = 0; fw < ker_size; ++fw) {
                    update(oh, ow, fh, fw, val);
                }
        }
        for (size_t iw = iw_end; iw < IW; ++iw) {
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < ker_size; ++fh)
                for (int fw = 0; fw < ker_size; ++fw) {
                    if (ow + fw >= 0 && ow + fw < (int)OW)
                        update(oh, ow, fh, fw, val);
                }
        }
    }
    for (size_t ih = ih_end; ih < IH; ++ih) {
        for (size_t iw = 0; iw < IW; ++iw) {
            int oh = ih * SH - PH;
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < ker_size; ++fh)
                if (oh + fh >= 0 && oh + fh < (int)OH)
                    for (int fw = 0; fw < ker_size; ++fw)
                        if (ow + fw >= 0 && ow + fw < (int)OW) {
                            update(oh, ow, fh, fw, val);
                        }
        }
    }
#undef divup
#undef update
}

void conv_backdata_single_channel_templated(
        const float* src, const float* filter, float* dst, size_t IH, size_t IW,
        size_t FH, size_t FW, size_t OH, size_t OW, size_t PH, size_t PW, size_t SH,
        size_t SW) {
    megdnn_ignore(FW);
#define DISPATCH(ker_size)                                             \
    if (FH == ker_size) {                                              \
        MIDOUT_BEGIN(megdnn_fallback_conv, ker_size) {                 \
            conv_backdata_single_channel_templated_impl<ker_size>(     \
                    src, filter, dst, IH, IW, OH, OW, PH, PW, SH, SW); \
        }                                                              \
        MIDOUT_END();                                                  \
        return;                                                        \
    }
    DISPATCH(1)
    DISPATCH(2)
    DISPATCH(3)
    DISPATCH(4)
    DISPATCH(5)
    DISPATCH(6)
    DISPATCH(7)
#undef DISPATCH
    megdnn_throw("internal error in conv_backdata template dispatching: impossible");
}

void conv_backdata_single_channel_nontemplated(
        const float* diff, const float* filter, float* grad, size_t IH, size_t IW,
        size_t FH_, size_t FW_, size_t OH, size_t OW, size_t PH, size_t PW, size_t SH,
        size_t SW) {
#define divup(x, y) (((x) + (y)-1) / (y))
#define update(oh, ow, fh, fw, val) \
    grad[(oh + fh) * OW + (ow + fw)] += filter[(fh)*FW + (fw)] * val
    int FH = FH_, FW = FW_;
    size_t ih_start = divup(PH, SH);
    size_t ih_end = OH + PH >= FH_ ? (OH + PH - FH) / SH + 1 : 0;
    size_t iw_start = divup(PW, SW);
    size_t iw_end = OW + PW >= FW_ ? (OW + PW - FW) / SW + 1 : 0;
    if (ih_start > ih_end)
        ih_start = ih_end = 0;
    if (iw_start > iw_end)
        iw_start = iw_end = 0;
    for (size_t ih = 0; ih < ih_start; ++ih)
        for (size_t iw = 0; iw < IW; ++iw) {
            int oh = ih * SH - PH;
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < FH; ++fh)
                if (oh + fh >= 0 && oh + fh < (int)OH)
                    for (int fw = 0; fw < FW; ++fw)
                        if (ow + fw >= 0 && ow + fw < (int)OW) {
                            update(oh, ow, fh, fw, val);
                        }
        }
    for (size_t ih = ih_start; ih < ih_end; ++ih) {
        int oh = ih * SH - PH;
        for (size_t iw = 0; iw < iw_start; ++iw) {
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < FH; ++fh)
                for (int fw = 0; fw < FW; ++fw) {
                    if (ow + fw >= 0 && ow + fw < (int)OW)
                        update(oh, ow, fh, fw, val);
                }
        }
        for (size_t iw = iw_start; iw < iw_end; ++iw) {
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < FH; ++fh)
                for (int fw = 0; fw < FW; ++fw) {
                    update(oh, ow, fh, fw, val);
                }
        }
        for (size_t iw = iw_end; iw < IW; ++iw) {
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < FH; ++fh)
                for (int fw = 0; fw < FW; ++fw) {
                    if (ow + fw >= 0 && ow + fw < (int)OW)
                        update(oh, ow, fh, fw, val);
                }
        }
    }
    for (size_t ih = ih_end; ih < IH; ++ih) {
        for (size_t iw = 0; iw < IW; ++iw) {
            int oh = ih * SH - PH;
            int ow = iw * SW - PW;
            float val = diff[ih * IW + iw];
            for (int fh = 0; fh < FH; ++fh)
                if (oh + fh >= 0 && oh + fh < (int)OH)
                    for (int fw = 0; fw < FW; ++fw)
                        if (ow + fw >= 0 && ow + fw < (int)OW) {
                            update(oh, ow, fh, fw, val);
                        }
        }
    }
#undef divup
#undef update
}

void conv_backdata_single_channel(
        const float* diff, const float* filter, float* grad, size_t IH, size_t IW,
        size_t FH, size_t FW, size_t OH, size_t OW, size_t PH, size_t PW, size_t SH,
        size_t SW) {
    if (can_run_xcorr_single_channel_templated(
                IH, IW, FH, FW, OH, OW, PH, PW, SH, SW)) {
        conv_backdata_single_channel_templated(
                diff, filter, grad, IH, IW, FH, FW, OH, OW, PH, PW, SH, SW);
    } else {
        MIDOUT_BEGIN(megdnn_fallback_conv, void) {
            conv_backdata_single_channel_nontemplated(
                    diff, filter, grad, IH, IW, FH, FW, OH, OW, PH, PW, SH, SW);
        }
        MIDOUT_END();
    }
}

}  // anonymous namespace

namespace megdnn {
namespace fallback {
namespace convolution {

void run_conv(
        const float* src, const float* filter, float* dst, void* workspace, size_t IH,
        size_t IW, size_t IC, size_t FH, size_t FW, size_t OH, size_t OW, size_t OC,
        size_t PH, size_t PW, size_t SH, size_t SW, bool xcorr) {
    for (size_t oc = 0; oc < OC; ++oc)
        for (size_t ic = 0; ic < IC; ++ic) {
            // ut for untransposed
            const float* fut = filter + oc * IC * FH * FW + ic * FH * FW;
            const float* f;
            if (!xcorr) {
                // need transpose
                f = (float*)workspace;
                for (size_t fh = 0; fh < FH; ++fh)
                    for (size_t fw = 0; fw < FW; ++fw) {
                        ((float*)f)[fh * FW + fw] =
                                fut[(FH - fh - 1) * FW + (FW - fw - 1)];
                    }
            } else {
                // do not need transpose
                f = fut;
            }
            run_xcorr_single_channel(
                    src + ic * IH * IW, f, dst + oc * OH * OW, IH, IW, FH, FW, OH, OW,
                    PH, PW, SH, SW, ic > 0);
        }
}

void run_conv_backward_data(
        const float* diff, const float* filter, float* grad, void* workspace, size_t IH,
        size_t IW, size_t IC, size_t FH, size_t FW, size_t OH, size_t OW, size_t OC,
        size_t PH, size_t PW, size_t SH, size_t SW, bool xcorr) {
    std::memset(grad, 0, sizeof(float) * IC * OH * OW);
    for (size_t oc = 0; oc < OC; ++oc)
        for (size_t ic = 0; ic < IC; ++ic) {
            // ut for untransposed
            const float* fut = filter + oc * IC * FH * FW + ic * FH * FW;
            const float* f;
            if (!xcorr) {
                // need transpose
                f = (float*)workspace;
                for (size_t fh = 0; fh < FH; ++fh)
                    for (size_t fw = 0; fw < FW; ++fw) {
                        ((float*)f)[fh * FW + fw] =
                                fut[(FH - fh - 1) * FW + (FW - fw - 1)];
                    }
            } else {
                // do not need transpose
                f = fut;
            }
            conv_backdata_single_channel(
                    diff + oc * IH * IW, f, grad + ic * OH * OW, IH, IW, FH, FW, OH, OW,
                    PH, PW, SH, SW);
        }
}

}  // namespace convolution
}  // namespace fallback
}  // namespace megdnn

// vim: syntax=cpp.doxygen