Много параллельных приложений последовательного преобразования в repa

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

В общем, такое преобразование может быть выражено как матрица M, и каждая запись M*v является просто точечным произведением соответствующей строки M с v. Поэтому я мог бы просто использовать traverse с функцией, которая вычисляет соответствующий точечный продукт. Это стоит d^2.

Однако мой M является особенным: он допускает последовательный алгоритм линейной работы. Например, M может быть нижней треугольной матрицей с 1 по всему нижнему треугольнику. Тогда M*v является просто вектором частичных сумм v (a.k.a. "scan" ). Эти суммы могут быть вычислены последовательно очевидным образом, но нужно, чтобы элемент (i-1) st был получен для эффективного вычисления i -й записи. (У меня есть несколько таких M, все из которых могут быть вычислены так или иначе в линейном последовательном времени.)

Я не вижу очевидного способа использовать traverse (или любые другие функции Repa), чтобы воспользоваться этим свойством M. Это можно сделать? Будет довольно расточительно использовать алгоритм d^2 -work (даже с обильным parallelism), когда имеется такой быстрый алгоритм линейной работы.

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

UPDATE

По запросу здесь приведен пример иллюстративного кода для M, который вычисляет частичные суммы (как описано выше). Как я и ожидал, среда выполнения (работа) растет сверхлинейно в d, второй аргумент степени массива (ext). Это происходит из-за того, что mulM' указывает только, как вычислить запись i th, независимо от всех других записей. Несмотря на то, что в общем размере массива существует алгоритм с линейным временем, я не знаю, как выразить его в Repa.

Интересно, что если я удалю строку, которая определяет манифест array' из main, тогда время выполнения масштабируется только линейно в общем размере массива! Поэтому, когда массивы откладываются "полностью вниз", fusion/optimization должно каким-то образом извлекать алгоритм линейной работы, но без какой-либо явной помощи от меня. Это удивительно, но и не очень полезно для меня, потому что на самом деле мне нужно будет называть mulM на манифестных массивах.

{-# LANGUAGE TypeOperators, ScopedTypeVariables, FlexibleContexts #-}

module Main where

import Data.Array.Repa as R

-- multiplication by M across innermost dimension
mulM arr = traverse arr id mulM'
    where mulM' _ [email protected](i' :. i) =
              sumAllS $ extract (Z:.0) (Z:.(i+1)) $ slice arr (i' :. All)

ext = Z :. (1000000::Int) :. (10::Int) -- super-linear runtime in 2nd arg
--ext = Z :. (10::Int) :. (1000000::Int) -- takes forever

array = fromFunction ext (\(Z:.j:.i) -> j+i)

main :: IO ()
main = do
  -- apply mulM to a manifest array
  array' :: Array U DIM2 Int <- computeP $ array
  ans :: Array U DIM2 Int <- computeP $ mulM array'
  print "done"

Ответы