Implementation of irregular flattening

futhark.cabal

@@ -390,6 +390,9 @@ library
       Futhark.Pass.ExtractKernels.ToGPU
       Futhark.Pass.ExtractMulticore
       Futhark.Pass.FirstOrderTransform
+      Futhark.Pass.Flatten
+      Futhark.Pass.Flatten.Builtins
+      Futhark.Pass.Flatten.Distribute
       Futhark.Pass.KernelBabysitting
       Futhark.Pass.LiftAllocations
       Futhark.Pass.LowerAllocations
@@ -491,6 +494,7 @@ library
     , lsp >= 1.5.0
     , mainland-pretty >=0.7.1
     , cmark-gfm >=0.2.1
+    , OneTuple
     , megaparsec >=9.0.0
     , mtl >=2.2.1
     , neat-interpolation >=0.3

src/Futhark/AD/Rev/Scan.hs

@@ -80,7 +80,7 @@ mkScanFusedMapLam ops w scn_lam xs ys ys_adj s d = do
         ( buildBody_ $ do
             j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))
             y_s <- forM ys_adj $ \y_ ->
-              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ (eSubExp j)
+              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ [eSubExp j]
             let zso = orderArgs s y_s
             let ido = orderArgs s $ case_jac k sc idmat
             pure $ subExpsRes $ concat $ zipWith (++) zso $ fmap concat ido
@@ -89,10 +89,10 @@ mkScanFusedMapLam ops w scn_lam xs ys ys_adj s d = do
             j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))
             j1 <- letSubExp "j1" =<< toExp (pe64 w - le64 i)
             y_s <- forM ys_adj $ \y_ ->
-              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ (eSubExp j)
+              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ [eSubExp j]
 
             let args =
-                  map (`eIndex` eSubExp j) ys ++ map (`eIndex` eSubExp j1) xs
+                  map (`eIndex` [eSubExp j]) ys ++ map (`eIndex` [eSubExp j1]) xs
             lam_rs <- traverse (`eLambda` args) lams
 
             let yso = orderArgs s $ subExpsRes y_s
@@ -171,7 +171,7 @@ mkScanFinalMap ops w scan_lam xs ys ds = do
 
       dj <-
         traverse
-          (\dd -> letExp (baseString dd ++ "_dj") =<< eIndex dd (eSubExp j))
+          (\dd -> letExp (baseString dd ++ "_dj") =<< eIndex dd [eSubExp j])
           ds
 
       fmap varsRes . letTupExp "scan_contribs"
@@ -182,9 +182,8 @@ mkScanFinalMap ops w scan_lam xs ys ds = do
               lam <- mkScanAdjointLam ops scan_lam WrtSecond $ fmap Var dj
 
               im1 <- letSubExp "im1" =<< toExp (le64 i - 1)
-              ys_im1 <- forM ys $ \y -> do
-                y_t <- lookupType y
-                letSubExp (baseString y ++ "_last") $ BasicOp $ Index y $ fullSlice y_t [DimFix im1]
+              ys_im1 <- forM ys $ \y ->
+                letSubExp (baseString y <> "_im1") =<< eIndex y [eSubExp im1]
 
               let args = map eSubExp $ ys_im1 ++ map (Var . paramName) par_x
               eLambda lam args
@@ -341,5 +340,5 @@ diffScanAdd _ops ys n lam' ne as = do
       mkLambda [par_i] $ do
         a <-
           letExp "ys_bar_rev"
-            =<< eIndex arr (toExp (pe64 n - le64 (paramName par_i) - 1))
+            =<< eIndex arr [toExp (pe64 n - le64 (paramName par_i) - 1)]
         pure [varRes a]

src/Futhark/CLI/Dev.hs

@@ -49,6 +49,7 @@ import Futhark.Pass.ExplicitAllocations.Seq qualified as Seq
 import Futhark.Pass.ExtractKernels
 import Futhark.Pass.ExtractMulticore
 import Futhark.Pass.FirstOrderTransform
