Yarr — dataflow-фреймворк (обработки изображений) на Хаскеле

Зондирование обстановки на Реддите показало, что едва ли хоть кто-то всерьез занимается обработкой изображений на Хаскеле, несмотря на то, что достаточно популярная библиотека Repa предполагает работу с изображениями как одно из основных приложений. Надеюсь, ситуацию сможет изменить библиотека Yarr (документация, гитхаб).

Я называю библиотеку dataflow-фреймворком, потому что она обобщена для обработки массивов (от одномерных до трехмерных) элементов любых типов, в том числе векторов чисел, например координат, комплексных чисел. Но основное предполагаемое применение — обработка двумерных массивов из векторов цветовых компонент, т. е. изображений. Фреймворк непосредственно не содержит алгоритмов обработки изображений, а предоставляет мощную инфраструктуру для их написания.

Суть

В фреймворке определено семейство типов — представлений массивов (UArray). Представления делятся на 2 рода:

• Явные, элементы массивов в таких представлениях действительно находятся где-то в памяти. Есть представления на основе низкоуровневых массивов GHC, на основе кусков сырой памяти из маллока — только для элементов из класса Storable.
• Неявные представления инкапсулируют функцию (индекс -> элемент).

Также есть огромное число функций трех основных направлений:

• Отображения (map), функции комбинирования (zipWith) и свертки по ядру (convolution) принимают на вход один или несколько массивов и возвращают неявный массив.
• Вычисляющие (Load) функции переводят неявный массив в явный, т. е. в каком-то виде записывают первый в память.
• Функции прохода (Walk) по массиву, в том числе с состоянием — свертки (fold).

Знакомые с библиотекой Repa могут заметить, что Yarr очень похож на нее. Основные отличия:

• Большинство функций имеют «покомпонентные» варианты, что удобно при работе с массивами векторов, например изображениями в формате RGB.
• Сильно расширены возможности сверток (fold).
• Зачастую Yarr быстрее Repa. Например, детектор границ Канни на Yarr работает примерно в 2 раза быстрее, чем детектор на Repa.

Примеры

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

Расчет гистограммы по светлоте — один из этапов эквализации (первая стрелка на верхней картинке).

import Data.Yarr
import Data.Yarr.Walk
import Data.Yarr.Shape as S
import Data.Yarr.Repr.Foreign
...

computeHist :: UArray FS L Dim2 Float -> IO (UArray F L Dim1 Int)
computeHist pixelValues =
    walkP caps (mutate (S.unrolledFill n8 noTouch) incHist)
               (newEmpty 256) joinHists pixelValues
  where
    incHist :: UArray F L Dim1 Int -> Float -> IO ()
    incHist hist value = do
        let level = truncate (value * 255.99)
        c <- index hist level
        write hist level (c + 1)

    joinHists h1 h2 = do
        loadS S.fill (dzip2 (+) h1 h2) h1
        return h1

В этой статье приведены замеры минимального времени выполнения из множества запусков, в расчете на пиксель изображения, на разном кол-ве потоков. Единица измерения — такт процессора. Версия Yarr — 1.2.3, тестовая конфигурация, общий Makefile, тестовое изображение (2560×1600, 6 MБ).

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

