http://stackoverflow.com/questions/15595140/repa-zipwithing-over-columns

compile

$ ghc -Odph -rtsopts -threaded -fno-liberate-case -funbox-strict-fields -funfolding-keeness-factor1000 -fllvm -optlo-O3 -fexpose-all-unfoldings -fsimpl-tick-factor=500 -ddump-simpl -dsuppress-all zip-columns.hs > z.hs

run

$ ./zip-columns 1000 1000

zip-columns.hs

{-# LANGUAGE ScopedTypeVariables, FlexibleContexts #-}

import System.Environment

import Data.Array.Repa as R
import Data.Array.Repa.Eval as R

import Data.Yarr as Y
import Debug.Yarr

import Data.Yarr.Benchmarking as B

main = do
    [columns, rows] <- fmap (fmap read) getArgs
    let sh = (rows, columns)
        repaSh = Z :. columns :. rows
    (repaArr :: Array R.U DIM2 Int) <-
        newMVec (rows * columns) >>= unsafeFreezeMVec repaSh
    (repaCol :: Array R.U DIM1 Int) <-
        newMVec rows >>= unsafeFreezeMVec (Z :. rows)
    let repaComp :: IO (Array R.U DIM2 Int)
        repaComp = computeP $ repaZipOverColumns (*) repaCol repaArr
        repaComp2 :: IO (Array R.U DIM2 Int)
        repaComp2 = computeP $ repaZipOverColumns2 (*) repaCol repaArr
    B.benchMin "repa" 100 sh (repaComp >> return ())
    B.benchMin "repa2" 100 sh (repaComp2 >> return ())

    (yarrArr :: UArray Y.F L Dim2 Int) <- new sh
    (yarrCol :: UArray Y.F L Dim1 Int) <- new rows
    let yarrComp :: IO (UArray Y.F L Dim2 Int)
        yarrComp =
            compute (Y.loadP (unrolledFill n8 noTouch) caps) $
                yarrZipOverColumns (*) yarrCol yarrArr
    B.benchMin "yarr" 100 sh (yarrComp >> return ())
    Y.touchArray yarrCol  -- seg faults without this


repaZipOverColumns
    :: (Source r1 a, Source r2 b)
    => (a -> b -> c)
    -> Array r1 DIM1 a
    -> Array r2 DIM2 b
    -> Array R.D DIM2 c
repaZipOverColumns f col arr = R.traverse arr id lookup
  where
    lookup get sh@(Z :. _ :. r) = f (col ! (Z :. r)) $ get sh

repaZipOverColumns2
    :: (Source r1 a, Source r2 b)
    => (a -> b -> c)
    -> Array r1 DIM1 a
    -> Array r2 DIM2 b
    -> Array R.D DIM2 c
repaZipOverColumns2 f col arr = R.zipWith f ccol arr
  where
    ccol = R.fromFunction (R.extent arr) (\(Z :. _ :. r) -> col ! (Z :. r))


yarrZipOverColumns
    :: (USource r1 l1 Dim1 a,
        USource r2 l2 Dim2 b, DefaultIFusion r2 l2 fr fl Dim2, USource fr fl Dim2 c)
    => (a -> b -> c)
    -> UArray r1 l1 Dim1 a
    -> UArray r2 l2 Dim2 b
    -> UArray fr fl Dim2 c
yarrZipOverColumns f col arr = imapM mapF arr
  where
    mapF (r, _) b = do
        a <- col `Y.index` r
        return $ f a b