inconsistent valuation @ shared 'Acc' when trying to lift non-Acc function to Acc

[justinlovinger]

Description
I am working on a machine learning/mathematical optimization library with accelerate for array computation. An optimization algorithm typically takes an objective function, like A.Acc (A.Vector Bool) -> A.Acc (A.Scalar b), as an argument. Not all objective functions can be described in terms of Acc, so it is important that a user can "lift" a non-Acc function to Acc, like A.use . f . A.run.

To ensure accelerate is capable of this, I put together a trivial lifted objective function:

liftedSumBools :: A.Acc (A.Vector Bool) -> A.Acc (A.Scalar Double)
liftedSumBools = A.use . A.fromList A.Z . (: []) . sumBools . A.toList . A.run

sumBools :: [Bool] -> Double
sumBools = sum . fmap (\b -> if b then 1 else 0)

However, when I tried to run an optimizer on this, I got an error:

*** Internal error in package accelerate ***
*** Please submit a bug report at https://github.com/AccelerateHS/accelerate/issues

inconsistent valuation @ shared 'Acc' tree with stable name 224;
  aenv = [296]

CallStack (from HasCallStack):
  internalError: Data.Array.Accelerate.Trafo.Sharing:267:5
  convertSharingAcc: Data.Array.Accelerate.Trafo.Sharing:292:16
  convertSharingAcc: Data.Array.Accelerate.Trafo.Sharing:292:16
  convertSharingAcc: Data.Array.Accelerate.Trafo.Sharing:292:16
  convertSharingAcc: Data.Array.Accelerate.Trafo.Sharing:292:16
  convertSharingAcc: Data.Array.Accelerate.Trafo.Sharing:285:13
  convertSharingAcc: Data.Array.Accelerate.Trafo.Sharing:282:14
  convertSharingAcc: Data.Array.Accelerate.Trafo.Sharing:292:16
  convertSharingAcc: Data.Array.Accelerate.Trafo.Sharing:292:16
  convertSharingAcc: Data.Array.Accelerate.Trafo.Sharing:243:3
  convertOpenAcc: Data.Array.Accelerate.Trafo.Sharing:161:35
  convertAccWith: Data.Array.Accelerate.Trafo:69:37

If you replace liftedSumBools with

sumBoolsAcc :: A.Acc (A.Vector Bool) -> A.Acc (A.Scalar Double)
sumBoolsAcc = A.sum . A.map (\b -> A.cond b 1 0)

there is no error.

Steps to reproduce
Run this program: https://gist.github.com/JustinLovinger/49b81dc83284732c05e4b657670b57c0.

Expected behaviour
Program runs without error.

Your environment

• Accelerate: 1.3
• Accelerate backend(s): Reference interpreter
• GHC: 8.10.3
• OS: NixOS 20.09

Additional context
While trying to create a minimal reproduction, I ran into a different error, derivative-free-comparison: Cyclic definition of a value of type 'Acc' (sa = 26):

module Main where

import qualified Data.Array.Accelerate         as A
import qualified Data.Array.Accelerate.Interpreter
                                               as A

main :: IO ()
main = do
  print $ A.run $ aiterate 2 (step liftedSumBools) $ A.use $ A.fromList
    (A.Z A.:. 2)
    [False, False]

step
  :: (A.Acc (A.Vector Bool) -> A.Acc (A.Scalar Double))
  -> A.Acc (A.Vector Bool)
  -> A.Acc (A.Vector Bool)
step f xs = A.acond (A.the (f xs) A.> 1) xs (A.map A.not xs)

liftedSumBools :: A.Acc (A.Vector Bool) -> A.Acc (A.Scalar Double)
liftedSumBools = A.use . A.fromList A.Z . (: []) . sumBools . A.toList . A.run

sumBools :: [Bool] -> Double
sumBools = sum . fmap (\b -> if b then 1 else 0)

-- | Repeatedly apply a function a fixed number of times.
aiterate
  :: (A.Arrays a)
  => A.Exp Int -- ^ number of times to apply function
  -> (A.Acc a -> A.Acc a) -- ^ function to apply
  -> A.Acc a -- ^ initial value
  -> A.Acc a
aiterate n f xs0 = A.asnd $ A.awhile
  (A.unit . (A.< n) . A.the . A.afst)
  (\(A.T2 i xs) -> A.T2 (A.map (+ 1) i) (f xs))
  (A.lift (A.unit $ A.constant (0 :: Int), xs0))

This program doesn't give an error if you replace liftedSumBools with

sumBoolsAcc :: A.Acc (A.Vector Bool) -> A.Acc (A.Scalar Double)
sumBoolsAcc = A.sum . A.map (\b -> A.cond b 1 0)

or aiterate 2 (step liftedSumBools) with step liftedSumBools $ step liftedSumBools.

[dpvanbalen]

As far as I can see (but it's late, I'll check again tomorrow) this behaviour can't be supported by Accelerate. You can 'compose' Acc programs using run and use, but you can't use such a "lifted" function within Accelerate-level control flow (such as acond or awhile), I think. The cleanest way to avoid doing so, is to really divide your program such that each part is a single Accelerate program, using only 1 call to 'run'.
We could certainly do a better job reporting this; "Cyclic definition of a value of type 'Acc'" hints in the right direction but the first error is really meant as an internal sharing recovery error.

[justinlovinger]

@dpvanbalen I don't see a good way to separate objective functions from optimizers in this situation. Different optimizers call objective functions in different ways, and some may call an objective function multiple times or an indeterminate number of times.

Maybe I don't understand something about Accelerate, but couldn't a function like liftAcc wrap a non-Acc function so that it takes an Acc value, evaluates it if necessary, loads it into host memory, runs the given function, and loads the result back into device memory? I understand if you can't implement this liftAcc in terms of run and use, but is it possible for Accelerate to support such functionality?

[tmcdonell]

It sounds like your liftAcc is foreignAcc: call some foreign code and return the result back into accelerate?

[justinlovinger]

@tmcdonell That looks close to what I'm looking for. My liftAcc would look something like

liftAcc :: (Arrays as, Arrays bs) => (as -> bs) -> Acc as -> Acc bs

Unlike foreignAcc, it would not depend on choice of backend or have a fallback written in Acc.

[tmcdonell]

The fallback choice allows you to chain multiple implementations, presumably ending with one written in pure accelerate (or, error). I don't think we can really entirely remove the dependency on the particular backend, however, because e.g. a GPU needs to copy the results back to the host, but the CPU does not.

[justinlovinger]

@tmcdonell It sounds like you could implement liftAcc with a chain of foreignAcc, one for each backend, followed by an error. The main downside I see is you would need to update liftAcc for every new backend. It would be nice if this functionality lived in the accelerate package.

P.S. One of my motivations for wanting a function like liftAcc is convenience for users of my library that don't know Accelerate, so they can use it without learning Accelerate.