From 8204274eea4799ed637d4a94f91fed783901c7ab Mon Sep 17 00:00:00 2001
From: Troels Henriksen <athas@sigkill.dk>
Date: Mon, 11 Dec 2023 15:03:32 +0100
Subject: [PATCH] More fusion across slices.

---
 CHANGELOG.md                             |  2 ++
 src/Futhark/IR/SOACS/SOAC.hs             | 24 ++++++++++--------------
 src/Futhark/Optimise/Fusion.hs           |  3 +--
 src/Futhark/Optimise/Fusion/TryFusion.hs | 11 +++++------
 tests/fusion/slicemap4.fut               |  4 ++++
 5 files changed, 22 insertions(+), 22 deletions(-)
 create mode 100644 tests/fusion/slicemap4.fut

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f11befaba4..69ece62dc2 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs
index b6933bb435..1fe9791428 100644
--- a/src/Futhark/IR/SOACS/SOAC.hs
+++ b/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
diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs
index 15895743d1..6cb9943416 100644
--- a/src/Futhark/Optimise/Fusion.hs
+++ b/src/Futhark/Optimise/Fusion.hs
@@ -279,10 +279,9 @@ vFuseNodeT
 vFuseNodeT
   _
   infusible
-  (SoacNode ots1 pat1 (H.Screma w form inps) aux1, _, _)
+  (SoacNode ots1 pat1 (H.Screma _ form inps) aux1, _, _)
   (TransNode stm2_out (H.Index cs slice@(Slice (ds@(DimSlice _ w' _) : ds_rest))) _, _)
     | null infusible,
-      w /= w',
       ots1 == mempty,
       Just _ <- isMapSOAC form,
       [pe] <- patElems pat1 = do
diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs
index e5bdcc665d..b83d868ec3 100644
--- a/src/Futhark/Optimise/Fusion/TryFusion.hs
+++ b/src/Futhark/Optimise/Fusion/TryFusion.hs
@@ -673,13 +673,12 @@ 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.viewf ots = do
+      let lam = scremaLambda form
+          sliceInput inp =
             SOAC.addTransform
               (SOAC.Index cs (fullSlice (SOAC.inputType inp) [ds]))
               inp
@@ -699,7 +698,7 @@ pullIndex (SOAC.Screma w form inps) ots
           else
             runLambdaBuilder (lambdaParams lam) $
               mapM sliceRes =<< bodyBind (lambdaBody lam)
-      pure (SOAC.Screma w' (mapSOAC lam') (map sliceInput inps), ots')
+      pure (SOAC.Screma w' (form {scremaLambda = lam'}) (map sliceInput inps), ots')
 pullIndex _ _ = fail "Cannot pull index"
 
 pushRearrange ::
diff --git a/tests/fusion/slicemap4.fut b/tests/fusion/slicemap4.fut
new file mode 100644
index 0000000000..078f002fdb
--- /dev/null
+++ b/tests/fusion/slicemap4.fut
@@ -0,0 +1,4 @@
+-- ==
+-- structure { Screma 1 }
+
+def main (xs: []i32) = xs |> map (+2) |> reverse |> map (*3)