More fusion across slices.

CHANGELOG.md

@@ -18,6 +18,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
   used in all public interfaces. The OpenCL names are still supported
   for backwards compatibility.
 
+* More fusion across array slicing.
+
 ### Removed
 
 ### Changed

src/Futhark/IR/SOACS/SOAC.hs

@@ -31,7 +31,6 @@ module Futhark.IR.SOACS.SOAC
     isScanSOAC,
     isReduceSOAC,
     isMapSOAC,
-    scremaLambda,
     ppScrema,
     ppHist,
     ppStream,
@@ -143,11 +142,16 @@ data HistOp rep = HistOp
 
 -- | The essential parts of a 'Screma' factored out (everything
 -- except the input arrays).
-data ScremaForm rep
-  = ScremaForm
-      [Scan rep]
-      [Reduce rep]
-      (Lambda rep)
+data ScremaForm rep = ScremaForm
+  { scremaScans :: [Scan rep],
+    scremaReduces :: [Reduce rep],
+    -- | The "main" lambda of the Screma. For a map, this is
+    -- equivalent to 'isMapSOAC'. Note that the meaning of the return
+    -- value of this lambda depends crucially on exactly which Screma
+    -- this is. The parameters will correspond exactly to elements of
+    -- the input arrays, however.
+    scremaLambda :: Lambda rep
+  }
   deriving (Eq, Ord, Show)
 
 singleBinOp :: (Buildable rep) => [Lambda rep] -> Lambda rep
@@ -316,14 +320,6 @@ isMapSOAC (ScremaForm scans reds map_lam) = do
   guard $ null reds
   pure map_lam
 
--- | Return the "main" lambda of the Screma.  For a map, this is
--- equivalent to 'isMapSOAC'.  Note that the meaning of the return
--- value of this lambda depends crucially on exactly which Screma this
--- is.  The parameters will correspond exactly to elements of the
--- input arrays, however.
-scremaLambda :: ScremaForm rep -> Lambda rep
-scremaLambda (ScremaForm _ _ map_lam) = map_lam
-
 -- | @groupScatterResults <output specification> <results>@
 --
 -- Blocks the index values and result values of <results> according to the

src/Futhark/Optimise/Fusion/TryFusion.hs

@@ -673,11 +673,10 @@ pullIndex ::
   SOAC ->
   SOAC.ArrayTransforms ->
   TryFusion (SOAC, SOAC.ArrayTransforms)
-pullIndex (SOAC.Screma w form inps) ots
+pullIndex (SOAC.Screma _ form inps) ots
   | SOAC.Index cs slice@(Slice (ds@(DimSlice _ w' _) : inner_ds))
       SOAC.:< ots' <-
       SOAC.viewf ots,
-    w /= w',
     Just lam <- isMapSOAC form = do
       let sliceInput inp =
             SOAC.addTransform

tests/fusion/slicemap4.fut

@@ -0,0 +1,4 @@
+-- ==
+-- structure { Screma 1 }
+
+def main (xs: []i32) = xs |> map (+2) |> reverse |> map (*3)