【C++】SIMD+OpenMPによる2値化処理の高速化【OpenCV比較】

#include <iostream>

#ifdef _OPENMP
#include <omp.h>
#endif

int main(void){
#ifdef _OPENMP
#pragma omp parallel
    {
        int tid = omp_get_thread_num();
        int nthreads = omp_get_num_threads();

#pragma omp critical
        {
            std::cout << "Hello from thread "
                << tid << " / " << nthreads << std::endl;
        }
    }
#else
    std::cout << "OpenMP is NOT enabled" << std::endl;
#endif

    return 0;
}

Hello from thread 0 / 20
Hello from thread 1 / 20
Hello from thread 2 / 20
Hello from thread 9 / 20
Hello from thread 3 / 20
Hello from thread 4 / 20
Hello from thread 5 / 20
Hello from thread 6 / 20
Hello from thread 7 / 20
Hello from thread 8 / 20
Hello from thread 10 / 20
Hello from thread 11 / 20
Hello from thread 12 / 20
Hello from thread 13 / 20
Hello from thread 14 / 20
Hello from thread 15 / 20
Hello from thread 16 / 20
Hello from thread 17 / 20
Hello from thread 19 / 20
Hello from thread 18 / 20

#if defined(_MSC_VER)
#include <intrin.h>
#endif

#if defined(__AVX2__) || defined(_M_AVX2)
#include <immintrin.h>
#elif defined(__SSE2__) || defined(_M_X64) || (defined(_M_IX86_FP) && _M_IX86_FP >= 2)
#include <emmintrin.h>
#endif

#if defined(__ARM_NEON) || defined(_M_ARM64)
#include <arm_neon.h>
#endif

struct v_uint8
{
#if defined(__AVX2__) || defined(_M_AVX2)
    static constexpr int lanes = 32;
    __m256i val;
    v_uint8() = default;
    explicit v_uint8(__m256i v) : val(v) {}
#elif defined(__SSE2__) || defined(_M_X64) || (defined(_M_IX86_FP) && _M_IX86_FP >= 2)
    static constexpr int lanes = 16;
    __m128i val;
    v_uint8() = default;
    explicit v_uint8(__m128i v) : val(v) {}
#elif defined(__ARM_NEON) || defined(_M_ARM64)
    static constexpr int lanes = 16;
    uint8x16_t val;
    v_uint8() = default;
    explicit v_uint8(uint8x16_t v) : val(v) {}
#else
    static constexpr int lanes = 16;
    uchar val[lanes];
#endif
};

static inline v_uint8 vx_setall_u8(uchar x)
{
#if defined(__AVX2__) || defined(_M_AVX2)
    // 32 lanes に同じ 8bit 値をブロードキャスト
    return v_uint8(_mm256_set1_epi8((char)x));

#elif defined(__SSE2__) || defined(_M_X64) || (defined(_M_IX86_FP) && _M_IX86_FP >= 2)
    // 16 lanes
    return v_uint8(_mm_set1_epi8((char)x));

#elif defined(__ARM_NEON) || defined(_M_ARM64)
    // 16 lanes
    return v_uint8(vdupq_n_u8(x));

#else
    // Portable fallback
    v_uint8 r;
    for (int i = 0; i < v_uint8::lanes; ++i)
        r.val[i] = x;
    return r;
#endif
}

static inline v_uint8 v_lt(const v_uint8& a, const v_uint8& b)
{
#if defined(__AVX2__) || defined(_M_AVX2)
    // x86の8bit比較は signed なので、unsigned比較にするため 0x80 を XOR してレンジをずらす
    const __m256i bias = _mm256_set1_epi8((char)0x80);
    __m256i ax = _mm256_xor_si256(a.val, bias);
    __m256i bx = _mm256_xor_si256(b.val, bias);

    // (a < b) は (b > a) と同値。signed compare を使う
    __m256i mask = _mm256_cmpgt_epi8(bx, ax); // 0xFF or 0x00 per lane
    return v_uint8(mask);

#elif defined(__SSE2__) || defined(_M_X64) || (defined(_M_IX86_FP) && _M_IX86_FP >= 2)
    const __m128i bias = _mm_set1_epi8((char)0x80);
    __m128i ax = _mm_xor_si128(a.val, bias);
    __m128i bx = _mm_xor_si128(b.val, bias);

    __m128i mask = _mm_cmpgt_epi8(bx, ax); // 0xFF or 0x00 per lane
    return v_uint8(mask);

#elif defined(__ARM_NEON) || defined(_M_ARM64)
    // NEON は unsigned compare が素直に使える（a < b）
    // 環境によって vcltq_u8 が無い場合は vcgtq_u8(b,a) で代替可能
#if defined(__aarch64__) || defined(_M_ARM64)
    uint8x16_t mask = vcltq_u8(a.val, b.val); // 0xFF or 0x00
    return v_uint8(mask);
#else
    uint8x16_t mask = vcgtq_u8(b.val, a.val); // (a < b) == (b > a)
    return v_uint8(mask);
#endif

#else
    v_uint8 r;
    for (int i = 0; i < v_uint8::lanes; ++i)
        r.val[i] = (a.val[i] < b.val[i]) ? 0xFF : 0x00;
    return r;
#endif
}

static inline v_uint8 v_and(const v_uint8& a, const v_uint8& b)
{
#if defined(__AVX2__) || defined(_M_AVX2)
    return v_uint8(_mm256_and_si256(a.val, b.val));

#elif defined(__SSE2__) || defined(_M_X64) || (defined(_M_IX86_FP) && _M_IX86_FP >= 2)
    return v_uint8(_mm_and_si128(a.val, b.val));

#elif defined(__ARM_NEON) || defined(_M_ARM64)
    return v_uint8(vandq_u8(a.val, b.val));

#else
    v_uint8 r;
    for (int i = 0; i < v_uint8::lanes; ++i)
        r.val[i] = (uchar)(a.val[i] & b.val[i]);
    return r;
#endif
}

