Вглубь еще одной кроличьей норы оптимизации

Недавно Джейк Уортон заставил меня спуститься в очередную глупую кроличью нору оптимизации, когда он во время беседы в Slack об отсутствии в Kotlin тернарного оператора беспечно сослался на кусок кода, используемый для подсчета количества цифр в Long. Это, конечно, вызвало у таких людей, как Мадис Пинк и я, желание оптимизировать его…

Подсчет цифр

Самым простым способом подсчета количества цифр было бы вычисление log10(n).toInt() + 1, где n — наше входное число. К сожалению, логарифмические функции, такие как log10 в Kotlin, определены только для чисел с плавающей точкой. Если бы мы вводили Int или Long, мы могли бы сначала преобразовать их в Double, а затем вызвать log10, но не все значения Long могут быть сохранены в Double (любое значение больше 2^53 не пройдет), и нам пришлось бы использовать на этот случай 0. Поэтому мы должны найти другое решение.

Начнем с самого простого решения:

fun simpleCount(v: Long): Int {
    if (v < 10L) return 1
    if (v < 100L) return 2
    if (v < 1000L) return 3
    if (v < 10000L) return 4
    if (v < 100000L) return 5
    if (v < 1000000L) return 6
    if (v < 10000000L) return 7
    if (v < 100000000L) return 8
    if (v < 1000000000L) return 9
    if (v < 10000000000L) return 10
    if (v < 100000000000L) return 11
    if (v < 1000000000000L) return 12
    if (v < 10000000000000L) return 13
    if (v < 100000000000000L) return 14
    if (v < 1000000000000000L) return 15
    if (v < 10000000000000000L) return 16
    if (v < 100000000000000000L) return 17
    if (v < 1000000000000000000L) return 18
    return 19
}

Мне нравится это решение из-за его очевидности. Что мне в нем не нравится, так это большое количество ветвлений и инструкций, которые оно генерирует. Я избавлю вас от полного дампа ассемблера arm64, просто знайте, что код генерирует 111 инструкций и 32 ветвления. Хотя ранние возвраты, скорее всего, избавят вас от полного выполнения сгенерированного кода, это не лучшим образом скажется ни на кэше инструкций, ни на предсказателе ветвлений.

На Pixel 6 для выполнения этого кода на большом массиве случайных входных данных требуется 63.236 нс.

Использование двоичного поиска

Решение Джейка похоже на то, которое мы только что рассмотрели, но он переставляет тесты, оптимизируя их под вероятные значения, и эффективно выполняет двоичный поиск.

fun binarySearchCount(v: Long): Int {
    return if (v < 100000000L) {
        if (v < 10000L) {
            if (v < 100L) {
                if (v < 10L) {
                    1
                } else {
                    2
                }
            } else if (v < 1000L) {
                3
            } else {
                4
            }
        } else if (v < 1000000L) {
            if (v < 100000L) {
                5
            } else {
                6
            }
        } else if (v < 10000000L) {
            7
        } else {
            8
        }
    } else if (v < 1000000000000L) {
        if (v < 10000000000L) {
            if (v < 1000000000L) {
                9
            } else {
                10
            }
        } else if (v < 100000000000L) {
            11
        } else {
            12
        }
    } else if (v < 1000000000000000L) {
        if (v < 10000000000000L) {
            13
        } else if (v < 100000000000000L) {
            14
        } else {
            15
        }
    } else if (v < 100000000000000000L) {
        if (v < 10000000000000000L) {
            16
        } else {
            17
        }
    } else if (v < 1000000000000000000L) {
        18
    } else {
        19
    }
}

Это новое решение имеет почти такое же количество инструкций, как и предыдущее (108 против 111), но только четверть ветвлений (в общей сложности 8), по сравнению с 32. Кроме того, вы гарантированно никогда не выполните все ветвления.

С более разумным способом прохождения ветвлений двоичный поиск выполняется за 27.034 нс, что составляет всего ~40% от первоначального времени выполнения. Это уже большое улучшение.

Обоснованное предположение

Сложной эту задачу делает то, что мы пытаемся работать с основанием 10, в то время как компьютеры работают с основанием 2. Однако мы можем использовать это в своих интересах, чтобы оценить, сколько цифр имеет наше число. Для этого нам нужно подсчитать количество ведущих нулей в двоичном представлении числа. Используя эту формулу:

