Как LLVM оптимизирует суммы степеней

LLVM оптимизирует суммы степеней, например:

int sum(int count)
{
  int result = 0;

  for (int j = 0; j < count; ++j)
    result += j*j;

  return result;
}

генерируя код, вычисляющий результат без цикла (godbolt):

sum(int):
        test    edi, edi
        jle     .LBB0_1
        lea     eax, [rdi - 1]
        lea     ecx, [rdi - 2]
        imul    rcx, rax
        lea     eax, [rdi - 3]
        imul    rax, rcx
        shr     rax
        imul    eax, eax, 1431655766
        add     eax, edi
        shr     rcx
        lea     ecx, [rcx + 2*rcx]
        lea     eax, [rax + rcx]
        add     eax, -1
        ret
.LBB0_1:
        xor     eax, eax
        ret

Также обрабатываются более сложные случаи (godbolt) – то есть оптимизация здесь не просто сравнивает паттерны. В этом посте мы рассмотрим, как выполняется эта оптимизация.

Анализ циклов — скалярное развёртывание

Есть много случаев, когда компилятору нужно отслеживать, как значение изменяется внутри цикла. Например, векторизатор цикла должен проверить, что указатели перемещаются на следующий элемент на новой итерации, и проверяет, что никакой другой указатель не ссылается на векторизируемый диапазон.

И GCC, и LLVM делают это сходным способом, в проходах «scalar evolution» (я предпочел не переводить такие термины во избежание потери смысла — прим перев.), в которых каждая переменная на итерации i (мы начинаем отсчитывать итерации с 0) представлена как функция , представленная как линейная рекуррентная форма

где

Пример 1

Рассмотрим простейший пример цикла:

void foo(int m, int *p)
{
  for (int j = 0; j < m; j++)
    *p++ = j;
}

Цикл записывает 0 в *p++ на первой итерации, 1 на второй, и т. д. Итак, мы можем выразить значение, записанное на итерации i как

Пример 2

Полиномы также могут быть выражены в этой форме.

void foo(int m, int k, int *p)
{
  for (int j = 0; < m; j++)
    *p++ = j*j*j - 2*j*j + k*j + 7;
}

Мы увидим ниже, как построить функции, сейчас приведём результат построений для значения, сохранённого в цикле:

Одну оптимизацию мы можем видеть напрямую из этих функций, она заключается в том, что значение может быть вычислено за три сложения в цикле

void foo(int m, int k, int *p)
{
  int t0 = 7;
  int t1 = k-1;
  int t2 = 2;
  for (int j = 0; j < m; j++) {
    *p++ = t0;
    t0 = t0 + t1;
    t1 = t1 + t2;
    t2 = t2 + 6;
  }
}

, что является полезной оптимизацией для архитектур, в которых умножение является дорогостоящим. Код такого вида, однако, не является общепринятым, и большинство компиляторов не выполняет такую оптимизацию, но они делают её для более простых случаев, таких как

void foo(int m, int k, int *p)
{
  for (int j = 0; < m; j++)
    *p++ = k*j + 7;
}

так как конструкции вида k*j+7 являются распространёнными в вычислениях адреса.

Рекуррентные цепи

Громоздко каждый раз писать рекурсивные функции, поэтому функции обычно пишутся в форме . Например:

Эти функции можно объединить в цепочку, и может быть записана как рекуррентная цепь (chain of recurrences, CR) . Внутренние фигурные скобки избыточны, и CR обычно записывается как кортеж .

Построение реккурентных цепей

Рекуррентные цепи строятся путём итераций над кодом и вычисления результирующего CR для каждой операции (или маркирования неизвестным результатом, если мы не можем обработать операцию), используя правила упрощения:

Итак, для цикла в функции sum:

for (int j = 0; j < count; ++j)
  result += j*j;

мы начинаем с j для которой известна CR из примера 1. Затем она используется как j*j, когда мы вычисляем result, и мы можем вычислить CR для j*j, используя правила упрощения:

Сходные вычисления для result даёт нам CR после добавления j*j.

Выполняем оптимизации

Оптимизация выполняется как упрощение по индукции (induction variable simplification), и LLVM преобразует функцию в форму, удобную для анализа и оптимизации

int sum(int count)
{
  int result = 0;

  if (count > 0) {
    int j = 0;
    do {
      result = result + j*j;
      ++j;
    } while (j < count);
  }

  return result;
}

или, как это выглядит в LLVM IR:

