Search code examples
haskellrepaaccelerate-haskell

Converting accelerate's A array representation to repa's U array representation


I'd like to convert an accelerate array to a repa array, before then using writeImageToBMP from repa-io to write the array to a BMP file. Ignore the fact that there exists such a function in accelerate-io, I'm just using it as an example of the question. Take:

{-# LANGUAGE ScopedTypeVariables #-}
import qualified Data.Array.Accelerate.Interpreter as A (run)
import Data.Array.Accelerate
import qualified Data.Array.Accelerate as A
import qualified Data.Array.Accelerate.IO as A
import qualified Data.Array.Repa as Repa
import qualified Data.Array.Repa.IO.BMP as Repa
import GHC.Word

main :: IO ()
main = do
  let accelArr = A.fromList
                 (Z :. (2::Int) :. (1::Int))
                 ([(1,2,3),(4,5,6)] :: [(Word8,Word8,Word8)])
      computation :: A.Acc (A.Array A.DIM2 (Word8,Word8,Word8))
      computation = A.map
       (\triple ->
           let (r',g',b') = unlift triple :: (Exp Word8,Exp Word8,Exp Word8)
           in lift (constant 0,g',b')) (use accelArr)
      repaArr :: Repa.Array A.A Repa.DIM2 (Word8,Word8,Word8)
      repaArr = A.toRepa (A.run computation)
  Repa.writeImageToBMP "out_repa.bmp" repaArr

The packages accelerate, accelerate-io, repa and repa-io are reuired to compile this code.

This doesn't compile because the type for writeImageToBmp in repa-io is:

writeImageToBMP
  :: FilePath
  -> Repa.Array Repa.U Repa.DIM2 (Word8, Word8, Word8)
  -> IO ()

Where U represents a manifest array in repa. The GHC error is:

Couldn't match type ‘A.A’ with ‘Repa.U’
Expected type: Repa.Array Repa.U Repa.DIM2 (Word8, Word8, Word8)
  Actual type: Repa.Array A.A Repa.DIM2 (Word8, Word8, Word8)
In the second argument of ‘Repa.writeImageToBMP’, namely ‘repaArr’
In a stmt of a 'do' block:
  Repa.writeImageToBMP "out_repa.bmp" repaArr

The type of repaArr is:

Repa.Array A.A Repa.DIM2 (Word8,Word8,Word8)

Which is why GHC is complaining about it being an argument to writeImageToBMP.

Question: How do I convert the accelerate array representation A to the repa manifest array representation U, so that I can call writeImageToBMP on it?


Solution

  • Note that the functions computeS, computeP, copyS and copyP in Data.Array.Repa have this kind of signature:

    ...constraints... => Array r1 sh e -> m (Array r2 sh e)
    

    and therefore they are functions which allow you to convert between different Repa representations (i.e. convert r1 -> r2).

    With that in mind I got this code to type check:

    main :: IO ()
    main = do
      let accelArr = A.fromList
                     (Z :. (2::Int) :. (1::Int))
                     ([(1,2,3),(4,5,6)] :: [(Word8,Word8,Word8)])
          computation :: A.Acc (A.Array A.DIM2 (Word8,Word8,Word8))
          computation = A.map
           (\triple ->
               let (r',g',b') = unlift triple :: (Exp Word8,Exp Word8,Exp Word8)
               in lift (constant 0,g',b')) (use accelArr)
          repaArr :: Repa.Array A.A Repa.DIM2 (Word8,Word8,Word8)
          repaArr = A.toRepa (A.run computation)
      zzz <- Repa.copyP repaArr
      Repa.writeImageToBMP "out_repa.bmp" zzz
    

    You can also use Repa.copyS here.