static inline void v_store(uchar* p, const v_uint8& a)
{
#if defined(__AVX2__) || defined(_M_AVX2)
    _mm256_storeu_si256(reinterpret_cast<__m256i*>(p), a.val);

#elif defined(__SSE2__) || defined(_M_X64) || (defined(_M_IX86_FP) && _M_IX86_FP >= 2)
    _mm_storeu_si128(reinterpret_cast<__m128i*>(p), a.val);

#elif defined(__ARM_NEON) || defined(_M_ARM64)
    vst1q_u8(p, a.val);

#else
    std::memcpy(p, a.val, v_uint8::lanes);
#endif
}

template<typename V>
struct VTraits;   // primary template（未定義）

template<>
struct VTraits<v_uint8>
{
    using lane_type = unsigned char;

    static inline int vlanes()
    {
        return v_uint8::lanes;
    }
};

struct Range {
    int start;
    int end;
};

static inline Range computeRange(
    int id,          // 分割番号（sr.start）
    int id_end,      // sr.end（通常は id+1）
    int wholeStart,  // wholeRange.start
    int wholeEnd,    // wholeRange.end
    int nstripes     // 分割数
) {
    const int len = wholeEnd - wholeStart;

    Range r;
    r.start = wholeStart + (int)(((uint64_t)id * len + nstripes / 2) / nstripes);

    if (id_end >= nstripes) r.end = wholeEnd;
    else                    r.end = wholeStart + (int)(((uint64_t)id_end * len + nstripes / 2) / nstripes);

    return r;
}

void parallel_threshold(cv::Mat& mat_src, cv::Mat& mat_dst, uchar thresh, uchar maxval, int id, int nstripes) {
    Range r = computeRange(id, id + 1, 0, mat_src.rows, nstripes);

    cv::Mat srcStripe = mat_src.rowRange(r.start, r.end);
    cv::Mat dstStripe = mat_dst.rowRange(r.start, r.end);

    const int lanes = VTraits<v_uint8>::vlanes();

    v_uint8 thresh_u = vx_setall_u8(thresh);
    v_uint8 maxval_u = vx_setall_u8(maxval);

    if (srcStripe.isContinuous() && dstStripe.isContinuous())
    {
        const int total = (int)srcStripe.total(); // 画素数
        const uchar* src = srcStripe.ptr<uchar>(0);
        uchar* dst = dstStripe.ptr<uchar>(0);

        int j = 0;
        for (; j <= total - lanes; j += lanes)
        {
            v_uint8 v0 = vx_load(src + j);
            v0 = v_lt(thresh_u, v0);
            v0 = v_and(v0, maxval_u);
            v_store(dst + j, v0);
        }
        // 端数
        for (; j < total; ++j)
            dst[j] = (src[j] > thresh) ? maxval : 0;

        return;
    }

    for (int i = 0; i < srcStripe.rows; ++i)
    {
        const uchar* src = srcStripe.ptr<uchar>(i);
        uchar* dst = dstStripe.ptr<uchar>(i);

        int j = 0;
        for (; j <= srcStripe.cols - lanes; j += lanes)
        {
            v_uint8 v0 = vx_load(src + j);
            v0 = v_lt(thresh_u, v0);
            v0 = v_and(v0, maxval_u);
            v_store(dst + j, v0);
        }
        for (; j < srcStripe.cols; ++j)
            dst[j] = (src[j] > thresh) ? maxval : 0;
    }
}

void custom_threshold(cv::Mat& mat_src, cv::Mat& mat_dst, uchar thresh, uchar maxval) {
    mat_dst.create(mat_src.size(), mat_src.type());

    int nstripes = (int)std::max(1.0, mat_dst.total() / (double)(1 << 16));

#pragma omp parallel for schedule(static)
    for (int i = 0; i < nstripes; i++) parallel_threshold(mat_src, mat_dst, thresh, maxval, i, nstripes);
}

#include <iostream>
#include "main.h"
#include <opencv2/opencv.hpp>

int main(void) {
	cv::Mat mat_img = cv::imread("shimons_labo.png");
	cv::Mat mat_gray;
	cv::Mat mat_bin[2];
	uchar thresh = 128;
	uchar maxval = 255;

	// AVX2が有効か確認
#if defined(__AVX2__)
	std::cout << "AVX2 enabled\n";
#elif defined(__SSE2__)
	std::cout << "SSE2 enabled\n";
#else
	std::cout << "No SIMD\n";
#endif

	// グレースケール変換
	cv::cvtColor(mat_img, mat_gray, cv::COLOR_BGR2GRAY);

	// 2値化処理(OpenMP)
	custom_threshold(mat_gray, mat_bin[0], thresh, maxval);
	cv::threshold(mat_gray, mat_bin[1], thresh, maxval, cv::THRESH_BINARY);

	cv::Mat mat_diff;
	cv::compare(mat_bin[0], mat_bin[1], mat_diff, cv::CMP_NE);   // 不一致 = 255
	if (cv::countNonZero(mat_diff) == 0)	std::cout << "一致しています\n";
	else									std::cout << "一致していません\n";

	cv::imshow("OpenMP", mat_bin[0]);
	cv::imshow("OpenCV", mat_bin[1]);
	cv::waitKey(0);

	return 0;
}

AVX2 enabled
一致しています