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

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

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

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

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

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

ОБНОВЛЕНИЕ

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

Интересно, что если я уберу строку, определяющую манифест array' из main, то время выполнения масштабируется только линейно по общему размеру массива! Поэтому, когда массивы задерживаются «полностью вниз», слияние/оптимизация должны каким-то образом извлекать алгоритм линейной работы, но без какой-либо явной помощи с моей стороны. Это удивительно, но также не очень полезно для меня, потому что на самом деле мне нужно будет вызывать 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' _ idx@(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"