define i32 @sum(i32) {
%2 = icmp sgt i32 %0, 0
br i1 %2, label %3, label %6
; <label>:3:
br label %8
; <label>:4:
%5 = phi i32 [ %12, %8 ]
 br label %6
; <label>:6:
%7 = phi i32 [ 0, %1 ], [ %5, %4 ]
 ret i32 %7
; <label>:8:
%9 = phi i32 [ %13, %8 ], [ 0, %3 ]     ; {0,+,1}
%10 = phi i32 [ %12, %8 ], [ 0, %3 ]    ; {0,+,0,+,1,+,2}
%11 = mul nsw i32 %9, %9                ; {0,+,1,+,2}
%12 = add nuw nsw i32 %11, %10          ; {0,+,1,+,3,+,2}
%13 = add nuw nsw i32 %9, 1             ; {1,+,1}
%14 = icmp slt i32 %13, %0
br i1 %14, label %8, label %4
}

Компилятор может видеть, что функция возвращает 0, если count <= 0, иначе возвращает результат цикла loop iteration count-1.

Приятное свойство рекуррентной цепи состоит в том, что легко вычислить значение определённой итерации: если мы знаем CR:, тогда значение итерации может быть вычислено как:
\begin{align}f(i) & = \sum_{j=0}^{n}\phi_j{i \choose j} \\ & = \phi_0 + \phi_1i + \phi_2{i(i-1)\over 2!} + \ldots + \phi_n{i(i-1)\cdots(i-n+1)\over n!}\end{align}
Подставляя значения для CR , описывающие result, получаем

Компилятору сейчас нужно подставить код, который вычисляет значение для
count-1, после цикла

result = count-1 + 3*(count-1)*(count-2)/2 + (count-1)*(count-2)(count-3)/3;

но нужна некоторая осторожность, при вычислениях может потеряться точность (временные значения могут не помещаться в 32-битные целые). Деление целых — медленная операция, и мы делаем некоторый трюк с заменой деления на умножение и сдвиги. Результат в LLVM IR

%4 = add i32 %0, -1
  %5 = zext i32 %4 to i33
  %6 = add i32 %0, -2
  %7 = zext i32 %6 to i33
  %8 = mul i33 %5, %7
  %9 = add i32 %0, -3
  %10 = zext i32 %9 to i33
  %11 = mul i33 %8, %10
  %12 = lshr i33 %11, 1
  %13 = trunc i33 %12 to i32
  %14 = mul i32 %13, 1431655766
  %15 = add i32 %14, %0
  %16 = lshr i33 %8, 1
  %17 = trunc i33 %16 to i32
  %18 = mul i32 %17, 3
  %19 = add i32 %15, %18
  %20 = add i32 %19, -1

Вставка этого кода делает цикл мёртвым, и позже он удаляется проходом удаления мёртвого кода (dead code elimination), и мы, наконец, получаем код

sum(int):
        test    edi, edi
        jle     .LBB0_1
        lea     eax, [rdi - 1]
        lea     ecx, [rdi - 2]
        imul    rcx, rax
        lea     eax, [rdi - 3]
        imul    rax, rcx
        shr     rax
        imul    eax, eax, 1431655766
        add     eax, edi
        shr     rcx
        lea     ecx, [rcx + 2*rcx]
        lea     eax, [rax + rcx]
        add     eax, -1
        ret
.LBB0_1:
        xor     eax, eax
        ret

Производительность

Эта оптимизация не всегда выгодна. Например,

int sum(int count)
{
  int result = 0;

  for (int j = 0; j < count; ++j)
    result += j*j*j*j*j*j;

  return result;
}

вычисляет три 32-битных умножения и одно сложение за цикл, а оптимизированная версия требует шесть 64-битных умножений, пять 32-битных умножений, и другие инструкции (godbolt), и оптимизированная версия выполняется медленнее для малых значений цикла. На маленьких CPU с, например, более дорогостоящим 64-битным умножением, значение числа циклов, при которых оптимизация будет полезна, будет больше, чем на обычных CPU. Для CPU, которые не имеют инструкций для 64-битного умножения, это значение будет ещё больше (godbolt).

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

/* Do not emit expensive expressions.  The rationale is that
   when someone writes a code like

   while (n > 45) n -= 45;

   he probably knows that n is not large, and does not want it
   to be turned into n %= 45.  */
|| expression_expensive_p (def))

Если GCC не выполнил оптимизацию, это не баг, это фича.

Литература:

Рекуррентные цепи:

1. Olaf Bachmann, Paul S. Wang, Eugene V. Zima. “Chains of recurrences – a method to expedite the evaluation of closed-form functions”
2. Eugene V. Zima. “On computational properties of chains of recurrences”

Цикловые оптимизации, использующие рекуррентные цепи:

1. Robert A. van Engelen. “Symbolic Evaluation of Chains of Recurrences for Loop Optimization”
2. Robert A. van Engelen. “Efficient Symbolic Analysis for Optimizing Compilers”

Оптимизация деления с использованием инструкций умножения и сдвига:

1. Torbjörn Granlund, Peter L. Montgomery. “Division by Invariant Integers using Multiplication”