Better tracing of AD values in interpreter.

CHANGELOG.md

@@ -16,6 +16,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 * Sizes that go out of scope due to use of higher order functions will
   now work in more cases by adding existentials. (#2193)
 
+* Tracing inside AD operators with the interpreter now prints values
+  properly.
+
 ### Changed
 
 ## [0.25.24]

src/Language/Futhark/Interpreter.hs

@@ -264,8 +264,8 @@ asInteger :: Value -> Integer
 asInteger (ValuePrim (SignedValue v)) = P.valueIntegral v
 asInteger (ValuePrim (UnsignedValue v)) =
   toInteger (P.valueIntegral (P.doZExt v Int64) :: Word64)
-asInteger (ValueAD d v)
-  | P.IntValue v' <- AD.primitive $ AD.primal $ AD.Variable d v =
+asInteger (ValueAD _ v)
+  | P.IntValue v' <- AD.varPrimal v =
       P.valueIntegral v'
 asInteger v = error $ "Unexpectedly not an integer: " <> show v
 
@@ -274,17 +274,17 @@ asInt = fromIntegral . asInteger
 
 asSigned :: Value -> IntValue
 asSigned (ValuePrim (SignedValue v)) = v
-asSigned (ValueAD d v)
-  | P.IntValue v' <- AD.primitive $ AD.primal $ AD.Variable d v = v'
+asSigned (ValueAD _ v)
+  | P.IntValue v' <- AD.varPrimal v = v'
 asSigned v = error $ "Unexpectedly not a signed integer: " <> show v
 
 asInt64 :: Value -> Int64
 asInt64 = fromIntegral . asInteger
 
 asBool :: Value -> Bool
 asBool (ValuePrim (BoolValue x)) = x
-asBool (ValueAD d v)
-  | P.BoolValue v' <- AD.primitive $ AD.primal $ AD.Variable d v = v'
+asBool (ValueAD _ v)
+  | P.BoolValue v' <- AD.varPrimal v = v'
 asBool v = error $ "Unexpectedly not a boolean: " <> show v
 
 lookupInEnv ::

src/Language/Futhark/Interpreter/AD.hs

@@ -7,9 +7,9 @@ module Language.Futhark.Interpreter.AD
     JVPValue (..),
     doOp,
     addFor,
-    primal,
     tapePrimal,
     primitive,
+    varPrimal,
     deriveTape,
   )
 where
@@ -97,9 +97,13 @@ primal (Variable _ (JVP (JVPValue v _))) = primal v
 primal (Constant v) = Constant v
 
 primitive :: ADValue -> PrimValue
-primitive v@(Variable _ _) = primitive $ primal v
+primitive (Variable _ v) = varPrimal v
 primitive (Constant v) = v
 
+varPrimal :: ADVariable -> PrimValue
+varPrimal (VJP (VJPValue t)) = primitive $ tapePrimal t
+varPrimal (JVP (JVPValue v _)) = primitive $ primal v
+
 -- Evaluates a PrimExp using doOp
 evalPrimExp :: M.Map VName ADValue -> PrimExp VName -> Maybe ADValue
 evalPrimExp m (LeafExp n _) = M.lookup n m

src/Language/Futhark/Interpreter/Values.hs

@@ -152,11 +152,15 @@ prettyValueWith pprPrim = pprPrec 0
     pprPrec _ ValueAcc {} = "#<acc>"
     pprPrec p (ValueSum _ n vs) =
       parensIf (p > (0 :: Int)) $ "#" <> sep (pretty n : map (pprPrec 1) vs)
-    -- TODO: This could be prettier. Perhaps add pretty printing for ADVariable / ADValues
-    pprPrec _ (ValueAD d v) = pretty $ "d[" ++ show d ++ "]" ++ show v
+    pprPrec _ (ValueAD _ v) = pprPrim $ putV $ AD.varPrimal v
     pprElem v@ValueArray {} = pprPrec 0 v
     pprElem v = group $ pprPrec 0 v
 
+    putV (P.IntValue x) = SignedValue x
+    putV (P.FloatValue x) = FloatValue x
+    putV (P.BoolValue x) = BoolValue x
+    putV P.UnitValue = BoolValue True
+
 -- | Prettyprint value.
 prettyValue :: Value m -> Doc a
 prettyValue = prettyValueWith pprPrim