diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,27 @@ +Copyright (c) Henning Thielemann 2014 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. +3. Neither the name of the author nor the names of his contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS +OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY +OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF +SUCH DAMAGE. diff --git a/Setup.lhs b/Setup.lhs new file mode 100644 --- /dev/null +++ b/Setup.lhs @@ -0,0 +1,3 @@ +#! /usr/bin/env runhaskell +> import Distribution.Simple +> main = defaultMain diff --git a/knead.cabal b/knead.cabal new file mode 100644 --- /dev/null +++ b/knead.cabal @@ -0,0 +1,86 @@ +Name: knead +Version: 0.2 +License: BSD3 +License-File: LICENSE +Author: Henning Thielemann +Maintainer: Henning Thielemann +Homepage: http://hub.darcs.net/thielema/knead/ +Category: Data Structures +Synopsis: Repa array processing using LLVM JIT +Description: + This library processes arrays like @Repa@ and @Accelerate@, + but it uses the just-in-time compiler of @LLVM@ + for generating the machine code. + That is, you get very efficient vectorised code + that can be run without a GPU. + You do not need to care about inlining and strictness annotations, + because the LLVM code is by default inlined and strict. + The package is the basis for an LLVM backend for the @Accelerate@ framework. + . + Highlights: + . + * Very flexible index handling, + even more flexible than the one of 'Data.Array'. + It is much more expressive and type-safe than that of @repa@ and @array@. + . + * Extensible element types, e.g. complex numbers. + (Maybe this is also possible with accelerate, e.g. with RGB type.) + . + * Every compilable program also runs. + In contrast to that, @accelerate@ may accept a program + that cannot be run by a particular backend, like @accelerate-cuda@. + . + Known deficiencies: + . + * The functions do not check array bounds. + (Of course, we can think about temporary bound checking + for debugging purposes.) + . + * The package does not try to distribute work across multiple processors. + It is certainly simpler, more efficient and more reliable + if you do that at a higher level. + . + The name of the package is inspired by the visualization of typical operations + like reshaping, collapsing a dimension and extruding another one. +Tested-With: GHC==7.4.2, GHC==7.8.4, GHC==8.0.1 +Cabal-Version: >=1.6 +Build-Type: Simple + +Source-Repository this + Tag: 0.2 + Type: darcs + Location: http://hub.darcs.net/thielema/knead/ + +Source-Repository head + Type: darcs + Location: http://hub.darcs.net/thielema/knead/ + +Library + Build-Depends: + llvm-extra >=0.6 && <0.7, + llvm-tf >=3.0.3 && <3.0.4, + storable-tuple >=0.0 && <0.1, + storable-record >=0.0.3 && <0.1, + utility-ht >=0.0.1 && <0.1, + base >=4 && <5 + + GHC-Options: -Wall + Hs-Source-Dirs: src + Exposed-Modules: + Data.Array.Knead.Index.Linear.Int + Data.Array.Knead.Index.Linear + Data.Array.Knead.Index.Nested.Shape + Data.Array.Knead.Expression + Data.Array.Knead.Parameter + Data.Array.Knead.Simple.Symbolic + Data.Array.Knead.Simple.ShapeDependent + Data.Array.Knead.Simple.Physical + Data.Array.Knead.Simple.Slice + Data.Array.Knead.Simple.Fold + Data.Array.Knead.Parameterized.Symbolic + Data.Array.Knead.Parameterized.Physical + Data.Array.Knead.Parameterized.Slice + Other-Modules: + Data.Array.Knead.Simple.Private + Data.Array.Knead.Parameterized.Private + Data.Array.Knead.Code diff --git a/src/Data/Array/Knead/Code.hs b/src/Data/Array/Knead/Code.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Code.hs @@ -0,0 +1,46 @@ +{-# LANGUAGE TypeFamilies #-} +module Data.Array.Knead.Code where + +import qualified Data.Array.Knead.Index.Nested.Shape as Shape + +import qualified LLVM.Extra.Multi.Value as MultiValue +import qualified LLVM.Extra.Execution as Exec + +import qualified LLVM.ExecutionEngine as EE +import qualified LLVM.Util.Optimize as Opt +import qualified LLVM.Core as LLVM + +import Foreign.Ptr (Ptr, ) + +import Control.Monad (void, liftM2, when, ) + + +getElementPtr :: + (Shape.C sh, Shape.Index sh ~ ix) => + MultiValue.T sh -> LLVM.Value (Ptr a) -> + MultiValue.T ix -> + LLVM.CodeGenFunction r (LLVM.Value (Ptr a)) +getElementPtr sh ptr ix = do + n <- Shape.flattenIndex sh ix + LLVM.getElementPtr ptr (n, ()) + + +compile :: + (Exec.Compile funcs) => + String -> + LLVM.CodeGenModule (Exec.LLVMFunction funcs) -> + IO funcs +compile name bld = do + LLVM.initializeNativeTarget + m <- LLVM.newModule + (funcs, mappings) <- + LLVM.defineModule m $ + liftM2 (,) bld LLVM.getGlobalMappings + LLVM.writeBitcodeToFile (name ++ ".bc") m + when False $ do + void $ Opt.optimizeModule 3 m + LLVM.writeBitcodeToFile (name ++ "-opt.bc") m + EE.runEngineAccess $ + EE.addModule m >> + EE.addGlobalMappings mappings >> + Exec.compile funcs diff --git a/src/Data/Array/Knead/Expression.hs b/src/Data/Array/Knead/Expression.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Expression.hs @@ -0,0 +1,306 @@ +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +module Data.Array.Knead.Expression where + +import qualified LLVM.Extra.Multi.Value as MultiValue +import qualified LLVM.Extra.Arithmetic as A +import qualified LLVM.Extra.Monad as LMonad +import qualified LLVM.Core as LLVM +import LLVM.Extra.Multi.Value (PatternTuple, Decomposed, Atom, atom, ) + +import qualified Control.Monad as Monad +import qualified Data.Tuple.HT as Tuple + +import Prelude hiding (fst, snd, min, max, zip, unzip, zip3, unzip3, ) + + +newtype Exp a = Exp {unExp :: forall r. LLVM.CodeGenFunction r (MultiValue.T a)} + + +class Value val where + lift0 :: MultiValue.T a -> val a + lift1 :: + (MultiValue.T a -> MultiValue.T b) -> + val a -> val b + lift2 :: + (MultiValue.T a -> MultiValue.T b -> MultiValue.T c) -> + val a -> val b -> val c + lift3 :: + (MultiValue.T a -> MultiValue.T b -> MultiValue.T c -> MultiValue.T d) -> + val a -> val b -> val c -> val d + lift4 :: + (MultiValue.T a -> MultiValue.T b -> MultiValue.T c -> MultiValue.T d -> MultiValue.T e) -> + val a -> val b -> val c -> val d -> val e + +instance Value MultiValue.T where + lift0 = id + lift1 = id + lift2 = id + lift3 = id + lift4 = id + +instance Value Exp where + lift0 a = Exp (return a) + lift1 f (Exp a) = Exp (Monad.liftM f a) + lift2 f (Exp a) (Exp b) = Exp (Monad.liftM2 f a b) + lift3 f (Exp a) (Exp b) (Exp c) = Exp (Monad.liftM3 f a b c) + lift4 f (Exp a) (Exp b) (Exp c) (Exp d) = Exp (Monad.liftM4 f a b c d) + + +liftM :: + (forall r. + MultiValue.T a -> + LLVM.CodeGenFunction r (MultiValue.T b)) -> + (Exp a -> Exp b) +liftM f (Exp a) = Exp (f =<< a) + +liftM2 :: + (forall r. + MultiValue.T a -> MultiValue.T b -> + LLVM.CodeGenFunction r (MultiValue.T c)) -> + (Exp a -> Exp b -> Exp c) +liftM2 f (Exp a) (Exp b) = Exp (LMonad.liftR2 f a b) + +liftM3 :: + (forall r. + MultiValue.T a -> MultiValue.T b -> MultiValue.T c -> + LLVM.CodeGenFunction r (MultiValue.T d)) -> + (Exp a -> Exp b -> Exp c -> Exp d) +liftM3 f (Exp a) (Exp b) (Exp c) = Exp (LMonad.liftR3 f a b c) + + +unliftM1 :: + (Exp a -> Exp b) -> + MultiValue.T a -> LLVM.CodeGenFunction r (MultiValue.T b) +unliftM1 f ix = unExp (f (lift0 ix)) + +unliftM2 :: + (Exp a -> Exp b -> Exp c) -> + MultiValue.T a -> MultiValue.T b -> + LLVM.CodeGenFunction r (MultiValue.T c) +unliftM2 f ix jx = unExp (f (lift0 ix) (lift0 jx)) + +unliftM3 :: + (Exp a -> Exp b -> Exp c -> Exp d) -> + MultiValue.T a -> MultiValue.T b -> MultiValue.T c -> + LLVM.CodeGenFunction r (MultiValue.T d) +unliftM3 f ix jx kx = unExp (f (lift0 ix) (lift0 jx) (lift0 kx)) + + + +min :: (MultiValue.Real a) => Exp a -> Exp a -> Exp a +min = liftM2 A.min + +max :: (MultiValue.Real a) => Exp a -> Exp a -> Exp a +max = liftM2 A.max + + +zip :: (Value val) => val a -> val b -> val (a, b) +zip = lift2 MultiValue.zip + +zip3 :: (Value val) => val a -> val b -> val c -> val (a, b, c) +zip3 = lift3 MultiValue.zip3 + +zip4 :: (Value val) => val a -> val b -> val c -> val d -> val (a, b, c, d) +zip4 = lift4 MultiValue.zip4 + +unzip :: (Value val) => val (a, b) -> (val a, val b) +unzip ab = + (lift1 MultiValue.fst ab, lift1 MultiValue.snd ab) + +unzip3 :: (Value val) => val (a, b, c) -> (val a, val b, val c) +unzip3 abc = + (lift1 MultiValue.fst3 abc, + lift1 MultiValue.snd3 abc, + lift1 MultiValue.thd3 abc) + +unzip4 :: (Value val) => val (a, b, c, d) -> (val a, val b, val c, val d) +unzip4 abcd = + (lift1 (\(MultiValue.Cons (a,_,_,_)) -> MultiValue.Cons a) abcd, + lift1 (\(MultiValue.Cons (_,b,_,_)) -> MultiValue.Cons b) abcd, + lift1 (\(MultiValue.Cons (_,_,c,_)) -> MultiValue.Cons c) abcd, + lift1 (\(MultiValue.Cons (_,_,_,d)) -> MultiValue.Cons d) abcd) + +fst :: (Value val) => val (a, b) -> val a +fst = lift1 MultiValue.fst + +snd :: (Value val) => val (a, b) -> val b +snd = lift1 MultiValue.snd + +mapFst :: (Exp a -> Exp b) -> Exp (a, c) -> Exp (b, c) +mapFst f = modify (atom, atom) $ \(a,c) -> (f a, c) + +mapSnd :: (Exp b -> Exp c) -> Exp (a, b) -> Exp (a, c) +mapSnd f = modify (atom, atom) $ \(a,b) -> (a, f b) + +swap :: (Value val) => val (a, b) -> val (b, a) +swap = lift1 MultiValue.swap + + +modifyMultiValue :: + (Value val, + MultiValue.Compose a, + MultiValue.Decompose pattern, + MultiValue.PatternTuple pattern ~ tuple) => + pattern -> + (Decomposed MultiValue.T pattern -> a) -> + val tuple -> val (MultiValue.Composed a) +modifyMultiValue p f = lift1 $ MultiValue.modify p f + +modifyMultiValue2 :: + (Value val, + MultiValue.Compose a, + MultiValue.Decompose patternA, + MultiValue.Decompose patternB, + MultiValue.PatternTuple patternA ~ tupleA, + MultiValue.PatternTuple patternB ~ tupleB) => + patternA -> + patternB -> + (Decomposed MultiValue.T patternA -> + Decomposed MultiValue.T patternB -> a) -> + val tupleA -> val tupleB -> val (MultiValue.Composed a) +modifyMultiValue2 pa pb f = lift2 $ MultiValue.modify2 pa pb f + +modifyMultiValueM :: + (MultiValue.Compose a, + MultiValue.Decompose pattern, + MultiValue.PatternTuple pattern ~ tuple) => + pattern -> + (forall r. + Decomposed MultiValue.T pattern -> + LLVM.CodeGenFunction r a) -> + Exp tuple -> Exp (MultiValue.Composed a) +modifyMultiValueM p f = liftM (MultiValue.modifyF p f) + +modifyMultiValueM2 :: + (MultiValue.Compose a, + MultiValue.Decompose patternA, + MultiValue.Decompose patternB, + MultiValue.PatternTuple patternA ~ tupleA, + MultiValue.PatternTuple patternB ~ tupleB) => + patternA -> + patternB -> + (forall r. + Decomposed MultiValue.T patternA -> + Decomposed MultiValue.T patternB -> + LLVM.CodeGenFunction r a) -> + Exp tupleA -> Exp tupleB -> Exp (MultiValue.Composed a) +modifyMultiValueM2 pa pb f = liftM2 (MultiValue.modifyF2 pa pb f) + + +class Compose multituple where + type Composed multituple + {- | + A nested 'zip'. + -} + compose :: multituple -> Exp (Composed multituple) + +class + (Composed (Decomposed Exp pattern) ~ PatternTuple pattern) => + Decompose pattern where + {- | + Analogous to 'MultiValue.decompose'. + -} + decompose :: pattern -> Exp (PatternTuple pattern) -> Decomposed Exp pattern + + +{- | +Analogus to 'MultiValue.modifyMultiValue'. +-} +modify :: + (Compose a, Decompose pattern) => + pattern -> + (Decomposed Exp pattern -> a) -> + Exp (PatternTuple pattern) -> Exp (Composed a) +modify p f = compose . f . decompose p + +modify2 :: + (Compose a, Decompose patternA, Decompose patternB) => + patternA -> + patternB -> + (Decomposed Exp patternA -> Decomposed Exp patternB -> a) -> + Exp (PatternTuple patternA) -> Exp (PatternTuple patternB) -> Exp (Composed a) +modify2 pa pb f a b = compose $ f (decompose pa a) (decompose pb b) + + + +instance Compose (Exp a) where + type Composed (Exp a) = a + compose = id + +instance Decompose (Atom a) where + decompose _ = id + + + +instance (Compose a, Compose b) => Compose (a,b) where + type Composed (a,b) = (Composed a, Composed b) + compose = uncurry zip . Tuple.mapPair (compose, compose) + +instance (Decompose pa, Decompose pb) => Decompose (pa,pb) where + decompose (pa,pb) = + Tuple.mapPair (decompose pa, decompose pb) . unzip + + +instance (Compose a, Compose b, Compose c) => Compose (a,b,c) where + type Composed (a,b,c) = (Composed a, Composed b, Composed c) + compose = Tuple.uncurry3 zip3 . Tuple.mapTriple (compose, compose, compose) + +instance + (Decompose pa, Decompose pb, Decompose pc) => + Decompose (pa,pb,pc) where + decompose (pa,pb,pc) = + Tuple.mapTriple (decompose pa, decompose pb, decompose pc) . unzip3 + + +instance (Compose a, Compose b, Compose c, Compose d) => Compose (a,b,c,d) where + type Composed (a,b,c,d) = (Composed a, Composed b, Composed c, Composed d) + compose (a,b,c,d) = zip4 (compose a) (compose b) (compose c) (compose d) + +instance + (Decompose pa, Decompose pb, Decompose pc, Decompose pd) => + Decompose (pa,pb,pc,pd) where + decompose (pa,pb,pc,pd) x = + case unzip4 x of + (a,b,c,d) -> + (decompose pa a, decompose pb b, decompose pc c, decompose pd d) + + +unit :: Exp () +unit = lift0 $ MultiValue.cons () + +zero :: (MultiValue.C a) => Exp a +zero = lift0 MultiValue.zero + +add :: (MultiValue.Additive a) => Exp a -> Exp a -> Exp a +add = liftM2 MultiValue.add + +sub :: (MultiValue.Additive a) => Exp a -> Exp a -> Exp a +sub = liftM2 MultiValue.sub + +mul :: (MultiValue.PseudoRing a) => Exp a -> Exp a -> Exp a +mul = liftM2 MultiValue.mul + +idiv :: (MultiValue.Integral a) => Exp a -> Exp a -> Exp a +idiv = liftM2 MultiValue.idiv + +fromInteger' :: (MultiValue.IntegerConstant a) => Integer -> Exp a +fromInteger' = lift0 . MultiValue.fromInteger' + + +instance + (MultiValue.PseudoRing a, MultiValue.Real a, MultiValue.IntegerConstant a) => + Num (Exp a) where + fromInteger n = lift0 (MultiValue.fromInteger' n) + (+) = liftM2 MultiValue.add + (-) = liftM2 MultiValue.sub + negate = liftM MultiValue.neg + (*) = liftM2 MultiValue.mul + abs = liftM MultiValue.abs + signum = liftM MultiValue.signum + +instance (MultiValue.Field a, MultiValue.Real a, MultiValue.RationalConstant a) => + Fractional (Exp a) where + fromRational n = lift0 (MultiValue.fromRational' n) + (/) = liftM2 MultiValue.fdiv diff --git a/src/Data/Array/Knead/Index/Linear.hs b/src/Data/Array/Knead/Index/Linear.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Index/Linear.hs @@ -0,0 +1,553 @@ +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE EmptyDataDecls #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +module Data.Array.Knead.Index.Linear ( + C(switch), + switchInt, + intersect, + value, + constant, + paramWith, + tunnel, + flattenIndex, + peek, + poke, + computeSize, + + Struct, + T(..), + Z(Z), + (:.)((:.)), + Shape, shape, + Index, index, + cons, (#:.), + head, + tail, + switchR, + loadMultiValue, + storeMultiValue, + ) where + +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Index.Linear.Int as Index + +import qualified Data.Array.Knead.Parameter as Param +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp, ) + +import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory +import qualified LLVM.Extra.Multi.Value as MultiValue +import qualified LLVM.Extra.Arithmetic as A +import qualified LLVM.Extra.Control as C +import LLVM.Extra.Multi.Value (Atom, ) + +import qualified LLVM.Util.Loop as Loop +import qualified LLVM.Core as LLVM + +import qualified Foreign.Storable as St +import Foreign.Storable.FixedArray (sizeOfArray, ) +import Foreign.Marshal.Array (advancePtr, ) +import Foreign.Ptr (Ptr, castPtr, ) + +import Control.Monad (liftM2, ) +import Data.Word (Word32, ) + +import Prelude hiding (min, head, tail, ) + + +class C ix where + switch :: + f Z -> + (forall ix0 i. (C ix0, Index.Single i) => f (ix0 :. i)) -> + f ix + +instance C Z where + switch x _ = x + +instance (C ix0, Index.Single i) => C (ix0 :. i) where + switch _ x = x + + +newtype SwitchInt f ix i = SwitchInt {runSwitchInt :: f (ix :. i)} + +switchInt :: + (C ix) => + f Z -> + (forall ix0. (C ix0) => f (ix0 :. Index.Int)) -> + f ix +switchInt z0 cons0 = + switch z0 + (runSwitchInt $ Index.switchSingle (SwitchInt cons0)) + + +newtype Op2 tag sh = Op2 {runOp2 :: Exp (T tag sh) -> Exp (T tag sh) -> Exp (T tag sh)} + +intersect :: C sh => Exp (Shape sh) -> Exp (Shape sh) -> Exp (Shape sh) +intersect = + runOp2 $ + switchInt + (Op2 $ \z0 _ -> z0) + (Op2 $ + switchR $ \is i -> + switchR $ \js j -> + intersect is js #:. Expr.min i j) + + +_value :: (C sh, MultiValue.C sh) => sh -> Exp sh +_value = Expr.lift0 . MultiValue.cons + + +newtype MakeValue val tag sh = MakeValue {runMakeValue :: T tag sh -> val (T tag sh)} + +value :: (C sh, Expr.Value val) => T tag sh -> val (T tag sh) +value = + runMakeValue $ + switchInt + (MakeValue $ \(Cons Z) -> z) + (MakeValue $ \(Cons (t:.h)) -> + value (Cons t) #:. Expr.lift0 (MultiValue.cons h)) + +paramWith :: + (C sh, Expr.Value val) => + Param.T p (T tag sh) -> + (forall parameters. + (St.Storable parameters, + MultiValueMemory.C parameters) => + (p -> parameters) -> + (MultiValue.T parameters -> val (T tag sh)) -> + a) -> + a +paramWith p f = + case tunnel p of + Param.Tunnel get val -> f get (Expr.lift0 . val) + +tunnel :: (C sh) => Param.T p (T tag sh) -> Param.Tunnel p (T tag sh) +tunnel p = + case structFieldsPropF p of + StructFieldsProp -> Param.tunnel value p + + +data StructFieldsProp sh = LLVM.StructFields (Struct sh) => StructFieldsProp + +_structFieldsProp :: (C sh) => f sh -> StructFieldsProp sh +_structFieldsProp _p = structFieldsRec + +structFieldsPropF :: (C sh) => f (g sh) -> StructFieldsProp sh +structFieldsPropF _p = structFieldsRec + +withStructFieldsPropFF :: + (C sh) => (StructFieldsProp sh -> f (g (h sh))) -> f (g (h sh)) +withStructFieldsPropFF f = f structFieldsRec + +structFieldsRec :: (C sh) => StructFieldsProp sh +structFieldsRec = + switchInt + StructFieldsProp + (succStructFieldsProp structFieldsRec) + +succStructFieldsProp :: + StructFieldsProp sh -> StructFieldsProp (sh:.Index.Int) +succStructFieldsProp StructFieldsProp = StructFieldsProp + + +data Z = Z + deriving (Eq, Ord, Read, Show) + + +infixl 3 :., #:. + +data tail :. head = !tail :. !head + deriving (Eq, Ord, Read, Show) + + +newtype T tag sh = Cons {decons :: sh} + +data ShapeTag +data IndexTag + +type Shape = T ShapeTag +type Index = T IndexTag + +shape :: sh -> Shape sh +shape = Cons + +index :: ix -> Index ix +index = Cons + + +(#:.) :: (Expr.Value val) => val (T tag sh) -> val i -> val (T tag (sh:.i)) +(#:.) = cons + +cons :: (Expr.Value val) => val (T tag sh) -> val i -> val (T tag (sh:.i)) +cons = + Expr.lift2 $ + \(MultiValue.Cons t) (MultiValue.Cons h) -> + MultiValue.Cons (t,h) + +z :: (Expr.Value val) => val (T tag Z) +z = Expr.lift0 $ MultiValue.Cons () + +head :: (Expr.Value val) => val (T tag (sh:.i)) -> val i +head = Expr.lift1 $ \(MultiValue.Cons (_t,h)) -> MultiValue.Cons h + +tail :: (Expr.Value val) => val (T tag (sh:.i)) -> val (T tag sh) +tail = Expr.lift1 $ \(MultiValue.Cons (t,_h)) -> MultiValue.Cons t + +switchR :: + Expr.Value val => + (val (T tag sh) -> val i -> a) -> val (T tag (sh :. i)) -> a +switchR f ix = f (tail ix) (head ix) + + +type family PatternTuple pattern +type family Decomposed (f :: * -> *) tag pattern + +type instance PatternTuple (sh:.s) = + PatternTuple sh :. MultiValue.PatternTuple s + +type instance Decomposed f tag (sh:.s) = + Decomposed f tag sh :. MultiValue.Decomposed f s + +type instance PatternTuple (Atom sh) = sh + +type instance Decomposed f tag (Atom sh) = f (T tag sh) + + +class + (Expr.Composed (Decomposed Exp tag pattern) ~ T tag (PatternTuple pattern)) => + Decompose tag pattern where + decompose :: + T tag pattern -> Exp (T tag (PatternTuple pattern)) -> + Decomposed Exp tag pattern + +instance Decompose tag (Atom sh) where + decompose (Cons _atom) x = x + +instance (Decompose tag sh, Expr.Decompose s) => Decompose tag (sh :. s) where + decompose (Cons (psh:.ps)) x = + decompose (Cons psh) (tail x) :. Expr.decompose ps (head x) + + +type instance MultiValue.PatternTuple (T tag sh) = T tag (PatternTuple sh) + +type instance MultiValue.Decomposed f (T tag sh) = Decomposed f tag sh + + +type family Unwrap sh +type instance Unwrap (T tag sh) = sh + +type family Tag sh +type instance Tag (T tag sh) = tag + +instance + (Expr.Compose sh, + Expr.Composed sh ~ T (Tag (Expr.Composed sh)) (Unwrap (Expr.Composed sh)), + Expr.Compose s) => + Expr.Compose (sh :. s) where + type Composed (sh :. s) = + T (Tag (Expr.Composed sh)) + (Unwrap (Expr.Composed sh) :. Expr.Composed s) + compose (sh :. s) = cons (Expr.compose sh) (Expr.compose s) + +instance (Decompose tag sh) => Expr.Decompose (T tag sh) where + decompose = decompose + + + +instance (C sh) => St.Storable (T tag sh) where + sizeOf (Cons sh) = sizeOfArray (rank sh) (0::Word32) + alignment (Cons _sh) = St.alignment (0::Word32) + poke ptr = poke (castPtr ptr) . decons + peek = fmap Cons . peek . castPtr + + +type family Repr (f :: * -> *) sh +type instance Repr f Z = () +type instance Repr f (tail :. head) = (Repr f tail, MultiValue.Repr f head) + +instance (C sh) => MultiValue.C (T tag sh) where + type Repr f (T tag sh) = Repr f sh + cons = value + undef = constant $ MultiValue.undef + zero = constant $ MultiValue.zero + addPhis = addPhis + phis = phis + +instance (tag ~ ShapeTag, C sh) => Shape.C (T tag sh) where + type Index (T tag sh) = Index sh + size = fromIntegral . size . decons + sizeCode = computeSize + intersectCode = Expr.unliftM2 intersect + flattenIndexRec sh ix = + -- a joint implementation would not be more efficient + liftM2 (,) + (computeSize sh) + (flattenIndex sh ix) + loop = loop + + +type family Struct sh +type instance Struct Z = () +type instance Struct (sh :. Index.Int) = (Word32, Struct sh) + +instance + (C sh, LLVM.StructFields (Struct sh)) => + MultiValueMemory.C (T tag sh) where + type Struct (T tag sh) = LLVM.Struct (Struct sh) + load = loadMultiValue + store = storeMultiValue + +loadMultiValue :: + (C sh) => + LLVM.Value (Ptr (LLVM.Struct (Struct sh))) -> + LLVM.CodeGenFunction r (MultiValue.T (T tag sh)) +loadMultiValue ptr = + withStructFieldsPropFF $ \StructFieldsProp -> + load =<< castPtrValue ptr + +storeMultiValue :: + (C sh) => + MultiValue.T (T tag sh) -> + LLVM.Value (Ptr (LLVM.Struct (Struct sh))) -> LLVM.CodeGenFunction r () +storeMultiValue x ptr = + case structFieldsPropF x of + StructFieldsProp -> store x =<< castPtrValue ptr + + +newtype FlattenIndex r sh = + FlattenIndex { + runFlattenIndex :: + MultiValue.T (Shape sh) -> MultiValue.T (Index sh) -> + LLVM.CodeGenFunction r (LLVM.Value Word32) + } + +flattenIndex :: + (C sh) => + MultiValue.T (Shape sh) -> MultiValue.T (Index sh) -> + LLVM.CodeGenFunction r (LLVM.Value Word32) +flattenIndex = + runFlattenIndex $ + switchInt + (FlattenIndex $ \_zerosh _zeroix -> return A.zero) + (FlattenIndex $ + switchR $ \sh (MultiValue.Cons s) -> + switchR $ \ix (MultiValue.Cons i) -> + A.add i =<< A.mul s =<< flattenIndex sh ix) + + +newtype Rank sh = Rank {runRank :: sh -> Int} + +rank :: (C sh) => sh -> Int +rank = + runRank $ + switch + (Rank $ const 0) + (Rank $ succ . rank . (\(sh :. _s) -> sh)) + + +newtype Peek sh = Peek {runPeek :: Ptr Word32 -> IO sh} + +peek :: (C sh) => Ptr Word32 -> IO sh +peek = + runPeek $ + switchInt + (Peek $ const $ return Z) + (Peek $ \ptr -> do + h <- St.peek ptr + t <- peek $ advancePtr ptr 1 + return (t :. Index.Int h)) + + +newtype Poke sh = Poke {runPoke :: Ptr Word32 -> sh -> IO ()} + +poke :: (C sh) => Ptr Word32 -> sh -> IO () +poke = + runPoke $ + switchInt + (Poke $ const $ const $ return ()) + (Poke $ \ptr (sh :. Index.Int i) -> do + St.poke ptr i + poke (advancePtr ptr 1) sh) + + +castPtrValue :: + (LLVM.StructFields sh) => + LLVM.Value (Ptr (LLVM.Struct sh)) -> + LLVM.CodeGenFunction r (LLVM.Value (Ptr Word32)) +castPtrValue = LLVM.bitcast + +newtype Load r tag sh = + Load { + runLoad :: + LLVM.Value (Ptr Word32) -> + LLVM.CodeGenFunction r (MultiValue.T (T tag sh)) + } + +load :: + (C sh) => + LLVM.Value (Ptr Word32) -> + LLVM.CodeGenFunction r (MultiValue.T (T tag sh)) +load = + runLoad $ + switchInt + (Load $ const $ return z) + (Load $ \ptr -> do + h <- LLVM.load ptr + t <- load =<< A.advanceArrayElementPtr ptr + return (t #:. MultiValue.Cons h)) + + +newtype Store r tag sh = + Store { + runStore :: + MultiValue.T (T tag sh) -> + LLVM.Value (Ptr Word32) -> + LLVM.CodeGenFunction r () + } + +store :: + (C sh) => + MultiValue.T (T tag sh) -> + LLVM.Value (Ptr Word32) -> + LLVM.CodeGenFunction r () +store = + runStore $ + switchInt + (Store $ \_z _ptr -> return ()) + (Store $ switchR $ \sh (MultiValue.Cons k) ptr -> do + LLVM.store k ptr + store sh =<< A.advanceArrayElementPtr ptr) + + +newtype Size sh = + Size { + runSize :: sh -> Word32 + } + +size :: (C sh) => sh -> Word32 +size = + runSize $ + switchInt + (Size $ \_z -> 1) + (Size $ \(sh :. Index.Int k) -> k * size sh) + + +newtype ComputeSize r sh = + ComputeSize { + runComputeSize :: + MultiValue.T (Shape sh) -> + LLVM.CodeGenFunction r (LLVM.Value Word32) + } + +computeSize :: + (C sh) => + MultiValue.T (Shape sh) -> + LLVM.CodeGenFunction r (LLVM.Value Word32) +computeSize = + runComputeSize $ + switchInt + (ComputeSize $ \_z -> return A.one) + (ComputeSize $ switchR $ \sh (MultiValue.Cons k) -> + A.mul k =<< computeSize sh) + + +newtype + Constant val tag sh = + Constant {getConstant :: val Index.Int -> val (T tag sh)} + +constant :: (C sh, Expr.Value val) => val Index.Int -> val (T tag sh) +constant = + getConstant $ + switchInt + (Constant $ const z) + (Constant $ \x -> constant x #:. x) + + +newtype AddPhis r tag sh = + AddPhis { + runAddPhis :: + LLVM.BasicBlock -> + MultiValue.T (T tag sh) -> + MultiValue.T (T tag sh) -> + LLVM.CodeGenFunction r () + } + +addPhis :: + (C sh) => + LLVM.BasicBlock -> + MultiValue.T (T tag sh) -> + MultiValue.T (T tag sh) -> + LLVM.CodeGenFunction r () +addPhis = + runAddPhis $ + switchInt + (AddPhis $ \_ _ _ -> return ()) + (AddPhis $ \bb -> + switchR $ \hx tx -> + switchR $ \hy ty -> + MultiValue.addPhis bb tx ty >> + addPhis bb hx hy) + + +newtype Phis r tag sh = + Phis { + runPhis :: + LLVM.BasicBlock -> + MultiValue.T (T tag sh) -> + LLVM.CodeGenFunction r (MultiValue.T (T tag sh)) + } + +phis :: + (C sh) => + LLVM.BasicBlock -> + MultiValue.T (T tag sh) -> + LLVM.CodeGenFunction r (MultiValue.T (T tag sh)) +phis = + runPhis $ + switchInt + (Phis $ \_ -> return) + (Phis $ \bb -> + switchR $ \h t -> + liftM2 (#:.) + (phis bb h) + (MultiValue.phis bb t)) + + +newtype Loop r state sh = + Loop { + runLoop :: + (MultiValue.T (Index sh) -> + state -> + LLVM.CodeGenFunction r state) -> + MultiValue.T (Shape sh) -> + state -> + LLVM.CodeGenFunction r state + } + +loop :: + (C sh, Loop.Phi state) => + (MultiValue.T (Index sh) -> + state -> + LLVM.CodeGenFunction r state) -> + MultiValue.T (Shape sh) -> + state -> + LLVM.CodeGenFunction r state +loop = + runLoop $ + switchInt + (Loop $ \code _z -> code z) + (Loop $ \code -> switchR $ \sh (MultiValue.Cons n) -> + loop + (\ix ptrStart -> + fmap fst $ + C.fixedLengthLoop n (ptrStart, A.zero) $ \(ptr, k) -> + liftM2 (,) + (code (ix #:. MultiValue.Cons k) ptr) + (A.inc k)) + sh) diff --git a/src/Data/Array/Knead/Index/Linear/Int.hs b/src/Data/Array/Knead/Index/Linear/Int.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Index/Linear/Int.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE TypeFamilies #-} +module Data.Array.Knead.Index.Linear.Int ( + Single(..), + Int(Int), cons, decons, + ) where + +import qualified Data.Array.Knead.Expression as Expr + +import qualified LLVM.Extra.Multi.Value as MultiValue +import qualified LLVM.Extra.Arithmetic as A + +import Data.Word (Word32, ) + +import Prelude hiding (Int, head, tail, ) + + +newtype Int = Int Word32 + +cons :: (Expr.Value val) => val Word32 -> val Int +cons = Expr.lift1 $ \(MultiValue.Cons x) -> MultiValue.Cons x + +decons :: (Expr.Value val) => val Int -> val Word32 +decons = Expr.lift1 $ \(MultiValue.Cons x) -> MultiValue.Cons x + + +class Single ix where + switchSingle :: f Int -> f ix + +instance Single Int where + switchSingle x = x + + +instance MultiValue.C Int where + type Repr f Int = f Word32 + cons (Int x) = MultiValue.consPrimitive x + undef = MultiValue.undefPrimitive + zero = MultiValue.zeroPrimitive + phis = MultiValue.phisPrimitive + addPhis = MultiValue.addPhisPrimitive + +instance MultiValue.Additive Int where + add = MultiValue.liftM2 A.add + sub = MultiValue.liftM2 A.sub + neg = MultiValue.liftM A.neg + +instance MultiValue.Real Int where + min = MultiValue.liftM2 A.min + max = MultiValue.liftM2 A.max + abs = MultiValue.liftM A.abs + signum = MultiValue.liftM A.signum diff --git a/src/Data/Array/Knead/Index/Nested/Shape.hs b/src/Data/Array/Knead/Index/Nested/Shape.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Index/Nested/Shape.hs @@ -0,0 +1,188 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE Rank2Types #-} +module Data.Array.Knead.Index.Nested.Shape where + +import qualified Data.Array.Knead.Expression as Expr +import qualified Data.Array.Knead.Parameter as Param +import Data.Array.Knead.Expression (Exp, ) + +import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory +import qualified LLVM.Extra.Multi.Value as MultiValue +import qualified LLVM.Extra.Arithmetic as A +import qualified LLVM.Extra.Control as C +import LLVM.Extra.Monad (liftR2) + +import qualified LLVM.Util.Loop as Loop +import qualified LLVM.Core as LLVM + +import Foreign.Storable (Storable, ) +import Foreign.Ptr (Ptr, ) + +import Data.Word (Word32, Word64) + +import qualified Control.Monad.HT as Monad + + +value :: (C sh, Expr.Value val) => sh -> val sh +value = Expr.lift0 . MultiValue.cons + +paramWith :: + (Storable b, MultiValueMemory.C b, Expr.Value val) => + Param.T p b -> + (forall parameters. + (Storable parameters, + MultiValueMemory.C parameters) => + (p -> parameters) -> + (MultiValue.T parameters -> val b) -> + a) -> + a +paramWith p f = + Param.withMulti p (\get val -> f get (Expr.lift0 . val)) + +load :: + (MultiValueMemory.C sh) => + f sh -> LLVM.Value (Ptr (MultiValueMemory.Struct sh)) -> + LLVM.CodeGenFunction r (MultiValue.T sh) +load _ = MultiValueMemory.load + +intersect :: (C sh) => Exp sh -> Exp sh -> Exp sh +intersect = Expr.liftM2 intersectCode + +flattenIndex :: + (C sh) => + MultiValue.T sh -> MultiValue.T (Index sh) -> + LLVM.CodeGenFunction r (LLVM.Value Word32) +flattenIndex sh ix = + fmap snd $ flattenIndexRec sh ix + +class (MultiValue.C sh) => C sh where + type Index sh :: * + {- + It would be better to restrict zipWith to matching shapes + and turn shape intersection into a bound check. + -} + intersectCode :: + MultiValue.T sh -> MultiValue.T sh -> + LLVM.CodeGenFunction r (MultiValue.T sh) + sizeCode :: + MultiValue.T sh -> + LLVM.CodeGenFunction r (LLVM.Value Word32) + size :: sh -> Int + {- | + Result is @(size, flattenedIndex)@. + @size@ must equal the result of 'sizeCode'. + We use this for sharing intermediate results. + -} + flattenIndexRec :: + MultiValue.T sh -> MultiValue.T (Index sh) -> + LLVM.CodeGenFunction r (LLVM.Value Word32, LLVM.Value Word32) + loop :: + (Index sh ~ ix, Loop.Phi state) => + (MultiValue.T ix -> state -> LLVM.CodeGenFunction r state) -> + MultiValue.T sh -> state -> LLVM.CodeGenFunction r state + + +instance C () where + type Index () = () + intersectCode _ _ = return $ MultiValue.cons () + sizeCode _ = return A.one + size _ = 1 + flattenIndexRec _ _ = return (A.one, A.zero) + loop = id + + +class C sh => Scalar sh where + scalar :: (Expr.Value val) => val sh + zeroIndex :: (Expr.Value val) => f sh -> val (Index sh) + +instance Scalar () where + scalar = Expr.lift0 $ MultiValue.Cons () + zeroIndex _ = Expr.lift0 $ MultiValue.Cons () + + +loopPrimitive :: + (MultiValue.Repr LLVM.Value i ~ LLVM.Value i, + Num i, LLVM.IsConst i, LLVM.IsInteger i, + LLVM.CmpRet i, LLVM.CmpResult i ~ Bool, Loop.Phi state) => + (MultiValue.T i -> state -> LLVM.CodeGenFunction r state) -> + MultiValue.T i -> state -> LLVM.CodeGenFunction r state +loopPrimitive code (MultiValue.Cons n) ptrStart = + fmap fst $ + C.fixedLengthLoop n (ptrStart, A.zero) $ \(ptr, k) -> + Monad.lift2 (,) + (code (MultiValue.Cons k) ptr) + (A.inc k) + +instance C Word32 where + type Index Word32 = Word32 + intersectCode = MultiValue.min + sizeCode (MultiValue.Cons n) = return n + size = fromIntegral + flattenIndexRec (MultiValue.Cons n) (MultiValue.Cons i) = return (n, i) + loop = loopPrimitive + +instance C Word64 where + type Index Word64 = Word64 + intersectCode = MultiValue.min + sizeCode (MultiValue.Cons n) = LLVM.trunc n + size = fromIntegral + flattenIndexRec (MultiValue.Cons n) (MultiValue.Cons i) = + Monad.lift2 (,) (LLVM.trunc n) (LLVM.trunc i) + loop = loopPrimitive + + +instance (C n, C m) => C (n,m) where + type Index (n,m) = (Index n, Index m) + intersectCode a b = + case (MultiValue.unzip a, MultiValue.unzip b) of + ((an,am), (bn,bm)) -> + Monad.lift2 MultiValue.zip + (intersectCode an bn) + (intersectCode am bm) + sizeCode nm = + case MultiValue.unzip nm of + (n,m) -> liftR2 A.mul (sizeCode n) (sizeCode m) + size (n,m) = size n * size m + flattenIndexRec nm ij = + case (MultiValue.unzip nm, MultiValue.unzip ij) of + ((n,m), (i,j)) -> do + (ns, il) <- flattenIndexRec n i + (ms, jl) <- flattenIndexRec m j + Monad.lift2 (,) + (A.mul ns ms) + (A.add jl =<< A.mul ms il) + loop code nm = + case MultiValue.unzip nm of + (n,m) -> loop (\i -> loop (\j -> code (MultiValue.zip i j)) m) n + +instance (C n, C m, C l) => C (n,m,l) where + type Index (n,m,l) = (Index n, Index m, Index l) + intersectCode a b = + case (MultiValue.unzip3 a, MultiValue.unzip3 b) of + ((ai,aj,ak), (bi,bj,bk)) -> + Monad.lift3 MultiValue.zip3 + (intersectCode ai bi) + (intersectCode aj bj) + (intersectCode ak bk) + sizeCode nml = + case MultiValue.unzip3 nml of + (n,m,l) -> + liftR2 A.mul (sizeCode n) $ + liftR2 A.mul (sizeCode m) (sizeCode l) + size (n,m,l) = size n * size m * size l + flattenIndexRec nml ijk = + case (MultiValue.unzip3 nml, MultiValue.unzip3 ijk) of + ((n,m,l), (i,j,k)) -> do + (ns, il) <- flattenIndexRec n i + (ms, jl) <- flattenIndexRec m j + x0 <- A.add jl =<< A.mul ms il + (ls, kl) <- flattenIndexRec l k + x1 <- A.add kl =<< A.mul ls x0 + sz <- A.mul ns =<< A.mul ms ls + return (sz, x1) + loop code nml = + case MultiValue.unzip3 nml of + (n,m,l) -> + loop (\i -> loop (\j -> loop (\k -> + code (MultiValue.zip3 i j k)) + l) m) n diff --git a/src/Data/Array/Knead/Parameter.hs b/src/Data/Array/Knead/Parameter.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Parameter.hs @@ -0,0 +1,218 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ExistentialQuantification #-} +module Data.Array.Knead.Parameter where + +import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory +import qualified LLVM.Extra.Multi.Value as MultiValue +import qualified LLVM.Extra.Class as Class +import qualified LLVM.Extra.Memory as Memory +import Foreign.Storable.Tuple () +import Foreign.Storable (Storable, ) + +import qualified Control.Category as Cat +import qualified Control.Arrow as Arr +import qualified Control.Applicative as App +import Control.Applicative (pure, liftA2, ) + +import Data.Tuple.HT (mapFst, ) +import Data.Word (Word32, ) + + +{- | +This data type is for parameters of parameterized signal generators and causal processes. +It is better than using plain functions of type @p -> a@ +since it allows for numeric instances +and we can make explicit, +whether a parameter is constant. + +We recommend to use parameters for atomic types. +Although a parameter of type @T p (a,b)@ is possible, +it means that the whole parameter is variable +if only one of the pair elements is variable. +This way you may miss optimizations. +-} +data T p a = + Constant a | + Variable (p -> a) + + +get :: T p a -> (p -> a) +get (Constant a) = const a +get (Variable f) = f + + +{- | +The call @value param v@ requires +that @v@ represents the same value as @valueTupleOf (get param p)@ for some @p@. +However @v@ might be the result of a load operation +and @param@ might be a constant. +In this case it is more efficient to use @valueTupleOf (get param undefined)@ +since the constant is translated to an LLVM constant +that allows for certain optimizations. + +This is the main function for taking advantage of a constant parameter +in low-level implementations. +For simplicity we do not omit constant parameters in the parameter struct +since this would mean to construct types at runtime and might become ugly. +Instead we just check using 'value' at the according places in LLVM code +whether a parameter is constant +and ignore the parameter from the struct in this case. +In many cases there will be no speed benefit +because the parameter will be loaded to a register anyway. +It can only lead to speed-up if subsequent optimizations +can precompute constant expressions. +Another example is 'drop' where a loop with constant loop count can be generated. +For small loop counts and simple loop bodies the loop might get unrolled. +-} +valueTuple :: + (Class.MakeValueTuple tuple, Class.ValueTuple tuple ~ value) => + T p tuple -> value -> value +valueTuple = genericValue Class.valueTupleOf + +multiValue :: + (MultiValue.C a) => + T p a -> MultiValue.T a -> MultiValue.T a +multiValue = genericValue MultiValue.cons + +genericValue :: + (a -> value) -> + T p a -> value -> value +genericValue cons p v = + case p of + Constant a -> cons a + Variable _ -> v + + +{- | +This function provides specialised variants of 'get' and 'value', +that use the unit type for constants +and thus save space in parameter structures. +-} +withTuple :: + (Storable tuple, Class.MakeValueTuple tuple, + Class.ValueTuple tuple ~ value, Memory.C value) => + T p tuple -> + (forall parameters. + (Storable parameters, + Class.MakeValueTuple parameters, + Memory.C (Class.ValueTuple parameters)) => + (p -> parameters) -> + (Class.ValueTuple parameters -> value) -> + a) -> + a +withTuple (Constant a) f = f (const ()) (\() -> Class.valueTupleOf a) +withTuple (Variable v) f = f v id + +withMulti :: + (Storable b, MultiValueMemory.C b) => + T p b -> + (forall parameters. + (Storable parameters, + MultiValueMemory.C parameters) => + (p -> parameters) -> + (MultiValue.T parameters -> MultiValue.T b) -> + a) -> + a +withMulti = with MultiValue.cons + +with :: + (Storable b, MultiValueMemory.C b) => + (b -> MultiValue.T b) -> + T p b -> + (forall parameters. + (Storable parameters, + MultiValueMemory.C parameters) => + (p -> parameters) -> + (MultiValue.T parameters -> MultiValue.T b) -> + a) -> + a +with cons p f = + case p of + Constant b -> f (const ()) (\_ -> cons b) + Variable v -> f v id + + +data Tunnel p a = + forall t. + (Storable t, MultiValueMemory.C t) => + Tunnel (p -> t) (MultiValue.T t -> MultiValue.T a) + +tunnel :: + (Storable a, MultiValueMemory.C a) => + (a -> MultiValue.T a) -> T p a -> Tunnel p a +tunnel cons p = + case p of + Constant b -> Tunnel (const ()) (\_ -> cons b) + Variable v -> Tunnel v id + + +word32 :: T p Int -> T p Word32 +word32 = fmap fromIntegral + + +infixl 0 $# + +($#) :: (T p a -> b) -> (a -> b) +($#) f a = f (pure a) + + +{- | +@.@ can be used for fetching a parameter from a super-parameter. +-} +instance Cat.Category T where + id = Variable id + Constant f . _ = Constant f + Variable f . Constant a = Constant (f a) + Variable f . Variable g = Variable (f . g) + +{- | +@arr@ is useful for lifting parameter selectors to our parameter type +without relying on the constructor. +-} +instance Arr.Arrow T where + arr = Variable + first f = Variable (mapFst (get f)) + + + +{- | +Useful for splitting @T p (a,b)@ into @T p a@ and @T p b@ +using @fmap fst@ and @fmap snd@. +-} +instance Functor (T p) where + fmap f (Constant a) = Constant (f a) + fmap f (Variable g) = Variable (f . g) + +{- | +Useful for combining @T p a@ and @T p b@ to @T p (a,b)@ +using @liftA2 (,)@. +However, we do not recommend to do so +because the result parameter can only be constant +if both operands are constant. +-} +instance App.Applicative (T p) where + pure a = Constant a + Constant f <*> Constant a = Constant (f a) + f <*> a = Variable (\p -> get f p (get a p)) + +instance Monad (T p) where + return = pure + Constant x >>= f = f x + Variable x >>= f = + Variable (\p -> get (f (x p)) p) + + +instance Num a => Num (T p a) where + (+) = liftA2 (+) + (-) = liftA2 (-) + (*) = liftA2 (*) + negate = fmap negate + abs = fmap abs + signum = fmap signum + fromInteger = pure . fromInteger + +instance Fractional a => Fractional (T p a) where + (/) = liftA2 (/) + fromRational = pure . fromRational diff --git a/src/Data/Array/Knead/Parameterized/Physical.hs b/src/Data/Array/Knead/Parameterized/Physical.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Parameterized/Physical.hs @@ -0,0 +1,242 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ForeignFunctionInterface #-} +module Data.Array.Knead.Parameterized.Physical ( + Phys.Array, + Phys.shape, + Phys.fromList, + feed, + the, + render, + renderShape, + scatter, + permute, + ) where + +import qualified Data.Array.Knead.Parameterized.Private as Sym +import qualified Data.Array.Knead.Simple.Physical as Phys +import qualified Data.Array.Knead.Simple.Private as Core +import qualified Data.Array.Knead.Parameter as Param +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp, unExp, ) +import Data.Array.Knead.Code (getElementPtr, compile, ) + +import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory +import qualified LLVM.Extra.Multi.Value as MultiValue +import qualified LLVM.Extra.Arithmetic as A +import qualified LLVM.Extra.Memory as Memory + +import qualified LLVM.Core as LLVM + +import Foreign.Marshal.Utils (with, ) +import Foreign.Marshal.Alloc (alloca, ) +import Foreign.Storable (Storable, peek, ) +import Foreign.ForeignPtr (withForeignPtr, mallocForeignPtrArray, touchForeignPtr, ) +import Foreign.Ptr (FunPtr, Ptr, ) + +import Control.Exception (bracket, ) +import Control.Monad.HT (void, (<=<), ) +import Control.Monad (liftM2, ) +import Data.Tuple.HT (mapFst, ) +import Data.Word (Word32, ) + +import Prelude hiding (scanl1, ) + + +feed :: + (Shape.C sh, Storable sh, MultiValueMemory.C sh, + MultiValueMemory.C a) => + Param.T p (Phys.Array sh a) -> Sym.Array p sh a +feed arr = + Param.withMulti (fmap Phys.shape arr) $ \getShape valueShape -> + Sym.Array + (\p -> + case mapFst valueShape $ MultiValue.unzip p of + (sh, MultiValue.Cons ptr) -> + Core.Array (Expr.lift0 sh) $ + Memory.load <=< getElementPtr sh ptr) + (\p -> + case Phys.buffer $ Param.get arr p of + fptr -> + withForeignPtr fptr $ \ptr -> + return (fptr, (getShape p, MultiValueMemory.castStructPtr ptr))) + touchForeignPtr + + +type Importer f = FunPtr f -> f + +foreign import ccall safe "dynamic" callThe :: + Importer (Ptr param -> Ptr am -> IO ()) + + +the :: + (Shape.Scalar z, MultiValueMemory.C a, Storable a) => + Sym.Array p z a -> IO (p -> IO a) +the (Sym.Array arr create delete) = do + func <- + compile "the" $ + LLVM.createNamedFunction LLVM.ExternalLinkage "eval" $ + \paramPtr resultPtr -> do + param <- Memory.load paramPtr + case arr param of + Core.Array z code -> + code (Shape.zeroIndex z) >>= flip Memory.store resultPtr + return $ \p -> + bracket (create p) (delete . fst) $ \(_ctx, param) -> + with param $ \pptr -> + alloca $ \aptr -> + callThe func (MultiValueMemory.castStructPtr pptr) (MultiValueMemory.castStructPtr aptr) >> + peek aptr + + +foreign import ccall safe "dynamic" callShaper :: + Importer (Ptr param -> Ptr shape -> IO Word32) + +foreign import ccall safe "dynamic" callRenderer :: + Importer (Ptr param -> Ptr shape -> Ptr am -> IO ()) + + +renderShape :: + (Shape.C sh, Storable sh, MultiValueMemory.C sh, + Storable a, MultiValueMemory.C a) => + Sym.Array p sh a -> IO (p -> IO (sh, Word32)) +renderShape (Sym.Array arr create delete) = do + fsh <- + compile "renderShape" $ + LLVM.createNamedFunction LLVM.ExternalLinkage "shape" $ + \paramPtr resultPtr -> do + param <- Memory.load paramPtr + case arr param of + Core.Array esh _code -> do + sh <- unExp esh + MultiValueMemory.store sh resultPtr + Shape.sizeCode sh >>= LLVM.ret + return $ \p -> + bracket (create p) (delete . fst) $ \(_ctx, param) -> + alloca $ \shptr -> + with param $ \pptr -> do + let lpptr = MultiValueMemory.castStructPtr pptr + let lshptr = MultiValueMemory.castStructPtr shptr + n <- callShaper fsh lpptr lshptr + sh <- peek shptr + return (sh, n) + + +render :: + (Shape.C sh, Storable sh, MultiValueMemory.C sh, + Storable a, MultiValueMemory.C a) => + Sym.Array p sh a -> IO (p -> IO (Phys.Array sh a)) +render (Sym.Array arr create delete) = do + (fsh, farr) <- + compile "render" $ + liftM2 (,) + (LLVM.createNamedFunction LLVM.ExternalLinkage "shape" $ + \paramPtr resultPtr -> do + param <- Memory.load paramPtr + case arr param of + Core.Array esh _code -> do + sh <- unExp esh + MultiValueMemory.store sh resultPtr + Shape.sizeCode sh >>= LLVM.ret) + (LLVM.createNamedFunction LLVM.ExternalLinkage "fill" $ + \paramPtr shapePtr bufferPtr -> do + param <- Memory.load paramPtr + case arr param of + Core.Array esh code -> do + let step ix p = do + flip Memory.store p =<< code ix + A.advanceArrayElementPtr p + sh <- Shape.load esh shapePtr + void $ Shape.loop step sh bufferPtr) + return $ \p -> + bracket (create p) (delete . fst) $ \(_ctx, param) -> + alloca $ \shptr -> + with param $ \pptr -> do + let lpptr = MultiValueMemory.castStructPtr pptr + let lshptr = MultiValueMemory.castStructPtr shptr + n <- callShaper fsh lpptr lshptr + fptr <- mallocForeignPtrArray (fromIntegral n) + withForeignPtr fptr $ + callRenderer farr lpptr lshptr . MultiValueMemory.castStructPtr + sh <- peek shptr + return (Phys.Array sh fptr) + + + +foreign import ccall safe "dynamic" callScatterer :: + Importer (Ptr paramBase -> Ptr paramMap -> Ptr shape -> Ptr am -> IO ()) + +scatter :: + (Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + Storable sh1, MultiValueMemory.C sh1, + Storable a, MultiValueMemory.C a) => + (Exp a -> Exp a -> Exp a) -> + Sym.Array p sh1 a -> + Sym.Array p sh0 (ix1, a) -> IO (p -> IO (Phys.Array sh1 a)) +scatter accum + (Sym.Array arrBase createBase deleteBase) + (Sym.Array arrMap createMap deleteMap) = do + + (fsh, farr) <- + compile "scatter" $ + liftM2 (,) + (LLVM.createNamedFunction LLVM.ExternalLinkage "shape" $ + \paramPtr resultPtr -> do + param <- Memory.load paramPtr + case arrBase param of + Core.Array esh _code -> do + sh <- unExp esh + MultiValueMemory.store sh resultPtr + Shape.sizeCode sh >>= LLVM.ret) + (LLVM.createNamedFunction LLVM.ExternalLinkage "fill" $ + \paramBasePtr paramMapPtr shapePtr bufferPtr -> do + paramBase <- Memory.load paramBasePtr + paramMap <- Memory.load paramMapPtr + case (arrBase paramBase, arrMap paramMap) of + (Core.Array esh codeBase, Core.Array eish codeMap) -> do + let clear ix p = do + flip Memory.store p =<< codeBase ix + A.advanceArrayElementPtr p + sh <- Shape.load esh shapePtr + void $ Shape.loop clear sh bufferPtr + + ish <- unExp eish + let fill ix () = do + (jx, a) <- fmap MultiValue.unzip $ codeMap ix + p <- getElementPtr sh bufferPtr jx + flip Memory.store p + =<< Expr.unliftM2 (flip accum) a + =<< Memory.load p + Shape.loop fill ish ()) + + return $ \p -> + bracket (createBase p) (deleteBase . fst) $ \(_ctxBase, paramBase) -> + bracket (createMap p) (deleteMap . fst) $ \(_ctxMap, paramMap) -> + alloca $ \shptr -> + with paramBase $ \paramBasePtr -> do + with paramMap $ \paramMapPtr -> do + let paramBaseMVPtr = MultiValueMemory.castStructPtr paramBasePtr + let paramMapMVPtr = MultiValueMemory.castStructPtr paramMapPtr + let shapeMVPtr = MultiValueMemory.castStructPtr shptr + n <- callShaper fsh paramBaseMVPtr shapeMVPtr + fptr <- mallocForeignPtrArray (fromIntegral n) + withForeignPtr fptr $ + callScatterer farr paramBaseMVPtr paramMapMVPtr shapeMVPtr . + MultiValueMemory.castStructPtr + sh <- peek shptr + return (Phys.Array sh fptr) + +permute :: + (Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + Storable sh1, MultiValueMemory.C sh1, + Storable a, MultiValueMemory.C a) => + (Exp a -> Exp a -> Exp a) -> + Sym.Array p sh1 a -> + (Exp ix0 -> Exp ix1) -> + Sym.Array p sh0 a -> + IO (p -> IO (Phys.Array sh1 a)) +permute accum deflt ixmap input = + scatter accum deflt + (Core.mapWithIndex (Expr.lift2 MultiValue.zip . ixmap) input) diff --git a/src/Data/Array/Knead/Parameterized/Private.hs b/src/Data/Array/Knead/Parameterized/Private.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Parameterized/Private.hs @@ -0,0 +1,194 @@ +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeOperators #-} +module Data.Array.Knead.Parameterized.Private where + +import qualified Data.Array.Knead.Simple.Symbolic as Core + +import qualified Data.Array.Knead.Parameter as Param +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp, ) + +import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory +import qualified LLVM.Extra.Multi.Value as MultiValue + +import Foreign.Storable (Storable, ) + +import Prelude hiding (id, map, zipWith, replicate, ) + + +data Array p sh a = + forall parameter context. + (Storable parameter, MultiValueMemory.C parameter) => + Array { + core :: MultiValue.T parameter -> Core.Array sh a, + createContext :: p -> IO (context, parameter), + deleteContext :: context -> IO () + } + +instance Core.C (Array p) where + lift0 arr = Array (const arr) (createPlain (const ())) deletePlain + lift1 f (Array arr create delete) = Array (f . arr) create delete + lift2 f (Array arrA createA deleteA) (Array arrB createB deleteB) = + Array + (\p -> + case MultiValue.unzip p of + (paramA, paramB) -> f (arrA paramA) (arrB paramB)) + (combineCreate createA createB) + (combineDelete deleteA deleteB) + + +{- +(!) :: (Shape.C sh) => Array p sh a -> Param.T p sh -> Array p z a +(!) arr pix = + paramArray + (\ix carr -> Core.fromScalar $ carr Core.! ix) + (Shape.tunnel pix) + arr +-} +(!) :: + (Shape.C sh, Shape.Index sh ~ ix, MultiValue.C ix, + Storable ix, MultiValueMemory.C ix, + Shape.Scalar z) => + Array p sh a -> Param.T p ix -> Array p z a +(!) arr pix = + paramArray + (\ix carr -> Core.fromScalar $ carr Core.! ix) + (Param.tunnel MultiValue.cons pix) + arr + + +fill :: + (Shape.C sh, Storable sh, MultiValueMemory.C sh, + MultiValue.C a, Storable a, MultiValueMemory.C a) => + Param.T p sh -> Param.T p a -> Array p sh a +fill sh a = + Shape.paramWith sh $ \getSh valueSh -> + Param.withMulti a $ \getA valueA -> + Array + (\p -> + case MultiValue.unzip p of + (vsh, va) -> + Core.fill (valueSh vsh) (Expr.lift0 $ valueA va)) + (createPlain $ \p -> (getSh p, getA p)) + deletePlain + +gather :: + (Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, MultiValue.C a) => + Array p sh1 ix0 -> + Array p sh0 a -> + Array p sh1 a +gather = Core.gather + + +id :: + (Shape.C sh, Storable sh, MultiValueMemory.C sh, Shape.Index sh ~ ix) => + Param.T p sh -> Array p sh ix +id sh = + Shape.paramWith sh $ \getSh valueSh -> + Array + (Core.id . valueSh) + (createPlain getSh) + deletePlain + +map :: + (Shape.C sh, MultiValueMemory.C c, Storable c) => + (Exp c -> Exp a -> Exp b) -> + Param.T p c -> Array p sh a -> Array p sh b +map = lift Core.map + +mapWithIndex :: + (Shape.C sh, MultiValueMemory.C c, Storable c, Shape.Index sh ~ ix) => + (Exp c -> Exp ix -> Exp a -> Exp b) -> + Param.T p c -> Array p sh a -> Array p sh b +mapWithIndex = lift Core.mapWithIndex + + +fold1 :: + (Shape.C sh0, Shape.C sh1, + MultiValueMemory.C c, Storable c, MultiValue.C a) => + (Exp c -> Exp a -> Exp a -> Exp a) -> + Param.T p c -> Array p (sh0, sh1) a -> Array p sh0 a +fold1 = lift Core.fold1 + +fold1All :: + (Shape.C sh, Shape.Scalar z, + MultiValueMemory.C c, Storable c, MultiValue.C a) => + (Exp c -> Exp a -> Exp a -> Exp a) -> + Param.T p c -> Array p sh a -> Array p z a +fold1All = lift Core.fold1All + +lift :: + (Shape.C sh0, Shape.C sh1, + MultiValueMemory.C c, Storable c) => + (f -> Core.Array sh0 a -> Core.Array sh1 b) -> + (Exp c -> f) -> + Param.T p c -> Array p sh0 a -> Array p sh1 b +lift g f c arr = + paramArray + (\cexp -> g (f cexp)) + (Param.tunnel MultiValue.cons c) + arr + +{- +Could be generalized to nested indices. + +foldSelected1 :: + (Fold.C sl, MultiValue.C a) => + (Exp a -> Exp a -> Exp a) -> + Param.T p (Linear.Shape sl) -> + Array p (Linear.Shape (Fold.FullShape sl)) a -> + Array p (Linear.Shape (Fold.FoldShape sl)) a +foldSelected1 f esl arr = + paramArray (Core.foldSelected1 f) (Fold.tunnel esl) arr +-} + + +paramArray :: + (Exp sl -> Core.Array shb b -> Core.Array sha a) -> + Param.Tunnel p sl -> + Array p shb b -> Array p sha a +paramArray f tunnel (Array arr create delete) = + case tunnel of + Param.Tunnel getSl valueSl -> + Array + (\p -> + case MultiValue.unzip p of + (arrp, sl) -> + f (Expr.lift0 $ valueSl sl) $ arr arrp) + (\p -> do + (ctx, param) <- create p + return (ctx, (param, getSl p))) + delete + + +createPlain :: (Monad m) => (p -> pl) -> p -> m ((), pl) +createPlain f p = return ((), f p) + +deletePlain :: (Monad m) => () -> m () +deletePlain () = return () + + +combineCreate :: + Monad m => + (p -> m (ctxA, paramA)) -> (p -> m (ctxB, paramB)) -> + p -> m ((ctxA, ctxB), (paramA, paramB)) +combineCreate createA createB p = do + (ctxA, paramA) <- createA p + (ctxB, paramB) <- createB p + return ((ctxA, ctxB), (paramA, paramB)) + +combineDelete :: + Monad m => + (ctxA -> m ()) -> (ctxB -> m ()) -> (ctxA, ctxB) -> m () +combineDelete deleteA deleteB (ctxA, ctxB) = do + deleteA ctxA + deleteB ctxB + + +extendParameter :: + (q -> p) -> Array p sh a -> Array q sh a +extendParameter f (Array arr create delete) = + Array arr (create . f) delete diff --git a/src/Data/Array/Knead/Parameterized/Slice.hs b/src/Data/Array/Knead/Parameterized/Slice.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Parameterized/Slice.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE TypeOperators #-} +module Data.Array.Knead.Parameterized.Slice ( + T, + apply, + Linear, + passAny, + pass, + pick, + extrude, + (Core.$:.), + ) where + +import qualified Data.Array.Knead.Parameterized.Private as Priv +import Data.Array.Knead.Parameterized.Private (Array(Array), ) + +import qualified Data.Array.Knead.Simple.Slice as Slice +import qualified Data.Array.Knead.Simple.Private as Core + +import qualified Data.Array.Knead.Index.Linear as Linear +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Parameter as Param +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp, ) +import Data.Array.Knead.Index.Linear ((:.), ) + +import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory +import qualified LLVM.Extra.Multi.Value as MultiValue + +import Foreign.Storable (Storable, ) + + +{- +This wrapper data type is pretty much the same as Parameterized.Array +but there seems to be no benefit from using the same data structure for it. +-} +data T p sh0 sh1 = + forall parameter context. + (Storable parameter, MultiValueMemory.C parameter) => + Cons { + _core :: MultiValue.T parameter -> Slice.T sh0 sh1, + _createContext :: p -> IO (context, parameter), + _deleteContext :: context -> IO () + } + +apply :: + (Shape.C sh0, Shape.C sh1, MultiValue.C a) => + T p sh0 sh1 -> + Array p sh0 a -> + Array p sh1 a +apply (Cons slice createSlice deleteSlice) (Array arr createArr deleteArr) = + Array + (\p -> + case MultiValue.unzip p of + (paramSlice, paramArr) -> + Slice.apply (slice paramSlice) (arr paramArr)) + (Priv.combineCreate createSlice createArr) + (Priv.combineDelete deleteSlice deleteArr) + + +type Linear p sh0 sh1 = T p (Linear.Shape sh0) (Linear.Shape sh1) + + +passAny :: Linear p sh sh +passAny = + Cons (const Slice.passAny) (Priv.createPlain $ const ()) Priv.deletePlain + +pass :: + Linear p sh0 sh1 -> + Linear p (sh0:.i) (sh1:.i) +pass (Cons slice create delete) = Cons (Slice.pass . slice) create delete + +pick :: + (MultiValueMemory.C i, Storable i) => + Param.T p i -> + Linear p sh0 sh1 -> + Linear p (sh0:.i) sh1 +pick = lift Slice.pick + +extrude :: + (MultiValueMemory.C i, Storable i) => + Param.T p i -> + Linear p sh0 sh1 -> + Linear p sh0 (sh1:.i) +extrude = lift Slice.extrude + +lift :: + (MultiValueMemory.C i, Storable i) => + (Exp i -> Slice.Linear sh0 sh1 -> Slice.Linear sh2 sh3) -> + Param.T p i -> + Linear p sh0 sh1 -> Linear p sh2 sh3 +lift f i (Cons slice create delete) = + Param.withMulti i $ \getI valueI -> + Cons + (\p -> + case MultiValue.unzip p of + (slicep, ip) -> + f (Expr.lift0 (valueI ip)) (slice slicep)) + (\p -> do + (ctx, param) <- create p + return (ctx, (param, getI p))) + delete + +instance Core.Process (T p sh0 sh1) where diff --git a/src/Data/Array/Knead/Parameterized/Symbolic.hs b/src/Data/Array/Knead/Parameterized/Symbolic.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Parameterized/Symbolic.hs @@ -0,0 +1,66 @@ +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeOperators #-} +module Data.Array.Knead.Parameterized.Symbolic ( + Array, + Exp, + Sym.extendParameter, + (Sym.!), + Sym.fill, + gather, + backpermute, + Sym.id, + Sym.map, + zipWith, + Sym.fold1, + Sym.fold1All, + ) where + +import qualified Data.Array.Knead.Parameterized.Private as Sym +import qualified Data.Array.Knead.Simple.Symbolic as Core +import Data.Array.Knead.Parameterized.Private (Array, gather, ) + +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Parameter as Param +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp, ) + +import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory +import qualified LLVM.Extra.Multi.Value as MultiValue + +import Foreign.Storable (Storable, ) + +import Prelude (uncurry, ($), ) + + +{- +fromScalar :: + (Storable a, MultiValueMemory.C a, MultiValue.C a) => + Param.T p a -> Array p Z a +fromScalar = Sym.fill (return Z) +-} + + +backpermute :: + (Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + Storable sh1, MultiValueMemory.C sh1, + MultiValue.C a) => + Param.T p sh1 -> + (Exp ix1 -> Exp ix0) -> + Array p sh0 a -> + Array p sh1 a +backpermute sh1 f = gather (Core.map f (Sym.id sh1)) + +{- +_backpermute sh1 f = + paramArray (flip Core.backpermute f) (Shape.tunnel sh1) +-} + + +zipWith :: + (Shape.C sh, MultiValueMemory.C d, Storable d) => + (Exp d -> Exp a -> Exp b -> Exp c) -> + Param.T p d -> Array p sh a -> Array p sh b -> Array p sh c +zipWith f d a b = + Sym.map (\di ab -> uncurry (f di) $ Expr.unzip ab) d $ Core.zip a b diff --git a/src/Data/Array/Knead/Simple/Fold.hs b/src/Data/Array/Knead/Simple/Fold.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Simple/Fold.hs @@ -0,0 +1,94 @@ +{- | +Reduce selected dimensions. +Alternatively you may reorder dimensions with 'ShapeDep.backpermute' +and fold once along a multiple dimensions. +-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +module Data.Array.Knead.Simple.Fold ( + T, + Linear, + apply, + passAny, + pass, + fold, + (Core.$:.), + ) where + +import qualified Data.Array.Knead.Simple.Private as Core +import Data.Array.Knead.Simple.Private (Array(Array), Code, Val, ) + +import qualified Data.Array.Knead.Index.Linear as Linear +import qualified Data.Array.Knead.Index.Linear.Int as IndexInt +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp, unExp, ) +import Data.Array.Knead.Index.Linear ((#:.), (:.)((:.)), ) + +import qualified LLVM.Extra.Multi.Value as MultiValue +import LLVM.Extra.Multi.Value (atom, ) + +import Prelude hiding (zipWith, zipWith3, zip, zip3, replicate, ) + + +data T sh0 sh1 a = + forall ix0 ix1. + (Shape.Index sh0 ~ ix0, Shape.Index sh1 ~ ix1) => + Cons + (Exp sh0 -> Exp sh1) + (forall r. Val sh0 -> (Val ix0 -> Code r a) -> (Val ix1 -> Code r a)) + + +apply :: + (Core.C array, Shape.C sh0, Shape.C sh1, MultiValue.C a) => + T sh0 sh1 a -> + array sh0 a -> + array sh1 a +apply (Cons fsh reduce) = + Core.lift1 $ \(Array sh code) -> + Array (fsh sh) (\ix -> do sh0 <- unExp sh; reduce sh0 code ix) + + +type Linear sh0 sh1 = T (Linear.Shape sh0) (Linear.Shape sh1) + +passAny :: Linear sh sh a +passAny = Cons id (const id) + +pass :: + Linear sh0 sh1 a -> + Linear (sh0:.i) (sh1:.i) a +pass (Cons fsh reduce) = + Cons + (Expr.modify (Linear.shape (atom:.atom)) $ \(sh:.s) -> fsh sh :. s) + (\sh code -> + Linear.switchR $ \jx j -> + reduce (Linear.tail sh) (\kx -> code (kx #:. j)) jx) + + +fold1CodeLinear :: + (MultiValue.C a) => + (Exp a -> Exp a -> Exp a) -> + Exp IndexInt.Int -> + (Val (Linear.Index (sh :. IndexInt.Int)) -> Code r a) -> + (Val (Linear.Index sh) -> Code r a) +fold1CodeLinear f nc code ix = + Core.fold1Code f (IndexInt.decons nc) + (\jx j -> code (jx #:. IndexInt.cons j)) + ix + +fold :: + (MultiValue.C a) => + (Exp a -> Exp a -> Exp a) -> + Linear sh0 sh1 a -> + Linear (sh0:.IndexInt.Int) sh1 a +fold f (Cons fsh reduce) = + Cons + (fsh . Linear.tail) + (\sh code jx -> + reduce (Linear.tail sh) + (fold1CodeLinear f (Expr.lift0 (Linear.head sh)) code) jx) + + +instance Core.Process (T sh0 sh1 a) where diff --git a/src/Data/Array/Knead/Simple/Physical.hs b/src/Data/Array/Knead/Simple/Physical.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Simple/Physical.hs @@ -0,0 +1,198 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ForeignFunctionInterface #-} +module Data.Array.Knead.Simple.Physical ( + Array(Array, shape, buffer), -- data constructor intended for PhysicalParameterized + fromList, + vectorFromList, + with, + render, + scanl1, + scatter, + permute, + ) where + +import qualified Data.Array.Knead.Simple.Private as Sym +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp, unExp, ) +import Data.Array.Knead.Code (getElementPtr, compile, ) + +import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory +import qualified LLVM.Extra.Multi.Value as MultiValue +import qualified LLVM.Extra.Arithmetic as A +import qualified LLVM.Extra.Control as C +import qualified LLVM.Extra.Memory as Memory +import qualified LLVM.Extra.Maybe as Maybe + +import qualified LLVM.Core as LLVM + +import Foreign.Marshal.Array (pokeArray, ) +import Foreign.Marshal.Alloc (alloca, ) +import Foreign.Storable (Storable, peek, ) +import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, mallocForeignPtrArray, ) +import Foreign.Ptr (FunPtr, Ptr, ) + +import Control.Monad.HT (void, ) +import Control.Monad (liftM2, ) +import Data.Word (Word32, ) + +import Prelude hiding (scanl1, ) + + +data Array sh a = + Array { + shape :: sh, + buffer :: ForeignPtr a + } + + +fromList :: + (Shape.C sh, Storable a) => + sh -> [a] -> IO (Array sh a) +fromList sh xs = do + let size = Shape.size sh + fptr <- mallocForeignPtrArray size + withForeignPtr fptr $ + \ptr -> + pokeArray ptr $ + take size $ + xs ++ repeat (error "Array.Knead.Physical.fromList: list too short for shape") + return (Array sh fptr) + +vectorFromList :: + (Shape.C sh, Num sh, Storable a) => + [a] -> IO (Array sh a) +vectorFromList xs = do + let size = length xs + fptr <- mallocForeignPtrArray size + withForeignPtr fptr $ flip pokeArray xs + return (Array (fromIntegral size) fptr) + + +{- | +The symbolic array is only valid inside the enclosed action. +-} +with :: + (Shape.C sh, MultiValueMemory.C a) => + (Sym.Array sh a -> IO b) -> + Array sh a -> IO b +with f (Array sh fptr) = + withForeignPtr fptr $ \ptr -> + f $ + Sym.Array + (Shape.value sh) + (\ix -> + Memory.load =<< + getElementPtr (Shape.value sh) + (LLVM.valueOf (MultiValueMemory.castStructPtr ptr)) ix) + + +type Importer f = FunPtr f -> f + +foreign import ccall safe "dynamic" callShaper :: + Importer (Ptr sh -> IO Word32) + +foreign import ccall safe "dynamic" callRenderer :: + Importer (Ptr sh -> Ptr am -> IO ()) + + +materialize :: + (Shape.C sh, Storable sh, MultiValueMemory.C sh, + Storable a, MultiValueMemory.C a, MultiValueMemory.Struct a ~ am) => + String -> + Exp sh -> + (LLVM.Value (Ptr (MultiValueMemory.Struct sh)) -> + LLVM.Value (Ptr am) -> LLVM.CodeGenFunction () ()) -> + IO (Array sh a) +materialize name esh code = + alloca $ \shptr -> do + (fsh, farr) <- + compile name $ + liftM2 (,) + (LLVM.createNamedFunction LLVM.ExternalLinkage "shape" $ \ptr -> do + sh <- unExp esh + MultiValueMemory.store sh ptr + Shape.sizeCode sh >>= LLVM.ret) + (LLVM.createNamedFunction LLVM.ExternalLinkage "fill" code) + let lshptr = MultiValueMemory.castStructPtr shptr + n <- callShaper fsh lshptr + fptr <- mallocForeignPtrArray (fromIntegral n) + withForeignPtr fptr $ + callRenderer farr lshptr . MultiValueMemory.castStructPtr + sh <- peek shptr + return (Array sh fptr) + +render :: + (Shape.C sh, Storable sh, MultiValueMemory.C sh, + Storable a, MultiValueMemory.C a) => + Sym.Array sh a -> IO (Array sh a) +render (Sym.Array esh code) = + materialize "render" esh $ \sptr ptr -> do + let step ix p = do + flip Memory.store p =<< code ix + A.advanceArrayElementPtr p + sh <- Shape.load esh sptr + void $ Shape.loop step sh ptr + +scanl1 :: + (Shape.C sh, Storable sh, MultiValueMemory.C sh, + Storable a, MultiValueMemory.C a) => + (Exp a -> Exp a -> Exp a) -> + Sym.Array (sh, Word32) a -> IO (Array (sh, Word32) a) +scanl1 f (Sym.Array esh code) = + materialize "scanl1" esh $ \sptr ptr -> do + (sh, MultiValue.Cons n) <- + fmap MultiValue.unzip $ Shape.load esh sptr + let step ix ptrStart = + fmap (fst.fst) $ + C.fixedLengthLoop n ((ptrStart, A.zero), Maybe.nothing) $ \((ptr0, k0), macc0) -> do + a <- code (MultiValue.zip ix $ MultiValue.Cons k0) + acc1 <- Maybe.run macc0 (return a) (flip (Expr.unliftM2 f) a) + flip Memory.store ptr0 acc1 + ptrK1 <- + liftM2 (,) + (A.advanceArrayElementPtr ptr0) + (A.inc k0) + return (ptrK1, Maybe.just acc1) + void $ Shape.loop step sh ptr + +scatter :: + (Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + Storable sh1, MultiValueMemory.C sh1, + Storable a, MultiValueMemory.C a) => + (Exp a -> Exp a -> Exp a) -> + Sym.Array sh1 a -> + Sym.Array sh0 (ix1, a) -> IO (Array sh1 a) +scatter accum (Sym.Array esh defltCode) (Sym.Array eish code) = + materialize "scatter" esh $ \sptr ptr -> do + let clear ix p = do + flip Memory.store p =<< defltCode ix + A.advanceArrayElementPtr p + sh <- Shape.load esh sptr + void $ Shape.loop clear sh ptr + + ish <- unExp eish + let fill ix () = do + (jx, a) <- fmap MultiValue.unzip $ code ix + p <- getElementPtr sh ptr jx + flip Memory.store p + =<< Expr.unliftM2 (flip accum) a + =<< Memory.load p + void $ Shape.loop fill ish () + +permute :: + (Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + Storable sh1, MultiValueMemory.C sh1, + Storable a, MultiValueMemory.C a) => + (Exp a -> Exp a -> Exp a) -> + Sym.Array sh1 a -> + (Exp ix0 -> Exp ix1) -> + Sym.Array sh0 a -> + IO (Array sh1 a) +permute accum deflt ixmap input = + scatter accum deflt + (Sym.mapWithIndex (Expr.lift2 MultiValue.zip . ixmap) input) diff --git a/src/Data/Array/Knead/Simple/Private.hs b/src/Data/Array/Knead/Simple/Private.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Simple/Private.hs @@ -0,0 +1,172 @@ +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +module Data.Array.Knead.Simple.Private where + +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp(Exp), ) + +import qualified LLVM.Extra.Multi.Value as MultiValue +import qualified LLVM.Extra.Monad as Monad +import qualified LLVM.Extra.Maybe as Maybe +import qualified LLVM.Core as LLVM + +import qualified Control.Category as Cat +import Control.Monad ((<=<), ) + +import Prelude hiding (id, map, zipWith, replicate, ) + + +type Val = MultiValue.T +type Code r a = LLVM.CodeGenFunction r (Val a) + +data Array sh a = + Array (Exp sh) (forall r. Val (Shape.Index sh) -> Code r a) + +shape :: Array sh a -> Exp sh +shape (Array sh _) = sh + +(!) :: + (Shape.C sh, Shape.Index sh ~ ix) => + Array sh a -> Exp ix -> Exp a +(!) (Array _ code) (Exp ix) = Exp (code =<< ix) + +the :: (Shape.Scalar sh) => Array sh a -> Exp a +the (Array z code) = Exp (code $ Shape.zeroIndex z) + +fromScalar :: (Shape.Scalar sh) => Exp a -> Array sh a +fromScalar = fill Shape.scalar + + +fill :: Exp sh -> Exp a -> Array sh a +fill sh (Exp code) = Array sh (\_z -> code) + + +{- | +This class allows to implement functions without parameters +for both simple and parameterized arrays. +-} +class C array where + lift0 :: Array sh a -> array sh a + lift1 :: (Array sha a -> Array shb b) -> array sha a -> array shb b + lift2 :: + (Array sha a -> Array shb b -> Array shc c) -> + array sha a -> array shb b -> array shc c + +instance C Array where + lift0 = Cat.id + lift1 = Cat.id + lift2 = Cat.id + + +gather :: + (C array, + Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + MultiValue.C a) => + array sh1 ix0 -> + array sh0 a -> + array sh1 a +gather = + lift2 $ \(Array sh1 f) (Array _sh0 code) -> + Array sh1 (code <=< f) + +backpermute2 :: + (C array, + Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + Shape.C sh, Shape.Index sh ~ ix) => + Exp sh -> + (Exp ix -> Exp ix0) -> + (Exp ix -> Exp ix1) -> + (Exp a -> Exp b -> Exp c) -> + array sh0 a -> array sh1 b -> array sh c +backpermute2 sh projectIndex0 projectIndex1 f = + lift2 $ \(Array _sha codeA) (Array _shb codeB) -> + Array sh + (\ix -> + Monad.liftR2 (Expr.unliftM2 f) + (codeA =<< Expr.unliftM1 projectIndex0 ix) + (codeB =<< Expr.unliftM1 projectIndex1 ix)) + + +id :: + (Shape.C sh, Shape.Index sh ~ ix) => + Exp sh -> Array sh ix +id sh = Array sh return + +map :: + (C array, Shape.C sh) => + (Exp a -> Exp b) -> + array sh a -> array sh b +map f = + lift1 $ \(Array sh code) -> + Array sh (Expr.unliftM1 f <=< code) + +mapWithIndex :: + (C array, Shape.C sh, Shape.Index sh ~ ix) => + (Exp ix -> Exp a -> Exp b) -> + array sh a -> array sh b +mapWithIndex f = + lift1 $ \(Array sh code) -> + Array sh (\ix -> Expr.unliftM2 f ix =<< code ix) + + +fold1Code :: + (Shape.C sh1, Shape.Index sh1 ~ ix1, MultiValue.C a) => + (Exp a -> Exp a -> Exp a) -> + Exp sh1 -> + (Val ix0 -> Val ix1 -> Code r a) -> + (Val ix0 -> Code r a) +fold1Code f (Exp nc) code ix = do + n <- nc + fmap Maybe.fromJust $ + Shape.loop + (\i0 macc0 -> do + a <- code ix i0 + acc1 <- Maybe.run macc0 (return a) (flip (Expr.unliftM2 f) a) + return $ Maybe.just acc1) + n Maybe.nothing + +fold1 :: + (C array, Shape.C sh0, Shape.C sh1, MultiValue.C a) => + (Exp a -> Exp a -> Exp a) -> + array (sh0, sh1) a -> array sh0 a +fold1 f = + lift1 $ \(Array shs code) -> + case Expr.unzip shs of + (sh, s) -> Array sh $ fold1Code f s $ MultiValue.curry code + + +fold1All :: + (Shape.C sh, MultiValue.C a) => + (Exp a -> Exp a -> Exp a) -> + Array sh a -> Array () a +fold1All f (Array esh code) = + fold1 f $ + Array + (Expr.lift1 (MultiValue.zip (MultiValue.Cons ())) esh) + (code . MultiValue.snd) + + + +class Process proc where + + +infixl 3 $:. + +{- | +Use this for combining several dimension manipulators. +E.g. + +> apply (passAny $:. pick 3 $:. pass $:. replicate 10) array + +The constraint @(Process proc0, Process proc1)@ is a bit weak. +We like to enforce that the type constructor like @Slice.T@ +is the same in @proc0@ and @proc1@, and only the parameters differ. +Currently this coherence is achieved, +because we only provide functions of type @proc0 -> proc1@ with this condition. +-} +($:.) :: (Process proc0, Process proc1) => proc0 -> (proc0 -> proc1) -> proc1 +($:.) = flip ($) diff --git a/src/Data/Array/Knead/Simple/ShapeDependent.hs b/src/Data/Array/Knead/Simple/ShapeDependent.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Simple/ShapeDependent.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE TypeFamilies #-} +module Data.Array.Knead.Simple.ShapeDependent where + +import qualified Data.Array.Knead.Simple.Private as Core +import Data.Array.Knead.Simple.Private (Array(Array), ) + +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp, ) + +import qualified LLVM.Extra.Monad as Monad + +import Control.Monad ((<=<), ) + + +shape :: (Core.C array, Shape.C sh, Shape.Scalar z) => array sh a -> array z sh +shape = Core.lift1 $ Core.fromScalar . Core.shape + +backpermute :: + (Core.C array, + Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1) => + (Exp sh0 -> Exp sh1) -> + (Exp ix1 -> Exp ix0) -> + array sh0 a -> + array sh1 a +backpermute createShape projectIndex = + Core.lift1 $ \(Array sh code) -> + Array (createShape sh) + (code <=< Expr.unliftM1 projectIndex) + +{- | +This is between 'backpermute' and 'backpermute2'. +You can access the shapes of two arrays, +but only the content of one of them. +This is necessary if the second array contributes only a virtual dimension. +-} +backpermuteExtra :: + (Core.C array, + Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + Shape.C sh, Shape.Index sh ~ ix) => + (Exp sh0 -> Exp sh1 -> Exp sh) -> + (Exp ix -> Exp ix0) -> + array sh0 a -> array sh1 b -> array sh a +backpermuteExtra newShape projectIndex = + Core.lift2 $ \(Array sh0 code) (Array sh1 _code) -> + Array (newShape sh0 sh1) + (\ix -> code =<< Expr.unliftM1 projectIndex ix) + +backpermute2 :: + (Core.C array, + Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + Shape.C sh, Shape.Index sh ~ ix) => + (Exp sh0 -> Exp sh1 -> Exp sh) -> + (Exp ix -> Exp ix0) -> + (Exp ix -> Exp ix1) -> + (Exp a -> Exp b -> Exp c) -> + array sh0 a -> array sh1 b -> array sh c +backpermute2 combineShape projectIndex0 projectIndex1 f = + Core.lift2 $ \(Array sha codeA) (Array shb codeB) -> + Array (combineShape sha shb) + (\ix -> + Monad.liftR2 (Expr.unliftM2 f) + (codeA =<< Expr.unliftM1 projectIndex0 ix) + (codeB =<< Expr.unliftM1 projectIndex1 ix)) + +fill :: + (Core.C array) => + (Exp sh0 -> Exp sh1) -> Exp b -> + array sh0 a -> array sh1 b +fill fsh a = + Core.lift1 $ \arr -> + Core.fill (fsh $ Core.shape arr) a diff --git a/src/Data/Array/Knead/Simple/Slice.hs b/src/Data/Array/Knead/Simple/Slice.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Simple/Slice.hs @@ -0,0 +1,143 @@ +{- | +Generate and apply index maps. +This unifies the @replicate@ and @slice@ functions of the @accelerate@ package. +However the structure of slicing and replicating cannot depend on parameters. +If you need that, you must use 'ShapeDep.backpermute' and friends. +-} +{- +Some notes on the design choice: + +Instead of the shallow embedding implemented by the 'T' type, +we could maintain a symbolic representation of the Slice and Replicate pattern, +like the accelerate package does. +We actually used that representation in former versions. +It has however some drawbacks: + +* We need additional type functions that map from the pattern + to the source and the target shape and we need a proof, + that the images of these type functions are actually shapes. + This worked already, but was rather cumbersome. + +* We need a way to store and pass this pattern through the Parameter handler. + This yields new problems: + We need a wrapper type for wrapping Index, Shape, Slice, Replicate, Fold patterns. + Then the question is whether we use one Wrap type with a phantom parameter + or whether we define a Wrap type for every pattern type. + That is, the options are to write either + + > Wrap Shape (Z:.Int:.Int) + + or + + > Shape (Z:.Int:.Int) + + The first one seems to save us many duplicate instances of + Storable, MultiValue etc. + and it allows us easily to reuse the (:.) for all kinds of patterns. + However, we need a way to restrict the element type of the (:.)-list elements. + We can define that using variable ConstraintKinds, + but e.g. we are not able to add a Storable superclass constraint + to the instance Storable (Wrap constr). + That is, we are left with the second option + and had to define a lot of similar Storable, MultiValue instances. +-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +module Data.Array.Knead.Simple.Slice ( + T, + Linear, + apply, + passAny, + pass, + pick, + extrude, + (Core.$:.), + ) where + +import qualified Data.Array.Knead.Simple.ShapeDependent as ShapeDep +import qualified Data.Array.Knead.Simple.Private as Core + +import qualified Data.Array.Knead.Index.Linear as Linear +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Index.Linear ((#:.), (:.)((:.)), ) +import Data.Array.Knead.Expression (Exp, ) + +import qualified LLVM.Extra.Multi.Value as MultiValue +import LLVM.Extra.Multi.Value (atom, ) + +import Prelude hiding (zipWith, zipWith3, zip, zip3, replicate, ) + + + +{- +This data type is almost identical to Core.Array. +The only difference is, +that the shape @sh1@ in T can depend on another shape @sh0@. +-} +data T sh0 sh1 = + forall ix0 ix1. + (Shape.Index sh0 ~ ix0, Shape.Index sh1 ~ ix1) => + Cons + (Exp sh0 -> Exp sh1) + (Exp ix1 -> Exp ix0) + +{- | +This is essentially a 'ShapeDep.backpermute'. +-} +apply :: + (Core.C array, Shape.C sh0, Shape.C sh1, MultiValue.C a) => + T sh0 sh1 -> + array sh0 a -> + array sh1 a +apply (Cons fsh fix) = + ShapeDep.backpermute fsh fix + + +type Linear sh0 sh1 = T (Linear.Shape sh0) (Linear.Shape sh1) + +{- | +Like @Any@ in @accelerate@. +-} +passAny :: Linear sh sh +passAny = Cons id id + +{- | +Like @All@ in @accelerate@. +-} +pass :: + Linear sh0 sh1 -> + Linear (sh0:.i) (sh1:.i) +pass (Cons fsh fix) = + Cons + (Expr.modify (Linear.shape (atom:.atom)) $ \(sh:.s) -> fsh sh :. s) + (Expr.modify (Linear.index (atom:.atom)) $ \(ix:.i) -> fix ix :. i) + +{- | +Like @Int@ in @accelerate/slice@. +-} +pick :: + Exp i -> + Linear sh0 sh1 -> + Linear (sh0:.i) sh1 +pick i (Cons fsh fix) = + Cons + (fsh . Linear.tail) + (\ix -> fix ix #:. i) + +{- | +Like @Int@ in @accelerate/replicate@. +-} +extrude :: + Exp i -> + Linear sh0 sh1 -> + Linear sh0 (sh1:.i) +extrude n (Cons fsh fix) = + Cons + (\sh -> fsh sh #:. n) + (fix . Linear.tail) + + +instance Core.Process (T sh0 sh1) where diff --git a/src/Data/Array/Knead/Simple/Symbolic.hs b/src/Data/Array/Knead/Simple/Symbolic.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Knead/Simple/Symbolic.hs @@ -0,0 +1,100 @@ +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeOperators #-} +module Data.Array.Knead.Simple.Symbolic ( + Core.Array, + Core.C(..), + Exp, + shape, + (Core.!), + Core.the, + Core.fromScalar, + Core.fill, + gather, + backpermute, + Core.backpermute2, + Core.id, + Core.map, + Core.mapWithIndex, + zipWith, + zipWith3, + zipWith4, + zip, + zip3, + zip4, + fold1, + fold1All, + ) where + +import qualified Data.Array.Knead.Simple.ShapeDependent as ShapeDep +import qualified Data.Array.Knead.Simple.Private as Core +import Data.Array.Knead.Simple.Private (Array, shape, fold1, gather, ) + +import qualified Data.Array.Knead.Index.Nested.Shape as Shape +import qualified Data.Array.Knead.Expression as Expr +import Data.Array.Knead.Expression (Exp, ) + +import qualified LLVM.Extra.Multi.Value as MultiValue + +import Prelude hiding (zipWith, zipWith3, zip, zip3, replicate, ) + + +backpermute :: + (Shape.C sh0, Shape.Index sh0 ~ ix0, + Shape.C sh1, Shape.Index sh1 ~ ix1, + MultiValue.C a) => + Exp sh1 -> + (Exp ix1 -> Exp ix0) -> + Array sh0 a -> + Array sh1 a +backpermute sh1 f = gather (Core.map f (Core.id sh1)) + +zipWith :: + (Core.C array, Shape.C sh) => + (Exp a -> Exp b -> Exp c) -> + array sh a -> array sh b -> array sh c +zipWith = ShapeDep.backpermute2 Shape.intersect id id + +zipWith3 :: + (Core.C array, Shape.C sh) => + (Exp a -> Exp b -> Exp c -> Exp d) -> + array sh a -> array sh b -> array sh c -> array sh d +zipWith3 f a b c = + zipWith (\ab -> uncurry f (Expr.unzip ab)) (zipWith Expr.zip a b) c + +zipWith4 :: + (Core.C array, Shape.C sh) => + (Exp a -> Exp b -> Exp c -> Exp d -> Exp e) -> + array sh a -> array sh b -> array sh c -> array sh d -> array sh e +zipWith4 f a b c d = + zipWith3 (\ab -> uncurry f (Expr.unzip ab)) (zipWith Expr.zip a b) c d + + +zip :: + (Core.C array, Shape.C sh) => + array sh a -> array sh b -> array sh (a,b) +zip = zipWith (Expr.lift2 MultiValue.zip) + +zip3 :: + (Core.C array, Shape.C sh) => + array sh a -> array sh b -> array sh c -> array sh (a,b,c) +zip3 = zipWith3 (Expr.lift3 MultiValue.zip3) + +zip4 :: + (Core.C array, Shape.C sh) => + array sh a -> array sh b -> array sh c -> array sh d -> + array sh (a,b,c,d) +zip4 = zipWith4 (Expr.lift4 MultiValue.zip4) + + +fold1All :: + (Core.C array, Shape.C sh, Shape.Scalar z, MultiValue.C a) => + (Exp a -> Exp a -> Exp a) -> + array sh a -> array z a +fold1All f = + Core.lift1 $ \arr -> + fold1 f $ + backpermute + (Expr.lift2 MultiValue.zip Shape.scalar (shape arr)) + (Expr.lift1 MultiValue.snd) + arr