Поговорим об оптимизирующих компиляторах. Сказ восьмой: размотка циклов

Это цикл статей об оптимизирующих компиляторах вообще и LLVM в частности. Смотри все статьи данного цикла:

1. SSA форма

2. Доминирование

3. Неопределённое поведение

4. Циклы

5. CSE

6. Цикловые инварианты

7. Проверки диапазонов

8. Размотка циклов

В этой статье мы будем много разматывать циклы. Не так весело и задорно, как Виталий Наливкин, но тоже неплохо.

Не все фломастеры одинаково полезны

Есть оптимизации, польза от которых очевидна всегда или почти всегда. Например, не делать лишнюю проверку лучше, чем делать. Не считать два раза одно и то же обычно лучше, чем считать (если только мы не упёрлись в нехватку регистров или имеем другие подобные проблемы на нижнем уровне). Вычислять выражения вне цикла выгоднее, чем в цикле. И так далее.

Но есть оптимизации, применение которых имеет как плюсы, так и минусы. Выиграв в одном месте, мы можем получить отрицательные эффекты в другом. Например, сэкономив на количестве проверок, мы можем раздуть общий объём кода и поломать микрооптимизации. Каноничным примером такой оптимизации, решение вопроса об использовании которой больше похоже на искусство, чем на науку, является размотка циклов (Loop Unrolling), о которой мы сегодня поговорим. В статье я попробую осветить как можно больше (хотя, наверное, и не все) соображений о том, почему эту оптимизацию может быть нужно или не нужно применять.

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

Полная размотка циклов

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

int foo_before_full_unroll(int *a, int x) {
  int sum = 0;
  for (int i = 0; i < 5; i++) {
    sum += a[i];
    sum += 2 * x;
    if (i == 3) {
      sum *= 2;
    }
  }
  return sum;
}

int foo_after_full_unroll(int *a, int x) {
  int sum = 0;
  sum += a[0];
  sum += 2 * x;
  if (0 == 3) {
    sum *= 2;
  }
  sum += a[1];
  sum += 2 * x;
  if (1 == 3) {
    sum *= 2;
  }
  sum += a[2];
  sum += 2 * x;
  if (2 == 3) {
    sum *= 2;
  }
  sum += a[3];
  sum += 2 * x;
  if (3 == 3) {
    sum *= 2;
  }
  sum += a[4];
  sum += 2 * x;
  if (4 == 3) {
    sum *= 2;
  }
  return sum;
}

Мы видим, что код стал сильно раздутым, и ни один нормальный человек бы так писать не стал. Однако уже сама по себе такая трансформация принесла нам некоторую пользу. Если мы будем исполнять данный код "в лоб", то мы сэкономили целых 2 операции на каждой итерации цикла: нам больше не нужно всякий раз увеличивать переменную i и сравнивать её значение с 5. Казалось бы, это немного, но если тело цикла и так было маленькое, то увеличение счётчика и сравнение могли занимать какой-то заметный процент общего времени.

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

• Выражение 2 * x можно вычислить всего один раз, а не на каждой итерации цикла. Для этого нужно провести удаление общих подвыражений. Альтернативным способом добиться того же результата мог бы быть вынос инвариантов. То, что два разных набора оптимизаций приводят к одному результату, при построении компиляторов бывает часто, и это считается нормальным: всякая оптимизация имеет свои ограничения, и некоторая избыточность позволяет повысить наши шансы на успех.

• Проверка if (i == 3), которую раньше приходилось исполнять на каждой итерации, теперь вообще может быть удалена. У нас больше нет i, а сравнения констант между собой тривиальным образом превращаются в true или false, после чего от ветвления можно избавиться.

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

int foo_after_full_unroll_optimized_1(int *a, int x) {
  int sum = 0;
  int x2 = 2 * x;
  sum += a[0];
  sum += x2;
  sum += a[1];
  sum += x2;
  sum += a[2];
  sum += x2;
  sum += a[3];
  sum += x2;
  sum *= 2;
  sum += a[4];
  sum += x2
  return sum;
}

Теперь код стал уж совсем простым, и теперь мы можем сделать ещё несколько интересных оптимизаций:

• Вместо того, чтобы вначале добавлять x2 четыре раза, можно умножить его на четыре и добавить один раз (умножение делается сдвигом, а потому получается дешевле). Замечание: мы не можем "в лоб" поступить так с последним добавлением x2, поскольку перед ним значение sum умножается на 2. Наверное, можно учесть и это (например, добавить сразу x2 * 9), но мы таких сложных преобразований делать не будем.

