Check whether original is consumed before simplifying.

CKuke · Nov 8, 2023 · 4575c8f · 4575c8f
1 parent 1bb9b00
commit 4575c8f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/src/Futhark/Optimise/Simplify/Rules.hs b/src/Futhark/Optimise/Simplify/Rules.hs
@@ -125,7 +125,7 @@ emptyArrayToScratch _ _ = Skip
 
 simplifyIndex :: (BuilderOps rep) => BottomUpRuleBasicOp rep
 simplifyIndex (vtable, used) pat@(Pat [pe]) (StmAux cs attrs _) (Index idd inds)
-  | Just m <- simplifyIndexing vtable seType idd inds consumed =
+  | Just m <- simplifyIndexing vtable seType idd inds consumed consuming =
       Simplify $ certifying cs $ do
         res <- m
         attributing attrs $ case res of
@@ -134,7 +134,8 @@ simplifyIndex (vtable, used) pat@(Pat [pe]) (StmAux cs attrs _) (Index idd inds)
           IndexResult extra_cs idd' inds' ->
             certifying extra_cs $ letBindNames (patNames pat) $ BasicOp $ Index idd' inds'
   where
-    consumed = patElemName pe `UT.isConsumed` used
+    consuming = (`UT.isConsumed` used)
+    consumed = consuming $ patElemName pe
     seType (Var v) = ST.lookupType v vtable
     seType (Constant v) = Just $ Prim $ primValueType v
 simplifyIndex _ _ _ _ = Skip

diff --git a/src/Futhark/Optimise/Simplify/Rules/Index.hs b/src/Futhark/Optimise/Simplify/Rules/Index.hs
@@ -39,8 +39,9 @@ simplifyIndexing ::
   VName ->
   Slice SubExp ->
   Bool ->
+  (VName -> Bool) ->
   Maybe (m IndexResult)
-simplifyIndexing vtable seType idd (Slice inds) consuming =
+simplifyIndexing vtable seType idd (Slice inds) consuming consumed =
   case defOf idd of
     _
       | Just t <- seType (Var idd),
@@ -100,11 +101,11 @@ simplifyIndexing vtable seType idd (Slice inds) consuming =
           <$> subExpSlice (sliceSlice (primExpSlice ais) (primExpSlice (Slice inds)))
     Just (Replicate (Shape [_]) (Var vv), cs)
       | [DimFix {}] <- inds,
-        not consuming,
         ST.available vv vtable ->
           Just $ pure $ SubExpResult cs $ Var vv
       | DimFix {} : is' <- inds,
         not consuming,
+        not $ consumed vv,
         ST.available vv vtable ->
           Just $ pure $ IndexResult cs vv $ Slice is'
     Just (Replicate (Shape [_]) val@(Constant _), cs)
@@ -130,6 +131,7 @@ simplifyIndexing vtable seType idd (Slice inds) consuming =
     Just (Replicate (Shape []) (Var src), cs)
       | Just dims <- arrayDims <$> seType (Var src),
         length inds == length dims,
+        not $ consumed src,
         -- It is generally not safe to simplify a slice of a copy,
         -- because the result may be used in an in-place update of the
         -- original.  But we know this can only happen if the original