Improve simplification of reshape/rearrange combinations (#2255)

This can be done with a fairly straightforward composition algebra. Currently restricted to rearranges that are (mapped) transpositions, but can be extended easily enough I think.

Author: athas

src/Futhark/IR/Prop/Reshape.hs

@@ -9,6 +9,7 @@ module Futhark.IR.Prop.Reshape
     reshapeInner,
 
     -- * Simplification
+    flipReshapeRearrange,
 
     -- * Shape calculations
     reshapeIndex,
@@ -18,8 +19,11 @@ module Futhark.IR.Prop.Reshape
   )
 where
 
+import Control.Monad (guard, mplus)
 import Data.Foldable
+import Futhark.IR.Prop.Rearrange (isMapTranspose)
 import Futhark.IR.Syntax
+import Futhark.Util (takeLast)
 import Futhark.Util.IntegralExp
 import Prelude hiding (product, quot, sum)
 
@@ -101,3 +105,53 @@ sliceSizes (n : ns) =
   product (n : ns) : sliceSizes ns
 
 {- HLINT ignore sliceSizes -}
+
+-- | Interchange a reshape and rearrange. Essentially, rewrite composition
+--
+-- @
+-- let v1 = reshape(v1_shape, v0)
+-- let v2 = rearrange(perm, v1)
+-- @
+--
+-- into
+--
+-- @
+-- let v1' = rearrange(perm', v0)
+-- let v2' = reshape(v1_shape', v1')
+--
+-- The function is given the shape of @v0@, @v1@, and the @perm@, and returns
+-- @perm'@. This is a meaningful operation when @v2@ is itself reshaped, as the
+-- reshape-reshape can be fused. This can significantly simplify long chains of
+-- reshapes and rearranges.
+flipReshapeRearrange ::
+  (Eq d) =>
+  [d] ->
+  [d] ->
+  [Int] ->
+  Maybe [Int]
+flipReshapeRearrange v0_shape v1_shape perm = do
+  (num_map_dims, num_a_dims, num_b_dims) <- isMapTranspose perm
+  guard $ num_a_dims == 1
+  guard $ num_b_dims == 1
+  let map_dims = take num_map_dims v0_shape
+      num_b_dims_expanded = length v0_shape - num_map_dims - num_a_dims
+      num_a_dims_expanded = length v0_shape - num_map_dims - num_b_dims
+      caseA = do
+        guard $ take num_a_dims v0_shape == take num_b_dims v1_shape
+        let perm' =
+              [0 .. num_map_dims - 1]
+                ++ map (+ num_map_dims) ([1 .. num_b_dims_expanded] ++ [0])
+        Just perm'
+      caseB = do
+        guard $ takeLast num_b_dims v0_shape == takeLast num_b_dims v1_shape
+        let perm' =
+              [0 .. num_map_dims - 1]
+                ++ map
+                  (+ num_map_dims)
+                  (num_a_dims_expanded : [0 .. num_a_dims_expanded - 1])
+        Just perm'
+
+  guard $ map_dims == take num_map_dims v1_shape
+
+  caseA `mplus` caseB
+{-# NOINLINE flipReshapeRearrange #-}

src/Futhark/Optimise/Simplify/Rules/BasicOp.hs

@@ -366,6 +366,16 @@ ruleBasicOp vtable pat aux (Manifest perm v1)
     ST.available v2 vtable =
       Simplify . auxing aux . certifying cs . letBind pat . BasicOp $
         Manifest perm v2
+ruleBasicOp vtable pat aux (Reshape ReshapeArbitrary v3_shape v2)
+  | Just (Rearrange perm v1, v2_cs) <- ST.lookupBasicOp v2 vtable,
+    Just (Reshape ReshapeArbitrary v1_shape v0, v1_cs) <- ST.lookupBasicOp v1 vtable,
+    Just v0_shape <- arrayShape <$> ST.lookupType v0 vtable,
+    Just perm' <-
+      flipReshapeRearrange (shapeDims v0_shape) (shapeDims v1_shape) perm =
+      Simplify $ do
+        v1' <- letExp (baseString v1) $ BasicOp $ Rearrange perm' v0
+        auxing aux . certifying (v1_cs <> v2_cs) . letBind pat $
+          BasicOp (Reshape ReshapeArbitrary v3_shape v1')
 ruleBasicOp _ _ _ _ =
   Skip
 

unittests/Futhark/IR/Prop/ReshapeTests.hs

@@ -11,6 +11,9 @@ import Futhark.IR.Syntax
 import Test.Tasty
 import Test.Tasty.HUnit
 
+intShape :: [Int] -> Shape
+intShape = Shape . map (intConst Int32 . toInteger)
+
 reshapeOuterTests :: [TestTree]
 reshapeOuterTests =
   [ testCase (unwords ["reshapeOuter", show sc, show n, show shape, "==", show sc_res]) $
@@ -35,10 +38,47 @@ reshapeInnerTests =
         ]
   ]
 
-intShape :: [Int] -> Shape
-intShape = Shape . map (intConst Int32 . toInteger)
+flipTests :: [TestTree]
+flipTests =
+  [ testCase
+      ( unwords
+          [ "flipReshapeRearrange",
+            show v0_shape,
+            show v1_shape,
+            show perm
+          ]
+      )
+      $ flipReshapeRearrange v0_shape v1_shape perm @?= res
+    | (v0_shape :: [String], v1_shape, perm, res) <-
+        [ ( ["A", "B", "C"],
+            ["A", "BC"],
+            [1, 0],
+            Just [1, 2, 0]
+          ),
+          ( ["A", "B", "C", "D"],
+            ["A", "BCD"],
+            [1, 0],
+            Just [1, 2, 3, 0]
+          ),
+          ( ["A"],
+            ["B", "C"],
+            [1, 0],
+            Nothing
+          ),
+          ( ["A", "B", "C"],
+            ["AB", "C"],
+            [1, 0],
+            Just [2, 0, 1]
+          ),
+          ( ["A", "B", "C", "D"],
+            ["ABC", "D"],
+            [1, 0],
+            Just [3, 0, 1, 2]
+          )
+        ]
+  ]
 
 tests :: TestTree
 tests =
   testGroup "ReshapeTests" $
-    reshapeOuterTests ++ reshapeInnerTests
+    reshapeOuterTests ++ reshapeInnerTests ++ flipTests