Merge branch 'master' into flattening

diku-dk · Dec 6, 2024 · f71c51e · f71c51e
2 parents be27284 + daa7181
commit f71c51e
Show file tree

Hide file tree

Showing 17 changed files with 145 additions and 154 deletions.
diff --git a/src/Futhark/AD/Fwd.hs b/src/Futhark/AD/Fwd.hs
@@ -264,13 +264,13 @@ zeroFromSubExp (Var v) = do
   letExp "zero" $ zeroExp t
 
 fwdSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM ()
-fwdSOAC pat aux (Screma size xs (ScremaForm scs reds f)) = do
+fwdSOAC pat aux (Screma size xs (ScremaForm f scs reds)) = do
   pat' <- bundleNewPat pat
   xs' <- bundleTangents xs
+  f' <- fwdLambda f
   scs' <- mapM fwdScan scs
   reds' <- mapM fwdRed reds
-  f' <- fwdLambda f
-  addStm $ Let pat' aux $ Op $ Screma size xs' $ ScremaForm scs' reds' f'
+  addStm $ Let pat' aux $ Op $ Screma size xs' $ ScremaForm f' scs' reds'
   where
     fwdScan :: Scan SOACS -> ADM (Scan SOACS)
     fwdScan sc = do

diff --git a/src/Futhark/AD/Rev/SOAC.hs b/src/Futhark/AD/Rev/SOAC.hs
@@ -227,7 +227,7 @@ mapOp (Lambda [pa1, pa2] _ lam_body)
     cs == mempty,
     [map_stm] <- stmsToList (bodyStms lam_body),
     (Let (Pat [pe]) _ (Op scrm)) <- map_stm,
-    (Screma _ [a1, a2] (ScremaForm [] [] map_lam)) <- scrm,
+    (Screma _ [a1, a2] (ScremaForm map_lam [] [])) <- scrm,
     (a1 == paramName pa1 && a2 == paramName pa2) || (a1 == paramName pa2 && a2 == paramName pa1),
     r == Var (patElemName pe) =
       Just map_lam

diff --git a/src/Futhark/Analysis/HORep/MapNest.hs b/src/Futhark/Analysis/HORep/MapNest.hs
@@ -76,7 +76,7 @@ fromSOAC' ::
   [Ident] ->
   SOAC rep ->
   m (Maybe (MapNest rep))
-fromSOAC' bound (SOAC.Screma w inps (SOAC.ScremaForm [] [] lam)) = do
+fromSOAC' bound (SOAC.Screma w inps (SOAC.ScremaForm lam [] [])) = do
   maybenest <- case ( stmsToList $ bodyStms $ lambdaBody lam,
                       bodyResult $ lambdaBody lam
                     ) of

diff --git a/src/Futhark/Analysis/HORep/SOAC.hs b/src/Futhark/Analysis/HORep/SOAC.hs
@@ -435,7 +435,7 @@ newWidth (inp : _) _ = arraySize 0 $ inputType inp
 lambda :: SOAC rep -> Lambda rep
 lambda (Stream _ _ _ lam) = lam
 lambda (Scatter _len _ivs _spec lam) = lam
-lambda (Screma _ _ (ScremaForm _ _ lam)) = lam
+lambda (Screma _ _ (ScremaForm lam _ _)) = lam
 lambda (Hist _ _ _ lam) = lam
 
 -- | Set the lambda used in the SOAC.
@@ -444,8 +444,8 @@ setLambda lam (Stream w arrs nes _) =
   Stream w arrs nes lam
 setLambda lam (Scatter len arrs spec _lam) =
   Scatter len arrs spec lam
-setLambda lam (Screma w arrs (ScremaForm scan red _)) =
-  Screma w arrs (ScremaForm scan red lam)
+setLambda lam (Screma w arrs (ScremaForm _ scan red)) =
+  Screma w arrs (ScremaForm lam scan red)
 setLambda lam (Hist w ops inps _) =
   Hist w ops inps lam
 

diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs
@@ -727,24 +727,25 @@ pSOAC pr =
           <*> p
     pScremaForm =
       SOAC.ScremaForm
-        <$> braces (pScan pr `sepBy` pComma)
+        <$> pLambda pr
         <* pComma
-        <*> braces (pReduce pr `sepBy` pComma)
+        <*> braces (pScan pr `sepBy` pComma)
         <* pComma
-        <*> pLambda pr
+        <*> braces (pReduce pr `sepBy` pComma)
     pRedomapForm =
-      SOAC.ScremaForm mempty
-        <$> braces (pReduce pr `sepBy` pComma)
+      SOAC.ScremaForm
+        <$> pLambda pr
+        <*> pure []
         <* pComma
-        <*> pLambda pr
+        <*> braces (pReduce pr `sepBy` pComma)
     pScanomapForm =
       SOAC.ScremaForm
-        <$> braces (pScan pr `sepBy` pComma)
+        <$> pLambda pr
         <* pComma
-        <*> pure mempty
-        <*> pLambda pr
+        <*> braces (pScan pr `sepBy` pComma)
+        <*> pure []
     pMapForm =
-      SOAC.ScremaForm mempty mempty <$> pLambda pr
+      SOAC.ScremaForm <$> pLambda pr <*> pure mempty <*> pure mempty
     pScatter =
       keyword "scatter"
         *> parens
@@ -797,19 +798,19 @@ pSOAC pr =
     pVJP =
       parens $
         SOAC.VJP
-          <$> pLambda pr
+          <$> braces (pSubExp `sepBy` pComma)
           <* pComma
           <*> braces (pSubExp `sepBy` pComma)
           <* pComma
-          <*> braces (pSubExp `sepBy` pComma)
+          <*> pLambda pr
     pJVP =
       parens $
         SOAC.JVP
-          <$> pLambda pr
+          <$> braces (pSubExp `sepBy` pComma)
           <* pComma
           <*> braces (pSubExp `sepBy` pComma)
           <* pComma
-          <*> braces (pSubExp `sepBy` pComma)
+          <*> pLambda pr
 
 pSizeClass :: Parser GPU.SizeClass
 pSizeClass =

diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs
@@ -55,7 +55,7 @@ usesAD prog = any stmUsesAD (progConsts prog) || any funUsesAD (progFuns prog)
     expUsesAD (Op JVP {}) = True
     expUsesAD (Op VJP {}) = True
     expUsesAD (Op (Stream _ _ _ lam)) = lamUsesAD lam
-    expUsesAD (Op (Screma _ _ (ScremaForm scans reds lam))) =
+    expUsesAD (Op (Screma _ _ (ScremaForm lam scans reds))) =
       lamUsesAD lam
         || any (lamUsesAD . scanLambda) scans
         || any (lamUsesAD . redLambda) reds