Compiler Explorer

Source code

#include <cstdint>
#include <string_view>
#include <immintrin.h>

std::uint32_t levenshtein_diagonal_avx512(std::string_view str0, std::string_view str1) {
    // Ensure that str0 is the longest of the strings. Allow logic to be simplified elsewhere
    if (str0.size() < str1.size()) {
        std::swap(str0, str1);
    }

/*
    if (255 < str1.size()) {
        std::cerr << "Inputs strings cannot be longer than 255 characters." << std::endl;
        return -1;
    }
    */

std::int32_t diagonal_count = 1 + str0.size() + str1.size();

std::int32_t max_diagonal_width = 1 + str1.size();

/*
    if (255 < max_diagonal_width) {
        std::cerr << "Input strings are collectively too long." << std::endl;
        return -1;
    }
    */

// Under the constraints on the lengths of the input strings we know that this will always suffice
    alignas(512) std::uint8_t diagonal_buffers[2][256]{};
    alignas(512) std::uint8_t* diagonal0 = diagonal_buffers[0];
    alignas(512) std::uint8_t* diagonal1 = diagonal_buffers[1];

std::uint8_t debug_buffer[256]{};

// Table containing indices to reverse bytes in vector
    alignas(64) constexpr std::uint8_t reverse_indices_data[64] {
        0x3f, 0x3e, 0x3d, 0x3c, 0x3b, 0x3a, 0x39, 0x38,
        0x37, 0x36, 0x35, 0x34, 0x33, 0x32, 0x31, 0x30,
        0x2f, 0x2e, 0x2d, 0x2c, 0x2b, 0x2a, 0x29, 0x28,
        0x27, 0x26, 0x25, 0x24, 0x23, 0x22, 0x21, 0x20,
        0x1f, 0x1e, 0x1d, 0x1c, 0x1b, 0x1a, 0x19, 0x18,
        0x17, 0x16, 0x15, 0x14, 0x13, 0x12, 0x11, 0x10,
        0x0f, 0x0e, 0x0d, 0x0c, 0x0b, 0x0a, 0x09, 0x08,
        0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01, 0x00
    };

// Vector containing indices to reverse bytes in vector
    __m512i reverse_indices = _mm512_load_si512(reverse_indices_data);

constexpr std::int32_t vector_width = 64;

// Return the maximum length of the input arrays if at least one string is empty
    // Necessary for subsequent optimization
    if (str1.empty()) {
        return str0.size();
    }

// Handle base cases for second iteration of exterior loop in scalar fashion to reduce overall iteration count by 1
    diagonal0[max_diagonal_width - 2] = 1;
    diagonal0[max_diagonal_width - 1] = 1;

// First iteration can be skipped since the first entry is always zero
    for (std::int32_t i = 2; i < diagonal_count; ++i) {
        // Index of first non-base-case in diagonal
        std::int32_t j_begin = std::max(max_diagonal_width - i, 0);

// Index one-past the last non-base-case in diagonal
        std::int32_t j_end = std::min(max_diagonal_width - 1, diagonal_count - i);

//Handle base case, first entry in diagonal
        if (i <= str1.size()) {
            diagonal1[j_begin - 1] = i;
        }

std::int32_t diagonal_width = (j_end - j_begin);

std::int32_t whole_iterations = (diagonal_width / vector_width) * vector_width;
        std::int32_t iteration_remainder = diagonal_width % vector_width;

std::int32_t j = 0;
        for (; j < whole_iterations; j += vector_width) {
            // Read values of other cells
            __m512i A = _mm512_loadu_epi8(diagonal0 + j_begin + j + 1);
            __m512i B = _mm512_loadu_epi8(diagonal1 + j_begin + j + 1);
            __m512i C = _mm512_loadu_epi8(diagonal0 + j_begin + j + 0);

// Compute cost for current cell under assumption that characters aren't equal
            __m512i construction_cost = _mm512_add_epi8(_mm512_set1_epi8(1), _mm512_min_epu8(A, _mm512_min_epu8(B, C)));

// Load vector of characters from first string
            std::int64_t char_load_offset0 = i - max_diagonal_width + j + j_begin;
            __m512i ch0 = _mm512_loadu_epi8(str0.data() + char_load_offset0);

// Load vectors of characters from second string
            std::int64_t char_load_offset1 = std::min(i - 1, std::int32_t(str1.size())) - j - vector_width;
            __m512i ch1_data = _mm512_loadu_epi8(str1.data() + char_load_offset1);
            __m512i ch1 = _mm512_permutexvar_epi8(reverse_indices, ch1_data);

// Compute final cost, factoring in equality of characters
            __mmask64 are_characters_equal = _mm512_cmpeq_epi8_mask(ch0, ch1);
            __m512i cost = _mm512_mask_blend_epi8(are_characters_equal, construction_cost, B);

// Write current vector results out
            _mm512_storeu_epi8(diagonal1 + j + j_begin, cost);
        }

{
            // Indicates which lanes are used by the current vector
            // Ternary used to avoid UB with shifting by an amount equal to integer width
            std::uint64_t active_lane_mask = (std::uint64_t(1) << iteration_remainder) - 1;

// Read values of other cells
            __m512i A = _mm512_maskz_loadu_epi8(active_lane_mask, diagonal0 + j_begin + j + 1);
            __m512i B = _mm512_maskz_loadu_epi8(active_lane_mask, diagonal1 + j_begin + j + 1);
            __m512i C = _mm512_maskz_loadu_epi8(active_lane_mask, diagonal0 + j_begin + j + 0);

// Compute cost for current cell under assumption that characters aren't equal
            __m512i construction_cost = _mm512_add_epi8(_mm512_set1_epi8(1), _mm512_min_epu8(A, _mm512_min_epu8(B, C)));

// Load vector of characters from first string
            std::int64_t char_load_offset0 = i - max_diagonal_width + j + j_begin;
            __m512i ch0 = _mm512_maskz_loadu_epi8(active_lane_mask, str0.data() + char_load_offset0);

// Load vectors of characters from second string
            __mmask64 ch1_load_mask = active_lane_mask << (vector_width - iteration_remainder);
            std::int64_t char_load_offset1 = std::min(i - 1, std::int32_t(str1.size())) - j - vector_width;
            __m512i ch1_data = _mm512_maskz_loadu_epi8(ch1_load_mask, str1.data() + char_load_offset1);
            __m512i ch1 = _mm512_permutexvar_epi8(reverse_indices, ch1_data);

// Write current vector results out
            _mm512_mask_storeu_epi8(diagonal1 + j + j_begin, active_lane_mask, cost);
        }

//Handle base case, last entry in diagonal
        if (i <= str0.size()) {
            diagonal1[j_end] = i;
        }

std::swap(diagonal0, diagonal1);
    }

return diagonal0[0];
}