• Добавление элементов a[0]-a[3] -- это довольно классическая векторная операция (редукция), и на векторных регистрах также может быть выполнена быстро (если соответствующая инструкция поддерживается вашим процессором).

Итого имеем следующее:

int foo_after_full_unroll_optimized_2(int *a, int x) {
  int sum = 0;
  int x2 = 2 * x;
  int x8 = x2 * 4;
  sum += x8;
  sum += vector_add_reduce(a[0], a[3]);
  sum *= 2;
  sum += a[4];
  sum += x2;
  return sum;
}

Итак, плюсы от применения полной размотки циклов:

• Общее количество исполненных операций уменьшается за счёт удаления цикловых инкрементов и операций сравнения;

• После размотки могут образоваться общие подвыражения, которые можно сократить;

• Некоторые вариантные условия становятся тривиальными, за счёт чего упростится CFG;

• Последовательные доступы к памяти могут быть выполнены быстрее за счёт векторизации;

• За счёт общего сокращения числа ветвлений упростится работа предсказателя переходов;

Но есть ли минусы? Увы, но да. Давайте для справедливости перечислим и их.

• Увеличивается общий объём кода, и как следствие - время компиляции для последующих оптимизаций. Это особенно критично в огромных проектах или в JIT-компилируемых языках. Также вероятно увеличение размера бинарника (насколько это критично -- тут по-разному; на телефонах его стараются экономить);

• Выигрыш от удаления ветвлений может оказаться маленьким или вообще нулевым, если предсказатель перехода в современном процессоре и так справляется с этим паттерном;

• Может возрасти общая стоимость декодирования инструкций процессором при исполнении. Например, если раньше цикл был маленький и состоял всего из нескольких инструкций, современные процессоры могли декодировать его всего один раз и сразу исполнять, не тратя время на декодирование каждой итерации. В размотанном цикле эти инструкции реально становятся разными, и их все придётся декодировать;

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

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

Частичная размотка циклов

Гораздо чаще возникают ситуации, когда точное число итераций в цикле неизвестно, либо слишком велико, чтобы применять полную размотку. В этом случае размотка может быть частичной. Для этого выбирается некоторое число N, которое называется фактором размотки (обычно, но не обязательно, в качестве N выступает какая-нибудь небольшая степень двойки, вроде 4 или 32). Для частичной размотки нужно, чтобы число итераций в цикле делилось на N. Тогда можно раскопировать тело цикла N раз и сократить число итераций в N раз. Приведём пример размотки с фактором 4:

int foo_before_partial_unroll(int *a) {
  int sum = 0;
  for (int i = 0; i < 1024; i++) {
    sum += a[i];
    sum += i / 4;
    if (i % 16 == 0) {
      sum *= 2;
    }
  }
  return sum;
}

int foo_after_partial_unroll_factor_4(int *a) {
  int sum = 0;
  for (int i = 0; i < 1024; i += 4) {
    sum += a[i];
    sum += i / 4;
    if (i % 16 == 0) {
      sum *= 2;
    }
    sum += a[i + 1];
    sum += (i + 1) / 4;
    if ((i + 1) % 16 == 0) {
      sum *= 2;
    }
    sum += a[i + 2];
    sum += (i + 2) / 4;
    if ((i + 2) % 16 == 0) {
      sum *= 2;
    }
    sum += a[i + 3];
    sum += (i + 3) / 4;
    if ((i + 3) % 16 == 0) {
      sum *= 2;
    }
  }
  return sum;
}

Итак, мы чисто механически раскопировали тело цикла. Как видно, теперь полностью избавиться от счётчика i не удалось, однако количество его сравнений с пороговым значением сократилось вчетверо. Также тут появились общие подвыражения, хоть это не так и очевидно. Поскольку i делится на 4, то i / 4 == (i + 1) / 4 == (i + 2) / 4 == (i + 3) / 4, и компилятор в теории может это доказать. Интересна также судьба условия if. Оно здесь не такое тривиальное, как в прошлом примере, однако можно заметить следующее: поскольку i всегда делится на 4, то выражения i + 1, i + 2 и i + 3 не делятся ни на 4, ни тем более на 16. Таким образом, первые 3 проверки можно не выполнять, т.к. их условия всегда false. С учётом всех оптимизаций (вынос инварианта специально не делаем) получается следующее:

