onednn-src 0.1.13

/*******************************************************************************
* Copyright 2018 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/
#include <cmath>

#include "common/dnnl_thread.hpp"
#include "common/utils.hpp"

#include "cpu/gemm/f32/gemm_utils_f32.hpp"

namespace dnnl {
namespace impl {
namespace cpu {
namespace gemm_utils {
#define BM_NOCOPY_AVX 64
#define BN_NOCOPY_AVX 48
#define BK_NOCOPY_AVX 384
#define BN_LARGE_NOCOPY_AVX 192
#define BM_SMALL_NOCOPY_AVX 16
#define BN_SMALL_NOCOPY_AVX 1
#define BK_SMALL_NOCOPY_AVX 4
// Determine number of threads for each dimension of a 3-D partitioning
// algorithm based on input parameters
// m/n/k - First/second/third parameter for GEMM
// nthrs - total available number of threads
// nthrs_m/nthrs_n/nthrs_k - number of threads to use in each dimension
// BM/BN/BK - blocking values
void calc_nthr_nocopy_avx(dim_t m, dim_t n, dim_t k, int nthrs, int *nthrs_m,
        int *nthrs_n, int *nthrs_k, dim_t *BM, dim_t *BN, dim_t *BK) {

    // Quick exit for single thread.
    if (nthrs == 1) {
        *nthrs_m = 1;
        *nthrs_n = 1;
        *nthrs_k = 1;

        *BM = m;
        *BN = n;
        *BK = k;
        return;
    }

    int nthr, nthr_m, nthr_n, nthr_k;
    dim_t MB, NB, KB;

    nthr = nthrs;
    nthr_m = static_cast<int>((m + BM_NOCOPY_AVX - 1) / BM_NOCOPY_AVX);
    nthr_n = static_cast<int>((n + BN_NOCOPY_AVX - 1) / BN_NOCOPY_AVX);
    nthr_k = 1;

    // Partition along K dimension
    //  - if threading allows having barriers (e.g. OMP)
    //  - if there is not enough parallelism along M or N
    if (dnnl_thr_syncable()) {
        int nthr_other = nthr_k = 1;
        while ((nthr_m * nthr_n * nthr_other < nthr)
                && (k / (nthr_other + 1) > BK_NOCOPY_AVX)) {
            nthr_other++;
            if ((nthr / nthr_other) * nthr_other > 0.9 * nthr)
                nthr_k = nthr_other;
        }
    }
    nthr /= nthr_k;

    if (nthr_m == 1) nthr_n = nthr;
    if (nthr_n == 1) nthr_m = nthr;

    // Simple partition reduction
    while (nthr_m * nthr_n > nthr)
        if (nthr_m > nthr_n)
            nthr_m--;
        else
            nthr_n--;
    while (nthr_m * nthr_n < nthr)
        if (nthr_m < nthr_n)
            nthr_m++;
        else
            nthr_n++;

    if ((nthr_m * nthr_n > nthr) && (nthr_m > 1) && (nthr_n > 1)) {

        if (nthr_m <= nthr_n) {
            nthr_m = (int)sqrt((double)nthr);
            if (nthr_m > (m + BM_SMALL_NOCOPY_AVX - 1) / BM_SMALL_NOCOPY_AVX)
                nthr_m = static_cast<int>(
                        (m + BM_SMALL_NOCOPY_AVX - 1) / BM_SMALL_NOCOPY_AVX);
            nthr_n = nthr / nthr_m;

            while ((nthr_m > 1) && (nthr_m * nthr_n != nthr)) {
                nthr_m--;
                nthr_n = nthr / nthr_m;
            }
        } else {
            nthr_n = (int)sqrt((double)nthr);
            if (nthr_n > (n + BN_SMALL_NOCOPY_AVX - 1) / BN_SMALL_NOCOPY_AVX)
                nthr_n = static_cast<int>(
                        (n + BN_SMALL_NOCOPY_AVX - 1) / BN_SMALL_NOCOPY_AVX);
            nthr_m = nthr / nthr_n;

            while ((nthr_n > 1) && (nthr_m * nthr_n != nthr)) {
                nthr_n--;
                nthr_m = nthr / nthr_n;
            }
        }
    }

    MB = (m + nthr_m - 1) / nthr_m + BM_SMALL_NOCOPY_AVX - 1;
    MB -= MB % BM_SMALL_NOCOPY_AVX;
    NB = (n + nthr_n - 1) / nthr_n + BN_SMALL_NOCOPY_AVX - 1;
    NB -= NB % BN_SMALL_NOCOPY_AVX;
    KB = (k + nthr_k - 1) / nthr_k + BK_SMALL_NOCOPY_AVX - 1;
    KB -= KB % BK_SMALL_NOCOPY_AVX;

    if (MB * nthr_m > m) nthr_m = static_cast<int>((m + MB - 1) / MB);
    if (NB * nthr_n > n) nthr_n = static_cast<int>((n + NB - 1) / NB);
    if (KB * nthr_k > k) nthr_k = static_cast<int>((k + KB - 1) / KB);

    *nthrs_m = nthr_m;
    *nthrs_n = nthr_n;
    *nthrs_k = nthr_k;

    *BM = MB;
    *BN = NB;
    *BK = KB;
}
#undef BM_NOCOPY_AVX
#undef BN_NOCOPY_AVX
#undef BK_NOCOPY_AVX
#undef BN_LARGE_NOCOPY_AVX
#undef BM_SMALL_NOCOPY_AVX
#undef BN_SMALL_NOCOPY_AVX
#undef BK_SMALL_NOCOPY_AVX

#define BM_NOCOPY_AVX512_COMMON 32
#define BN_NOCOPY_AVX512_COMMON 64
#define BK_NOCOPY_AVX512_COMMON 192
#define BN_LARGE_NOCOPY_AVX512_COMMON 192
#define BM_SMALL_NOCOPY_AVX512_COMMON 16
#define BN_SMALL_NOCOPY_AVX512_COMMON 1
#define BK_SMALL_NOCOPY_AVX512_COMMON 4
// Determine number of threads for each dimension of a 3-D partitioning
// algorithm based on input parameters
// m/n/k - First/second/third parameter for GEMM
// nthrs - total available number of threads
// nthrs_m/nthrs_n/nthrs_k - number of threads to use in each dimension
// BM/BN/BK - blocking values
void calc_nthr_nocopy_avx512_common(dim_t m, dim_t n, dim_t k, int nthrs,
        int *nthrs_m, int *nthrs_n, int *nthrs_k, dim_t *BM, dim_t *BN,
        dim_t *BK) {

    // Quick exit for single thread.
    if (nthrs == 1) {
        *nthrs_m = 1;
        *nthrs_n = 1;
        *nthrs_k = 1;

        *BM = m;
        *BN = n;
        *BK = k;
        return;
    }

    int nthr, nthr_m, nthr_n, nthr_k = 1;
    dim_t MB, NB, KB;
    nthr = nthrs;

    int counter = 0;
    float ratio_float = 1.f;
    int ratio = 1;
    nthr = nthrs;
    int nthr_m_gt_n;

    // Partition along K dimension
    //  - if threading allows having barriers (e.g. OMP)
    //  - if there is not enough parallelism along M or N
    if (dnnl_thr_syncable()) {
        if (n <= 2 * BN_NOCOPY_AVX512_COMMON
                && m <= 2 * BM_NOCOPY_AVX512_COMMON * nthr && k > m && k > n) {
            nthr_k = static_cast<int>(k / BK_NOCOPY_AVX512_COMMON);
            if (nthr_k > nthr / 4) nthr_k = nthr / 4;
            if (nthr_k < 1) nthr_k = 1;

            while ((nthr_k > 1) && (nthr % nthr_k)) {
                nthr_k--;
            }
            nthr /= nthr_k;
        } else {
            nthr_k = 1;
        }
    }
    nthr_m = static_cast<int>(
            (m + BM_NOCOPY_AVX512_COMMON - 1) / BM_NOCOPY_AVX512_COMMON);
    nthr_n = static_cast<int>(
            (n + BN_NOCOPY_AVX512_COMMON - 1) / BN_NOCOPY_AVX512_COMMON);

    if (nthr_m < 1) nthr_m = 1;
    if (nthr_n < 1) nthr_n = 1;

    nthr_m_gt_n = nthr_m > nthr_n ? 1 : 0;
    ratio_float = static_cast<float>(nthr_m) / static_cast<float>(nthr_n);

    if (nthr_m_gt_n)
        ratio = (int)ratio_float;
    else
        ratio = (int)(1.f / ratio_float);

    // scale down nthr_m and nthr_n if they are too large
    while (nthr_m * nthr_n > 4 * nthr) {
        nthr_m /= 2;
        nthr_n /= 2;
    }

    if (nthr_m < 1) nthr_m = 1;
    if (nthr_n < 1) nthr_n = 1;

    // Simple partition reduction
    counter = 0;
    while (nthr_m * nthr_n > nthr) {
        if (nthr_m > nthr_n) {
            if (counter < ratio)
                nthr_m--;
            else {
                nthr_n--;
                counter = -1;
            }
        } else {
            if (counter < ratio)
                nthr_n--;
            else {
                nthr_m--;
                counter = -1;
            }
        }
        counter++;
    }

    // Simple partition increment
    counter = 0;
    while (nthr_m * nthr_n < 0.95 * nthr) {
        if (nthr_m > nthr_n) {
            if (counter < ratio)
                nthr_m++;
            else {
                nthr_n++;
                counter = -1;
            }
        } else {
            if (counter < ratio)
                nthr_n++;
            else {
                nthr_m++;
                counter = -1;
            }
        }
        counter++;
    }

    // if nothing works out, then this should work
    if ((nthr_m * nthr_n > nthr)) {

        if (nthr_m <= nthr_n) {
            nthr_m = (int)sqrt((double)nthr);
            if (nthr_m > (m + BM_SMALL_NOCOPY_AVX512_COMMON - 1)
                            / BM_SMALL_NOCOPY_AVX512_COMMON)
                nthr_m = static_cast<int>(
                        (m + BM_SMALL_NOCOPY_AVX512_COMMON - 1)
                        / BM_SMALL_NOCOPY_AVX512_COMMON);
            nthr_n = nthr / nthr_m;

            while ((nthr_m > 1) && (nthr_m * nthr_n != nthr)) {
                nthr_m--;
                nthr_n = nthr / nthr_m;
            }
        } else {
            nthr_n = (int)sqrt((double)nthr);
            if (nthr_n > (n + BN_SMALL_NOCOPY_AVX512_COMMON - 1)
                            / BN_SMALL_NOCOPY_AVX512_COMMON)
                nthr_n = static_cast<int>(
                        (n + BN_SMALL_NOCOPY_AVX512_COMMON - 1)
                        / BN_SMALL_NOCOPY_AVX512_COMMON);
            nthr_m = nthr / nthr_n;

            while ((nthr_n > 1) && (nthr_m * nthr_n != nthr)) {
                nthr_n--;
                nthr_m = nthr / nthr_n;
            }
        }
    }

    MB = (m + nthr_m - 1) / nthr_m + BM_SMALL_NOCOPY_AVX512_COMMON - 1;
    MB -= MB % BM_SMALL_NOCOPY_AVX512_COMMON;
    NB = (n + nthr_n - 1) / nthr_n + BN_SMALL_NOCOPY_AVX512_COMMON - 1;
    NB -= NB % BN_SMALL_NOCOPY_AVX512_COMMON;
    KB = (k + nthr_k - 1) / nthr_k + BK_SMALL_NOCOPY_AVX512_COMMON - 1;
    KB -= KB % BK_SMALL_NOCOPY_AVX512_COMMON;

    if (MB * nthr_m > m) nthr_m = static_cast<int>((m + MB - 1) / MB);
    if (NB * nthr_n > n) nthr_n = static_cast<int>((n + NB - 1) / NB);
    if (KB * nthr_k > k) nthr_k = static_cast<int>((k + KB - 1) / KB);

    *nthrs_m = nthr_m;
    *nthrs_n = nthr_n;
    *nthrs_k = nthr_k;

    *BM = MB;
    *BN = NB;
    *BK = KB;
}
#undef BM_NOCOPY_AVX512_COMMON
#undef BN_NOCOPY_AVX512_COMMON
#undef BK_NOCOPY_AVX512_COMMON
#undef BN_LARGE_NOCOPY_AVX512_COMMON
#undef BM_SMALL_NOCOPY_AVX512_COMMON
#undef BN_SMALL_NOCOPY_AVX512_COMMON
#undef BK_SMALL_NOCOPY_AVX512_COMMON

// Partition n values as equally as possible among nthr threads
// and set the offset (t_offset) and number of values (t_block) for ithr
// Assumption: 0 <= ithr < nthr
void partition_unit_diff(
        int ithr, int nthr, dim_t n, dim_t *t_offset, dim_t *t_block) {

    dim_t band = n / nthr;
    if (band == 0) band = 1;
    dim_t tail = n - band * nthr;
    if (tail < 0) tail = 0;

    if (ithr < tail) {
        band++;
        *t_offset = band * ithr;
        *t_block = band;
    } else {
        *t_offset = band * ithr + tail;
        *t_block = band;
    }

    if (*t_offset >= n) {
        *t_offset = 0;
        *t_block = 0;
    }

    if (*t_offset + *t_block > n) { *t_block = n - *t_offset; }
}

// Sum the m*n values from p_src into p_dst, assuming the two-dimensional
// arrays have leading dimensions ld_src and ld_dst, respectively
template <typename data_t>
void sum_two_matrices(dim_t m, dim_t n, data_t *__restrict p_src, dim_t ld_src,
        data_t *__restrict p_dst, dim_t ld_dst) {

    for (dim_t j = 0; j < n; j++) {
        for (dim_t i = 0; i < m; i++) {
            p_dst[i + j * ld_dst] += p_src[i + j * ld_src];
        }
    }
}

template void sum_two_matrices<float>(dim_t m, dim_t n, float *__restrict p_src,
        dim_t ld_src, float *__restrict p_dst, dim_t ld_dst);

template void sum_two_matrices<double>(dim_t m, dim_t n,
        double *__restrict p_src, dim_t ld_src, double *__restrict p_dst,
        dim_t ld_dst);
} // namespace gemm_utils
} // namespace cpu
} // namespace impl
} // namespace dnnl