+import Futhark.Pass.Flatten (flattenSOACs)
 import Futhark.Pass.KernelBabysitting
 import Futhark.Pass.LiftAllocations as LiftAllocations
 import Futhark.Pass.LowerAllocations as LowerAllocations
@@ -605,6 +606,7 @@ commandLineOptions =
     kernelsPassOption sinkGPU [],
     kernelsPassOption reduceDeviceSyncs [],
     typedPassOption soacsProg GPU extractKernels [],
+    typedPassOption soacsProg GPU flattenSOACs [],
     typedPassOption soacsProg MC extractMulticore [],
     iplOption [],
     allocateOption "a",

src/Futhark/Construct.hs

@@ -417,11 +417,11 @@ eOutOfBounds arr is = do
   foldBinOp LogOr (constant False) =<< zipWithM checkDim ws is'
 
 -- | The array element at this index.
-eIndex :: MonadBuilder m => VName -> m (Exp (Rep m)) -> m (Exp (Rep m))
-eIndex arr i = do
-  i' <- letSubExp "i" =<< i
+eIndex :: MonadBuilder m => VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
+eIndex arr is = do
+  is' <- mapM (letSubExp "i" =<<) is
   arr_t <- lookupType arr
-  pure $ BasicOp $ Index arr $ fullSlice arr_t [DimFix i']
+  pure $ BasicOp $ Index arr $ fullSlice arr_t $ map DimFix is'
 
 -- | The last element of the given array.
 eLast :: MonadBuilder m => VName -> m (Exp (Rep m))
@@ -430,7 +430,7 @@ eLast arr = do
   nm1 <-
     letSubExp "nm1" . BasicOp $
       BinOp (Sub Int64 OverflowUndef) n (intConst Int64 1)
-  eIndex arr (eSubExp nm1)
+  eIndex arr [eSubExp nm1]
 
 -- | Construct an unspecified value of the given type.
 eBlank :: MonadBuilder m => Type -> m (Exp (Rep m))

src/Futhark/IR/TypeCheck.hs

@@ -723,7 +723,7 @@ checkSubExp (Var ident) = context ("In subexp " <> prettyText ident) $ do
   lookupType ident
 
 checkCerts :: Checkable rep => Certs -> TypeM rep ()
-checkCerts (Certs cs) = mapM_ (requireI [Prim Unit]) cs
+checkCerts = mapM_ lookupType . unCerts
 
 checkSubExpRes :: Checkable rep => SubExpRes -> TypeM rep Type
 checkSubExpRes (SubExpRes cs se) = do
@@ -1276,8 +1276,8 @@ checkStm ::
   Stm (Aliases rep) ->
   TypeM rep a ->
   TypeM rep a
-checkStm stm@(Let pat (StmAux (Certs cs) _ (_, dec)) e) m = do
-  context "When checking certificates" $ mapM_ (requireI [Prim Unit]) cs
+checkStm stm@(Let pat (StmAux cs _ (_, dec)) e) m = do
+  context "When checking certificates" $ checkCerts cs
   context "When checking expression annotation" $ checkExpDec dec
   context ("When matching\n" <> message "  " pat <> "\nwith\n" <> message "  " e) $
     matchPat pat e

src/Futhark/Pass/ExtractKernels/ToGPU.hs

@@ -5,6 +5,7 @@ module Futhark.Pass.ExtractKernels.ToGPU
     segThread,
     soacsLambdaToGPU,
     soacsStmToGPU,
+    soacsExpToGPU,
     scopeForGPU,
     scopeForSOACs,
     injectSOACS,
@@ -74,6 +75,9 @@ injectSOACS f =
 soacsStmToGPU :: Stm SOACS -> Stm GPU
 soacsStmToGPU = runIdentity . rephraseStm (injectSOACS OtherOp)
 
+soacsExpToGPU :: Exp SOACS -> Exp GPU
+soacsExpToGPU = runIdentity . rephraseExp (injectSOACS OtherOp)
+
 soacsLambdaToGPU :: Lambda SOACS -> Lambda GPU
 soacsLambdaToGPU = runIdentity . rephraseLambda (injectSOACS OtherOp)