int foo_after_partial_unroll_factor_4_optimized_1(int *a) {
  int sum = 0;
  for (int i = 0; i < 1024; i += 4) {
    int tmp = i / 4; // = также (i + 1) / 4, (i + 2) / 4, (i + 3) / 4.
    sum += a[i];
    sum += tmp;
    if (i % 16 == 0) {
      sum *= 2;
    }
    sum += a[i + 1];
    sum += tmp;
    if (false) {
      sum *= 2;
    }
    sum += a[i + 2];
    sum += tmp;
    if (false) {
      sum *= 2;
    }
    sum += a[i + 3];
    sum += tmp;
    if (false) {
      sum *= 2;
    }
  }
  return sum;
}

Или после упрощений

int foo_after_partial_unroll_factor_4_optimized_2(int *a) {
  int sum = 0;
  for (int i = 0; i < 1024; i += 4) {
    int tmp = i / 4; // = также (i + 1) / 4, (i + 2) / 4, (i + 3) / 4.
    sum += a[i];
    sum += tmp;
    if (i % 16 == 0) {
      sum *= 2;
    }
    sum += vector_add_reduce(a[i + 1], a[i + 3]);
    sum += tmp * 3;
  }
  return sum;
}

Заметим, что тут, к сожалению, мы не смогли склеить операции до и после sum *= 2, поскольку это умножение происходит не всегда. Тем не менее, наблюдательный читатель может заметить, что новый цикл опять можно размотать по 4, и тогда проверку if (i % 16 == 0) можно будет убрать. Стоит ли это делать исходя из других факторов (рост объёма кода) -- опять же, вопрос стоимостной модели.

Минусы частичной размотки, в целом, те же, что и у полной размотки, и по примерно тем же причинам. Приведём сравнительную таблицу этих оптимизаций:

В сухом остатке имеем, что частичная размотка менее требовательна, но её плюсы более скромные при потенциально тех же минусах.

Частичная размотка против лишних свопов

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

int swap_before_partial_unroll(int a, int b, int c) {
  int sum = 0;
  for (int i = 0; i < 3000; i++) {
    sum += a;
    int swap = a;
    a = b;
    b = c;
    c = swap;
  }
  return sum;
}

Объясняем, что происходит. Здесь на каждой итерации цикла переменные a, b и с циклически меняются значениями. Чтобы понять, как оптимизировать такую функцию, есть смысл посмотреть на неё в SSA-форме:

define i32 @swap_before_partial_unroll(i32 noundef %a, i32 noundef %b, i32 noundef %c) {
entry:
  br label %loop

exit:                                              ; preds = %loop
  ret i32 %add

loop:                                              ; preds = %loop, %entry
  %a_phi = phi i32 [ %a, %entry ], [ %b_phi, %loop ]
  %i = phi i32 [ 0, %entry ], [ %i.next, %loop ]
  %sum = phi i32 [ 0, %entry ], [ %add, %loop ]
  %c_phi = phi i32 [ %c, %entry ], [ %a_phi, %loop ]
  %b_phi = phi i32 [ %b, %entry ], [ %c_phi, %loop ]
  %add = add nsw i32 %a_phi, %sum
  %i.next = add nuw nsw i32 %i, 1
  %icmp = icmp eq i32 %i.next, 3000
  br i1 %icmp, label %exit, label %loop
}

Здесь финоды %a_phi, %b_phi и %c_phi циклически переходят друг в друга при переходе по обратной дуге. Поскольку длина цикла равна 3, то каждые 3 итерации все они возвращают свои изначальные значения. Таким образом, если частично размотать этот цикл по 3 итерации, то в новом цикле этих финод попросту не будет, т.к. они будут принимать те же значения, что у них были 3 итерации назад. Размотанный цикл на С++ будет выглядеть примерно так:

int swap_after_partial_unroll_factor_3(int a0, int b0, int c0) {
  int sum = 0;
  int swap;
  int a = a0;
  int b = b0;
  int c = c0;
  for (int i = 0; i < 3000; i += 3) {
    sum += a; // sum += a0;
    swap = a; // swap = a0
    a = b;    // a = b0
    b = c;    // b = c0
    c = swap; // c = a0
    sum += a; // sum += b0
    swap = a; // swap = b0
    a = b;    // a = c0
    b = c;    // b = a0
    c = swap; // c = b0
    sum += a; // sum += c0
    swap = a; // swap = c0
    a = b;    // a = a0
    b = c;    // b = b0
    c = swap; // c = c0
  }
  return sum;
}

