diff --git a/examples/categorifier-c-examples.cabal b/examples/categorifier-c-examples.cabal index 037ec0c..10e0ed0 100644 --- a/examples/categorifier-c-examples.cabal +++ b/examples/categorifier-c-examples.cabal @@ -278,3 +278,46 @@ executable th-compile , categorifier-unconcat-category , categorifier-unconcat-integration default-language: Haskell2010 + +executable grad-descent + hs-source-dirs: grad-descent + main-is: Main.hs + other-modules: + F + ghc-options: + -O0 + -ddump-splices + -fexpose-all-unfoldings + -fmax-simplifier-iterations=0 + -fno-ignore-interface-pragmas + -fno-omit-interface-pragmas + -Wall + -fplugin Categorifier + -fplugin-opt Categorifier:hierarchy:Categorifier.Hierarchy.UnconCat.hierarchy + -fplugin-opt Categorifier:hierarchy:Categorifier.Hierarchy.ConCat.functionHierarchy + -fplugin-opt Categorifier:hierarchy:Categorifier.Hierarchy.ConCatExtensions.hierarchy + -fplugin-opt Categorifier:maker-map:Categorifier.Core.MakerMap.baseMakerMapFun + -fplugin-opt Categorifier:maker-map:Categorifier.C.MakerMap.cMakerMapFun + -fplugin-opt Categorifier:autointerpreter:Categorifier.C.UnconCat.tryAutoInterpret + -optc-Wno-implicit-function-declaration + -fprint-potential-instances + build-depends: + base + , ghc-prim + , concat-classes + , categorifier-c + , categorifier-c-hk-classes + , categorifier-c-maker-map + , categorifier-c-test-lib + , categorifier-c-unconcat + , categorifier-category + , categorifier-client + , categorifier-concat-extensions-category + , categorifier-concat-extensions-integration + , categorifier-concat-integration + , categorifier-plugin + , categorifier-unconcat-category + , categorifier-unconcat-integration + , ad + , reflection + default-language: Haskell2010 diff --git a/examples/grad-descent/F.hs b/examples/grad-descent/F.hs new file mode 100644 index 0000000..001ab5f --- /dev/null +++ b/examples/grad-descent/F.hs @@ -0,0 +1,132 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +module F + ( Input (..), + Output (..), + Param (..), + XY (..), + rosenbrock, + dRosenbrock, + wrap_rosenbrockF, + wrap_dRosenbrockF, + ) +where + +import qualified Categorifier.C.CExpr.Cat as C +import Categorifier.C.CExpr.Cat.TargetOb (TargetOb, TargetObTC1) +import Categorifier.C.CExpr.Types.Core (CExpr) +import Categorifier.C.CTypes.CGeneric (CGeneric) +import qualified Categorifier.C.CTypes.CGeneric as CG +import Categorifier.C.CTypes.GArrays (GArrays) +import Categorifier.C.KTypes.C (C) +import Categorifier.C.KTypes.Function (kFunctionCall) +import Categorifier.C.KTypes.KType1 (KType1) +import qualified Categorifier.Categorify as Categorify +import qualified Categorifier.Category as Category +import Categorifier.Client (deriveHasRep) +import Data.Int (Int32) +import Data.Proxy (Proxy (..)) +import Data.Reflection (Reifies) +import Data.Word (Word64) +import GHC.Generics (Generic) +import Numeric.AD (grad) +import Numeric.AD.Internal.Reverse (Reverse (Lift), Tape) + +data Param f = Param + { paramA :: f Double, + paramB :: f Double + } + deriving (Generic) + +deriving instance Show (Param C) + +deriveHasRep ''Param + +instance CGeneric (Param f) + +instance GArrays f (Param f) + +type instance TargetOb (Param f) = Param (TargetObTC1 f) + +data XY f = XY + { xyX :: f Double, + xyY :: f Double + } + deriving (Generic) + +deriving instance Show (XY C) + +deriveHasRep ''XY + +instance CGeneric (XY f) + +instance GArrays f (XY f) + +type instance TargetOb (XY f) = XY (TargetObTC1 f) + +data Input f = Input + { iParam :: Param f, + iCoord :: XY f + } + deriving (Generic) + +deriving instance Show (Input C) + +deriveHasRep ''Input + +instance CGeneric (Input f) + +instance GArrays f (Input f) + +type instance TargetOb (Input f) = Input (TargetObTC1 f) + +newtype Output f = Output + {oF :: f Double} + deriving (Generic) + +deriving instance Show (Output C) + +deriveHasRep ''Output + +instance CGeneric (Output f) + +instance GArrays f (Output f) + +type instance TargetOb (Output f) = Output (TargetObTC1 f) + +rosenbrock :: Num a => (a, a) -> (a, a) -> a +rosenbrock (a, b) (x, y) = (a - x) ^ 2 + b * (y - x ^ 2) ^ 2 + +dRosenbrock :: forall a. Num a => (a, a) -> (a, a) -> (a, a) +dRosenbrock (a, b) (x, y) = + let rosenbrock' :: forall s. Reifies s Tape => [Reverse s a] -> Reverse s a + rosenbrock' [x', y'] = + let a' = Lift a + b' = Lift b + in rosenbrock (a', b') (x', y') + [dfdx, dfdy] = grad rosenbrock' [x, y] + in (dfdx, dfdy) + +rosenbrockF :: KType1 f => Input f -> Output f +rosenbrockF (Input (Param a b) (XY x y)) = Output $ rosenbrock (a, b) (x, y) + +dRosenbrockF :: forall f. (KType1 f) => Input f -> XY f +dRosenbrockF (Input (Param a b) (XY x y)) = + let (dfdx, dfdy) = dRosenbrock (a, b) (x, y) + in XY dfdx dfdy + +$(Categorify.separately 'rosenbrockF [t|C.Cat|] [pure [t|C|]]) + +$(Categorify.separately 'dRosenbrockF [t|C.Cat|] [pure [t|C|]]) diff --git a/examples/grad-descent/Main.hs b/examples/grad-descent/Main.hs new file mode 100644 index 0000000..b88e548 --- /dev/null +++ b/examples/grad-descent/Main.hs @@ -0,0 +1,68 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} + +module Main where + +import Categorifier.C.Codegen.FFI.TH (embedFunction) +import Categorifier.C.KTypes.C (C (unsafeC)) +import Categorifier.C.KTypes.KLiteral (kliteral) +import Data.Foldable (traverse_) +import F + ( Input (..), + Output (..), + Param (..), + XY (..), + dRosenbrock, + rosenbrock, + wrap_dRosenbrockF, + wrap_rosenbrockF, + ) + +$(embedFunction "rosenbrockF" wrap_rosenbrockF) + +$(embedFunction "dRosenbrockF" wrap_dRosenbrockF) + +gamma :: Double +gamma = 0.01 + +step :: + ((Double, Double) -> IO Double) -> + ((Double, Double) -> IO (Double, Double)) -> + (Double, Double) -> + IO (Double, Double) +step _f df (x0, y0) = do + (dfdx, dfdy) <- df (x0, y0) + let (x1, y1) = (x0 - gamma * dfdx, y0 - gamma * dfdy) + pure (x1, y1) + +iterateNM :: (Monad m) => Int -> (a -> m a) -> a -> m [a] +iterateNM n f x0 = go n x0 id + where + go k x acc + | k > 0 = do + y <- f x + go (k - 1) y (acc . (y :)) + | otherwise = pure (acc []) + +main :: IO () +main = do + let (x0, y0) = (0.1, 0.4) + -- pure haskell + putStrLn "pure haskell" + let f = pure . rosenbrock (1, 10) + df = pure . dRosenbrock (1, 10) + histH <- iterateNM 10 (step f df) (x0, y0) + traverse_ print histH + + -- C + putStrLn "codegen C" + let g (x, y) = do + Output z <- hs_rosenbrockF (Input (Param 1 10) (XY (kliteral x) (kliteral y))) + pure (unsafeC z) + dg (x, y) = do + XY x' y' <- hs_dRosenbrockF (Input (Param 1 10) (XY (kliteral x) (kliteral y))) + pure (unsafeC x', unsafeC y') + histC <- iterateNM 10 (step g dg) (x0, y0) + traverse_ print histC diff --git a/examples/separate-categorification/Main.hs b/examples/separate-categorification/Main.hs index 80ea0d5..f0fdd80 100644 --- a/examples/separate-categorification/Main.hs +++ b/examples/separate-categorification/Main.hs @@ -7,4 +7,4 @@ import F (fCategorified) -- This generates /tmp/separate_categorification.c main :: IO () -main = writeCFiles "/tmp" "separate_categorification" fCategorified +main = writeCFiles "." "separate_categorification" fCategorified