Debug printing of expression tree (default off)

3 files changed, 111 insertions, 22 deletions

diff --git a/Debug.hs b/Debug.hs
new file mode 100644
index 0000000..bb670d9
--- /dev/null
+++ b/Debug.hs
@@ -0,0 +1,69 @@
+{-# LANGUAGE GADTs #-}
+module Debug where
+
+import Data.Array.Accelerate.AST
+
+
+showPrimFun :: PrimFun a -> String
+showPrimFun PrimAdd{}                = "PrimAdd"
+showPrimFun PrimSub{}                = "PrimSub"
+showPrimFun PrimMul{}                = "PrimMul"
+showPrimFun PrimNeg{}                = "PrimNeg"
+showPrimFun PrimAbs{}                = "PrimAbs"
+showPrimFun PrimSig{}                = "PrimSig"
+showPrimFun PrimQuot{}               = "PrimQuot"
+showPrimFun PrimRem{}                = "PrimRem"
+showPrimFun PrimQuotRem{}            = "PrimQuotRem"
+showPrimFun PrimIDiv{}               = "PrimIDiv"
+showPrimFun PrimMod{}                = "PrimMod"
+showPrimFun PrimDivMod{}             = "PrimDivMod"
+showPrimFun PrimBAnd{}               = "PrimBAnd"
+showPrimFun PrimBOr{}                = "PrimBOr"
+showPrimFun PrimBXor{}               = "PrimBXor"
+showPrimFun PrimBNot{}               = "PrimBNot"
+showPrimFun PrimBShiftL{}            = "PrimBShiftL"
+showPrimFun PrimBShiftR{}            = "PrimBShiftR"
+showPrimFun PrimBRotateL{}           = "PrimBRotateL"
+showPrimFun PrimBRotateR{}           = "PrimBRotateR"
+showPrimFun PrimPopCount{}           = "PrimPopCount"
+showPrimFun PrimCountLeadingZeros{}  = "PrimCountLeadingZeros"
+showPrimFun PrimCountTrailingZeros{} = "PrimCountTrailingZeros"
+showPrimFun PrimFDiv{}               = "PrimFDiv"
+showPrimFun PrimRecip{}              = "PrimRecip"
+showPrimFun PrimSin{}                = "PrimSin"
+showPrimFun PrimCos{}                = "PrimCos"
+showPrimFun PrimTan{}                = "PrimTan"
+showPrimFun PrimAsin{}               = "PrimAsin"
+showPrimFun PrimAcos{}               = "PrimAcos"
+showPrimFun PrimAtan{}               = "PrimAtan"
+showPrimFun PrimSinh{}               = "PrimSinh"
+showPrimFun PrimCosh{}               = "PrimCosh"
+showPrimFun PrimTanh{}               = "PrimTanh"
+showPrimFun PrimAsinh{}              = "PrimAsinh"
+showPrimFun PrimAcosh{}              = "PrimAcosh"
+showPrimFun PrimAtanh{}              = "PrimAtanh"
+showPrimFun PrimExpFloating{}        = "PrimExpFloating"
+showPrimFun PrimSqrt{}               = "PrimSqrt"
+showPrimFun PrimLog{}                = "PrimLog"
+showPrimFun PrimFPow{}               = "PrimFPow"
+showPrimFun PrimLogBase{}            = "PrimLogBase"
+showPrimFun PrimTruncate{}           = "PrimTruncate"
+showPrimFun PrimRound{}              = "PrimRound"
+showPrimFun PrimFloor{}              = "PrimFloor"
+showPrimFun PrimCeiling{}            = "PrimCeiling"
+showPrimFun PrimAtan2{}              = "PrimAtan2"
+showPrimFun PrimIsNaN{}              = "PrimIsNaN"
+showPrimFun PrimIsInfinite{}         = "PrimIsInfinite"
+showPrimFun PrimLt{}                 = "PrimLt"
+showPrimFun PrimGt{}                 = "PrimGt"
+showPrimFun PrimLtEq{}               = "PrimLtEq"
+showPrimFun PrimGtEq{}               = "PrimGtEq"
+showPrimFun PrimEq{}                 = "PrimEq"
+showPrimFun PrimNEq{}                = "PrimNEq"
+showPrimFun PrimMax{}                = "PrimMax"
+showPrimFun PrimMin{}                = "PrimMin"
+showPrimFun PrimLAnd                 = "PrimLAnd"
+showPrimFun PrimLOr                  = "PrimLOr"
+showPrimFun PrimLNot                 = "PrimLNot"
+showPrimFun PrimFromIntegral{}       = "PrimFromIntegral"
+showPrimFun PrimToFloating{}         = "PrimToFloating"
diff --git a/SC/Exp.hs b/SC/Exp.hs
index ef18a7f..cf4e096 100644
--- a/SC/Exp.hs
+++ b/SC/Exp.hs
@@ -4,12 +4,16 @@ module SC.Exp where
 
 import qualified Data.Array.Accelerate.AST as A
 import Data.Array.Accelerate.AST.LeftHandSide
+import Data.Array.Accelerate.AST.Idx (idxToInt)
 import Data.Array.Accelerate.AST.Var
 import Data.Array.Accelerate.Representation.Array
 import Data.Array.Accelerate.Representation.Shape
 import Data.Array.Accelerate.Representation.Type
 import Data.Array.Accelerate.Type
 
+-- import Debug.Trace
+
+import Debug
 import qualified Language.C as C
 import SC.Defs
 import SC.Monad
@@ -41,7 +45,8 @@ compileFun aenv (A.Lam lhs (A.Body body)) = do
     (argnames, env) <- genVarsEnv lhs VENil
     outnames <- itupmap (\(TypedName t n) -> TypedName (C.TPtr t) n)
                     <$> genVars (A.expType body)
-    (usedA, res) <- compileExp' aenv env body
+    ((_tree, usedA), res) <- compileExp' aenv env body
+    -- traceM ("Compiled expression:\n" ++ prettyTree "  " "  " tree)
     (sts1, retexprs) <- toStExprs (A.expType body) res
     let sts2 = genoutstores outnames retexprs
         arrayarguments =
@@ -71,27 +76,38 @@ compileFun _ _ = error "compileFun: Not single-argument function"
 compileExp :: AVarEnv aenv -> A.Exp aenv t -> SC (CompiledFun aenv () t)
 compileExp aenv expr = compileFun aenv (A.Lam (LeftHandSideWildcard TupRunit) (A.Body expr))
 
+data Tree = Node String [Tree] | Leaf String
+
+prettyTree :: String -> String -> Tree -> String
+prettyTree pre _    (Leaf s) = pre ++ s ++ "\n"
+prettyTree pre pre2 (Node s []) = prettyTree pre pre2 (Leaf s)
+prettyTree pre pre2 (Node s ts) =
+    let (ts1, t2) = (init ts, last ts)
+    in pre ++ s ++ "\n" ++ concatMap (prettyTree (pre2 ++ "├─") (pre2 ++ "│ ")) ts1 ++ prettyTree (pre2 ++ "└─") (pre2 ++ "  ") t2
+
 compileExp' :: AVarEnv aenv -> VarEnv env -> A.OpenExp env aenv t
-            -> SC ([SomeArray], Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
+            -> SC ((Tree, [SomeArray]), Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
 compileExp' aenv env = \case
+    -- Foreign, IndexSlice, IndexFull, FromIndex, Case, Cond, While, PrimConst, ShapeSize, Undef, Coerce
+
     A.Let lhs rhs body -> do
         (names, env') <- genVarsEnv lhs env
         let sts1 = [C.SDecl t n Nothing | TypedName t n <- itupList names]
-        (usedA2, sts2) <- fmap (`toStoring` names) <$> compileExp' aenv env rhs
-        (usedA3, res3) <- compileExp' aenv env' body
-        return (usedA2 ++ usedA3
+        ((tree2, usedA2), sts2) <- fmap (`toStoring` names) <$> compileExp' aenv env rhs
+        ((tree3, usedA3), res3) <- compileExp' aenv env' body
+        return ((Node ("Let [" ++ show (length (itupList names)) ++ " vars]") [tree2, tree3], usedA2 ++ usedA3)
                ,fmap (\(sts, exs) -> (sts1 ++ sts2 ++ sts, exs)) res3)
 
     A.Evar (Var _ idx) ->
-        return ([], Right ([], ITupSingle (C.EVar (veprj env idx))))
+        return ((Leaf ("Evar " ++ show (idxToInt idx)), []), Right ([], ITupSingle (C.EVar (veprj env idx))))
 
     A.Nil ->
-        return ([], Right ([], ITupIgnore))
+        return ((Leaf "Nil", []), Right ([], ITupIgnore))
 
     A.Pair a b -> do
-        (usedA1, res1) <- compileExp' aenv env a
-        (usedA2, res2) <- compileExp' aenv env b
-        return (usedA1 ++ usedA2
+        ((tree1, usedA1), res1) <- compileExp' aenv env a
+        ((tree2, usedA2), res2) <- compileExp' aenv env b
+        return ((Node "Pair" [tree1, tree2], usedA1 ++ usedA2)
                ,case (res1, res2) of
                   (Right (sts1, exp1), Right (sts2, exp2)) ->
                       Right (sts1 ++ sts2, ITupPair exp1 exp2)
@@ -102,20 +118,21 @@ compileExp' aenv env = \case
 
     A.Const ty x
       | Just str <- showExpConst ty x
-      -> return ([], Right ([], ITupSingle (C.ELit str)))
+      -> return ((Leaf ("Const (" ++ str ++ ")"), []), Right ([], ITupSingle (C.ELit str)))
 
     A.PrimApp (A.PrimAdd _) e -> binary aenv env "+" e
     A.PrimApp (A.PrimSub _) e -> binary aenv env "-" e
     A.PrimApp (A.PrimMul _) e -> binary aenv env "*" e
     A.PrimApp (A.PrimQuot _) e -> binary aenv env "/" e
     A.PrimApp (A.PrimRem _) e -> binary aenv env "%" e
+    A.PrimApp op _ -> throw $ "Unsupported Exp primitive operator: " ++ showPrimFun op
 
     A.Shape (Var _ idx) ->
         let (shnames, _) = aveprj aenv idx
             buildExprs :: ShNames sh -> Exprs sh
             buildExprs ShZ = ITupIgnore
             buildExprs (ShS names n) = ITupPair (buildExprs names) (ITupSingle (C.EVar n))
-        in return ([], Right ([], buildExprs shnames))
+        in return ((Leaf ("Shape a" ++ show (idxToInt idx)), []), Right ([], buildExprs shnames))
 
     A.ToIndex shr she idxe -> do
         let build :: ShapeR sh -> Exprs sh -> Exprs sh -> C.Expr
@@ -125,11 +142,11 @@ compileExp' aenv env = \case
                                     (ITupPair idxes' (ITupSingle idxe')) =
                 C.EOp (C.EOp (build shr' shes' idxes') "*" she') "+" idxe'
             build _ _ _ = error "wat"
-        (usedA1, res1) <- compileExp' aenv env she
+        ((tree1, usedA1), res1) <- compileExp' aenv env she
         (sts1, shes) <- toStExprs (shapeType shr) res1
-        (usedA2, res2) <- compileExp' aenv env idxe
+        ((tree2, usedA2), res2) <- compileExp' aenv env idxe
         (sts2, idxes) <- toStExprs (shapeType shr) res2
-        return (usedA1 ++ usedA2, Right (sts1 ++ sts2, ITupSingle (build shr shes idxes)))
+        return ((Node "ToIndex" [tree1, tree2], usedA1 ++ usedA2), Right (sts1 ++ sts2, ITupSingle (build shr shes idxes)))
 
     A.Index avar@(Var (ArrayR shr _) _) she ->
         compileExp' aenv env $
@@ -138,22 +155,24 @@ compileExp' aenv env = \case
     A.LinearIndex (Var _ idx) e -> do
         temp <- genName "i"
         let sts0 = [C.SDecl (C.TInt C.B64) temp Nothing]
-        (usedA1, sts1) <- fmap (`toStoring` ITupSingle (TypedName (C.TInt C.B64) temp))
-                              <$> compileExp' aenv env e
+        ((tree1, usedA1), sts1) <-
+            fmap (`toStoring` ITupSingle (TypedName (C.TInt C.B64) temp))
+              <$> compileExp' aenv env e
         let (shnames, anames) = aveprj aenv idx
             usedA = SomeArray shnames anames : usedA1
-        return (usedA, Right (sts0 ++ sts1
-                             ,itupmap (\(TypedAName _ name) -> C.EIndex name (C.EVar temp)) anames))
+        return ((Node ("LinearIndex a" ++ show (idxToInt idx)) [tree1], usedA)
+               ,Right (sts0 ++ sts1
+                      ,itupmap (\(TypedAName _ name) -> C.EIndex name (C.EVar temp)) anames))
 
     e -> throw $ "Unsupported Exp constructor: " ++ A.showExpOp e
   where
     binary :: AVarEnv aenv -> VarEnv env -> String -> A.OpenExp env aenv (a, b)
-           -> SC ([SomeArray], Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
+           -> SC ((Tree, [SomeArray]), Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t))
     binary aenv' env' op e' = do
-        (usedA, res) <- compileExp' aenv' env' e'
+        ((tree, usedA), res) <- compileExp' aenv' env' e'
         (sts, ITupPair (ITupSingle e1) (ITupSingle e2)) <-
             toStExprs (A.expType e') res
-        return (usedA, Right (sts, ITupSingle (C.EOp e1 op e2)))
+        return ((Node ("binary " ++ show op) [tree], usedA), Right (sts, ITupSingle (C.EOp e1 op e2)))
 
 toStExprs :: TypeR t -> Either (Names t -> [C.Stmt]) ([C.Stmt], Exprs t) -> SC ([C.Stmt], Exprs t)
 toStExprs ty (Left fun) = do
diff --git a/accelerate-sc.cabal b/accelerate-sc.cabal
index 6283c53..378bb2a 100644
--- a/accelerate-sc.cabal
+++ b/accelerate-sc.cabal
@@ -12,6 +12,7 @@ library
     Data.Array.Accelerate.C
   other-modules:
     Data.Array.Accelerate.Trafo.UnDelayed
+    Debug
     Language.C
     Language.C.Print
     SC.Acc