Здесь в комментариях a0, b0 и с0 -- изначальные значения соответствующих переменных. Все эти рассуждения делаются автоматически, если мы будем строить SSA-форму данной функции. Теперь мы заметили интересный факт: перед каждой новой итерацией a, b и с имеют те же значения, что и перед всем циклом. Это позволяет избавиться от всех этих присваиваний по кругу, заменив только sum += a на те значения, которые реально добавляются:

int swap_after_partial_unroll_factor_3_optimized_1(int a0, int b0, int c0) {
  int sum = 0;
  for (int i = 0; i < 3000; i += 3) {
    sum += a0;
    sum += b0;
    sum += c0;
  }
  return sum;
}

Конечно, умный компилятор этим заниматься не станет, и сразу вычислит

int swap_after_partial_unroll_factor_3_optimized_2(int a0, int b0, int c0) {
  return (a0 + b0 + c0) * 1000;
}

Итак, удивительным образом частичная размотка позволила выбросить затейливый цикл, который в противном случае занимался бы перекладыванием значений из регистра в регистр. Это хоть и частный случай, но тем не менее достаточно интересный и важный. В LLVM он формализован в терминах нахождения цикла из взаимозависимых финод.

Избавляемся от требований по делимости

Частичная размотка всё ещё требует, чтобы число итераций делилось на N. Но что если это не так? Число итераций вполне может оказаться простым, или вообще не быть известным в момент компиляции. Значит ли это, что в этом случае частичная размотка невозможна? К счастью, нет.

Проще всего, когда число итераций известно, но не обладает нужными свойствами делимости. В этом случае, выбрав фактор размотки, мы можем провести частичную размотку, а все "лишние" итерации выполнить либо в самом начале (в виде пролога), либо в конце (в виде эпилога). Пример:

int foo_before_partial_unroll(int *a) {
  int sum = 0;
  for (int i = 0; i < 1027; i++) {
    <тело(i)>
  }
  return sum;
}

int foo_after_partial_unroll_with_prologue(int *a) {
  int sum = 0;
  for (int i = 0; i < 3; i++) {
    <тело(i)>
  }
  for (int i = 3; i < 1027; i += 4) {
    <тело(i)>
    <тело(i + 1)>
    <тело(i + 2)>
    <тело(i + 3)>
  }
  return sum;
}

int foo_after_partial_unroll_with_epilogue(int *a) {
  int sum = 0;
  for (int i = 0; i < 1024; i += 4) {
    <тело(i)>
    <тело(i + 1)>
    <тело(i + 2)>
    <тело(i + 3)>
  }
  for (int i = 1024; i < 1027; i++) {
    <тело(i)>
  }
  return sum;
}

Выбор между прологом и эпилогом, если он вообще имеет значение, обычно определяется выравниванием доступов к массивам. Например, некоторые векторные операции могут требовать выравненности адресов памяти, с которой они работают, по 4 или 8 байт. В этом случае выгодно "подогнать" первый индекс так, чтобы он был выравнен так, как нужно. За исключением таких случаев, выбор между прологом и эпилогом особого значения не имеет. Заметим также, что поскольку фактор размотки обычно небольшой, то пролог и эпилог можно подвергнуть полной размотке.

Ситуация усложняется, хотя и не сильно, если границы цикла неизвестны, но цикл всё ещё зависит от индукционной переменной (т.е. счётчика). Пример такого цикла:

int foo_before_partial_unroll(int *a, int len) {
  int sum = 0;
  for (int i = 0; i < len; i++) {
    <тело(i)>
  }
  return sum;
}

Мы не знаем ни чему равен len, ни тем более делится ли он на что-нибудь. Но мы можем заставить его делиться на что нам нужно с помощью пролога или эпилога:

int foo_after_partial_unroll_prologue(int *a, int len) {
  int sum = 0;
  int start = 0;
  while (len - start % 4 != 0) {
    <тело(start)>
    start++;
  }
  for (int i = start; i < len; i  += 4) {
    <тело(i)>
    <тело(i + 1)>
    <тело(i + 2)>
    <тело(i + 3)>
  }
  return sum;
}

int foo_after_partial_epilogue(int *a, int len) {
  int sum = 0;
  int end = len / 4 * 4;
  for (int i = 0; i < end; i += 4) {
    <тело(i)>
    <тело(i + 1)>
    <тело(i + 2)>
    <тело(i + 3)>
  }
  for (int i = end; i < len; i++) {
    <тело(i)>
  }
  return sum;
}

Конечно, тут нужно аккуратно работать с переполнениями и отрицательными значениями len, но я эти проверки специально опустил для понятности изложения.

Динамическая размотка циклов

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

