Ускоряем цикл foreach до for

Привет! Хочу рассказать об интересном опыте, как я писал енумератор для типа Range, который был бы таким же дешевым, как цикл for.

Что мы хотим?

У System.Range, как известно, очень красивый синтаксис создания:

var a = 1..2; // эквивалент new Range(1, 2)

Поэтому, как эксперимент, я решил абузить его для цикла foreach, как замена for:

foreach (var i in 1..5)
    Console.Write(i);

(выводит 12345)

Для foreach Розлин требует метод GetEnumerator, который возвращал бы что-то, у чего есть bool MoveNext() и T Current. Этот метод можно добавить как метод расширения:

public static class Extensions
{
		public ... GetEnumerator(this System.Range range) => ...
}

Baseline (на что равняемся)

Вообще-то for это куда больше, чем пройтись по последовательности. Но так как очень часто мы видим

for (int i = 0; i < n; i++)

то это стоит своего кейса использования for. Сравнивать мы будем с циклом for у которого в качестве итерации стоит i += inc:

for (int i = 0; i < n; i += inc)

По производительности он идентичен версии с i++:

Код метода с for-ом:

[Benchmark]
public int ForWithCustomIncrement()
{
    var iters = Iterations;
    var a = 0;
    var inc = Inc;
    for (int i = 0; i < iters; i += inc)
        a += i;
    return a;
}

(мы выгрузили все проперти в локальные переменные чтобы исключить чтение из памяти при итерации)

Кодген for-а:

    L000f: add edx, r8d
    L0012: add r8d, ecx
    L0015: cmp r8d, eax
    L0018: jl short L000f

Мы будем смотреть только на код одной итерации цикла, начальный оверхед и окружающий мусор нас не интересует.

Эксперимент 1. Enumerable.Range

Конечно, сначала стоит попробовать вариант из BCL: Enumerable.Range, исходный код которого можно поизучать здесь.

Так выглядит наш бенчмарк:

[Benchmark]
public int ForeachWithEnumerableRange()
{
    var a = 0;
    foreach (var i in Enumerable.Range(0, Iterations))
        a += i;
    return a;
}

Этот код Розлином раскрывается как

public int ForeachWithEnumerableRange()
{
    int num = 0;
    IEnumerator<int> enumerator = Enumerable.Range(0, Iterations).GetEnumerator();
    try
    {
        while (enumerator.MoveNext())
        {
            int current = enumerator.Current;
            num += current;
        }
        return num;
    }
    finally
    {
        if (enumerator != null)
        {
            enumerator.Dispose();
        }
    }
}

Изучать ассемблерный кодген смысла нет, достаточно взглянуть на бенчмарк:

Эксперимент 2. yield return

Второй по простоте способ - написать метод-генератор, и вызывать его енумератор:

private static IEnumerable<int> MyRange(int from, int to, int inc)
{
    for (int i = from; i <= to; i += inc)
        yield return i;
}

[Benchmark]
public int ForeachWithYieldReturn()
{
    var a = 0;
    foreach (var i in MyRange(0, Iterations - 1, 1))
        a += i;
    return a;
}

Принципиально нового мы ничего не получаем, наш енумератор это все еще огромный референс тип с кучей полей. Бенчмарки сообщают, что с yield-ом даже хуже, чем с Enumerable.Range:

Эксперимент 3. Собственный енумератор

Теперь будем писать собственный енумератор. Чтобы оно работало с foreach-ем, нужно, чтобы был метод bool MoveNext() и свойство T Current (в нашем случае int Current).

public struct RangeEnumerator
{
    private readonly int to;
    private readonly int step;

    private int curr;

    public RangeEnumerator(int @from, int to, int step)
    {
        this.to = to + step;
        this.step = step;
        this.curr = @from - step;
    }

    public bool MoveNext()
    {
        curr += step;
        return curr != to;
    }
    
    public int Current => curr;
}

Чтобы Розлин увидел, что по System.Range можно форычиться, нужно сделать метод расширения:

public static class Extensions
{
    public static RangeEnumerator GetEnumerator(this Range @this)
        => (@this.Start, @this.End) switch
        {
            ({ IsFromEnd: true, Value: 0 }, { IsFromEnd: true, Value: 0 }) => new RangeEnumerator(0, int.MaxValue, 1),
            ({ IsFromEnd: true, Value: 0 }, { IsFromEnd: false, Value: var to }) => new RangeEnumerator(0, to + 1, 1),
            ({ IsFromEnd: false, Value: var from }, { IsFromEnd: true, Value: 0 }) => new RangeEnumerator(from, int.MaxValue, 1),
            ({ IsFromEnd: false, Value: var from }, { IsFromEnd: false, Value: var to })
                => (from < to) switch
                {
                    true => new RangeEnumerator(from, to, 1),
                    false => new RangeEnumerator(from, to, -1),
                },
            _ => throw new InvalidOperationException("Invalid range")
        };
}