{-# LANGUAGE QuasiQuotes #-}
...
import Data.Yarr.Convolution
...

noOrient = 0 :: Word8
posDiag  = 64 :: Word8
vert     = 128 :: Word8
negDiag  = 192 :: Word8
horiz    = 255 :: Word8

gradientMagOrient
    :: Word8
    -> UArray F L Dim2 Float
    -> FImage (VecTuple N2 Word8)
    -> IO ()
gradientMagOrient !threshLow image target =
    loadP S.fill caps delayedMagOrient target
  where
    delayedMagOrient = dzip2 magnitudeAndOrient gradientX gradientY

    -- Sobel-X operator application
    gradientX =
        dConvolveLinearDim2WithStaticStencil
            [dim2St| -1  0  1
                     -2  0  2
                     -1  0  1 |]
            image

    -- Sobel-Y
    gradientY =
        dConvolveLinearDim2WithStaticStencil
            [dim2St|  1   2   1
                      0   0   0
                     -1  -2  -1 |]
            image

    magnitudeAndOrient :: Float -> Float -> VecTuple N2 Word8
    magnitudeAndOrient gX gY =
        VT_2 (mag, if mag < threshLow then noOrient else orient gX gY)
      where
        mag = floatToWord8 $ sqrt (gX * gX + gY * gY)

        orient :: Float -> Float -> Word8
        orient gX gY
            | atan_1q < 0.33 = horiz
            | atan_1q > 0.66 = vert
            | otherwise      = posDiag + diagInv
          where
            rr = gY / gX
            (r, diagInv) =
                if rr < 0 then (negate rr, negDiag - posDiag) else (rr, 0)
            -- 2nd order Taylor series of atan,
            -- see http://stackoverflow.com/a/14101194/648955
            br = 0.596227 * r
            num = br + (r * r)
            atan_1q = num / (1.0 + br + num)

Вот сравнение производительности детекторов границ целиком, на Yarr и OpenCV, самой популярной библиотеке обработки изображений. Код на Yarr, код на OpenCV. OpenCV не поддерживает параллельную обработку одного изображения.

Вывод: выделить границы на большом кол-ве обычных изображений быстрее получится на OpenCV (если раскидать изображения по потокам), а на одной гигантской, в пару гигапикселей, фотографии — быстрее на Хаскеле.

Самоценные оптимизации

Автоматическое распараллеливание

В примерах выше вы могли обратить внимание на функции с суффиксом -P и первым аргументом caps: walkP, loadP. Это они отвечают за «магическое» ускорение программ при задании через консоль разного кол-ва доступных процессу ядер. caps это сокращение для getNumCapabilities — функции из стандартного модуля Control.Concurrent. Вместо caps можно написать, например, (threads 2), с понятным результатом.

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

Если вычисление на элементе достаточно простое, а конвейер у целевого процессора широкий, можно развернуть циклы на несколько итераций, чтобы забить конвейер получше. Пример — в 1-й строчке функции computeHist (расчет гистограммы развернут на 8 итераций). Производительность этого вычисления при других параметрах разворачивания, и совсем без него (шкала слегка логарифмическая):

Как видно, в данном случае разворачивание позволяет ускорить программу примерно втрое. Интуитивно это объясняется тем, что в конвейере моего процессора 3 пайплайна. При длинном разворачивании они работают на полную, а без разворачивания — 2 из 3 простаивают.

Да, к сожалению, ни собственный бекенд GHC, ни LLVM (предполагается использовать с Yarr именно его) фактически не разворачивают циклы сами, хотя как минимум LLVM теоретически должен.

Напротив, при обработке массивов векторов, если вычисления отдельных компонент независимы и при этом содержат много арифметики, можно попробовать вычислить компоненты раздельно во времени (в параллельном варианте они также будут распределены по разным потокам), чтобы разгрузить конвейер и регистровый файл. Например, на моем процессоре это дает приращение производительности на 15% при Гауссовом сглаживании цветного изображения с ядром 5×5 (blur.hs). Покомпонентные варианты функций вычисления и итерации имеют инфикс -Slices- или -SlicesSeparate-.

2D-разворачивание циклов вычисления сверток по ядру для сокращения операций чтения

Для примера возьмем Гауссово сглаживание черно-белого изображения с ядром 3×3:

Если подряд вычислить значения 4 соседних пикселей, LLVM может доказать, что при повторных чтениях исходные значения пикселей не могут поменяться (так как между ними нет операций записи) и повторно использовать последние.

Выше изображен исходный двумерный массив интенсивностей пикселей, красный, синий, желтый и зеленый прямоугольники — ядра для вычисления сглаженных значений пикселей по координатам (y=2, x=2), (2, 3), (3, 2) и (3, 3) соответственно. Множители в ядре написаны только для первого пикселя.

Если в процессоре хватает регистров общего назначения, LLVM компилирует такое вычисление с 16 операциями чтения вместо 36 — срабатывает оптимизация GVN (Global Value Numbering).

В примере используется разворачивание на 2 индекса по горизонтали и на 2 по вертикали. Конечно, это не всегда оптимально. Скорость такого сглаживания в зависимости от шагов разворачивания:

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

Недостатки

Циклы не векторизуются. Как и в случае с разворачиванием, LLVM обещает, но фактически ничего не векторизует. В LLVM 3.3 анонсировано улучшение в этом направлении, поэтому в обозримом будущем можно надеяться на еще большее ускорение фреймворка. В данный момент стабильным бекендом GHC 7.6.1 считается LLVM 3.1.

Нет поддержки поточной фильтрации и объединения массивов (streaming). Эта возможность будет в Repa 4 (версия в разработке). Реализована в Yarr не будет, потому что для этого придется переписать библиотеку с нуля, на что желающих не найдется. Прикрутить сбоку при текущей архитектуре не получится. Впрочем, Ben Lippmeier тоже переписывает Repa 3 с нуля, значит мотивация нашлась :)

Реализация

Большинство функций отображения/комбинирования (maps, zips) и функции, переводящие массив из неявного представления в явное (Load) определены в классах с несколькими параметрами. Упрощенно, в случае отображений, 1-й параметр — тип исходного массива, 2-й — тип выходного неявного массива. В случае вычислений, 1-й параметр — тип входного неявного массива, 2-й — тип результирующего вычисленного явного массива. В итоге при отображениях в неявных массивах сохраняется информация о «нижележащих» явных источниках. При вычислениях знание о структуре входного и выходного массивов позволяет выбрать оптимальный по скорости способ итерирования.

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

Отдельный пост — о борьбе с медлительностью GHC.

Функции с разворачиванием циклов, например unrolledFill, принимают в качестве шага разворачивания число не обычное (Int), а поднятое на уровень типов. В функции computeHist (из примера выше) n8 — это значение типа N8, который принадлежит классу Arity из библиотеки fixed-vector. Вообще, фреймворк фактически основан на этой библиотеке, она используется очень активно.

Внутри функций вроде unrolledFill формируется вектор действий заданной арности. GHC полностью редуцирует вектор, после компиляции остается плоская последовательность действий, т. е. как раз разворачивание. Поразительно, GHC справляется даже с редукцией векторов двойной вложенности, с помощью которых реализована функция dim2BlockFill!

Разработка

Самым сложным при разработке фреймворка было придумать архитектуру, систему типов. Я кардинально менял ее около 8 раз, посматривая на свою заметку о проектировании на Хаскеле сквозь слезы самоиронии. Абстракции текли, даже скорее прорывались, архитектура не позволяла реализовать нужные функции или была несовместима с оптимизатором GHC, и я создавал новую папку и начинал писать с начала.

Первых 3-4 подхода были попытками прикрутить покомпонентные операции к Repa. Когда я наконец понял, что это невозможно, специально отошел от архитектуры Repa как можно дальше. Оставшиеся «итерации» снова приближали фреймворк по дизайну к Repa, в итоге снаружи он снова напоминает просто расширение Repa.

Забавно, что Алексей Худяков выложил библиотеку fixed-vector только в ноябре прошлого года, т. е. через 2 месяца после того, как я решил писать фреймворк (взял тему курсача). Но т. к. приступил в декабре, оказалось без особой разницы. Можно сказать, мне сильно повезло :)

Автор: leventov

Источник