val bitCount = 64 - n.countLeadingZeroBits()

Мы знаем, сколько битов требуется для кодирования числа. Зная это, мы можем использовать таблицу поиска, чтобы сопоставить это количество бит с возможным количеством цифр. Например, если на входе 237, формула выше говорит нам, что для кодирования значения нужно 8 бит. Отсюда следует, что нам нужно преобразовать значение в ближайшую степень 10. Это можно сделать либо с помощью другой таблицы поиска (как я сделал изначально), либо с помощью решения Мадиса:

val powerOfTen = bitCount * 3 / 10

В нашем примере с 237 возвращаемое значение равно 2, потому что ближайшая степень десяти равна 100, или 10^2. Если входные данные меньше, чем соответствующая степень десяти, мы можем вернуть powerOfTen в качестве количества цифр, а если входные данные больше, нам нужно скорректировать результат и добавить 1:

private val PowersOfTen = longArrayOf(
    0,
    10,
    100,
    1000,
    10000,
    100000,
    1000000,
    10000000,
    100000000,
    1000000000,
    10000000000,
    100000000000,
    1000000000000,
    10000000000000,
    100000000000000,
    1000000000000000,
    10000000000000000,
    100000000000000000,
    1000000000000000000
)

fun log2Count(v: Long): Int {
    val guess = ((64 - v.countLeadingZeroBits()) * 3) / 10
    return guess + (if (v >= PowersOfTen[guess]) 1 else 0)
}

Дорогостоящее целочисленное деление можно устранить с помощью еще одного трюка от Madis (который похож на то, что, как я видел, компиляторы делают автоматически):

fun log2Count(v: Long): Int {
    val guess = ((64 - v.countLeadingZeroBits()) * 10) ushr 5
    return guess + (if (v >= PowersOfTen[guess]) 1 else 0)
}

При таком решении мы получаем всего 24 инструкции и 2 ветви (одна ветвь всегда выполняется для проверки границ массива, другая — в случае провала проверки, чего в нашем случае произойти не может):

str x0, [sp, #-32]!
str lr, [sp, #24]
mov w3, #0x40
mov w2, #0xa
clz x4, x1
sub w3, w3, w4
mul w2, w3, w2
lsr w2, w2, #5
ldr w3, [x0]
ldr w0, [x3, #224]
ldr w3, [x0, #8]
cmp w2, w3
b.hs #+0x24
add w0, w0, #0x10 (16)
ldr x0, [x0, x2, lsl #3]
cmp x1, x0
cset w0, ge
add w0, w2, w0
ldr lr, [sp, #24]
add sp, sp, #0x20 (32)
ret
mov x0, x2
mov x1, x3
bl #+0x144

Это решение выполняется за 16.396 нс на Pixel 6, занимая лишь 60% времени, затраченного на двоичный поиск, и лишь 25% времени наивной реализации.

Консервативное приближение

Если вам нужно знать не точное количество цифр, а близкое приближение, которое гарантированно всегда будет больше или равно реальному ответу (например, для определения размера буфера), вы можете отказаться от всей таблицы поиска, что значительно упрощает генерируемый ассемблерный код arm64.

fun conservativeLog2Count(v: Long): Int {
    val guess = ((64 - v.countLeadingZeroBits()) * 10) ushr 5
    return guess + 1
}

Полученный код aarch64 содержит всего 8 инструкций и не имеет ветвлений:

mov w3, #0x40
mov w2, #0x4d1
mov w0, #0x1
clz x1, x1
sub w1, w3, w1
mul w1, w1, w2
add w0, w0, w1, lsr #12
ret

Подсчет

Вот окончательное сравнение между различными реализациями:

Удручающе — но неудивительно — реализация log2Count на C++ выдает 12 инструкций вместо 24 в Kotlin, и никаких ветвлений. Это, конечно, связано с тем, что C++ не делает для нас никакой проверки границ массива. Я не пробовал тестировать это решение с помощью JNI и @CriticalNative, но сомневаюсь, что экономия будет больше, чем накладные расходы на JNI.

Источник