Код бенчмарка выглядит так:

[Benchmark]
public int ForeachWithRangeEnumerator()
{
    var a = 0;
    foreach (var i in 0..(Iterations - 1))
        a += i;
    return a;
}

И раскрывается Розлином так:

public int ForeachWithRangeEnumerator()
{
    int num = 0;
    RangeEnumerator enumerator = Extensions.GetEnumerator(new Range(0, Iterations - 1));
    while (enumerator.MoveNext())
    {
        int current = enumerator.Current;
        num += current;
    }
    return num;
}

На этот раз нет референс типов, а так как RangeEnumerator не имплементит IDisposable, то и нет try-catch блоков вокруг. Взглянем на бенчмарк:

18 микросекунд на 10000 операций. Намного лучше, чем аналоги до этого, но сильно хуже, чем for. Время взглянуть на кодген!

Одна итерация цикла выглядит так:

    L003e: add esi, eax
    L0040: mov eax, [rsp+0x38]
    L0044: add eax, [rsp+0x34]
    L0048: mov [rsp+0x38], eax
    L004c: mov eax, [rsp+0x38]
    L0050: cmp eax, [rsp+0x30]
    L0054: jne short L003a

Да, у нас референс-типов нет, но все данные-то все равно в памяти! Хоть и на стеке.

Эксперимент 4. Собственный енумератор прячем в локальную функцию

Что за магия? Дело в том, что чтобы получить наш енумератор таким же быстрым, как for, следует поместить его поля в регистры. Но компилятор сам этого не сделал, слишком много локальных переменных в методе было, и он не положил нужные. Давайте облегчим задачу.

Из эксперимента 3:

[Benchmark]
public int ForeachWithRangeEnumeratorRaw()
{
    var a = 0;
    var enumerator = (0..(Iterations - 1)).GetEnumerator();
    while (enumerator.MoveNext())
        a += enumerator.Current;
    return a;
}

(я заменил foreach на точно то же самое с постфиксом Raw, только написанное ручками, для наглядности. См. эксперимент 3, там показано, как Розлин раскрывает foreach).

Чтобы уменьшить количество локальных переменных, спрячем цикл в локальную функцию:

[Benchmark]
public int ForeachWithRangeEnumeratorRawWithLocal()
{
    var enumerator = (0..(Iterations - 1)).GetEnumerator();        
    return EnumerateItAll(enumerator);

    static int EnumerateItAll(RangeEnumerator enumerator)
    {
        var a = 0;
        while (enumerator.MoveNext())
            a += enumerator.Current;
        return a;
    }
}

И побенчим:

ForeachWithRangeEnumeratorRawWithLocal - это как раз вариант с локальной функцией. И... он почти в шесть раз быстрее, чем заинлайненный! Как же так вышло-то?

Смотрим кодген:

В нашем методе локальная функция не заинлайнилась:

Benchmarks.ForeachWithRangeEnumeratorRawWithLocal()
    ...
    L0044: call Benchmarks.<ForeachWithRangeEnumeratorRawWithLocal>g__EnumerateItAll|8_0(RangeEnumerator)

И это нам на руку. Смотрим эту локальную функцию:

    L0000: mov eax, [rcx]
    L0002: mov edx, [rcx+4]
    L0005: mov ecx, [rcx+8]
    L0008: xor r8d, r8d
    L000b: jmp short L0010
    L000d: add r8d, ecx         |
    L0010: add ecx, edx         |
    L0012: cmp ecx, eax         |
    L0014: jne short L000d      |
    L0016: mov eax, r8d
    L0019: ret

Отмеченные четыре строчки,

    L000d: add r8d, ecx
    L0010: add ecx, edx
    L0012: cmp ecx, eax
    L0014: jne short L000d

это и есть наша победа, так как этот код идентичен итерации for-а.

Первые три строчки

    L0000: mov eax, [rcx]
    L0002: mov edx, [rcx+4]
    L0005: mov ecx, [rcx+8]

Выгружают все, что нужно, в регистры, и поэтому так быстро работает наш метод.

Вывод

Во-первых, хоть и желаемое получилось, в реальной жизни проще написать for, чем проверять, догадался ли компилятор положить ваш енумератор в регистры.

Во-вторых, это очень показательный пример, что инлайнинг может сделать не только не лучше, но и сильно хуже (в данном случае он увеличил время работы с 3.5 до 18 us на 10000 итераций).

Спасибо за внимание!

Ссылки

Весь код, который я писал, размещен в репозитории на github.

• Бенчмарки

• Код енумератора

• Тесты, если есть подозрение на неправильный код

• BenchmarkDotNet - утилита/библиотека для бенчинга