repa 中顺序变换的许多并行应用

Question

在 Repa 中，我想在数组的最内维度（即所有“列”向量）上并行应用某个

d

维线性变换。

一般来说，这样的变换可以表示为矩阵

M

，

M*v

的每个条目只是

M

的相应行与

v

的点积。所以我可以将

traverse

与计算适当点积的函数一起使用。这需要

d^2

工作。

但是，我的

M

很特别：它承认线性工作顺序算法。例如，

M

可能是一个下三角矩阵，其中

1

遍布整个下三角。那么

M*v

只是

v

的部分和的向量（又名“扫描”）。这些总和可以以一种明显的方式顺序计算，但是需要结果的第

(i-1)

条目才能有效地计算第

i

条目。 （我有几个这样的

M

，所有这些都可以在线性顺序时间内以一种方式或另一种方式计算。）

我没有看到任何明显的方法来使用

traverse

（或任何其他 Repa 函数）来利用

M

的这一属性。能做到吗？当有如此快速的线性工作算法可用时，使用

d^2

工作算法（即使具有丰富的并行性）将非常浪费。

（我见过一些旧的SO帖子（例如，here）提出类似的问题，但没有什么与我的情况完全相符。）

更新

根据要求，这里有一些说明性代码，用于计算部分和的

M

（如上所述）。正如我所预期的，运行时间（工作）在数组范围的第二个参数（

d

）中超线性增长。这是因为

ext

仅指定如何计算输出的第

mulM'

条目，独立于所有其他条目。尽管数组的总大小有一个线性时间算法，但我不知道如何在Repa中表达它。有趣的是，如果我从

i

中删除定义清单

array'

的行，那么运行时仅在数组的总大小中线性缩放！因此，当数组“一直向下”延迟时，融合/优化必须以某种方式提取线性工作算法，但没有我的任何明确帮助。这很神奇，但对我来说也不是很有用，因为实际上，我需要在清单数组上调用

main

。mulM

Answer 1

{-# 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"

请注意，这是一个简化的示例，您可能需要进一步调整和试验 Repa 的功能，以完全捕获特定矩阵 M 所需的线性功算法。请务必参阅 Repa 的文档以获取有关其功能的更多信息以及如何最好地表达你的计算。