don't SLV unit arrays

dpvanbalen · May 16, 2024 · 6149f49 · 6149f49
1 parent 54a1a86
commit 6149f49
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/src/Data/Array/Accelerate/Pretty/Partitioned.hs b/src/Data/Array/Accelerate/Pretty/Partitioned.hs
@@ -50,7 +50,7 @@ instance PrettyOp op => PrettyOp (Cluster op) where
   prettyOp (Op (SLV (SOp (SOAOp op _) _) _) _) = prettyOp op
   prettyOpWithArgs env (Fused f l r) args = "Fused" -- (" <> prettyOpWithArgs env l (left f args) <> ", " <> prettyOpWithArgs env r (right f args)
   prettyOpWithArgs env (Op (SLV (SOp (SOAOp op soa) (SA _ unsort)) subargs) _) args = 
-    prettyOpWithArgs env op (soaShrink combine soa . unsort . slv' varout subargs $ args)
+    prettyOp op -- WithArgs env op (soaShrink combine soa . unsort . slv' varout subargs $ args)
 
   -- prettyOpWithArgs :: forall env t. Val env -> Cluster op t -> Args env t -> Adoc
   -- prettyOpWithArgs env (Op (SLVOp (SOp (SOAOp op soa) (SA _ unsort)) sa) _) args = prettyOpWithArgs env op (soaShrink combine soa . unsort . slv' varToOut sa $ args)

diff --git a/src/Data/Array/Accelerate/Trafo.hs b/src/Data/Array/Accelerate/Trafo.hs
@@ -100,8 +100,8 @@ testWithObjective obj f
   ++ Pretty.renderForTerminal (Pretty.prettyAfun operation)
   ++ "\n\nPartitionedAcc:\n"
   ++ Pretty.renderForTerminal (Pretty.prettyAfun partitioned)
-  -- ++ "\nSLV'd PartitionedAcc:\n"
-  -- ++ Pretty.renderForTerminal (Pretty.prettyAfun slvpartitioned)
+  ++ "\nSLV'd PartitionedAcc:\n"
+  ++ Pretty.renderForTerminal (Pretty.prettyAfun slvpartitioned)
   ++ "\n\nSchedule:\n"
   ++ Pretty.renderForTerminal (Pretty.prettySchedule schedule)
   where

diff --git a/src/Data/Array/Accelerate/Trafo/Operation/LiveVars.hs b/src/Data/Array/Accelerate/Trafo/Operation/LiveVars.hs
@@ -338,7 +338,10 @@ reEnvSubArgs re (a :>: as)
   | ReEnvSubArgs subs as' <- reEnvSubArgs re as =
     case a of
       ArgArray Out (ArrayR shr tp) sh buffers -> case reEnvSubBuffers re tp buffers of
-        ReEnvSubBuffers SubTupRskip _        -> ReEnvSubArgs (SubArgsDead subs) (ArgVar (fromGrounds $ expectJust $ reEnvVars re sh) :>: as')
+        ReEnvSubBuffers SubTupRskip _
+          -- don't kill unit outputs
+          | TupRunit <- tp                   -> ReEnvSubArgs (SubArgsLive SubArgKeep subs) (ArgArray Out (ArrayR shr tp) (expectJust $ reEnvVars re sh) TupRunit :>: as')
+          | otherwise                        -> ReEnvSubArgs (SubArgsDead subs) (ArgVar (fromGrounds $ expectJust $ reEnvVars re sh) :>: as')
         ReEnvSubBuffers SubTupRkeep buffers' -> ReEnvSubArgs (SubArgsLive SubArgKeep subs) (ArgArray Out (ArrayR shr tp) (expectJust $ reEnvVars re sh) buffers' :>: as')
         ReEnvSubBuffers sub         buffers' -> ReEnvSubArgs (SubArgsLive (SubArgOut sub) subs) (ArgArray Out (ArrayR shr $ subTupR sub tp) (expectJust $ reEnvVars re sh) buffers' :>: as')
       _ -> ReEnvSubArgs (SubArgsLive SubArgKeep subs) (reEnvArg re a :>: as')
@@ -374,4 +377,4 @@ subTupUniqueness :: SubTupR t t' -> Uniquenesses t -> Uniquenesses t'
 subTupUniqueness SubTupRskip         _                = TupRunit
 subTupUniqueness SubTupRkeep         t                = t
 subTupUniqueness (SubTupRpair s1 s2) (TupRpair t1 t2) = subTupUniqueness s1 t1 `TupRpair` subTupUniqueness s2 t2
-subTupUniqueness (SubTupRpair s1 s2) (TupRsingle Shared) = TupRsingle $ Shared
+subTupUniqueness (SubTupRpair s1 s2) (TupRsingle Shared) = TupRpair (subTupUniqueness s1 (TupRsingle Shared)) (subTupUniqueness s2 (TupRsingle Shared)) --TupRsingle Shared