int example_before_runtime_unrolling(int *a, int x) {
  int idx = 0;
  while (a[idx / 7] != x) {
    x++;
    idx++;
  }
  return idx;
}

Заметим, что тут переменные x и idx меняются на каждой итерации, однако выражение idx / 7 меняется только 1 раз на 7 итераций. Деление -- довольно тяжёлая операция в сравнении со всеми остальными операциями в данном цикле. Поэтому на ней выгодно экономить. Добиться этого можно при помощи динамической размотки (runtime unrolling):

int example_after_runtime_unrolling(int *a, int x) {
  int idx = 0;
  while (true) {
    if (a[idx / 7] == x) break;
    if (a[(idx + 1) / 7] == x + 1) { idx += 1; x += 1; break; }
    if (a[(idx + 2) / 7] == x + 2) { idx += 2; x += 2; break; }
    if (a[(idx + 3) / 7] == x + 3) { idx += 3; x += 3; break; }
    if (a[(idx + 4) / 7] == x + 4) { idx += 4; x += 4; break; }
    if (a[(idx + 5) / 7] == x + 5) { idx += 5; x += 5; break; }
    if (a[(idx + 6) / 7] == x + 6) { idx += 6; x += 6; break; }
    x += 7;
    idx += 7;
  }
  return idx;
}

Здесь вместо того, чтобы менять переменную idx несколько раз в теле цикла, используем временные значения. Вычисления перед break нужны, чтобы вернуть правильное значение. В SSA-форме операции на выходе из цикла -- бесплатные и реализуются через Phi-ноды. Я специально оставил этот громоздкий псевдокод для понятности изложения.

Далее замечаем, что idx всегда делится на 7, и поэтому все индексы на самом деле равны между собой. Более того, достаточно всего один раз загрузить элемент по этому индексу из памяти для работы с ним:

int example_after_runtime_unrolling_optimized_1(int *a, int x) {
  int idx = 0;
  while (true) {
    int v = a[idx / 7];
    if (v == x) break;
    if (v == x + 1) { idx += 1; x += 1; break; }
    if (v == x + 2) { idx += 2; x += 2; break; }
    if (v == x + 3) { idx += 3; x += 3; break; }
    if (v == x + 4) { idx += 4; x += 4; break; }
    if (v == x + 5) { idx += 5; x += 5; break; }
    if (v == x + 6) { idx += 6; x += 6; break; }
    x += 7;
    idx += 7;
  }
  return idx;
}

Если оптимизатор достаточно умный, он может сэкономить на сравнениях: если v изначально не попадает в диапазон от x до x + 6, то ни одна из этих проверок не может выстрелить. В итоге получится что-то вроде:

int example_after_runtime_unrolling_optimized_2(int *a, int x) {
  int idx = 0;
  while (true) {
    int v = a[idx / 7];
    if (x <= v && v <= x + 6) {
      if (v == x) break;
      if (v == x + 1) { idx += 1; x += 1; break; }
      if (v == x + 2) { idx += 2; x += 2; break; }
      if (v == x + 3) { idx += 3; x += 3; break; }
      if (v == x + 4) { idx += 4; x += 4; break; }
      if (v == x + 5) { idx += 5; x += 5; break; }
      if (v == x + 6) { idx += 6; x += 6; break; }
    }
    x += 7;
    idx += 7;
  }
  return idx;
}

Теперь, если предположить, что в среднем цикл выполнялся достаточно долго, то на горячем пути мы сэкономили на делениях и чтениях элементов массива (их стало в 7 раз меньше), а также теперь выполняем 2 сравнения на 7 итераций вместо 7. Таким образом, динамическая размотка позволила неплохо ускорить цикл. Правда, если раньше цикл в среднем раньше занимал меньше 7 итераций, то мы сделали много лишней работы и, скорее всего, замедлили код. Поэтому и динамическая размотка не является волшебным лекарством от всех проблем.

Вместо заключения

В этой статье я попытался коротко описать, что такое размотка циклов, для чего она применяется, какая может быть польза и каков вред. Желающих получше разобраться с тем, как компилятор принимает решение о целесообразности размотки, отсылаю читать пасс Loop Unroll в LLVM.

Размотка циклов бывает настолько не полезна, что к ней существует обратная трансформация, называемая Loop Rerolling. Причём они могут применяться вместе по схеме "размотали -- удалили всё лишнее -- смотали обратно в цикл". Все трейдоффы в этой сфере -- настолько тонкий лёд и зыбкое болото, что без реальных замеров рассуждать о нём можно только теоретически.