futhark 0.22.3 → 0.22.4

raw patch · 45 files changed

+5096/−283 lines, 45 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• futhark.cabal +13/−1
• src/Futhark/Analysis/AlgSimplify.hs +262/−0
• src/Futhark/Analysis/Interference.hs +8/−10
• src/Futhark/Analysis/LastUse.hs +8/−6
• src/Futhark/CLI/Dev.hs +23/−0
• src/Futhark/CodeGen/Backends/GenericC/Fun.hs +38/−14
• src/Futhark/CodeGen/ImpCode.hs +25/−11
• src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +3/−0
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +106/−42
• src/Futhark/IR/GPU/Op.hs +1/−1
• src/Futhark/IR/Mem/Interval.hs +139/−0
• src/Futhark/IR/Mem/IxFun.hs +280/−15
• src/Futhark/IR/SOACS/Simplify.hs +8/−4
• src/Futhark/Internalise/Defunctionalise.hs +10/−10
• src/Futhark/Internalise/Defunctorise.hs +1/−1
• src/Futhark/Optimise/ArrayShortCircuiting.hs +175/−0
• src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs +1548/−0
• src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs +420/−0
• src/Futhark/Optimise/ArrayShortCircuiting/LastUse.hs +347/−0
• src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs +464/−0
• src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs +301/−0
• src/Futhark/Optimise/CSE.hs +11/−10
• src/Futhark/Optimise/Simplify/Rule.hs +4/−3
• src/Futhark/Optimise/Simplify/Rules.hs +2/−2
• src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +6/−6
• src/Futhark/Pass/LiftAllocations.hs +120/−0
• src/Futhark/Pass/LowerAllocations.hs +114/−0
• src/Futhark/Passes.hs +21/−0
• src/Futhark/Script.hs +13/−1
• src/Futhark/Util.hs +28/−3
• src/Futhark/Util/IntegralExp.hs +4/−0
• src/Futhark/Version.hs +5/−0
• src/Language/Futhark/Pretty.hs +14/−24
• src/Language/Futhark/Prop.hs +3/−2
• src/Language/Futhark/TypeChecker/Terms.hs +0/−49
• src/Language/Futhark/TypeChecker/Terms/DoLoop.hs +21/−9
• src/Language/Futhark/TypeChecker/Terms/Monad.hs +2/−14
• src/Language/Futhark/TypeChecker/Terms/Pat.hs +1/−1
• src/Language/Futhark/TypeChecker/Types.hs +54/−10
• src/Language/Futhark/TypeChecker/Unify.hs +55/−32
• unittests/Futhark/Analysis/AlgSimplifyTests.hs +90/−0
• unittests/Futhark/IR/Mem/IntervalTests.hs +68/−0
• unittests/Futhark/IR/Mem/IxFun/Alg.hs +31/−2
• unittests/Futhark/IR/Mem/IxFunTests.hs +245/−0
• unittests/futhark_tests.hs +4/−0

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.22.3+version:        0.22.4 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -156,6 +156,7 @@       Futhark.AD.Rev.Scan       Futhark.AD.Rev.Scatter       Futhark.AD.Rev.SOAC+      Futhark.Analysis.AlgSimplify       Futhark.Analysis.Alias       Futhark.Analysis.CallGraph       Futhark.Analysis.DataDependencies@@ -281,6 +282,7 @@       Futhark.IR.MC.Op       Futhark.IR.MCMem       Futhark.IR.Mem+      Futhark.IR.Mem.Interval       Futhark.IR.Mem.IxFun       Futhark.IR.Mem.Simplify       Futhark.IR.Parse@@ -342,6 +344,12 @@       Futhark.Optimise.InliningDeadFun       Futhark.Optimise.MemoryBlockMerging       Futhark.Optimise.MemoryBlockMerging.GreedyColoring+      Futhark.Optimise.ArrayShortCircuiting+      Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing+      Futhark.Optimise.ArrayShortCircuiting.DataStructs+      Futhark.Optimise.ArrayShortCircuiting.LastUse+      Futhark.Optimise.ArrayShortCircuiting.MemRefAggreg+      Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis       Futhark.Optimise.MergeGPUBodies       Futhark.Optimise.ReduceDeviceSyncs       Futhark.Optimise.ReduceDeviceSyncs.MigrationTable@@ -381,6 +389,8 @@       Futhark.Pass.ExtractMulticore       Futhark.Pass.FirstOrderTransform       Futhark.Pass.KernelBabysitting+      Futhark.Pass.LiftAllocations+      Futhark.Pass.LowerAllocations       Futhark.Pass.Simplify       Futhark.Passes       Futhark.Pipeline@@ -518,11 +528,13 @@       Futhark.AD.DerivativesTests       Futhark.BenchTests       Futhark.Pkg.SolveTests+      Futhark.Analysis.AlgSimplifyTests       Futhark.IR.Prop.RearrangeTests       Futhark.IR.Prop.ReshapeTests       Futhark.IR.PropTests       Futhark.IR.Syntax.CoreTests       Futhark.IR.SyntaxTests+      Futhark.IR.Mem.IntervalTests       Futhark.IR.Mem.IxFun.Alg       Futhark.IR.Mem.IxFunTests       Futhark.IR.Mem.IxFunWrapper

@@ -0,0 +1,262 @@+module Futhark.Analysis.AlgSimplify+  ( Prod (..),+    SofP,+    simplify0,+    simplify,+    simplify',+    simplifySofP,+    simplifySofP',+    sumOfProducts,+    sumToExp,+    prodToExp,+    add,+    sub,+    negate,+    isMultipleOf,+    maybeDivide,+    removeLessThans,+    lessThanish,+    compareComplexity,+  )+where++import Data.Bits (xor)+import Data.Function ((&))+import Data.List (findIndex, intersect, partition, sort, (\\))+import Data.Maybe (mapMaybe)+import Futhark.Analysis.PrimExp+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.Prop.Names+import Futhark.IR.Syntax.Core+import Futhark.Util+import Futhark.Util.Pretty+import Prelude hiding (negate)++type Exp = PrimExp VName++type TExp = TPrimExp Int64 VName++data Prod = Prod+  { negated :: Bool,+    atoms :: [Exp]+  }+  deriving (Show, Eq, Ord)++type SofP = [Prod]++sumOfProducts :: Exp -> SofP+sumOfProducts = map sortProduct . sumOfProducts'++sortProduct :: Prod -> Prod+sortProduct (Prod n as) = Prod n $ sort as++sumOfProducts' :: Exp -> SofP+sumOfProducts' (BinOpExp (Add Int64 _) e1 e2) =+  sumOfProducts' e1 <> sumOfProducts' e2+sumOfProducts' (BinOpExp (Sub Int64 _) (ValueExp (IntValue (Int64Value 0))) e) =+  map negate $ sumOfProducts' e+sumOfProducts' (BinOpExp (Sub Int64 _) e1 e2) =+  sumOfProducts' e1 <> map negate (sumOfProducts' e2)+sumOfProducts' (BinOpExp (Mul Int64 _) e1 e2) =+  sumOfProducts' e1 `mult` sumOfProducts' e2+sumOfProducts' (ValueExp (IntValue (Int64Value i))) =+  [Prod (i < 0) [ValueExp $ IntValue $ Int64Value $ abs i]]+sumOfProducts' e = [Prod False [e]]++mult :: SofP -> SofP -> SofP+mult xs ys = [Prod (b `xor` b') (x <> y) | Prod b x <- xs, Prod b' y <- ys]++negate :: Prod -> Prod+negate p = p {negated = not $ negated p}++sumToExp :: SofP -> Exp+sumToExp [] = val 0+sumToExp [x] = prodToExp x+sumToExp (x : xs) =+  foldl (BinOpExp $ Add Int64 OverflowUndef) (prodToExp x) $+    map prodToExp xs++prodToExp :: Prod -> Exp+prodToExp (Prod _ []) = val 1+prodToExp (Prod True [ValueExp (IntValue (Int64Value i))]) = ValueExp $ IntValue $ Int64Value (-i)+prodToExp (Prod True as) =+  foldl (BinOpExp $ Mul Int64 OverflowUndef) (val (-1)) as+prodToExp (Prod False (a : as)) =+  foldl (BinOpExp $ Mul Int64 OverflowUndef) a as++simplifySofP :: SofP -> SofP+simplifySofP =+  -- TODO: Maybe 'constFoldValueExps' is not necessary after adding scaleConsts+  fixPoint (mapMaybe (applyZero . removeOnes) . scaleConsts . constFoldValueExps . removeNegations)++simplifySofP' :: SofP -> SofP+simplifySofP' = fixPoint (mapMaybe (applyZero . removeOnes) . scaleConsts . removeNegations)++simplify0 :: Exp -> SofP+simplify0 = simplifySofP . sumOfProducts++simplify :: Exp -> Exp+simplify = constFoldPrimExp . sumToExp . simplify0++simplify' :: TExp -> TExp+simplify' = TPrimExp . simplify . untyped++applyZero :: Prod -> Maybe Prod+applyZero p@(Prod _ as)+  | val 0 `elem` as = Nothing+  | otherwise = Just p++removeOnes :: Prod -> Prod+removeOnes (Prod neg as) =+  let as' = filter (/= val 1) as+   in Prod neg $ if null as' then [ValueExp $ IntValue $ Int64Value 1] else as'++removeNegations :: SofP -> SofP+removeNegations [] = []+removeNegations (t : ts) =+  case break (== negate t) ts of+    (start, _ : rest) -> removeNegations $ start <> rest+    _ -> t : removeNegations ts++constFoldValueExps :: SofP -> SofP+constFoldValueExps prods =+  let (value_exps, others) = partition (all isPrimValue . atoms) prods+      value_exps' = sumOfProducts $ constFoldPrimExp $ sumToExp value_exps+   in value_exps' <> others++intFromExp :: Exp -> Maybe Int64+intFromExp (ValueExp (IntValue x)) = Just $ valueIntegral x+intFromExp _ = Nothing++-- | Given @-[2, x]@ returns @(-2, [x])@+prodToScale :: Prod -> (Int64, [Exp])+prodToScale (Prod b exps) =+  let (scalars, exps') = partitionMaybe intFromExp exps+   in if b+        then (-(product scalars), exps')+        else (product scalars, exps')++-- | Given @(-2, [x])@ returns @-[1, 2, x]@+scaleToProd :: (Int64, [Exp]) -> Prod+scaleToProd (i, exps) =+  Prod (i < 0) $ ValueExp (IntValue $ Int64Value $ abs i) : exps++-- | Given @[[2, x], -[x]]@ returns @[[x]]@+scaleConsts :: SofP -> SofP+scaleConsts =+  helper [] . map prodToScale+  where+    helper :: [Prod] -> [(Int64, [Exp])] -> [Prod]+    helper acc [] = reverse acc+    helper acc ((scale, exps) : rest) =+      case flip focusNth rest =<< findIndex ((==) exps . snd) rest of+        Nothing -> helper (scaleToProd (scale, exps) : acc) rest+        Just (before, (scale', _), after) ->+          helper acc $ (scale + scale', exps) : (before <> after)++isPrimValue :: Exp -> Bool+isPrimValue (ValueExp _) = True+isPrimValue _ = False++val :: Int64 -> Exp+val = ValueExp . IntValue . Int64Value++add :: SofP -> SofP -> SofP+add ps1 ps2 = simplifySofP $ ps1 <> ps2++sub :: SofP -> SofP -> SofP+sub ps1 ps2 = add ps1 $ map negate ps2++isMultipleOf :: Prod -> [Exp] -> Bool+isMultipleOf (Prod _ as) term =+  let quotient = as \\ term+   in sort (quotient <> term) == sort as++maybeDivide :: Prod -> Prod -> Maybe Prod+maybeDivide dividend divisor+  | Prod dividend_b dividend_factors <- dividend,+    Prod divisor_b divisor_factors <- divisor,+    quotient <- dividend_factors \\ divisor_factors,+    sort (quotient <> divisor_factors) == sort dividend_factors =+      Just $ Prod (dividend_b `xor` divisor_b) quotient+  | (dividend_scale, dividend_rest) <- prodToScale dividend,+    (divisor_scale, divisor_rest) <- prodToScale divisor,+    dividend_scale `mod` divisor_scale == 0,+    null $ divisor_rest \\ dividend_rest =+      Just $+        Prod+          (signum (dividend_scale `div` divisor_scale) < 0)+          ( ValueExp (IntValue $ Int64Value $ dividend_scale `div` divisor_scale)+              : (dividend_rest \\ divisor_rest)+          )+  | otherwise = Nothing++-- | Given a list of 'Names' that we know are non-negative (>= 0), determine+-- whether we can say for sure that the given 'AlgSimplify.SofP' is+-- non-negative. Conservatively returns 'False' if there is any doubt.+--+-- TODO: We need to expand this to be able to handle cases such as @i*n + g < (i+-- + 1) * n@, if it is known that @g < n@, eg. from a 'SegSpace' or a loop form.+nonNegativeish :: Names -> SofP -> Bool+nonNegativeish non_negatives = all (nonNegativeishProd non_negatives)++nonNegativeishProd :: Names -> Prod -> Bool+nonNegativeishProd _ (Prod True _) = False+nonNegativeishProd non_negatives (Prod False as) =+  all (nonNegativeishExp non_negatives) as++nonNegativeishExp :: Names -> PrimExp VName -> Bool+nonNegativeishExp _ (ValueExp v) = not $ negativeIsh v+nonNegativeishExp non_negatives (LeafExp vname _) = vname `nameIn` non_negatives+nonNegativeishExp _ _ = False++-- | Is e1 symbolically less than or equal to e2?+lessThanOrEqualish :: [(VName, PrimExp VName)] -> Names -> TPrimExp Int64 VName -> TPrimExp Int64 VName -> Bool+lessThanOrEqualish less_thans0 non_negatives e1 e2 =+  case e2 - e1 & untyped & simplify0 of+    [] -> True+    simplified ->+      nonNegativeish non_negatives $+        fixPoint (`removeLessThans` less_thans) simplified+  where+    less_thans =+      concatMap+        (\(i, bound) -> [(Var i, bound), (Constant $ IntValue $ Int64Value 0, bound)])+        less_thans0++lessThanish :: [(VName, PrimExp VName)] -> Names -> TPrimExp Int64 VName -> TPrimExp Int64 VName -> Bool+lessThanish less_thans non_negatives e1 =+  lessThanOrEqualish less_thans non_negatives (e1 + 1)++removeLessThans :: SofP -> [(SubExp, PrimExp VName)] -> SofP+removeLessThans =+  foldl+    ( \sofp (i, bound) ->+        let to_remove =+              simplifySofP $+                Prod True [primExpFromSubExp (IntType Int64) i]+                  : simplify0 bound+         in case to_remove `intersect` sofp of+              to_remove' | to_remove' == to_remove -> sofp \\ to_remove+              _ -> sofp+    )++compareComplexity :: SofP -> SofP -> Ordering+compareComplexity xs0 ys0 =+  case length xs0 `compare` length ys0 of+    EQ -> helper xs0 ys0+    c -> c+  where+    helper [] [] = EQ+    helper [] _ = LT+    helper _ [] = GT+    helper (px : xs) (py : ys) =+      case (prodToScale px, prodToScale py) of+        ((ix, []), (iy, [])) -> case ix `compare` iy of+          EQ -> helper xs ys+          c -> c+        ((_, []), (_, _)) -> LT+        ((_, _), (_, [])) -> GT+        ((_, x), (_, y)) -> case length x `compare` length y of+          EQ -> helper xs ys+          c -> c

@@ -211,16 +211,14 @@  analyseProgGPU :: Prog GPUMem -> Graph VName analyseProgGPU prog =-  let (lumap, _) = LastUse.analyseGPUMem prog-      graph =-        foldMap-          ( \f ->-              runReader (analyseGPU lumap $ bodyStms $ funDefBody f) $-                scopeOf f-          )-          $ progFuns prog-      graph' = applyAliases (MemAlias.analyzeGPUMem prog) graph-   in graph'+  applyAliases (MemAlias.analyzeGPUMem prog) $+    onConsts (progConsts prog) <> foldMap onFun (progFuns prog)+  where+    (lumap, _) = LastUse.analyseGPUMem prog+    onFun f =+      runReader (analyseGPU lumap $ bodyStms $ funDefBody f) $ scopeOf f+    onConsts stms =+      runReader (analyseGPU lumap stms) (mempty :: Scope GPUMem)  applyAliases :: MemAlias.MemAliases -> Graph VName -> Graph VName applyAliases aliases =

@@ -12,7 +12,6 @@  import Control.Monad.Reader import Data.Bifunctor (bimap, first)-import Data.Foldable import Data.Function ((&)) import Data.Map (Map) import Data.Map qualified as M@@ -106,13 +105,16 @@   runReader helper (Env onOp)   where     helper = do-      let consts =+      let bound_in_consts =             progConsts prog-              & concatMap (toList . fmap patElemName . patElems . stmPat)+              & concatMap (patNames . stmPat)               & namesFromList-          funs = progFuns $ aliasAnalysis prog-      (lus, used) <- mconcat <$> mapM (analyseFun mempty consts) funs-      pure (flipMap lus, used)+          prog_alias = aliasAnalysis prog+          consts = progConsts prog_alias+          funs = progFuns prog_alias+      (consts_lu, consts_used) <- analyseStms mempty mempty consts+      (lus, used) <- mconcat <$> mapM (analyseFun mempty bound_in_consts) funs+      pure (flipMap $ consts_lu <> lus, consts_used <> used)  analyseFun :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> FunDef (Aliases rep) -> LastUseM rep analyseFun lumap used fun = do

@@ -28,6 +28,7 @@ import Futhark.Internalise.Defunctorise as Defunctorise import Futhark.Internalise.LiftLambdas as LiftLambdas import Futhark.Internalise.Monomorphise as Monomorphise+import Futhark.Optimise.ArrayShortCircuiting qualified as ArrayShortCircuiting import Futhark.Optimise.CSE import Futhark.Optimise.DoubleBuffer import Futhark.Optimise.Fusion@@ -48,6 +49,8 @@ import Futhark.Pass.ExtractMulticore import Futhark.Pass.FirstOrderTransform import Futhark.Pass.KernelBabysitting+import Futhark.Pass.LiftAllocations as LiftAllocations+import Futhark.Pass.LowerAllocations as LowerAllocations import Futhark.Pass.Simplify import Futhark.Passes import Futhark.Util.Log@@ -219,6 +222,13 @@   externalErrorS $     "Pass " ++ name ++ " expects GPU representation, but got " ++ representation rep +seqMemProg :: String -> UntypedPassState -> FutharkM (Prog SeqMem.SeqMem)+seqMemProg _ (SeqMem prog) =+  pure prog+seqMemProg name rep =+  externalErrorS $+    "Pass " ++ name ++ " expects SeqMem representation, but got " ++ representation rep+ typedPassOption ::   Checkable torep =>   (String -> UntypedPassState -> FutharkM (Prog fromrep)) ->@@ -246,6 +256,13 @@ kernelsPassOption =   typedPassOption kernelsProg GPU +seqMemPassOption ::+  Pass SeqMem.SeqMem SeqMem.SeqMem ->+  String ->+  FutharkOption+seqMemPassOption =+  typedPassOption seqMemProg SeqMem+ kernelsMemPassOption ::   Pass GPUMem.GPUMem GPUMem.GPUMem ->   String ->@@ -577,6 +594,12 @@     kernelsMemPassOption doubleBufferGPU [],     kernelsMemPassOption expandAllocations [],     kernelsMemPassOption MemoryBlockMerging.optimise [],+    seqMemPassOption LiftAllocations.liftAllocationsSeqMem [],+    kernelsMemPassOption LiftAllocations.liftAllocationsGPUMem [],+    seqMemPassOption LowerAllocations.lowerAllocationsSeqMem [],+    kernelsMemPassOption LowerAllocations.lowerAllocationsGPUMem [],+    seqMemPassOption ArrayShortCircuiting.optimiseSeqMem [],+    kernelsMemPassOption ArrayShortCircuiting.optimiseGPUMem [],     cseOption [],     simplifyOption "e",     soacsPipelineOption

@@ -3,6 +3,7 @@ -- | C code generation for functions. module Futhark.CodeGen.Backends.GenericC.Fun   ( compileFun,+    compileVoidFun,     module Futhark.CodeGen.Backends.GenericC.Monad,     module Futhark.CodeGen.Backends.GenericC.Code,   )@@ -34,6 +35,24 @@     setRetVal' p (ScalarParam name _) =       stm [C.cstm|*$exp:p = $id:name;|] +compileInput :: Param -> CompilerM op s C.Param+compileInput (ScalarParam name bt) = do+  let ctp = primTypeToCType bt+  pure [C.cparam|$ty:ctp $id:name|]+compileInput (MemParam name space) = do+  ty <- memToCType name space+  pure [C.cparam|$ty:ty $id:name|]++compileOutput :: Param -> CompilerM op s (C.Param, C.Exp)+compileOutput (ScalarParam name bt) = do+  let ctp = primTypeToCType bt+  p_name <- newVName $ "out_" ++ baseString name+  pure ([C.cparam|$ty:ctp *$id:p_name|], [C.cexp|$id:p_name|])+compileOutput (MemParam name space) = do+  ty <- memToCType name space+  p_name <- newVName $ baseString name ++ "_p"+  pure ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|])+ compileFun :: [C.BlockItem] -> [C.Param] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func) compileFun get_constants extra (fname, func@(Function _ outputs inputs body)) = inNewFunction $ do   (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs@@ -66,18 +85,23 @@     -- actually need to use them.     ignores = [[C.cstm|(void)$id:p;|] | C.Param (Just p) _ _ _ <- extra] -    compileInput (ScalarParam name bt) = do-      let ctp = primTypeToCType bt-      pure [C.cparam|$ty:ctp $id:name|]-    compileInput (MemParam name space) = do-      ty <- memToCType name space-      pure [C.cparam|$ty:ty $id:name|]+-- | Generate code for a function that returns void (meaning it cannot+-- fail) and has no extra parameters (meaning it cannot allocate+-- memory non-lexxical or do anything fancy).+compileVoidFun :: [C.BlockItem] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func)+compileVoidFun get_constants (fname, func@(Function _ outputs inputs body)) = inNewFunction $ do+  (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs+  inparams <- mapM compileInput inputs -    compileOutput (ScalarParam name bt) = do-      let ctp = primTypeToCType bt-      p_name <- newVName $ "out_" ++ baseString name-      pure ([C.cparam|$ty:ctp *$id:p_name|], [C.cexp|$id:p_name|])-    compileOutput (MemParam name space) = do-      ty <- memToCType name space-      p_name <- newVName $ baseString name ++ "_p"-      pure ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|])+  cachingMemory (lexicalMemoryUsage func) $ \decl_cached free_cached -> do+    body' <- collect $ compileFunBody out_ptrs outputs body++    pure+      ( [C.cedecl|static void $id:(funName fname)($params:outparams, $params:inparams);|],+        [C.cfun|static void $id:(funName fname)($params:outparams, $params:inparams) {+               $items:decl_cached+               $items:get_constants+               $items:body'+               $stms:free_cached+               }|]+      )

@@ -74,6 +74,7 @@     declaredIn,     lexicalMemoryUsage,     calledFuncs,+    callGraph,      -- * Typed enumerations     Bytes,@@ -95,6 +96,7 @@   ) where +import Data.Bifunctor (second) import Data.List (intersperse) import Data.Map qualified as M import Data.Set qualified as S@@ -153,7 +155,7 @@     Definitions types (fmap f consts) (fmap f funs)  -- | A collection of imperative functions.-newtype Functions a = Functions [(Name, Function a)]+newtype Functions a = Functions {unFunctions :: [(Name, Function a)]}   deriving (Show)  instance Semigroup (Functions a) where@@ -374,16 +376,28 @@         onArg (MemArg x) = oneName x     set x = go set x --- | The set of functions that are called by this code.  Assumes there--- are no function calls in 'Op's.-calledFuncs :: Code a -> S.Set Name-calledFuncs (x :>>: y) = calledFuncs x <> calledFuncs y-calledFuncs (If _ x y) = calledFuncs x <> calledFuncs y-calledFuncs (For _ _ x) = calledFuncs x-calledFuncs (While _ x) = calledFuncs x-calledFuncs (Comment _ x) = calledFuncs x-calledFuncs (Call _ f _) = S.singleton f-calledFuncs _ = mempty+-- | The set of functions that are called by this code.  Accepts a+-- function for determing function calls in 'Op's.+calledFuncs :: (a -> S.Set Name) -> Code a -> S.Set Name+calledFuncs _ (Call _ v _) = S.singleton v+calledFuncs f (Op x) = f x+calledFuncs f (x :>>: y) = calledFuncs f x <> calledFuncs f y+calledFuncs f (If _ x y) = calledFuncs f x <> calledFuncs f y+calledFuncs f (For _ _ x) = calledFuncs f x+calledFuncs f (While _ x) = calledFuncs f x+calledFuncs f (Comment _ x) = calledFuncs f x+calledFuncs _ _ = mempty++-- | Compute call graph, as per 'calledFuncs', but also include+-- transitive calls.+callGraph :: (a -> S.Set Name) -> Functions a -> M.Map Name (S.Set Name)+callGraph f (Functions funs) =+  loop $ M.fromList $ map (second $ calledFuncs f . functionBody) funs+  where+    loop cur =+      let grow v = maybe (S.singleton v) (S.insert v) (M.lookup v cur)+          next = M.map (foldMap grow) cur+       in if next == cur then cur else loop next  -- | A side-effect free expression whose execution will produce a -- single primitive value.

@@ -841,6 +841,9 @@                 negate groups_per_segment        sLoopNest (slugShape slug) $ \vec_is -> do+        unless (null $ slugShape slug) $+          sOp (Imp.Barrier Imp.FenceLocal)+         -- There is no guarantee that the number of workgroups for the         -- segment is less than the workgroup size, so each thread may         -- have to read multiple elements.  We do this in a sequential

@@ -47,7 +47,7 @@   let ( prog',         ToOpenCL kernels device_funs used_types sizes failures         ) =-          (`runState` initialOpenCL) . (`runReaderT` defFuns prog) $ do+          (`runState` initialOpenCL) . (`runReaderT` envFromProg prog) $ do             let ImpGPU.Definitions                   types                   (ImpGPU.Constants ps consts)@@ -135,13 +135,51 @@ initialOpenCL :: ToOpenCL initialOpenCL = ToOpenCL mempty mempty mempty mempty mempty -type AllFunctions = ImpGPU.Functions ImpGPU.HostOp+data Env = Env+  { envFuns :: ImpGPU.Functions ImpGPU.HostOp,+    envFunsMayFail :: S.Set Name+  } -lookupFunction :: Name -> AllFunctions -> Maybe (ImpGPU.Function HostOp)-lookupFunction fname (ImpGPU.Functions fs) = lookup fname fs+codeMayFail :: (a -> Bool) -> ImpGPU.Code a -> Bool+codeMayFail _ (Assert {}) = True+codeMayFail f (Op x) = f x+codeMayFail f (x :>>: y) = codeMayFail f x || codeMayFail f y+codeMayFail f (For _ _ x) = codeMayFail f x+codeMayFail f (While _ x) = codeMayFail f x+codeMayFail f (If _ x y) = codeMayFail f x || codeMayFail f y+codeMayFail f (Comment _ x) = codeMayFail f x+codeMayFail _ _ = False -type OnKernelM = ReaderT AllFunctions (State ToOpenCL)+hostOpMayFail :: ImpGPU.HostOp -> Bool+hostOpMayFail (CallKernel k) = codeMayFail kernelOpMayFail $ kernelBody k+hostOpMayFail _ = False +kernelOpMayFail :: ImpGPU.KernelOp -> Bool+kernelOpMayFail = const False++funsMayFail :: M.Map Name (S.Set Name) -> ImpGPU.Functions ImpGPU.HostOp -> S.Set Name+funsMayFail cg (Functions funs) =+  S.fromList $ map fst $ filter mayFail funs+  where+    base_mayfail =+      map fst $ filter (codeMayFail hostOpMayFail . ImpGPU.functionBody . snd) funs+    mayFail (fname, _) =+      any (`elem` base_mayfail) $ fname : S.toList (M.findWithDefault mempty fname cg)++envFromProg :: ImpGPU.Program -> Env+envFromProg prog = Env funs (funsMayFail cg funs)+  where+    funs = defFuns prog+    cg = ImpGPU.callGraph calledInHostOp funs++lookupFunction :: Name -> Env -> Maybe (ImpGPU.Function HostOp)+lookupFunction fname = lookup fname . unFunctions . envFuns++functionMayFail :: Name -> Env -> Bool+functionMayFail fname = S.member fname . envFunsMayFail++type OnKernelM = ReaderT Env (State ToOpenCL)+ addSize :: Name -> SizeClass -> OnKernelM () addSize key sclass =   modify $ \s -> s {clSizes = M.insert key sclass $ clSizes s}@@ -158,14 +196,15 @@   pure $ ImpOpenCL.GetSizeMax v size_class  genGPUCode ::+  Env ->   OpsMode ->   KernelCode ->   [FailureMsg] ->   GC.CompilerM KernelOp KernelState a ->   (a, GC.CompilerState KernelState)-genGPUCode mode body failures =+genGPUCode env mode body failures =   GC.runCompilerM-    (inKernelOperations mode body)+    (inKernelOperations env mode body)     blankNameSource     (newKernelState failures) @@ -175,16 +214,25 @@ generateDeviceFun fname device_func = do   when (any memParam $ functionInput device_func) bad +  env <- ask   failures <- gets clFailures -  let params =-        [ [C.cparam|__global int *global_failure|],-          [C.cparam|__global typename int64_t *global_failure_args|]-        ]-      (func, cstate) =-        genGPUCode FunMode (functionBody device_func) failures $-          GC.compileFun mempty params (fname, device_func)-      kstate = GC.compUserState cstate+  let (func, kstate) =+        if functionMayFail fname env+          then+            let params =+                  [ [C.cparam|__global int *global_failure|],+                    [C.cparam|__global typename int64_t *global_failure_args|]+                  ]+                (f, cstate) =+                  genGPUCode env FunMode (functionBody device_func) failures $+                    GC.compileFun mempty params (fname, device_func)+             in (f, GC.compUserState cstate)+          else+            let (f, cstate) =+                  genGPUCode env FunMode (functionBody device_func) failures $+                    GC.compileVoidFun mempty (fname, device_func)+             in (f, GC.compUserState cstate)    modify $ \s ->     s@@ -209,22 +257,28 @@   exists <- gets $ M.member fname . clDevFuns   unless exists $ generateDeviceFun fname host_func +calledInHostOp :: HostOp -> S.Set Name+calledInHostOp (CallKernel k) = calledFuncs calledInKernelOp $ kernelBody k+calledInHostOp _ = mempty++calledInKernelOp :: KernelOp -> S.Set Name+calledInKernelOp = const mempty+ ensureDeviceFuns :: ImpGPU.KernelCode -> OnKernelM [Name] ensureDeviceFuns code = do-  let called = calledFuncs code-  fmap catMaybes $-    forM (S.toList called) $ \fname -> do-      def <- asks $ lookupFunction fname-      case def of-        Just host_func -> do-          -- Functions are a priori always considered host-level, so we have-          -- to convert them to device code.  This is where most of our-          -- limitations on device-side functions (no arrays, no parallelism)-          -- comes from.-          let device_func = fmap toDevice host_func-          ensureDeviceFun fname device_func-          pure $ Just fname-        Nothing -> pure Nothing+  let called = calledFuncs calledInKernelOp code+  fmap catMaybes . forM (S.toList called) $ \fname -> do+    def <- asks $ lookupFunction fname+    case def of+      Just host_func -> do+        -- Functions are a priori always considered host-level, so we have+        -- to convert them to device code.  This is where most of our+        -- limitations on device-side functions (no arrays, no parallelism)+        -- comes from.+        let device_func = fmap toDevice host_func+        ensureDeviceFun fname device_func+        pure $ Just fname+      Nothing -> pure Nothing   where     bad = compilerLimitationS "Cannot generate GPU functions that contain parallelism."     toDevice :: HostOp -> KernelOp@@ -237,9 +291,10 @@   -- Crucial that this is done after 'ensureDeviceFuns', as the device   -- functions may themselves define failure points.   failures <- gets clFailures+  env <- ask    let (kernel_body, cstate) =-        genGPUCode KernelMode (kernelBody kernel) failures . GC.collect $ do+        genGPUCode env KernelMode (kernelBody kernel) failures . GC.collect $ do           body <- GC.collect $ GC.compileCode $ kernelBody kernel           -- No need to free, as we cannot allocate memory in kernels.           mapM_ GC.item =<< GC.declAllocatedMem@@ -286,7 +341,11 @@   let (safety, error_init)         -- We conservatively assume that any called function can fail.         | not $ null called =-            (SafetyFull, [])+            ( SafetyFull,+              [C.citems|volatile __local bool local_failure;+                        // Harmless for all threads to write this.+                        local_failure = false;|]+            )         | length (kernelFailures kstate) == length failures =             if kernelFailureTolerant kernel               then (SafetyNone, [])@@ -611,10 +670,11 @@ data OpsMode = KernelMode | FunMode deriving (Eq)  inKernelOperations ::+  Env ->   OpsMode ->   ImpGPU.KernelCode ->   GC.Operations KernelOp KernelState-inKernelOperations mode body =+inKernelOperations env mode body =   GC.Operations     { GC.opsCompiler = kernelOps,       GC.opsMemoryType = kernelMemoryType,@@ -785,17 +845,21 @@               then [C.citems|return 1;|]               else [C.citems|return;|] -    callInKernel dests fname args = do-      let out_args = [[C.cexp|&$id:d|] | d <- dests]-          args' =-            [C.cexp|global_failure|]-              : [C.cexp|global_failure_args|]-              : out_args-              ++ args--      what_next <- whatNext+    callInKernel dests fname args+      | functionMayFail fname env = do+          let out_args = [[C.cexp|&$id:d|] | d <- dests]+              args' =+                [C.cexp|global_failure|]+                  : [C.cexp|global_failure_args|]+                  : out_args+                  ++ args -      GC.item [C.citem|if ($id:(funName fname)($args:args') != 0) { $items:what_next; }|]+          what_next <- whatNext+          GC.item [C.citem|if ($id:(funName fname)($args:args') != 0) { $items:what_next; }|]+      | otherwise = do+          let out_args = [[C.cexp|&$id:d|] | d <- dests]+              args' = out_args ++ args+          GC.item [C.citem|$id:(funName fname)($args:args');|]      errorInKernel msg@(ErrorMsg parts) backtrace = do       n <- length . kernelFailures <$> GC.getUserState

@@ -281,7 +281,7 @@   safeOp (SegOp op) = safeOp op   safeOp (OtherOp op) = safeOp op   safeOp (SizeOp op) = safeOp op-  safeOp GPUBody {} = True+  safeOp (GPUBody _ body) = all (safeExp . stmExp) $ bodyStms body    cheapOp (SegOp op) = cheapOp op   cheapOp (OtherOp op) = cheapOp op

@@ -0,0 +1,139 @@+{-# LANGUAGE OverloadedStrings #-}++module Futhark.IR.Mem.Interval+  ( Interval (..),+    distributeOffset,+    expandOffset,+    intervalOverlap,+    selfOverlap,+    primBool,+    intervalPairs,+    justLeafExp,+  )+where++import Data.Function (on)+import Data.List (maximumBy, minimumBy, (\\))+import Futhark.Analysis.AlgSimplify qualified as AlgSimplify+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.Prop+import Futhark.IR.Syntax hiding (Result)+import Futhark.Util++data Interval = Interval+  { lowerBound :: TPrimExp Int64 VName,+    numElements :: TPrimExp Int64 VName,+    stride :: TPrimExp Int64 VName+  }+  deriving (Show, Eq)++instance FreeIn Interval where+  freeIn' (Interval lb ne st) = freeIn' lb <> freeIn' ne <> freeIn' st++distributeOffset :: MonadFail m => AlgSimplify.SofP -> [Interval] -> m [Interval]+distributeOffset [] interval = pure interval+distributeOffset offset [] = fail $ "Cannot distribute offset " <> show offset <> " across empty interval"+distributeOffset offset [Interval lb ne 1] = pure [Interval (lb + TPrimExp (AlgSimplify.sumToExp offset)) ne 1]+distributeOffset offset (Interval lb ne st0 : is)+  | st <- AlgSimplify.Prod False [untyped st0],+    Just (before, quotient, after) <- focusMaybe (`AlgSimplify.maybeDivide` st) offset =+      distributeOffset (before <> after) $+        Interval (lb + TPrimExp (AlgSimplify.sumToExp [quotient])) ne st0 : is+  | [st] <- AlgSimplify.simplify0 $ untyped st0,+    Just (before, quotient, after) <- focusMaybe (`AlgSimplify.maybeDivide` st) offset =+      distributeOffset (before <> after) $+        Interval (lb + TPrimExp (AlgSimplify.sumToExp [quotient])) ne st0 : is+  | otherwise = do+      rest <- distributeOffset offset is+      pure $ Interval lb ne st0 : rest++findMostComplexTerm :: AlgSimplify.SofP -> (AlgSimplify.Prod, AlgSimplify.SofP)+findMostComplexTerm prods =+  let max_prod = maximumBy (compare `on` (length . AlgSimplify.atoms)) prods+   in (max_prod, prods \\ [max_prod])++findClosestStride :: [PrimExp VName] -> [Interval] -> (PrimExp VName, [PrimExp VName])+findClosestStride offset_term is =+  let strides = map (untyped . stride) is+      p =+        minimumBy+          ( compare+              `on` ( termDifferenceLength+                       . minimumBy (compare `on` \s -> length (offset_term \\ AlgSimplify.atoms s))+                       . AlgSimplify.simplify0+                   )+          )+          strides+   in ( p,+        (offset_term \\) $+          AlgSimplify.atoms $+            minimumBy (compare `on` \s -> length (offset_term \\ AlgSimplify.atoms s)) $+              AlgSimplify.simplify0 p+      )+  where+    termDifferenceLength (AlgSimplify.Prod _ xs) = length (offset_term \\ xs)++expandOffset :: AlgSimplify.SofP -> [Interval] -> Maybe AlgSimplify.SofP+expandOffset [] _ = Nothing+expandOffset offset i1+  | (AlgSimplify.Prod b term_to_add, offset_rest) <- findMostComplexTerm offset, -- Find gnb+    (closest_stride, first_term_divisor) <- findClosestStride term_to_add i1, -- find (nb-b, g)+    target <- [AlgSimplify.Prod b $ closest_stride : first_term_divisor], -- g(nb-b)+    diff <- AlgSimplify.sumOfProducts $ AlgSimplify.sumToExp $ AlgSimplify.Prod b term_to_add : map AlgSimplify.negate target, -- gnb - gnb + gb = gnb - g(nb-b)+    replacement <- target <> diff -- gnb = g(nb-b) + gnb - gnb + gb+    =+      Just (replacement <> offset_rest)++intervalOverlap :: [(VName, PrimExp VName)] -> Names -> Interval -> Interval -> Bool+intervalOverlap less_thans non_negatives (Interval lb1 ne1 st1) (Interval lb2 ne2 st2)+  | st1 == st2,+    AlgSimplify.lessThanish less_thans non_negatives lb1 lb2,+    AlgSimplify.lessThanish less_thans non_negatives (lb1 + ne1 - 1) lb2 =+      False+  | st1 == st2,+    AlgSimplify.lessThanish less_thans non_negatives lb2 lb1,+    AlgSimplify.lessThanish less_thans non_negatives (lb2 + ne2 - 1) lb1 =+      False+  | otherwise = True++primBool :: TPrimExp Bool VName -> Maybe Bool+primBool p+  | Just (BoolValue b) <- evalPrimExp (const Nothing) $ untyped p = Just b+  | otherwise = Nothing++intervalPairs :: [Interval] -> [Interval] -> [(Interval, Interval)]+intervalPairs = intervalPairs' []+  where+    intervalPairs' :: [(Interval, Interval)] -> [Interval] -> [Interval] -> [(Interval, Interval)]+    intervalPairs' acc [] [] = reverse acc+    intervalPairs' acc (i@(Interval lb _ st) : is) [] = intervalPairs' ((i, Interval lb 1 st) : acc) is []+    intervalPairs' acc [] (i@(Interval lb _ st) : is) = intervalPairs' ((Interval lb 1 st, i) : acc) [] is+    intervalPairs' acc (i1@(Interval lb1 _ st1) : is1) (i2@(Interval lb2 _ st2) : is2)+      | st1 == st2 = intervalPairs' ((i1, i2) : acc) is1 is2+      | otherwise =+          let res1 = intervalPairs' ((i1, Interval lb1 1 st1) : acc) is1 (i2 : is2)+              res2 = intervalPairs' ((Interval lb2 1 st2, i2) : acc) (i1 : is1) is2+           in if length res1 <= length res2+                then res1+                else res2++-- | Returns true if the intervals are self-overlapping, meaning that for a+-- given dimension d, the stride of d is larger than the aggregate spans of the+-- lower dimensions.+selfOverlap :: scope -> asserts -> [(VName, PrimExp VName)] -> [PrimExp VName] -> [Interval] -> Maybe Interval+selfOverlap _ _ _ _ [_] = Nothing+selfOverlap _ _ less_thans non_negatives' is+  | Just non_negatives <- namesFromList <$> mapM justLeafExp non_negatives' =+      -- TODO: Do we need to do something clever using some ranges of known values?+      let selfOverlap' acc (x : xs) =+            let interval_span = (lowerBound x + numElements x - 1) * stride x+                res = AlgSimplify.lessThanish less_thans non_negatives (AlgSimplify.simplify' acc) (AlgSimplify.simplify' $ stride x)+             in if res then selfOverlap' (acc + interval_span) xs else Just x+          selfOverlap' _ [] = Nothing+       in selfOverlap' 0 $ reverse is+selfOverlap _ _ _ _ (x : _) = Just x+selfOverlap _ _ _ _ [] = Nothing++justLeafExp :: PrimExp VName -> Maybe VName+justLeafExp (LeafExp v _) = Just v+justLeafExp _ = Nothing

@@ -19,15 +19,23 @@     flatSlice,     rebase,     shape,+    lmadShape,     rank,     linearWithOffset,     rearrangeWithOffset,     isDirect,     isLinear,     substituteInIxFun,+    substituteInLMAD,     existentialize,     closeEnough,     equivalent,+    hasOneLmad,+    permuteInv,+    conservativeFlatten,+    disjoint,+    disjoint2,+    disjoint3,     dynamicEqualsLMAD,   ) where@@ -37,30 +45,34 @@ import Control.Monad.State import Control.Monad.Writer import Data.Function (on, (&))-import Data.List (sort, sortBy, zip4, zipWith4)+import Data.List (elemIndex, partition, sort, sortBy, zip4, zipWith4) import Data.List.NonEmpty (NonEmpty (..)) import Data.List.NonEmpty qualified as NE import Data.Map.Strict qualified as M-import Data.Maybe (isJust)+import Data.Maybe (fromJust, isJust, isNothing)+import Futhark.Analysis.AlgSimplify qualified as AlgSimplify import Futhark.Analysis.PrimExp-import Futhark.Analysis.PrimExp.Convert (substituteInPrimExp)+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.Mem.Interval import Futhark.IR.Prop import Futhark.IR.Syntax   ( DimIndex (..),     FlatDimIndex (..),     FlatSlice (..),     Slice (..),+    Type,     dimFix,     flatSliceDims,     flatSliceStrides,     unitSlice,   )-import Futhark.IR.Syntax.Core (Ext (..))+import Futhark.IR.Syntax.Core (Ext (..), VName (..)) import Futhark.Transform.Rename import Futhark.Transform.Substitute+import Futhark.Util import Futhark.Util.IntegralExp import Futhark.Util.Pretty-import Prelude hiding (id, mod, (.))+import Prelude hiding (gcd, id, mod, (.))  -- | The shape of an index function. type Shape num = [num]@@ -89,6 +101,24 @@   }   deriving (Show, Eq) +instance Ord Monotonicity where+  (<=) _ Inc = True+  (<=) Unknown _ = True+  (<=) _ Unknown = False+  (<=) Inc Dec = False+  (<=) _ Dec = True++instance Ord num => Ord (LMADDim num) where+  (LMADDim s1 q1 p1 m1) <= (LMADDim s2 q2 p2 m2) =+    ([q1, s1] < [q2, s2])+      || ( ([q1, s1] == [q2, s2])+             && ( (p1 < p2)+                    || ( (p1 == p2)+                           && (m1 <= m2)+                       )+                )+         )+ -- | LMAD's representation consists of a general offset and for each dimension a -- stride, number of elements (or shape), permutation, and -- monotonicity. Note that the permutation is not strictly necessary in that the@@ -119,7 +149,7 @@   { lmadOffset :: num,     lmadDims :: [LMADDim num]   }-  deriving (Show, Eq)+  deriving (Show, Eq, Ord)  -- | An index function is a mapping from a multidimensional array -- index space (the domain) to a one-dimensional memory index space.@@ -131,6 +161,8 @@ -- An index function is represented as a sequence of 'LMAD's. data IxFun num = IxFun   { ixfunLMADs :: NonEmpty (LMAD num),+    -- | the shape of the support array, i.e., the original array+    --   that birthed (is the start point) of this index function.     base :: Shape num,     -- | ignoring permutations, is the index function contiguous?     contiguous :: Bool@@ -178,6 +210,9 @@ instance FreeIn num => FreeIn (IxFun num) where   freeIn' = foldMap freeIn' +instance FreeIn num => FreeIn (LMADDim num) where+  freeIn' (LMADDim s n _ _) = freeIn' s <> freeIn' n+ instance Functor LMAD where   fmap f = runIdentity . traverse (pure . f) @@ -228,22 +263,25 @@ -- | Substitute a name with a PrimExp in an LMAD. substituteInLMAD ::   Ord a =>-  M.Map a (PrimExp a) ->-  LMAD (PrimExp a) ->-  LMAD (PrimExp a)+  M.Map a (TPrimExp t a) ->+  LMAD (TPrimExp t a) ->+  LMAD (TPrimExp t a) substituteInLMAD tab (LMAD offset dims) =-  let offset' = substituteInPrimExp tab offset+  let offset' = sub offset       dims' =         map           ( \(LMADDim s n p m) ->               LMADDim-                (substituteInPrimExp tab s)-                (substituteInPrimExp tab n)+                (sub s)+                (sub n)                 p                 m           )           dims    in LMAD offset' dims'+  where+    tab' = fmap untyped tab+    sub = TPrimExp . substituteInPrimExp tab' . untyped  -- | Substitute a name with a PrimExp in an index function. substituteInIxFun ::@@ -253,7 +291,7 @@   IxFun (TPrimExp t a) substituteInIxFun tab (IxFun lmads oshp cg) =   IxFun-    (NE.map (fmap TPrimExp . substituteInLMAD tab' . fmap untyped) lmads)+    (NE.map (substituteInLMAD tab) lmads)     (map (TPrimExp . substituteInPrimExp tab' . untyped) oshp)     cg   where@@ -271,6 +309,11 @@           (zip4 dims [0 .. length dims - 1] oshp strides_expected) isDirect _ = False +-- | Is index function "analyzable", i.e., consists of one LMAD+hasOneLmad :: IxFun num -> Bool+hasOneLmad (IxFun (_ :| []) _ _) = True+hasOneLmad _ = False+ -- | Does the index function have an ascending permutation? hasContiguousPerm :: IxFun num -> Bool hasContiguousPerm (IxFun (lmad :| []) _ _) =@@ -612,7 +655,9 @@   Int rank (IxFun (LMAD _ sss :| _) _ _) = length sss --- | Handle the case where a rebase operation can stay within m + n - 1 LMADs,+-- | Essentially @rebase new_base ixfun = ixfun o new_base@+-- Core soundness condition: @base ixfun == shape new_base@+-- Handles the case where a rebase operation can stay within m + n - 1 LMADs, -- where m is the number of LMADs in the index function, and n is the number of -- LMADs in the new base.  If both index function have only on LMAD, this means -- that we stay within the single-LMAD domain.@@ -838,7 +883,9 @@     isMonDim mon (LMADDim s _ _ ldmon) =       s == 0 || mon == ldmon --- | Turn all the leaves of the index function into 'Ext's.+-- | Turn all the leaves of the index function into 'Ext's.  We+--  require that there's only one LMAD, that the index function is+--  contiguous, and the base shape has only one dimension. existentialize ::   IxFun (TPrimExp Int64 a) ->   IxFun (TPrimExp Int64 (Ext b))@@ -890,6 +937,224 @@           == lmadOffset lmad2         && map ldStride (lmadDims lmad1)           == map ldStride (lmadDims lmad2)++-- | Computes the maximum span of an 'LMAD'. The result is the lowest and+-- highest flat values representable by that 'LMAD'.+flatSpan :: LMAD (TPrimExp Int64 VName) -> TPrimExp Int64 VName+flatSpan (LMAD _ dims) =+  foldr+    ( \dim upper ->+        let spn = ldStride dim * (ldShape dim - 1)+         in -- If you've gotten this far, you've already lost+            spn + upper+    )+    0+    dims++-- | Conservatively flatten a list of LMAD dimensions+--+-- Since not all LMADs can actually be flattened, we try to overestimate the+-- flattened array instead. This means that any "holes" in betwen dimensions+-- will get filled out.+-- conservativeFlatten :: (IntegralExp e, Ord e, Pretty e) => LMAD e -> LMAD e+conservativeFlatten :: LMAD (TPrimExp Int64 VName) -> Maybe (LMAD (TPrimExp Int64 VName))+conservativeFlatten (LMAD offset []) =+  pure $ LMAD offset [LMADDim 1 1 0 Inc]+conservativeFlatten l@(LMAD _ [_]) =+  pure l+conservativeFlatten l@(LMAD offset dims) = do+  strd <-+    foldM+      gcd+      (ldStride $ head dims)+      $ map ldStride dims+  pure $ LMAD offset [LMADDim strd (shp + 1) 0 Unknown]+  where+    shp = flatSpan l++-- | Very conservative GCD calculation. Returns 'Nothing' if the result cannot+-- be immediately determined. Does not recurse at all.+gcd :: TPrimExp Int64 VName -> TPrimExp Int64 VName -> Maybe (TPrimExp Int64 VName)+gcd x y = gcd' (abs x) (abs y)+  where+    gcd' a b | a == b = Just a+    gcd' 1 _ = Just 1+    gcd' _ 1 = Just 1+    gcd' a 0 = Just a+    gcd' _ _ = Nothing -- gcd' b (a `Futhark.Util.IntegralExp.rem` b)++-- | Returns @True@ if the two 'LMAD's could be proven disjoint.+--+-- Uses some best-approximation heuristics to determine disjointness. For two+-- 1-dimensional arrays, we can guarantee whether or not they are disjoint, but+-- as soon as more than one dimension is involved, things get more+-- tricky. Currently, we try to 'conservativelyFlatten' any LMAD with more than+-- one dimension.+disjoint :: [(VName, PrimExp VName)] -> Names -> LMAD (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName) -> Bool+disjoint less_thans non_negatives (LMAD offset1 [dim1]) (LMAD offset2 [dim2]) =+  doesNotDivide (gcd (ldStride dim1) (ldStride dim2)) (offset1 - offset2)+    || AlgSimplify.lessThanish+      less_thans+      non_negatives+      (offset2 + (ldShape dim2 - 1) * ldStride dim2)+      offset1+    || AlgSimplify.lessThanish+      less_thans+      non_negatives+      (offset1 + (ldShape dim1 - 1) * ldStride dim1)+      offset2+  where+    doesNotDivide :: Maybe (TPrimExp Int64 VName) -> TPrimExp Int64 VName -> Bool+    doesNotDivide (Just x) y =+      Futhark.Util.IntegralExp.mod y x+        & untyped+        & constFoldPrimExp+        & TPrimExp+        & (.==.) (0 :: TPrimExp Int64 VName)+        & primBool+        & maybe False not+    doesNotDivide _ _ = False+disjoint less_thans non_negatives lmad1 lmad2 =+  case (conservativeFlatten lmad1, conservativeFlatten lmad2) of+    (Just lmad1', Just lmad2') -> disjoint less_thans non_negatives lmad1' lmad2'+    _ -> False++disjoint2 :: scope -> asserts -> [(VName, PrimExp VName)] -> Names -> LMAD (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName) -> Bool+disjoint2 _ _ less_thans non_negatives lmad1 lmad2 =+  let (offset1, interval1) = lmadToIntervals lmad1+      (offset2, interval2) = lmadToIntervals lmad2+      (neg_offset, pos_offset) =+        partition AlgSimplify.negated $+          offset1 `AlgSimplify.sub` offset2+      (interval1', interval2') =+        unzip $+          sortBy (flip AlgSimplify.compareComplexity `on` (AlgSimplify.simplify0 . untyped . stride . fst)) $+            intervalPairs interval1 interval2+   in case ( distributeOffset pos_offset interval1',+             distributeOffset (map AlgSimplify.negate neg_offset) interval2'+           ) of+        (Just interval1'', Just interval2'') ->+          isNothing+            ( selfOverlap () () less_thans (map (flip LeafExp $ IntType Int64) $ namesToList non_negatives) interval1''+            )+            && isNothing+              ( selfOverlap () () less_thans (map (flip LeafExp $ IntType Int64) $ namesToList non_negatives) interval2''+              )+            && any+              (not . uncurry (intervalOverlap less_thans non_negatives))+              (zip interval1'' interval2'')+        _ ->+          False++disjoint3 :: M.Map VName Type -> [PrimExp VName] -> [(VName, PrimExp VName)] -> [PrimExp VName] -> LMAD (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName) -> Bool+disjoint3 scope asserts less_thans non_negatives lmad1 lmad2 =+  let (offset1, interval1) = lmadToIntervals lmad1+      (offset2, interval2) = lmadToIntervals lmad2+      interval1' = fixPoint (mergeDims . joinDims) $ sortBy (flip AlgSimplify.compareComplexity `on` (AlgSimplify.simplify0 . untyped . stride)) interval1+      interval2' = fixPoint (mergeDims . joinDims) $ sortBy (flip AlgSimplify.compareComplexity `on` (AlgSimplify.simplify0 . untyped . stride)) interval2+      (interval1'', interval2'') =+        unzip $+          sortBy (flip AlgSimplify.compareComplexity `on` (AlgSimplify.simplify0 . untyped . stride . fst)) $+            intervalPairs interval1' interval2'+   in disjointHelper 4 interval1'' interval2'' $ offset1 `AlgSimplify.sub` offset2+  where+    disjointHelper :: Int -> [Interval] -> [Interval] -> AlgSimplify.SofP -> Bool+    disjointHelper 0 _ _ _ = False+    disjointHelper i is10 is20 offset =+      let (is1, is2) =+            unzip $+              sortBy (flip AlgSimplify.compareComplexity `on` (AlgSimplify.simplify0 . untyped . stride . fst)) $+                intervalPairs is10 is20+          (neg_offset, pos_offset) = partition AlgSimplify.negated offset+       in case ( distributeOffset pos_offset is1,+                 distributeOffset (map AlgSimplify.negate neg_offset) is2+               ) of+            (Just is1', Just is2') -> do+              let overlap1 = selfOverlap scope asserts less_thans non_negatives is1'+              let overlap2 = selfOverlap scope asserts less_thans non_negatives is2'+              case (overlap1, overlap2) of+                (Nothing, Nothing) ->+                  case namesFromList <$> mapM justLeafExp non_negatives of+                    Just non_negatives' ->+                      any+                        (not . uncurry (intervalOverlap less_thans non_negatives'))+                        (zip is1 is2)+                    _ -> False+                (Just overlapping_dim, _) ->+                  let expanded_offset = AlgSimplify.simplifySofP' <$> expandOffset offset is1+                      splits = splitDim overlapping_dim is1'+                   in all (\(new_offset, new_is1) -> disjointHelper (i - 1) (joinDims new_is1) (joinDims is2') new_offset) splits+                        || maybe False (disjointHelper (i - 1) is1 is2) expanded_offset+                (_, Just overlapping_dim) ->+                  let expanded_offset = AlgSimplify.simplifySofP' <$> expandOffset offset is2+                      splits = splitDim overlapping_dim is2'+                   in all+                        ( \(new_offset, new_is2) ->+                            disjointHelper (i - 1) (joinDims is1') (joinDims new_is2) $+                              map AlgSimplify.negate new_offset+                        )+                        splits+                        || maybe False (disjointHelper (i - 1) is1 is2) expanded_offset+            _ -> False++joinDims :: [Interval] -> [Interval]+joinDims = helper []+  where+    helper acc [] = reverse acc+    helper acc [x] = reverse $ x : acc+    helper acc (x : y : rest) =+      if stride x == stride y && lowerBound x == 0 && lowerBound y == 0+        then helper acc $ x {numElements = numElements x * numElements y} : rest+        else helper (x : acc) (y : rest)++mergeDims :: [Interval] -> [Interval]+mergeDims = helper [] . reverse+  where+    helper acc [] = acc+    helper acc [x] = x : acc+    helper acc (x : y : rest) =+      if stride x * numElements x == stride y && lowerBound x == 0 && lowerBound y == 0+        then helper acc $ x {numElements = numElements x * numElements y} : rest+        else helper (x : acc) (y : rest)++splitDim :: Interval -> [Interval] -> [(AlgSimplify.SofP, [Interval])]+splitDim overlapping_dim0 is+  | [st] <- AlgSimplify.simplify0 $ untyped $ stride overlapping_dim0,+    [st1] <- AlgSimplify.simplify0 $ untyped $ stride overlapping_dim,+    [spn] <- AlgSimplify.simplify0 $ untyped $ stride overlapping_dim * numElements overlapping_dim,+    lowerBound overlapping_dim == 0,+    Just big_dim_elems <- AlgSimplify.maybeDivide spn st,+    Just small_dim_elems <- AlgSimplify.maybeDivide st st1 =+      [ ( [],+          init before+            <> [ Interval 0 (isInt64 $ AlgSimplify.prodToExp big_dim_elems) (stride overlapping_dim0),+                 Interval 0 (isInt64 $ AlgSimplify.prodToExp small_dim_elems) (stride overlapping_dim)+               ]+            <> after+        )+      ]+  | otherwise =+      let shrunk_dim = overlapping_dim {numElements = numElements overlapping_dim - 1}+          point_offset = AlgSimplify.simplify0 $ untyped $ (numElements overlapping_dim - 1 + lowerBound overlapping_dim) * stride overlapping_dim+       in [ (point_offset, before <> after),+            ([], before <> [shrunk_dim] <> after)+          ]+  where+    (before, overlapping_dim, after) =+      fromJust $+        elemIndex overlapping_dim0 is+          >>= (flip focusNth is . (+ 1))++lmadToIntervals :: LMAD (TPrimExp Int64 VName) -> (AlgSimplify.SofP, [Interval])+lmadToIntervals (LMAD offset []) = (AlgSimplify.simplify0 $ untyped offset, [Interval 0 1 1])+lmadToIntervals lmad@(LMAD offset dims0) =+  (offset', map helper $ permuteInv (lmadPermutation lmad) dims0)+  where+    offset' = AlgSimplify.simplify0 $ untyped offset++    helper :: LMADDim (TPrimExp Int64 VName) -> Interval+    helper (LMADDim strd shp _ _) = do+      Interval 0 (AlgSimplify.simplify' shp) (AlgSimplify.simplify' strd)  -- | Dynamically determine if two 'LMADDim' are equal. --

@@ -747,10 +747,12 @@         Nothing -> execState (walkExpM walker e) mempty     onOp op       | Just soac <- asSOAC op =-          execWriter $-            mapSOACM-              identitySOACMapper {mapOnSOACLambda = onLambda}-              (soac :: SOAC rep)+          -- Copies are not safe to move out of nested ops (#1753).+          S.filter (notCopy . snd) $+            execWriter $+              mapSOACM+                identitySOACMapper {mapOnSOACLambda = onLambda}+                (soac :: SOAC rep)       | otherwise =           mempty     onLambda lam = do@@ -761,6 +763,8 @@         { walkOnBody = const $ modify . (<>) . arrayOps,           walkOnOp = modify . (<>) . onOp         }+    notCopy (ArrayCopy {}) = False+    notCopy _ = True  replaceArrayOps ::   forall rep.

@@ -40,7 +40,7 @@     -- holes.)     DynamicFun (Exp, StaticVal) StaticVal   | IntrinsicSV-  | HoleSV SrcLoc+  | HoleSV PatType SrcLoc   deriving (Show)  -- | The type is Just if this is a polymorphic binding that must be@@ -112,8 +112,8 @@         replaceStaticValSizes globals orig_substs sv2     IntrinsicSV ->       IntrinsicSV-    HoleSV loc ->-      HoleSV loc+    HoleSV t loc ->+      HoleSV t loc   where     tv substs =       identityMapper@@ -222,7 +222,7 @@     restrict' u (DynamicFun (e, sv1) sv2) =       DynamicFun (e, restrict' u sv1) $ restrict' u sv2     restrict' _ IntrinsicSV = IntrinsicSV-    restrict' _ (HoleSV loc) = HoleSV loc+    restrict' _ (HoleSV t loc) = HoleSV t loc     restrict'' u (Binding t sv) = Binding t $ restrict' u sv  -- | Defunctionalization monad.  The Reader environment tracks both@@ -504,13 +504,13 @@     IntrinsicSV -> do       (pats, body, tp) <- etaExpand (typeOf e) e       defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty-    HoleSV hole_loc ->-      pure (Hole (Info t) hole_loc, HoleSV hole_loc)+    HoleSV _ hole_loc ->+      pure (Hole (Info t) hole_loc, sv)     _ ->       let tp = typeFromSV sv        in pure (Var qn (Info tp) loc, sv) defuncExp (Hole (Info t) loc) =-  pure (Hole (Info t) loc, HoleSV loc)+  pure (Hole (Info t) loc, HoleSV t loc) defuncExp (Ascript e0 tydecl loc)   | orderZero (typeOf e0) = do       (e0', sv) <- defuncExp e0@@ -942,7 +942,7 @@     -- Propagate the 'IntrinsicsSV' until we reach the outermost application,     -- where we construct a dynamic static value with the appropriate type.     IntrinsicSV -> intrinsicOrHole argtypes e' sv1-    HoleSV _ -> intrinsicOrHole argtypes e' sv1+    HoleSV {} -> intrinsicOrHole argtypes e' sv1     _ ->       error $         "Application of an expression\n"@@ -1138,10 +1138,10 @@ typeFromSV (SumSV name svs fields) =   let svs' = map typeFromSV svs    in Scalar $ Sum $ M.insert name svs' $ M.fromList fields+typeFromSV (HoleSV t _) =+  t typeFromSV IntrinsicSV =   error "Tried to get the type from the static value of an intrinsic."-typeFromSV HoleSV {} =-  error "Tried to get the type from the static value of a hole."  -- | Construct the type for a fully-applied dynamic function from its -- static value and the original types of its arguments.

@@ -121,7 +121,7 @@ lookupImport :: String -> TransformM Scope lookupImport name = maybe bad pure =<< asks (M.lookup name . envImports)   where-    bad = error $ "Unknown import: " ++ name+    bad = error $ "Defunctorise: unknown import: " ++ name  lookupMod' :: QualName VName -> Scope -> Either String Mod lookupMod' mname scope =

@@ -0,0 +1,175 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE PartialTypeSignatures #-}+{-# LANGUAGE TypeFamilies #-}++-- | Perform array short circuiting+module Futhark.Optimise.ArrayShortCircuiting (optimiseSeqMem, optimiseGPUMem) where++import Control.Monad.Reader+import Data.Function ((&))+import Data.Map qualified as M+import Data.Maybe (fromMaybe)+import Futhark.Analysis.Alias qualified as AnlAls+import Futhark.IR.Aliases+import Futhark.IR.GPUMem+import Futhark.IR.Mem.IxFun (substituteInIxFun)+import Futhark.IR.SeqMem+import Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing+import Futhark.Optimise.ArrayShortCircuiting.DataStructs+import Futhark.Pass (Pass (..))+import Futhark.Pass qualified as Pass+import Futhark.Util++----------------------------------------------------------------+--- Printer/Tester Main Program+----------------------------------------------------------------++data Env inner = Env+  { envCoalesceTab :: M.Map VName Coalesced,+    onInner :: inner -> ReplaceM inner inner+  }++type ReplaceM inner a = Reader (Env inner) a++optimiseSeqMem :: Pass SeqMem SeqMem+optimiseSeqMem = pass "short-circuit" "Array Short-Circuiting" mkCoalsTab pure replaceInParams++optimiseGPUMem :: Pass GPUMem GPUMem+optimiseGPUMem = pass "short-circuit-gpu" "Array Short-Circuiting (GPU)" mkCoalsTabGPU replaceInHostOp replaceInParams++replaceInParams :: CoalsTab -> [Param FParamMem] -> (Names, [Param FParamMem])+replaceInParams coalstab fparams =+  let (mem_allocs_to_remove, fparams') =+        foldl+          replaceInParam+          (mempty, mempty)+          fparams+   in (mem_allocs_to_remove, reverse fparams')+  where+    replaceInParam (to_remove, acc) (Param attrs name dec) =+      case dec of+        MemMem DefaultSpace+          | Just entry <- M.lookup name coalstab ->+              (oneName (dstmem entry) <> to_remove, Param attrs (dstmem entry) dec : acc)+        MemArray pt shp u (ArrayIn m ixf)+          | Just entry <- M.lookup m coalstab ->+              (to_remove, Param attrs name (MemArray pt shp u $ ArrayIn (dstmem entry) ixf) : acc)+        _ -> (to_remove, Param attrs name dec : acc)++removeStms :: Names -> Body rep -> Body rep+removeStms to_remove (Body dec stms res) =+  Body dec (stmsFromList $ filter (not . flip nameIn to_remove . head . patNames . stmPat) $ stmsToList stms) res++pass ::+  (Mem rep inner, LetDec rep ~ LetDecMem, CanBeAliased inner) =>+  String ->+  String ->+  (FunDef (Aliases rep) -> Pass.PassM CoalsTab) ->+  (inner -> ReplaceM inner inner) ->+  (CoalsTab -> [FParam (Aliases rep)] -> (Names, [FParam (Aliases rep)])) ->+  Pass rep rep+pass flag desc mk on_inner on_fparams =+  Pass flag desc $+    Pass.intraproceduralTransformationWithConsts pure $ \_ f -> do+      coaltab <- mk (AnlAls.analyseFun f)+      let (mem_allocs_to_remove, new_fparams) = on_fparams coaltab $ funDefParams f+      pure $+        f+          { funDefBody =+              onBody (foldMap vartab $ M.elems coaltab) $+                removeStms mem_allocs_to_remove $+                  funDefBody f,+            funDefParams = new_fparams+          }+  where+    onBody coaltab body =+      body {bodyStms = runReader (mapM replaceInStm $ bodyStms body) (Env coaltab on_inner)}++replaceInStm :: (Mem rep inner, LetDec rep ~ LetDecMem) => Stm rep -> ReplaceM inner (Stm rep)+replaceInStm (Let (Pat elems) d e) = do+  elems' <- mapM replaceInPatElem elems+  e' <- replaceInExp elems' e+  pure $ Let (Pat elems') d e'+  where+    replaceInPatElem :: PatElem LetDecMem -> ReplaceM inner (PatElem LetDecMem)+    replaceInPatElem p@(PatElem vname (MemArray _ _ u _)) =+      fromMaybe p <$> lookupAndReplace vname PatElem u+    replaceInPatElem p = pure p++replaceInExp :: (Mem rep inner, LetDec rep ~ LetDecMem) => [PatElem LetDecMem] -> Exp rep -> ReplaceM inner (Exp rep)+replaceInExp _ e@(BasicOp _) = pure e+replaceInExp pat_elems (Match cond_ses cases defbody dec) = do+  defbody' <- replaceInIfBody defbody+  cases' <- mapM (\(Case p b) -> Case p <$> replaceInIfBody b) cases+  case_rets <- zipWithM (generalizeIxfun pat_elems) pat_elems $ matchReturns dec+  let dec' = dec {matchReturns = case_rets}+  pure $ Match cond_ses cases' defbody' dec'+replaceInExp _ (DoLoop loop_inits loop_form (Body dec stms res)) = do+  loop_inits' <- mapM (replaceInFParam . fst) loop_inits+  stms' <- mapM replaceInStm stms+  pure $ DoLoop (zip loop_inits' $ map snd loop_inits) loop_form $ Body dec stms' res+replaceInExp _ e@(Op (Alloc _ _)) = pure e+replaceInExp _ (Op (Inner i)) = do+  on_op <- asks onInner+  Op . Inner <$> on_op i+replaceInExp _ (Op _) = error "Unreachable" -- This shouldn't be possible?+replaceInExp _ e@WithAcc {} = pure e+replaceInExp _ e@Apply {} = pure e++replaceInHostOp :: HostOp GPUMem () -> ReplaceM (HostOp GPUMem ()) (HostOp GPUMem ())+replaceInHostOp (SegOp (SegMap lvl sp tps body)) = do+  stms <- mapM replaceInStm $ kernelBodyStms body+  pure $ SegOp $ SegMap lvl sp tps $ body {kernelBodyStms = stms}+replaceInHostOp (SegOp (SegRed lvl sp binops tps body)) = do+  stms <- mapM replaceInStm $ kernelBodyStms body+  pure $ SegOp $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}+replaceInHostOp (SegOp (SegScan lvl sp binops tps body)) = do+  stms <- mapM replaceInStm $ kernelBodyStms body+  pure $ SegOp $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}+replaceInHostOp (SegOp (SegHist lvl sp hist_ops tps body)) = do+  stms <- mapM replaceInStm $ kernelBodyStms body+  pure $ SegOp $ SegHist lvl sp hist_ops tps $ body {kernelBodyStms = stms}+replaceInHostOp op = pure op++generalizeIxfun :: [PatElem dec] -> PatElem LetDecMem -> BodyReturns -> ReplaceM inner BodyReturns+generalizeIxfun+  pat_elems+  (PatElem vname (MemArray _ _ _ (ArrayIn mem ixf)))+  m@(MemArray pt shp u _) = do+    coaltab <- asks envCoalesceTab+    if vname `M.member` coaltab+      then+        existentialiseIxFun (map patElemName pat_elems) ixf+          & ReturnsInBlock mem+          & MemArray pt shp u+          & pure+      else pure m+generalizeIxfun _ _ m = pure m++replaceInIfBody :: (Mem rep inner, LetDec rep ~ LetDecMem) => Body rep -> ReplaceM inner (Body rep)+replaceInIfBody b@(Body _ stms _) = do+  stms' <- mapM replaceInStm stms+  pure $ b {bodyStms = stms'}++replaceInFParam :: Param FParamMem -> ReplaceM inner (Param FParamMem)+replaceInFParam p@(Param _ vname (MemArray _ _ u _)) = do+  fromMaybe p <$> lookupAndReplace vname (Param mempty) u+replaceInFParam p = pure p++lookupAndReplace ::+  VName ->+  (VName -> MemBound u -> a) ->+  u ->+  ReplaceM inner (Maybe a)+lookupAndReplace vname f u = do+  coaltab <- asks envCoalesceTab+  case M.lookup vname coaltab of+    Just (Coalesced _ (MemBlock pt shp mem ixf) subs) ->+      ixf+        & fixPoint (substituteInIxFun subs)+        & ArrayIn mem+        & MemArray pt shp u+        & f vname+        & Just+        & pure+    Nothing -> pure Nothing

@@ -0,0 +1,1548 @@+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE TypeFamilies #-}++-- | The bulk of the short-circuiting implementation.+module Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing (mkCoalsTab, CoalsTab, mkCoalsTabGPU) where++import Control.Exception.Base qualified as Exc+import Control.Monad.Reader+import Control.Monad.State.Strict+import Data.Function ((&))+import Data.List qualified as L+import Data.List.NonEmpty (NonEmpty (..))+import Data.List.NonEmpty qualified as NE+import Data.Map.Strict qualified as M+import Data.Maybe+import Data.Sequence (Seq (..))+import Data.Set qualified as S+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.Aliases+import Futhark.IR.GPUMem+import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.SeqMem+import Futhark.MonadFreshNames+import Futhark.Optimise.ArrayShortCircuiting.DataStructs+import Futhark.Optimise.ArrayShortCircuiting.LastUse+import Futhark.Optimise.ArrayShortCircuiting.MemRefAggreg+import Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis+import Futhark.Util++-- | A helper type describing representations that can be short-circuited.+type Coalesceable rep inner =+  ( CreatesNewArrOp (OpWithAliases inner),+    ASTRep rep,+    CanBeAliased inner,+    Op rep ~ MemOp inner,+    HasMemBlock (Aliases rep),+    LetDec rep ~ LetDecMem,+    TopDownHelper (OpWithAliases inner)+  )++-- Helper type for computing scalar tables on ops.+newtype ComputeScalarTableOnOp rep = ComputeScalarTableOnOp+  { scalarTableOnOp :: ScopeTab rep -> Op (Aliases rep) -> ScalarTableM rep (M.Map VName (PrimExp VName))+  }++type ScalarTableM rep a = Reader (ComputeScalarTableOnOp rep) a++newtype ShortCircuitReader rep = ShortCircuitReader+  { onOp :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Op (Aliases rep) -> TopdownEnv rep -> BotUpEnv -> ShortCircuitM rep BotUpEnv+  }++newtype ShortCircuitM rep a = ShortCircuitM (ReaderT (ShortCircuitReader rep) (State VNameSource) a)+  deriving (Functor, Applicative, Monad, MonadReader (ShortCircuitReader rep), MonadState VNameSource)++instance MonadFreshNames (ShortCircuitM rep) where+  putNameSource = put+  getNameSource = get++emptyTopdownEnv :: TopdownEnv rep+emptyTopdownEnv =+  TopdownEnv+    { alloc = mempty,+      scope = mempty,+      inhibited = mempty,+      v_alias = mempty,+      m_alias = mempty,+      nonNegatives = mempty,+      scalarTable = mempty,+      knownLessThan = mempty,+      td_asserts = mempty+    }++emptyBotUpEnv :: BotUpEnv+emptyBotUpEnv =+  BotUpEnv+    { scals = mempty,+      activeCoals = mempty,+      successCoals = mempty,+      inhibit = mempty+    }++--------------------------------------------------------------------------------+--- Main Coalescing Transformation computes a successful coalescing table    ---+--------------------------------------------------------------------------------++-- | Given a 'FunDef' in 'SegMem' representation, compute the coalescing table+-- by folding over each function.+mkCoalsTab :: (MonadFreshNames m) => FunDef (Aliases SeqMem) -> m CoalsTab+mkCoalsTab =+  mkCoalsTabFun+    (snd . lastUseSeqMem)+    (ShortCircuitReader shortCircuitSeqMem)+    (ComputeScalarTableOnOp $ const $ const $ pure mempty)++-- | Given a 'FunDef' in 'GPUMem' representation, compute the coalescing table+-- by folding over each function.+mkCoalsTabGPU :: (MonadFreshNames m) => FunDef (Aliases GPUMem) -> m CoalsTab+mkCoalsTabGPU =+  mkCoalsTabFun+    (snd . lastUseGPUMem)+    (ShortCircuitReader shortCircuitGPUMem)+    (ComputeScalarTableOnOp computeScalarTableGPUMem)++-- | Given a function, compute the coalescing table+mkCoalsTabFun ::+  (MonadFreshNames m, Coalesceable rep inner, FParamInfo rep ~ FParamMem) =>+  (FunDef (Aliases rep) -> LUTabFun) ->+  ShortCircuitReader rep ->+  ComputeScalarTableOnOp rep ->+  FunDef (Aliases rep) ->+  m CoalsTab+mkCoalsTabFun lufun r computeScalarOnOp fun@(FunDef _ _ _ _ fpars body) = do+  -- First compute last-use information+  let lutab = lufun fun+      unique_mems = getUniqueMemFParam fpars+      scalar_table =+        runReader+          ( concatMapM+              (computeScalarTable $ scopeOf fun <> scopeOf (bodyStms body))+              (stmsToList $ bodyStms body)+          )+          computeScalarOnOp+      topenv =+        emptyTopdownEnv+          { scope = scopeOfFParams fpars,+            alloc = unique_mems,+            scalarTable = scalar_table,+            nonNegatives = foldMap paramSizes fpars+          }+      ShortCircuitM m = fixPointCoalesce lutab fpars body topenv+  modifyNameSource $ runState (runReaderT m r)++paramSizes :: Param FParamMem -> Names+paramSizes (Param _ _ (MemArray _ shp _ _)) = freeIn shp+paramSizes _ = mempty++-- | Short-circuit handler for a 'SeqMem' 'Op'.+--+-- Because 'SeqMem' don't have any special operation, simply return the input+-- 'BotUpEnv'.+shortCircuitSeqMem :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Op (Aliases SeqMem) -> TopdownEnv SeqMem -> BotUpEnv -> ShortCircuitM SeqMem BotUpEnv+shortCircuitSeqMem _ _ _ _ = pure++-- | Short-circuit handler for 'GPUMem' 'Op'.+--+-- When the 'Op' is a 'SegOp', we handle it accordingly, otherwise we do+-- nothing.+shortCircuitGPUMem ::+  LUTabFun ->+  Pat (VarAliases, LetDecMem) ->+  Op (Aliases GPUMem) ->+  TopdownEnv GPUMem ->+  BotUpEnv ->+  ShortCircuitM GPUMem BotUpEnv+shortCircuitGPUMem _ _ (Alloc _ _) _ bu_env = pure bu_env+shortCircuitGPUMem lutab pat (Inner (SegOp (SegMap lvl@SegThread {} space _ kernel_body))) td_env bu_env =+  -- No special handling necessary for 'SegMap'. Just call the helper-function.+  shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env+shortCircuitGPUMem lutab pat (Inner (SegOp (SegMap lvl@SegGroup {} space _ kernel_body))) td_env bu_env =+  -- No special handling necessary for 'SegMap'. Just call the helper-function.+  shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env+shortCircuitGPUMem lutab pat (Inner (SegOp (SegMap lvl@SegThreadInGroup {} space _ kernel_body))) td_env bu_env =+  -- No special handling necessary for 'SegMap'. Just call the helper-function.+  shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env+shortCircuitGPUMem lutab pat (Inner (SegOp (SegRed lvl space binops _ kernel_body))) td_env bu_env =+  -- When handling 'SegRed', we we first invalidate all active coalesce-entries+  -- where any of the variables in 'vartab' are also free in the list of+  -- 'SegBinOp'. In other words, anything that is used as part of the reduction+  -- step should probably not be coalesced.+  let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . segBinOpLambda) binops) $ activeCoals bu_env+      (active, inh) =+        foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail+      bu_env' = bu_env {activeCoals = active, inhibit = inh}+      num_reds = length red_ts+   in shortCircuitGPUMemHelper num_reds lvl lutab pat space kernel_body td_env bu_env'+  where+    segment_dims = init $ segSpaceDims space+    red_ts = do+      op <- binops+      let shp = Shape segment_dims <> segBinOpShape op+      map (`arrayOfShape` shp) (lambdaReturnType $ segBinOpLambda op)+shortCircuitGPUMem lutab pat (Inner (SegOp (SegScan lvl space binops _ kernel_body))) td_env bu_env =+  -- Like in the handling of 'SegRed', we do not want to coalesce anything that+  -- is used in the 'SegBinOp'+  let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . segBinOpLambda) binops) $ activeCoals bu_env+      (active, inh) = foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail+      bu_env' = bu_env {activeCoals = active, inhibit = inh}+   in shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env'+shortCircuitGPUMem lutab pat (Inner (SegOp (SegHist lvl space histops _ kernel_body))) td_env bu_env = do+  -- Need to take zipped patterns and histDest (flattened) and insert transitive coalesces+  let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . histOp) histops) $ activeCoals bu_env+      (active, inh) = foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail+      bu_env' = bu_env {activeCoals = active, inhibit = inh}+  bu_env'' <- shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env'+  pure $+    foldl insertHistCoals bu_env'' $+      zip (patElems pat) $+        concatMap histDest histops+  where+    insertHistCoals acc (PatElem p _, hist_dest) =+      case ( getScopeMemInfo p $ scope td_env,+             getScopeMemInfo hist_dest $ scope td_env+           ) of+        (Just (MemBlock _ _ p_mem _), Just (MemBlock _ _ dest_mem _)) ->+          case M.lookup p_mem $ successCoals acc of+            Just entry ->+              -- Update this entry with an optdep for the memory block of hist_dest+              let entry' = entry {optdeps = M.insert p p_mem $ optdeps entry}+               in acc+                    { successCoals = M.insert p_mem entry' $ successCoals acc,+                      activeCoals = M.insert dest_mem entry $ activeCoals acc+                    }+            Nothing -> acc+        _ -> acc+shortCircuitGPUMem lutab pat (Inner (GPUBody _ body)) td_env bu_env = do+  fresh1 <- newNameFromString "gpubody"+  fresh2 <- newNameFromString "gpubody"+  shortCircuitGPUMemHelper+    0+    -- Construct a 'SegLevel' corresponding to a single thread+    ( SegThread SegNoVirt $+        Just $+          KernelGrid+            (Count $ Constant $ IntValue $ Int64Value 1)+            (Count $ Constant $ IntValue $ Int64Value 1)+    )+    lutab+    pat+    (SegSpace fresh1 [(fresh2, Constant $ IntValue $ Int64Value 1)])+    (bodyToKernelBody body)+    td_env+    bu_env+shortCircuitGPUMem _ _ (Inner (SizeOp _)) _ bu_env = pure bu_env+shortCircuitGPUMem _ _ (Inner (OtherOp ())) _ bu_env = pure bu_env++dropLastSegSpace :: SegSpace -> SegSpace+dropLastSegSpace space = space {unSegSpace = init $ unSegSpace space}++isSegThread :: SegLevel -> Bool+isSegThread SegThread {} = True+isSegThread _ = False++-- | Computes the slice written at the end of a thread in a 'SegOp'.+threadSlice :: SegSpace -> KernelResult -> Maybe (Slice (TPrimExp Int64 VName))+threadSlice space Returns {} =+  Just $+    Slice $+      map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst) $+        unSegSpace space+threadSlice space (RegTileReturns _ dims _) =+  Just+    $ Slice+    $ zipWith+      ( \(_, block_tile_size0, reg_tile_size0) (x0, _) ->+          let x = pe64 $ Var x0+              block_tile_size = pe64 block_tile_size0+              reg_tile_size = pe64 reg_tile_size0+           in DimSlice (x * block_tile_size * reg_tile_size) (block_tile_size * reg_tile_size) 1+      )+      dims+    $ unSegSpace space+threadSlice _ _ = Nothing++bodyToKernelBody :: Body (Aliases GPUMem) -> KernelBody (Aliases GPUMem)+bodyToKernelBody (Body dec stms res) =+  KernelBody dec stms $ map (\(SubExpRes cert subexps) -> Returns ResultNoSimplify cert subexps) res++-- | A helper for all the different kinds of 'SegOp'.+--+-- Consists of four parts:+--+-- 1. Create coalescing relations between the pattern elements and the kernel+-- body results using 'makeSegMapCoals'.+--+-- 2. Process the statements of the 'KernelBody'.+--+-- 3. Check the overlap between the different threads.+--+-- 4. Mark active coalescings as finished, since a 'SegOp' is an array creation+-- point.+shortCircuitGPUMemHelper ::+  -- | The number of returns for which we should drop the last seg space+  Int ->+  SegLevel ->+  LUTabFun ->+  Pat (VarAliases, LetDecMem) ->+  SegSpace ->+  KernelBody (Aliases GPUMem) ->+  TopdownEnv GPUMem ->+  BotUpEnv ->+  ShortCircuitM GPUMem BotUpEnv+shortCircuitGPUMemHelper num_reds lvl lutab pat@(Pat ps0) space0 kernel_body td_env bu_env = do+  -- We need to drop the last element of the 'SegSpace' for pattern elements+  -- that correspond to reductions.+  let ps_space_and_res =+        zip3 ps0 (replicate num_reds (dropLastSegSpace space0) <> repeat space0) $+          kernelBodyResult kernel_body+  -- Create coalescing relations between pattern elements and kernel body+  -- results+  let (actv0, inhibit0) =+        filterSafetyCond2and5+          (activeCoals bu_env)+          (inhibit bu_env)+          (scals bu_env)+          td_env+          (patElems pat)+      (actv_return, inhibit_return) =+        if num_reds > 0+          then (actv0, inhibit0)+          else foldl (makeSegMapCoals lvl td_env kernel_body) (actv0, inhibit0) ps_space_and_res++  -- Start from empty references, we'll update with aggregates later.+  let actv0' = M.map (\etry -> etry {memrefs = mempty}) $ actv0 <> actv_return+  -- Process kernel body statements+  bu_env' <-+    mkCoalsTabStms lutab (kernelBodyStms kernel_body) td_env $+      bu_env {activeCoals = actv0', inhibit = inhibit_return}++  let actv_coals_after =+        M.mapWithKey+          ( \k etry ->+              etry+                { memrefs = memrefs etry <> maybe mempty memrefs (M.lookup k $ actv0 <> actv_return)+                }+          )+          $ activeCoals bu_env'++  -- Check partial overlap.+  let checkPartialOverlap bu_env_f (k, entry) = do+        let sliceThreadAccess (p, space, res) =+              case M.lookup (patElemName p) $ vartab entry of+                Just (Coalesced _ (MemBlock _ _ _ ixf) _) ->+                  maybe+                    Undeterminable+                    ( ixfunToAccessSummary+                        . IxFun.slice ixf+                        . fullSlice (IxFun.shape ixf)+                    )+                    $ threadSlice space res+                Nothing -> mempty+            thread_writes = foldMap sliceThreadAccess ps_space_and_res+            source_writes = srcwrts (memrefs entry) <> thread_writes+        destination_uses <-+          case dstrefs (memrefs entry)+            `accessSubtract` dstrefs (maybe mempty memrefs $ M.lookup k $ activeCoals bu_env) of+            Set s ->+              concatMapM+                (aggSummaryMapPartial (scalarTable td_env) $ unSegSpace space0)+                (S.toList s)+            Undeterminable -> pure Undeterminable+        let res = noMemOverlap td_env destination_uses source_writes+        if res+          then pure bu_env_f+          else do+            let (ac, inh) = markFailedCoal (activeCoals bu_env_f, inhibit bu_env_f) k+            pure $ bu_env_f {activeCoals = ac, inhibit = inh}++  bu_env'' <-+    foldM+      checkPartialOverlap+      (bu_env' {activeCoals = actv_coals_after})+      $ M.toList actv_coals_after++  let updateMemRefs entry = do+        wrts <- aggSummaryMapTotal (scalarTable td_env) (unSegSpace space0) $ srcwrts $ memrefs entry+        uses <- aggSummaryMapTotal (scalarTable td_env) (unSegSpace space0) $ dstrefs $ memrefs entry++        -- Add destination uses from the pattern+        let uses' =+              foldMap+                ( \case+                    PatElem _ (_, MemArray _ _ _ (ArrayIn p_mem p_ixf))+                      | p_mem `nameIn` alsmem entry ->+                          ixfunToAccessSummary p_ixf+                    _ -> mempty+                )+                ps0++        pure $ entry {memrefs = MemRefs (uses <> uses') wrts}++  actv <- mapM updateMemRefs $ activeCoals bu_env''+  let bu_env''' = bu_env'' {activeCoals = actv}++  -- Process pattern and return values+  let mergee_writes =+        mapMaybe+          ( \(p, _, _) ->+              fmap (p,) $+                getDirAliasedIxfn' td_env (activeCoals bu_env''') $+                  patElemName p+          )+          ps_space_and_res++  -- Now, for each mergee write, we need to check that it doesn't overlap with any previous uses of the destination.+  let checkMergeeOverlap bu_env_f (p, (m_b, _, ixf)) =+        let as = ixfunToAccessSummary ixf+         in -- Should be @bu_env@ here, because we need to check overlap+            -- against previous uses.+            case M.lookup m_b $ activeCoals bu_env of+              Just coal_entry -> do+                let mrefs =+                      memrefs coal_entry+                    res = noMemOverlap td_env as $ dstrefs mrefs+                    fail_res =+                      let (ac, inh) = markFailedCoal (activeCoals bu_env_f, inhibit bu_env_f) m_b+                       in bu_env_f {activeCoals = ac, inhibit = inh}++                if res+                  then case M.lookup (patElemName p) $ vartab coal_entry of+                    Nothing -> pure bu_env_f+                    Just (Coalesced knd mbd@(MemBlock _ _ _ ixfn) _) -> pure $+                      case freeVarSubstitutions (scope td_env) (scalarTable td_env) ixfn of+                        Just fv_subst ->+                          if ixfunPermutation ixfn+                            == ixfunPermutation (ixfun $ fromJust $ getScopeMemInfo (patElemName p) $ scope td_env)+                            then+                              let entry =+                                    coal_entry+                                      { vartab =+                                          M.insert+                                            (patElemName p)+                                            (Coalesced knd mbd fv_subst)+                                            (vartab coal_entry)+                                      }+                                  (ac, suc) =+                                    markSuccessCoal (activeCoals bu_env_f, successCoals bu_env_f) m_b entry+                               in bu_env_f {activeCoals = ac, successCoals = suc}+                            else fail_res+                        Nothing ->+                          fail_res+                  else pure fail_res+              _ -> pure bu_env_f++  foldM checkMergeeOverlap bu_env''' mergee_writes++ixfunPermutation :: IxFun -> [Int]+ixfunPermutation = map IxFun.ldPerm . IxFun.lmadDims . NE.head . IxFun.ixfunLMADs++-- | Given a pattern element and the corresponding kernel result, try to put the+-- kernel result directly in the memory block of pattern element+makeSegMapCoals :: SegLevel -> TopdownEnv GPUMem -> KernelBody (Aliases GPUMem) -> (CoalsTab, InhibitTab) -> (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult) -> (CoalsTab, InhibitTab)+makeSegMapCoals lvl td_env kernel_body (active, inhb) (PatElem pat_name (_, MemArray _ _ _ (ArrayIn pat_mem pat_ixf)), space, Returns _ _ (Var return_name))+  | Just mb@(MemBlock tp return_shp return_mem _) <-+      getScopeMemInfo return_name $ scope td_env <> scopeOf (kernelBodyStms kernel_body),+    isSegThread lvl,+    MemMem pat_space <- runReader (lookupMemInfo pat_mem) $ removeScopeAliases $ scope td_env,+    MemMem return_space <- runReader (lookupMemInfo return_mem) $ removeScopeAliases $ scope td_env <> scopeOf (kernelBodyStms kernel_body) <> scopeOfSegSpace space,+    pat_space == return_space =+      case M.lookup pat_mem active of+        Nothing ->+          -- We are not in a transitive case+          if IxFun.hasOneLmad pat_ixf+            then case ( maybe False (pat_mem `nameIn`) $ M.lookup return_mem inhb,+                        Coalesced InPlaceCoal mb mempty+                          & M.singleton return_name+                          & flip (addInvAliassesVarTab td_env) return_name+                          & fmap+                            ( M.adjust+                                ( \(Coalesced knd (MemBlock pt shp _ _) subst) ->+                                    Coalesced+                                      knd+                                      ( MemBlock pt shp pat_mem $+                                          IxFun.slice pat_ixf $+                                            fullSlice (IxFun.shape pat_ixf) $+                                              Slice $+                                                map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst) $+                                                  unSegSpace space+                                      )+                                      subst+                                )+                                return_name+                            )+                      ) of+              (False, Just vtab) ->+                (active <> M.singleton return_mem (CoalsEntry pat_mem pat_ixf (oneName pat_mem) vtab mempty mempty), inhb)+              _ -> (active, inhb)+            else (active, inhb)+        Just trans ->+          case ( maybe False (dstmem trans `nameIn`) $ M.lookup return_mem inhb,+                 Coalesced InPlaceCoal (MemBlock tp return_shp (dstmem trans) (dstind trans)) mempty+                   & M.singleton return_name+                   & flip (addInvAliassesVarTab td_env) return_name+                   & fmap+                     ( M.adjust+                         ( \(Coalesced knd (MemBlock pt shp mem ixf@(IxFun.IxFun _ base_shape _)) subst) ->+                             Coalesced+                               knd+                               ( MemBlock pt shp mem $+                                   IxFun.slice ixf $+                                     fullSlice base_shape $+                                       Slice $+                                         map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst) $+                                           unSegSpace space+                               )+                               subst+                         )+                         return_name+                     )+               ) of+            (False, Just vtab) ->+              let opts = if dstmem trans == pat_mem then mempty else M.insert pat_name pat_mem $ optdeps trans+               in ( M.insert+                      return_mem+                      ( CoalsEntry+                          (dstmem trans)+                          (dstind trans)+                          (oneName pat_mem <> alsmem trans)+                          vtab+                          opts+                          mempty+                      )+                      active,+                    inhb+                  )+            _ -> (active, inhb)+makeSegMapCoals _ td_env _ x (_, _, WriteReturns _ _ return_name _) =+  case getScopeMemInfo return_name $ scope td_env of+    Just (MemBlock _ _ return_mem _) -> markFailedCoal x return_mem+    Nothing -> error "Should not happen?"+makeSegMapCoals _ td_env _ x (_, _, result) =+  freeIn result+    & namesToList+    & mapMaybe (flip getScopeMemInfo $ scope td_env)+    & foldr (\(MemBlock _ _ mem _) -> flip markFailedCoal mem) x++fullSlice :: [TPrimExp Int64 VName] -> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)+fullSlice shp (Slice slc) =+  Slice $ slc ++ map (\d -> DimSlice 0 d 1) (drop (length slc) shp)++fixPointCoalesce ::+  (Coalesceable rep inner) =>+  LUTabFun ->+  [Param FParamMem] ->+  Body (Aliases rep) ->+  TopdownEnv rep ->+  ShortCircuitM rep CoalsTab+fixPointCoalesce lutab fpar bdy topenv = do+  buenv <- mkCoalsTabStms lutab (bodyStms bdy) topenv (emptyBotUpEnv {inhibit = inhibited topenv})+  let (succ_tab, actv_tab, inhb_tab) = (successCoals buenv, activeCoals buenv, inhibit buenv)+      -- Allow short-circuiting function parameters that are unique and have+      -- matching index functions, otherwise mark as failed+      handleFunctionParams (a, i, s) (_, u, MemBlock _ _ m ixf) =+        case (u, M.lookup m a) of+          (Unique, Just entry)+            | dstind entry == ixf ->+                let (a', s') = markSuccessCoal (a, s) m entry+                 in (a', i, s')+          _ ->+            let (a', i') = markFailedCoal (a, i) m+             in (a', i', s)+      (actv_tab', inhb_tab', succ_tab') =+        foldl+          handleFunctionParams+          (actv_tab, inhb_tab, succ_tab)+          $ getArrMemAssocFParam fpar++      (succ_tab'', failed_optdeps) = fixPointFilterDeps succ_tab' M.empty+      inhb_tab'' = M.unionWith (<>) failed_optdeps inhb_tab'+   in if not $ M.null actv_tab'+        then error ("COALESCING ROOT: BROKEN INV, active not empty: " ++ show (M.keys actv_tab'))+        else+          if M.null $ inhb_tab'' `M.difference` inhibited topenv+            then pure succ_tab''+            else fixPointCoalesce lutab fpar bdy (topenv {inhibited = inhb_tab''})+  where+    fixPointFilterDeps :: CoalsTab -> InhibitTab -> (CoalsTab, InhibitTab)+    fixPointFilterDeps coaltab inhbtab =+      let (coaltab', inhbtab') = foldl filterDeps (coaltab, inhbtab) (M.keys coaltab)+       in if length (M.keys coaltab) == length (M.keys coaltab')+            then (coaltab', inhbtab')+            else fixPointFilterDeps coaltab' inhbtab'++    filterDeps (coal, inhb) mb+      | not (M.member mb coal) = (coal, inhb)+    filterDeps (coal, inhb) mb+      | Just coal_etry <- M.lookup mb coal =+          let failed = M.filterWithKey (failedOptDep coal) (optdeps coal_etry)+           in if M.null failed+                then (coal, inhb) -- all ok+                else -- optimistic dependencies failed for the current+                -- memblock; extend inhibited mem-block mergings.+                  markFailedCoal (coal, inhb) mb+    filterDeps _ _ = error "In ArrayCoalescing.hs, fun filterDeps, impossible case reached!"+    failedOptDep coal _ mr+      | not (mr `M.member` coal) = True+    failedOptDep coal r mr+      | Just coal_etry <- M.lookup mr coal = not $ r `M.member` vartab coal_etry+    failedOptDep _ _ _ = error "In ArrayCoalescing.hs, fun failedOptDep, impossible case reached!"++-- | Perform short-circuiting on 'Stms'.+mkCoalsTabStms ::+  (Coalesceable rep inner) =>+  LUTabFun ->+  Stms (Aliases rep) ->+  TopdownEnv rep ->+  BotUpEnv ->+  ShortCircuitM rep BotUpEnv+mkCoalsTabStms lutab stms0 = traverseStms stms0+  where+    non_negs_in_pats = foldMap (nonNegativesInPat . stmPat) stms0+    traverseStms Empty _ bu_env = pure bu_env+    traverseStms (stm :<| stms) td_env bu_env = do+      -- Compute @td_env@ top down+      let td_env' = updateTopdownEnv td_env stm+      -- Compute @bu_env@ bottom up+      bu_env' <- traverseStms stms td_env' bu_env+      mkCoalsTabStm lutab stm (td_env' {nonNegatives = nonNegatives td_env' <> non_negs_in_pats}) bu_env'++-- | Array (register) coalescing can have one of three shapes:+--      a) @let y    = copy(b^{lu})@+--      b) @let y    = concat(a, b^{lu})@+--      c) @let y[i] = b^{lu}@+--   The intent is to use the memory block of the left-hand side+--     for the right-hand side variable, meaning to store @b@ in+--     @m_y@ (rather than @m_b@).+--   The following five safety conditions are necessary:+--      1. the right-hand side is lastly-used in the current statement+--      2. the allocation of @m_y@ dominates the creation of @b@+--         ^ relax it by hoisting the allocation of @m_y@+--      3. there is no use of the left-hand side memory block @m_y@+--           during the liveness of @b@, i.e., in between its last use+--           and its creation.+--         ^ relax it by pointwise/interval-based checking+--      4. @b@ is a newly created array, i.e., does not aliases anything+--         ^ relax it to support exitential memory blocks for if-then-else+--      5. the new index function of @b@ corresponding to memory block @m_y@+--           can be translated at the definition of @b@, and the+--           same for all variables aliasing @b@.+--   Observation: during the live range of @b@, @m_b@ can only be used by+--                variables aliased with @b@, because @b@ is newly created.+--                relax it: in case @m_b@ is existential due to an if-then-else+--                          then the checks should be extended to the actual+--                          array-creation points.+mkCoalsTabStm ::+  (Coalesceable rep inner) =>+  LUTabFun ->+  Stm (Aliases rep) ->+  TopdownEnv rep ->+  BotUpEnv ->+  ShortCircuitM rep BotUpEnv+mkCoalsTabStm _ (Let (Pat [pe]) _ e) td_env bu_env+  | Just primexp <- primExpFromExp (vnameToPrimExp (scope td_env) (scals bu_env)) e =+      pure $ bu_env {scals = M.insert (patElemName pe) primexp (scals bu_env)}+mkCoalsTabStm lutab (Let patt _ (Match _ cases defbody _)) td_env bu_env = do+  let pat_val_elms = patElems patt+      -- ToDo: 1. we need to record existential memory blocks in alias table on the top-down pass.+      --       2. need to extend the scope table++      --  i) Filter @activeCoals@ by the 2ND AND 5th safety conditions:+      (activeCoals0, inhibit0) =+        filterSafetyCond2and5+          (activeCoals bu_env)+          (inhibit bu_env)+          (scals bu_env)+          td_env+          pat_val_elms++      -- ii) extend @activeCoals@ by transfering the pattern-elements bindings existent+      --     in @activeCoals@ to the body results of the then and else branches, but only+      --     if the current pattern element can be potentially coalesced and also+      --     if the current pattern element satisfies safety conditions 2 & 5.+      res_mem_def = findMemBodyResult activeCoals0 (scope td_env) pat_val_elms defbody+      res_mem_cases = map (findMemBodyResult activeCoals0 (scope td_env) pat_val_elms . caseBody) cases++      subs_def = mkSubsTab patt $ map resSubExp $ bodyResult defbody+      subs_cases = map (mkSubsTab patt . map resSubExp . bodyResult . caseBody) cases++      actv_def_i = foldl (transferCoalsToBody subs_def) activeCoals0 res_mem_def+      actv_cases_i = zipWith (\subs res -> foldl (transferCoalsToBody subs) activeCoals0 res) subs_cases res_mem_cases++      -- eliminate the original pattern binding of the if statement,+      -- @let x = if y[0,0] > 0 then map (+y[0,0]) a else map (+1) b@+      -- @let y[0] = x@+      -- should succeed because @m_y@ is used before @x@ is created.+      aux ac (MemBodyResult m_b _ _ m_r) = if m_b == m_r then ac else M.delete m_b ac+      actv_def = foldl aux actv_def_i res_mem_def+      actv_cases = zipWith (foldl aux) actv_cases_i res_mem_cases++  -- iii) process the then and else bodies+  res_def <- mkCoalsTabStms lutab (bodyStms defbody) td_env (bu_env {activeCoals = actv_def})+  res_cases <- zipWithM (\c a -> mkCoalsTabStms lutab (bodyStms $ caseBody c) td_env (bu_env {activeCoals = a})) cases actv_cases+  let (actv_def0, succ_def0, inhb_def0) = (activeCoals res_def, successCoals res_def, inhibit res_def)++      -- iv) optimistically mark the pattern succesful:+      ((activeCoals1, inhibit1), successCoals1) =+        foldl+          ( foldfun+              ( (actv_def0, succ_def0)+                  : zip (map activeCoals res_cases) (map successCoals res_cases)+              )+          )+          ((activeCoals0, inhibit0), successCoals bu_env)+          (L.transpose $ res_mem_def : res_mem_cases)++      --  v) unify coalescing results of all branches by taking the union+      --     of all entries in the current/then/else success tables.++      actv_res = foldr (M.intersectionWith unionCoalsEntry) activeCoals1 $ actv_def0 : map activeCoals res_cases++      succ_res = foldr (M.unionWith unionCoalsEntry) successCoals1 $ succ_def0 : map successCoals res_cases++      -- vi) The step of filtering by 3rd safety condition is not+      --       necessary, because we perform index analysis of the+      --       source/destination uses, and they should have been+      --       filtered during the analysis of the then/else bodies.+      inhibit_res =+        M.unionsWith+          (<>)+          ( inhibit1+              : zipWith+                ( \actv inhb ->+                    let failed = M.difference actv $ M.intersectionWith unionCoalsEntry actv activeCoals0+                     in snd $ foldl markFailedCoal (failed, inhb) (M.keys failed)+                )+                (actv_def0 : map activeCoals res_cases)+                (inhb_def0 : map inhibit res_cases)+          )+  pure+    bu_env+      { activeCoals =+          actv_res,+        successCoals = succ_res,+        inhibit = inhibit_res+      }+  where+    foldfun _ _ [] =+      error "Imposible Case 1!!!"+    foldfun _ ((act, _), _) mem_body_results+      | Nothing <- M.lookup (patMem $ head mem_body_results) act =+          error "Imposible Case 2!!!"+    foldfun+      acc+      ((act, inhb), succc)+      mem_body_results@(MemBodyResult m_b _ _ _ : _)+        | Just info <- M.lookup m_b act,+          Just _ <- zipWithM (M.lookup . bodyMem) mem_body_results $ map snd acc =+            -- Optimistically promote to successful coalescing and append!+            let info' =+                  info+                    { optdeps =+                        foldr+                          (\mbr -> M.insert (bodyName mbr) (bodyMem mbr))+                          (optdeps info)+                          mem_body_results+                    }+                (act', succc') = markSuccessCoal (act, succc) m_b info'+             in ((act', inhb), succc')+    foldfun+      acc+      ((act, inhb), succc)+      mem_body_results@(MemBodyResult m_b _ _ _ : _)+        | Just info <- M.lookup m_b act,+          all ((==) m_b . bodyMem) mem_body_results,+          Just info' <- zipWithM (M.lookup . bodyMem) mem_body_results $ map fst acc =+            -- Treating special case resembling:+            -- @let x0 = map (+1) a                                  @+            -- @let x3 = if cond then let x1 = x0 with [0] <- 2 in x1@+            -- @                 else let x2 = x0 with [1] <- 3 in x2@+            -- @let z[1] = x3                                        @+            -- In this case the result active table should be the union+            -- of the @m_x@ entries of the then and else active tables.+            let info'' =+                  foldl unionCoalsEntry info info'+                act' = M.insert m_b info'' act+             in ((act', inhb), succc)+    foldfun _ ((act, inhb), succc) (mbr : _) =+      -- one of the branches has failed coalescing,+      -- hence remove the coalescing of the result.++      (markFailedCoal (act, inhb) (patMem mbr), succc)+mkCoalsTabStm lutab (Let pat _ (DoLoop arginis lform body)) td_env bu_env = do+  let pat_val_elms = patElems pat++      --  i) Filter @activeCoals@ by the 2nd, 3rd AND 5th safety conditions. In+      --  other words, for each active coalescing target, the creation of the+      --  array we're trying to merge should happen before the allocation of the+      --  merge target and the index function should be translateable.+      (actv0, inhibit0) =+        filterSafetyCond2and5+          (activeCoals bu_env)+          (inhibit bu_env)+          (scals bu_env)+          td_env+          pat_val_elms+      -- ii) Extend @activeCoals@ by transfering the pattern-elements bindings+      --     existent in @activeCoals@ to the loop-body results, but only if:+      --       (a) the pattern element is a candidate for coalescing,        &&+      --       (b) the pattern element satisfies safety conditions 2 & 5,+      --           (conditions (a) and (b) have already been checked above), &&+      --       (c) the memory block of the corresponding body result is+      --           allocated outside the loop, i.e., non-existential,        &&+      --       (d) the init name is lastly-used in the initialization+      --           of the loop variant.+      --     Otherwise fail and remove from active-coalescing table!+      bdy_ress = bodyResult body+      (patmems, argmems, inimems, resmems) =+        L.unzip4 $+          mapMaybe (mapmbFun actv0) (zip3 pat_val_elms arginis $ map resSubExp bdy_ress) -- td_env'++      -- remove the other pattern elements from the active coalescing table:+      coal_pat_names = namesFromList $ map fst patmems+      (actv1, inhibit1) =+        foldl+          ( \(act, inhb) (b, MemBlock _ _ m_b _) ->+              if b `nameIn` coal_pat_names+                then (act, inhb) -- ok+                else markFailedCoal (act, inhb) m_b -- remove from active+          )+          (actv0, inhibit0)+          (getArrMemAssoc pat)++      -- iii) Process the loop's body.+      --      If the memory blocks of the loop result and loop variant param differ+      --      then make the original memory block of the loop result conflict with+      --      the original memory block of the loop parameter. This is done in+      --      order to prevent the coalescing of @a1@, @a0@, @x@ and @db@ in the+      --      same memory block of @y@ in the example below:+      --      @loop(a1 = a0) = for i < n do @+      --      @    let x = map (stencil a1) (iota n)@+      --      @    let db = copy x          @+      --      @    in db                    @+      --      @let y[0] = a1                @+      --      Meaning the coalescing of @x@ in @let db = copy x@ should fail because+      --      @a1@ appears in the definition of @let x = map (stencil a1) (iota n)@.+      res_mem_bdy = zipWith (\(b, m_b) (r, m_r) -> MemBodyResult m_b b r m_r) patmems resmems+      res_mem_arg = zipWith (\(b, m_b) (r, m_r) -> MemBodyResult m_b b r m_r) patmems argmems+      res_mem_ini = zipWith (\(b, m_b) (r, m_r) -> MemBodyResult m_b b r m_r) patmems inimems++      actv2 =+        let subs_res = mkSubsTab pat $ map resSubExp $ bodyResult body+            actv11 = foldl (transferCoalsToBody subs_res) actv1 res_mem_bdy+            subs_arg = mkSubsTab pat $ map (Var . paramName . fst) arginis+            actv12 = foldl (transferCoalsToBody subs_arg) actv11 res_mem_arg+            subs_ini = mkSubsTab pat $ map snd arginis+         in foldl (transferCoalsToBody subs_ini) actv12 res_mem_ini++      -- The code below adds an aliasing relation to the loop-arg memory+      --   so that to prevent, e.g., the coalescing of an iterative stencil+      --   (you need a buffer for the result and a separate one for the stencil).+      -- @ let b =               @+      -- @    loop (a) for i<N do@+      -- @        stencil a      @+      -- @  ...                  @+      -- @  y[slc_y] = b         @+      -- This should fail coalescing because we are aliasing @m_a@ with+      --   the memory block of the result.+      insertMemAliases tab (MemBodyResult _ _ _ m_r, MemBodyResult _ _ _ m_a) =+        if m_r == m_a+          then tab+          else case M.lookup m_r tab of+            Nothing -> tab+            Just etry ->+              M.insert m_r (etry {alsmem = alsmem etry <> oneName m_a}) tab+      actv3 = foldl insertMemAliases actv2 (zip res_mem_bdy res_mem_arg)+      -- analysing the loop body starts from a null memory-reference set;+      --  the results of the loop body iteration are aggregated later+      actv4 = M.map (\etry -> etry {memrefs = mempty}) actv3+  res_env_body <-+    mkCoalsTabStms+      lutab+      (bodyStms body)+      td_env'+      ( bu_env+          { activeCoals = actv4,+            inhibit = inhibit1+          }+      )+  let scals_loop = scals res_env_body+      (res_actv0, res_succ0, res_inhb0) = (activeCoals res_env_body, successCoals res_env_body, inhibit res_env_body)+      -- iv) Aggregate memory references across loop and filter unsound coalescing+      -- a) Filter the active-table by the FIRST SOUNDNESS condition, namely:+      --     W_i does not overlap with Union_{j=i+1..n} U_j,+      --     where W_i corresponds to the Write set of src mem-block m_b,+      --     and U_j correspond to the uses of the destination+      --     mem-block m_y, in which m_b is coalesced into.+      --     W_i and U_j correspond to the accesses within the loop body.+      mb_loop_idx = mbLoopIndexRange lform+  res_actv1 <- filterMapM1 (loopSoundness1Entry scals_loop mb_loop_idx) res_actv0++  -- b) Update the memory-reference summaries across loop:+  --   W = Union_{i=0..n-1} W_i Union W_{before-loop}+  --   U = Union_{i=0..n-1} U_i Union U_{before-loop}+  res_actv2 <- mapM (aggAcrossLoopEntry (scope td_env' <> scopeOf (bodyStms body)) scals_loop mb_loop_idx) res_actv1++  -- c) check soundness of the successful promotions for:+  --      - the entries that have been promoted to success during the loop-body pass+  --      - for all the entries of active table+  --    Filter the entries by the SECOND SOUNDNESS CONDITION, namely:+  --      Union_{i=1..n-1} W_i does not overlap the before-the-loop uses+  --        of the destination memory block.+  let res_actv3 = M.filterWithKey (loopSoundness2Entry actv3) res_actv2++  let tmp_succ =+        M.filterWithKey (okLookup actv3) $+          M.difference res_succ0 (successCoals bu_env)+      ver_succ = M.filterWithKey (loopSoundness2Entry actv3) tmp_succ+  let suc_fail = M.difference tmp_succ ver_succ+      (res_succ, res_inhb1) = foldl markFailedCoal (res_succ0, res_inhb0) $ M.keys suc_fail+      --+      act_fail = M.difference res_actv0 res_actv3+      (_, res_inhb) = foldl markFailedCoal (res_actv0, res_inhb1) $ M.keys act_fail+      res_actv =+        M.mapWithKey (addBeforeLoop actv3) res_actv3++      -- v) optimistically mark the pattern succesful if there is any chance to succeed+      ((fin_actv1, fin_inhb1), fin_succ1) =+        foldl foldFunOptimPromotion ((res_actv, res_inhb), res_succ) $+          L.zip4 patmems argmems resmems inimems+      (fin_actv2, fin_inhb2) =+        M.foldlWithKey+          ( \acc k _ ->+              if k `nameIn` namesFromList (map (paramName . fst) arginis)+                then markFailedCoal acc k+                else acc+          )+          (fin_actv1, fin_inhb1)+          fin_actv1+  pure bu_env {activeCoals = fin_actv2, successCoals = fin_succ1, inhibit = fin_inhb2}+  where+    allocs_bdy = foldl getAllocs (alloc td_env') $ bodyStms body+    td_env_allocs = td_env' {alloc = allocs_bdy, scope = scope td_env' <> scopeOf (bodyStms body)}+    td_env' = updateTopdownEnvLoop td_env arginis lform+    getAllocs tab (Let (Pat [pe]) _ (Op (Alloc _ sp))) =+      M.insert (patElemName pe) sp tab+    getAllocs tab _ = tab+    okLookup tab m _+      | Just _ <- M.lookup m tab = True+    okLookup _ _ _ = False+    --+    mapmbFun actv0 (patel, (arg, ini), bdyres)+      | b <- patElemName patel,+        (_, MemArray _ _ _ (ArrayIn m_b _)) <- patElemDec patel,+        a <- paramName arg,+        Var a0 <- ini,+        Var r <- bdyres,+        Just coal_etry <- M.lookup m_b actv0,+        Just _ <- M.lookup b (vartab coal_etry),+        Just (MemBlock _ _ m_a _) <- getScopeMemInfo a (scope td_env_allocs),+        Just (MemBlock _ _ m_a0 _) <- getScopeMemInfo a0 (scope td_env_allocs),+        Just (MemBlock _ _ m_r _) <- getScopeMemInfo r (scope td_env_allocs),+        Just nms <- M.lookup a lutab,+        a0 `nameIn` nms,+        m_r `elem` M.keys (alloc td_env_allocs) =+          Just ((b, m_b), (a, m_a), (a0, m_a0), (r, m_r))+    mapmbFun _ (_patel, (_arg, _ini), _bdyres) = Nothing+    foldFunOptimPromotion ::+      ((CoalsTab, InhibitTab), CoalsTab) ->+      ((VName, VName), (VName, VName), (VName, VName), (VName, VName)) ->+      ((CoalsTab, InhibitTab), CoalsTab)+    foldFunOptimPromotion ((act, inhb), succc) ((b, m_b), (a, m_a), (_r, m_r), (b_i, m_i))+      | m_r == m_i,+        Just info <- M.lookup m_i act,+        Just vtab_i <- addInvAliassesVarTab td_env (vartab info) b_i =+          Exc.assert+            (m_r == m_b && m_a == m_b)+            ((M.insert m_b (info {vartab = vtab_i}) act, inhb), succc)+      | m_r == m_i =+          Exc.assert+            (m_r == m_b && m_a == m_b)+            (markFailedCoal (act, inhb) m_b, succc)+      | Just info_b0 <- M.lookup m_b act,+        Just info_a0 <- M.lookup m_a act,+        Just info_i <- M.lookup m_i act,+        M.member m_r succc,+        Just vtab_i <- addInvAliassesVarTab td_env (vartab info_i) b_i,+        [Just info_b, Just info_a] <- map translateIxFnInScope [(b, info_b0), (a, info_a0)] =+          let info_b' = info_b {optdeps = M.insert b_i m_i $ optdeps info_b}+              info_a' = info_a {optdeps = M.insert b_i m_i $ optdeps info_a}+              info_i' =+                info_i+                  { optdeps = M.insert b m_b $ optdeps info_i,+                    memrefs = mempty,+                    vartab = vtab_i+                  }+              act' = M.insert m_i info_i' act+              (act1, succc1) =+                foldl+                  (\acc (m, info) -> markSuccessCoal acc m info)+                  (act', succc)+                  [(m_b, info_b'), (m_a, info_a')]+           in -- ToDo: make sure that ixfun translates and update substitutions (?)+              ((act1, inhb), succc1)+    foldFunOptimPromotion ((act, inhb), succc) ((_, m_b), (_a, m_a), (_r, m_r), (_b_i, m_i)) =+      Exc.assert+        (m_r /= m_i)+        (foldl markFailedCoal (act, inhb) [m_b, m_a, m_r, m_i], succc)++    translateIxFnInScope (x, info)+      | Just (Coalesced knd mbd@(MemBlock _ _ _ ixfn) _subs0) <- M.lookup x (vartab info),+        isInScope td_env (dstmem info) =+          let scope_tab =+                scope td_env+                  <> scopeOfFParams (map fst arginis)+           in case freeVarSubstitutions scope_tab (scals bu_env) ixfn of+                Just fv_subst ->+                  Just $ info {vartab = M.insert x (Coalesced knd mbd fv_subst) (vartab info)}+                Nothing -> Nothing+    translateIxFnInScope _ = Nothing+    se0 = intConst Int64 0+    mbLoopIndexRange ::+      LoopForm (Aliases rep) ->+      Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))+    mbLoopIndexRange (WhileLoop _) = Nothing+    mbLoopIndexRange (ForLoop inm _inttp seN _) = Just (inm, (pe64 se0, pe64 seN))+    addBeforeLoop actv_bef m_b etry =+      case M.lookup m_b actv_bef of+        Nothing -> etry+        Just etry0 ->+          etry {memrefs = memrefs etry0 <> memrefs etry}+    aggAcrossLoopEntry scope_loop scal_tab idx etry = do+      wrts <-+        aggSummaryLoopTotal (scope td_env) scope_loop scal_tab idx $+          (srcwrts . memrefs) etry+      uses <-+        aggSummaryLoopTotal (scope td_env) scope_loop scal_tab idx $+          (dstrefs . memrefs) etry+      pure $ etry {memrefs = MemRefs uses wrts}+    loopSoundness1Entry scal_tab idx etry = do+      let wrt_i = (srcwrts . memrefs) etry+      use_p <-+        aggSummaryLoopPartial (scal_tab <> scalarTable td_env) idx $+          dstrefs $+            memrefs etry+      pure $ noMemOverlap td_env' wrt_i use_p+    loopSoundness2Entry :: CoalsTab -> VName -> CoalsEntry -> Bool+    loopSoundness2Entry old_actv m_b etry =+      case M.lookup m_b old_actv of+        Nothing -> True+        Just etry0 ->+          let uses_before = (dstrefs . memrefs) etry0+              write_loop = (srcwrts . memrefs) etry+           in noMemOverlap td_env write_loop uses_before++-- The case of in-place update:+--   @let x' = x with slice <- elm@+mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ e@(BasicOp (Update safety x _ _elm))) td_env bu_env+  | [(_, MemBlock _ _ m_x _)] <- getArrMemAssoc pat =+      do+        -- (a) filter by the 3rd safety for @elm@ and @x'@+        let (actv, inhbt) = recordMemRefUses td_env bu_env stm+            -- (b) if @x'@ is in active coalesced table, then add an entry for @x@ as well+            (actv', inhbt') =+              case M.lookup m_x actv of+                Nothing -> (actv, inhbt)+                Just info ->+                  case M.lookup (patElemName x') (vartab info) of+                    Nothing ->+                      markFailedCoal (actv, inhbt) m_x+                    Just (Coalesced k mblk@(MemBlock _ _ _ x_indfun) _) ->+                      case freeVarSubstitutions (scope td_env) (scals bu_env) x_indfun of+                        Just fv_subs+                          | isInScope td_env (dstmem info) ->+                              let coal_etry_x = Coalesced k mblk fv_subs+                                  info' =+                                    info+                                      { vartab =+                                          M.insert x coal_etry_x $+                                            M.insert (patElemName x') coal_etry_x (vartab info)+                                      }+                               in (M.insert m_x info' actv, inhbt)+                        _ ->+                          markFailedCoal (actv, inhbt) m_x++            -- (c) this stm is also a potential source for coalescing, so process it+            actv'' = if safety == Unsafe then mkCoalsHelper3PatternMatch pat e lutab td_env (successCoals bu_env) actv' inhbt' else actv'+        pure $+          bu_env {activeCoals = actv'', inhibit = inhbt'}++-- The case of flat in-place update:+--   @let x' = x with flat-slice <- elm@+mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ e@(BasicOp (FlatUpdate x _ _elm))) td_env bu_env+  | [(_, MemBlock _ _ m_x _)] <- getArrMemAssoc pat =+      do+        -- (a) filter by the 3rd safety for @elm@ and @x'@+        let (actv, inhbt) = recordMemRefUses td_env bu_env stm+            -- (b) if @x'@ is in active coalesced table, then add an entry for @x@ as well+            (actv', inhbt') =+              case M.lookup m_x actv of+                Nothing -> (actv, inhbt)+                Just info ->+                  case M.lookup (patElemName x') (vartab info) of+                    Nothing ->+                      -- error "In ArrayCoalescing.hs, fun mkCoalsTabStm, case in-place update!"+                      -- this case should not happen, but if it can that just fail conservatively+                      markFailedCoal (actv, inhbt) m_x+                    Just (Coalesced k mblk@(MemBlock _ _ _ x_indfun) _) ->+                      case freeVarSubstitutions (scope td_env) (scals bu_env) x_indfun of+                        Just fv_subs+                          | isInScope td_env (dstmem info) ->+                              let coal_etry_x = Coalesced k mblk fv_subs+                                  info' =+                                    info+                                      { vartab =+                                          M.insert x coal_etry_x $+                                            M.insert (patElemName x') coal_etry_x (vartab info)+                                      }+                               in (M.insert m_x info' actv, inhbt)+                        _ ->+                          markFailedCoal (actv, inhbt) m_x++            -- (c) this stm is also a potential source for coalescing, so process it+            actv'' = mkCoalsHelper3PatternMatch pat e lutab td_env (successCoals bu_env) actv' inhbt'+        pure $+          bu_env {activeCoals = actv'', inhibit = inhbt'}+--+mkCoalsTabStm _ (Let pat _ (BasicOp Update {})) _ _ =+  error $ "In ArrayCoalescing.hs, fun mkCoalsTabStm, illegal pattern for in-place update: " ++ show pat+-- default handling+mkCoalsTabStm lutab (Let pat _ (Op op)) td_env bu_env = do+  -- Process body+  on_op <- asks onOp+  on_op lutab pat op td_env bu_env+mkCoalsTabStm lutab stm@(Let pat _ e) td_env bu_env = do+  --   i) Filter @activeCoals@ by the 3rd safety condition:+  --      this is now relaxed by use of LMAD eqs:+  --      the memory referenced in stm are added to memrefs::dstrefs+  --      in corresponding coal-tab entries.+  let (activeCoals', inhibit') = recordMemRefUses td_env bu_env stm+      -- mkCoalsHelper1FilterActive pat (freeIn e) (scope td_env) (scals bu_env)+      --                           (activeCoals bu_env) (inhibit bu_env)++      --  ii) promote any of the entries in @activeCoals@ to @successCoals@ as long as+      --        - this statement defined a variable consumed in a coalesced statement+      --        - and safety conditions 2, 4, and 5 are satisfied.+      --      AND extend @activeCoals@ table for any definition of a variable that+      --      aliases a coalesced variable.+      safe_4 = createsNewArrOK e+      ((activeCoals'', inhibit''), successCoals') =+        foldl (foldfun safe_4) ((activeCoals', inhibit'), successCoals bu_env) (getArrMemAssoc pat)++      -- iii) record a potentially coalesced statement in @activeCoals@+      activeCoals''' = mkCoalsHelper3PatternMatch pat e lutab td_env successCoals' activeCoals'' (inhibited td_env)+  pure bu_env {activeCoals = activeCoals''', inhibit = inhibit'', successCoals = successCoals'}+  where+    foldfun safe_4 ((a_acc, inhb), s_acc) (b, MemBlock tp shp mb _b_indfun) =+      case M.lookup mb a_acc of+        Nothing -> ((a_acc, inhb), s_acc)+        Just info@(CoalsEntry x_mem _ _ vtab _ _) ->+          let failed = markFailedCoal (a_acc, inhb) mb+           in case M.lookup b vtab of+                Nothing ->+                  -- we hit the definition of some variable @b@ aliased with+                  --    the coalesced variable @x@, hence extend @activeCoals@, e.g.,+                  --       @let x = map f arr  @+                  --       @let b = alias x  @ <- current statement+                  --       @ ... use of b ...  @+                  --       @let c = alias b    @ <- currently fails+                  --       @let y[i] = x       @+                  -- where @alias@ can be @transpose@, @slice@, @rotate@, @reshape@.+                  -- We use getTransitiveAlias helper function to track the aliasing+                  --    through the td_env, and to find the updated ixfun of @b@:+                  case getDirAliasedIxfn td_env a_acc b of+                    Nothing -> (failed, s_acc)+                    Just (_, _, b_indfun') ->+                      case freeVarSubstitutions (scope td_env) (scals bu_env) b_indfun' of+                        Nothing -> (failed, s_acc)+                        Just fv_subst ->+                          let mem_info = Coalesced TransitiveCoal (MemBlock tp shp x_mem b_indfun') fv_subst+                              info' = info {vartab = M.insert b mem_info vtab}+                           in ((M.insert mb info' a_acc, inhb), s_acc)+                Just (Coalesced k mblk@(MemBlock _ _ _ new_indfun) _) ->+                  -- we are at the definition of the coalesced variable @b@+                  -- if 2,4,5 hold promote it to successful coalesced table,+                  -- or if e = transpose, etc. then postpone decision for later on+                  let safe_2 = isInScope td_env x_mem+                   in case freeVarSubstitutions (scope td_env) (scals bu_env) new_indfun of+                        Just fv_subst+                          | safe_2 ->+                              let mem_info = Coalesced k mblk fv_subst+                                  info' = info {vartab = M.insert b mem_info vtab}+                               in if safe_4+                                    then -- array creation point, successful coalescing verified!++                                      let (a_acc', s_acc') = markSuccessCoal (a_acc, s_acc) mb info'+                                       in ((a_acc', inhb), s_acc')+                                    else -- this is an invertible alias case of the kind+                                    -- @ let b    = alias a @+                                    -- @ let x[i] = b @+                                    -- do not promote, but update the index function++                                      ((M.insert mb info' a_acc, inhb), s_acc)+                        _ -> (failed, s_acc) -- fail!++ixfunToAccessSummary :: IxFun.IxFun (TPrimExp Int64 VName) -> AccessSummary+ixfunToAccessSummary (IxFun.IxFun (lmad NE.:| []) _ _) = Set $ S.singleton lmad+ixfunToAccessSummary _ = Undeterminable++-- | Check safety conditions 2 and 5 and update new substitutions:+-- called on the pat-elements of loop and if-then-else expressions.+--+-- The safety conditions are: The allocation of merge target should dominate the+-- creation of the array we're trying to merge and the new index function of the+-- array can be translated at the definition site of b. The latter requires that+-- any variables used in the index function of the target array are available at+-- the definition site of b.+filterSafetyCond2and5 ::+  HasMemBlock (Aliases rep) =>+  CoalsTab ->+  InhibitTab ->+  ScalarTab ->+  TopdownEnv rep ->+  [PatElem (VarAliases, LetDecMem)] ->+  (CoalsTab, InhibitTab)+filterSafetyCond2and5 act_coal inhb_coal scals_env td_env =+  foldl helper (act_coal, inhb_coal)+  where+    helper (acc, inhb) patel =+      -- For each pattern element in the input list+      case (patElemName patel, patElemDec patel) of+        (b, (_, MemArray tp0 shp0 _ (ArrayIn m_b _idxfn_b))) ->+          -- If it is an array in memory block m_b+          case M.lookup m_b acc of+            Nothing -> (acc, inhb)+            Just info@(CoalsEntry x_mem _ _ vtab _ _) ->+              -- And m_b we're trying to coalesce m_b+              let failed = markFailedCoal (acc, inhb) m_b+               in case M.lookup b vtab of+                    Nothing ->+                      case getDirAliasedIxfn td_env acc b of+                        Nothing -> failed+                        Just (_, _, b_indfun') ->+                          -- And we have the index function of b+                          case freeVarSubstitutions (scope td_env) scals_env b_indfun' of+                            Nothing -> failed+                            Just fv_subst ->+                              let mem_info = Coalesced TransitiveCoal (MemBlock tp0 shp0 x_mem b_indfun') fv_subst+                                  info' = info {vartab = M.insert b mem_info vtab}+                               in (M.insert m_b info' acc, inhb)+                    Just (Coalesced k (MemBlock pt shp _ new_indfun) _) ->+                      let safe_2 = isInScope td_env x_mem+                       in case freeVarSubstitutions (scope td_env) scals_env new_indfun of+                            Just fv_subst+                              | safe_2 ->+                                  let mem_info = Coalesced k (MemBlock pt shp x_mem new_indfun) fv_subst+                                      info' = info {vartab = M.insert b mem_info vtab}+                                   in (M.insert m_b info' acc, inhb)+                            _ -> failed+        _ -> (acc, inhb)++-- |   Pattern matches a potentially coalesced statement and+--     records a new association in @activeCoals@+mkCoalsHelper3PatternMatch ::+  HasMemBlock (Aliases rep) =>+  Pat (VarAliases, LetDecMem) ->+  Exp (Aliases rep) ->+  LUTabFun ->+  TopdownEnv rep ->+  CoalsTab ->+  CoalsTab ->+  InhibitTab ->+  CoalsTab+mkCoalsHelper3PatternMatch pat e lutab td_env _ activeCoals_tab _+  | Nothing <- genCoalStmtInfo lutab (scope td_env) pat e =+      activeCoals_tab+mkCoalsHelper3PatternMatch pat e lutab td_env successCoals_tab activeCoals_tab inhibit_tab+  | Just clst <- genCoalStmtInfo lutab (scope td_env) pat e =+      foldl processNewCoalesce activeCoals_tab clst+  where+    processNewCoalesce acc (knd, alias_fn, x, m_x, ind_x, b, m_b, _, tp_b, shp_b) =+      -- test whether we are in a transitive coalesced case, i.e.,+      --      @let b = scratch ...@+      --      @.....@+      --      @let x[j] = b@+      --      @let y[i] = x@+      -- and compose the index function of @x@ with that of @y@,+      -- and update aliasing of the @m_b@ entry to also contain @m_y@+      -- on top of @m_x@, i.e., transitively, any use of @m_y@ should+      -- be checked for the lifetime of @b@.+      let proper_coals_tab = case knd of+            InPlaceCoal -> activeCoals_tab+            _ -> successCoals_tab+          (m_yx, ind_yx, mem_yx_al, x_deps) =+            case M.lookup m_x proper_coals_tab of+              Nothing ->+                (m_x, alias_fn ind_x, oneName m_x, M.empty)+              Just (CoalsEntry m_y ind_y y_al vtab x_deps0 _) ->+                let ind = case M.lookup x vtab of+                      Just (Coalesced _ (MemBlock _ _ _ ixf) _) ->+                        ixf+                      Nothing ->+                        ind_y+                 in (m_y, alias_fn ind, oneName m_x <> y_al, x_deps0)+          success0 = IxFun.hasOneLmad ind_yx+          m_b_aliased_m_yx = areAnyAliased td_env m_b [m_yx] -- m_b \= m_yx+       in case (success0, not m_b_aliased_m_yx, isInScope td_env m_yx) of -- nameIn m_yx (alloc td_env)+            (True, True, True) ->+              -- Finally update the @activeCoals@ table with a fresh+              --   binding for @m_b@; if such one exists then overwrite.+              -- Also, add all variables from the alias chain of @b@ to+              --   @vartab@, for example, in the case of a sequence:+              --   @ b0 = if cond then ... else ... @+              --   @ b1 = alias0 b0 @+              --   @ b  = alias1 b1 @+              --   @ x[j] = b @+              -- Then @b1@ and @b0@ should also be added to @vartab@ if+              --   @alias1@ and @alias0@ are invertible, otherwise fail early!+              let mem_info = Coalesced knd (MemBlock tp_b shp_b m_yx ind_yx) M.empty+                  opts' =+                    if m_yx == m_x+                      then M.empty+                      else M.insert x m_x x_deps+                  vtab = M.singleton b mem_info+                  mvtab = addInvAliassesVarTab td_env vtab b++                  is_inhibited = case M.lookup m_b inhibit_tab of+                    Just nms -> m_yx `nameIn` nms+                    Nothing -> False+               in case (is_inhibited, mvtab) of+                    (True, _) -> acc -- fail due to inhibited+                    (_, Nothing) -> acc -- fail early due to non-invertible aliasing+                    (_, Just vtab') ->+                      -- successfully adding a new coalesced entry+                      let coal_etry =+                            CoalsEntry+                              m_yx+                              ind_yx+                              mem_yx_al+                              vtab'+                              opts'+                              mempty+                       in M.insert m_b coal_etry acc+            _ -> acc+mkCoalsHelper3PatternMatch _ _ _ _ _ _ _ =+  error "In ArrayCoalescing.hs, fun mkCoalsHelper3PatternMatch: Unreachable!!!"++genCoalStmtInfo ::+  HasMemBlock (Aliases rep) =>+  LUTabFun ->+  ScopeTab rep ->+  Pat (VarAliases, LetDecMem) ->+  Exp (Aliases rep) ->+  Maybe [(CoalescedKind, IxFun -> IxFun, VName, VName, IxFun, VName, VName, IxFun, PrimType, Shape)]+-- CASE a) @let x <- copy(b^{lu})@+genCoalStmtInfo lutab scopetab pat (BasicOp (Copy b))+  | Pat [PatElem x (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =+      case (M.lookup x lutab, getScopeMemInfo b scopetab) of+        (Just last_uses, Just (MemBlock tpb shpb m_b ind_b)) ->+          if b `notNameIn` last_uses+            then Nothing+            else Just [(CopyCoal, id, x, m_x, ind_x, b, m_b, ind_b, tpb, shpb)]+        _ -> Nothing+-- CASE c) @let x[i] = b^{lu}@+genCoalStmtInfo lutab scopetab pat (BasicOp (Update _ x slice_x (Var b)))+  | Pat [PatElem x' (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =+      case (M.lookup x' lutab, getScopeMemInfo b scopetab) of+        (Just last_uses, Just (MemBlock tpb shpb m_b ind_b)) ->+          if b `notNameIn` last_uses+            then Nothing+            else Just [(InPlaceCoal, (`updateIndFunSlice` slice_x), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb)]+        _ -> Nothing+  where+    updateIndFunSlice :: IxFun -> Slice SubExp -> IxFun+    updateIndFunSlice ind_fun slc_x =+      let slc_x' = map (fmap pe64) $ unSlice slc_x+       in IxFun.slice ind_fun $ Slice slc_x'+genCoalStmtInfo lutab scopetab pat (BasicOp (FlatUpdate x slice_x b))+  | Pat [PatElem x' (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =+      case (M.lookup x' lutab, getScopeMemInfo b scopetab) of+        (Just last_uses, Just (MemBlock tpb shpb m_b ind_b)) ->+          if b `notNameIn` last_uses+            then Nothing+            else Just [(InPlaceCoal, (`updateIndFunSlice` slice_x), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb)]+        _ -> Nothing+  where+    updateIndFunSlice :: IxFun -> FlatSlice SubExp -> IxFun+    updateIndFunSlice ind_fun (FlatSlice offset dims) =+      IxFun.flatSlice ind_fun $ FlatSlice (pe64 offset) $ map (fmap pe64) dims++-- CASE b) @let x = concat(a, b^{lu})@+genCoalStmtInfo lutab scopetab pat (BasicOp (Concat concat_dim (b0 :| bs) _))+  | Pat [PatElem x (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =+      case M.lookup x lutab of+        Nothing -> Nothing+        Just last_uses ->+          let zero = pe64 $ intConst Int64 0+              markConcatParts (acc, offs, succ0) b =+                if not succ0+                  then (acc, offs, succ0)+                  else case getScopeMemInfo b scopetab of+                    Just (MemBlock tpb shpb@(Shape dims@(_ : _)) m_b ind_b)+                      | Just d <- maybeNth concat_dim dims ->+                          let offs' = offs + pe64 d+                           in if b `nameIn` last_uses+                                then+                                  let slc =+                                        Slice $+                                          map (unitSlice zero . pe64) (take concat_dim dims)+                                            <> [unitSlice offs (pe64 d)]+                                            <> map (unitSlice zero . pe64) (drop (concat_dim + 1) dims)+                                   in ( acc ++ [(ConcatCoal, (`IxFun.slice` slc), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb)],+                                        offs',+                                        True+                                      )+                                else (acc, offs', True)+                    _ -> (acc, offs, False)+              (res, _, _) = foldl markConcatParts ([], zero, True) (b0 : bs)+           in if null res then Nothing else Just res+-- CASE other than a), b), or c) not supported+genCoalStmtInfo _ _ _ _ = Nothing++data MemBodyResult = MemBodyResult+  { patMem :: VName,+    _patName :: VName,+    bodyName :: VName,+    bodyMem :: VName+  }++-- | Results in pairs of pattern-blockresult pairs of (var name, mem block)+--   for those if-patterns that are candidates for coalescing.+findMemBodyResult ::+  (HasMemBlock (Aliases rep)) =>+  CoalsTab ->+  ScopeTab rep ->+  [PatElem (VarAliases, LetDecMem)] ->+  Body (Aliases rep) ->+  [MemBodyResult]+findMemBodyResult activeCoals_tab scope_env patelms bdy =+  mapMaybe+    findMemBodyResult'+    (zip patelms $ map resSubExp $ bodyResult bdy)+  where+    scope_env' = scope_env <> scopeOf (bodyStms bdy)+    findMemBodyResult' (patel, se_r) =+      case (patElemName patel, patElemDec patel, se_r) of+        (b, (_, MemArray _ _ _ (ArrayIn m_b _)), Var r) ->+          case getScopeMemInfo r scope_env' of+            Nothing -> Nothing+            Just (MemBlock _ _ m_r _) ->+              case M.lookup m_b activeCoals_tab of+                Nothing -> Nothing+                Just coal_etry ->+                  case M.lookup b (vartab coal_etry) of+                    Nothing -> Nothing+                    Just _ -> Just $ MemBodyResult m_b b r m_r+        _ -> Nothing++-- | transfers coalescing from if-pattern to then|else body result+--   in the active coalesced table. The transfer involves, among+--   others, inserting @(r,m_r)@ in the optimistically-dependency+--   set of @m_b@'s entry and inserting @(b,m_b)@ in the opt-deps+--   set of @m_r@'s entry. Meaning, ultimately, @m_b@ can be merged+--   if @m_r@ can be merged (and vice-versa). This is checked by a+--   fix point iteration at the function-definition level.+transferCoalsToBody ::+  M.Map VName (TPrimExp Int64 VName) -> -- (PrimExp VName)+  CoalsTab ->+  MemBodyResult ->+  CoalsTab+transferCoalsToBody exist_subs activeCoals_tab (MemBodyResult m_b b r m_r)+  | -- the @Nothing@ pattern for the two lookups cannot happen+    -- because they were already cheked in @findMemBodyResult@+    Just etry <- M.lookup m_b activeCoals_tab,+    Just (Coalesced knd (MemBlock btp shp _ ind_b) subst_b) <- M.lookup b $ vartab etry =+      -- by definition of if-stmt, r and b have the same basic type, shape and+      -- index function, hence, for example, do not need to rebase+      -- We will check whether it is translatable at the definition point of r.+      let ind_r = IxFun.substituteInIxFun exist_subs ind_b+          subst_r = M.union exist_subs subst_b+          mem_info = Coalesced knd (MemBlock btp shp (dstmem etry) ind_r) subst_r+       in if m_r == m_b -- already unified, just add binding for @r@+            then+              let etry' =+                    etry+                      { optdeps = M.insert b m_b (optdeps etry),+                        vartab = M.insert r mem_info (vartab etry)+                      }+               in M.insert m_r etry' activeCoals_tab+            else -- make them both optimistically depend on each other++              let opts_x_new = M.insert r m_r (optdeps etry)+                  -- Here we should translate the @ind_b@ field of @mem_info@+                  -- across the existential introduced by the if-then-else+                  coal_etry =+                    etry+                      { vartab = M.singleton r mem_info,+                        optdeps = M.insert b m_b (optdeps etry)+                      }+               in M.insert m_b (etry {optdeps = opts_x_new}) $+                    M.insert m_r coal_etry activeCoals_tab+  | otherwise = error "Impossible"++mkSubsTab ::+  Pat (aliases, LetDecMem) ->+  [SubExp] ->+  M.Map VName (TPrimExp Int64 VName)+mkSubsTab pat res =+  let pat_elms = patElems pat+   in M.fromList $ mapMaybe mki64subst $ zip pat_elms res+  where+    mki64subst (a, Var v)+      | (_, MemPrim (IntType Int64)) <- patElemDec a = Just (patElemName a, le64 v)+    mki64subst (a, se@(Constant (IntValue (Int64Value _)))) = Just (patElemName a, pe64 se)+    mki64subst _ = Nothing++computeScalarTable ::+  (Coalesceable rep inner) =>+  ScopeTab rep ->+  Stm (Aliases rep) ->+  ScalarTableM rep (M.Map VName (PrimExp VName))+computeScalarTable scope_table (Let (Pat [pe]) _ e)+  | Just primexp <- primExpFromExp (vnameToPrimExp scope_table mempty) e =+      pure $ M.singleton (patElemName pe) primexp+computeScalarTable scope_table (Let _ _ (DoLoop loop_inits loop_form body)) =+  concatMapM+    ( computeScalarTable $+        scope_table+          <> scopeOfFParams (map fst loop_inits)+          <> scopeOf loop_form+          <> scopeOf (bodyStms body)+    )+    (stmsToList $ bodyStms body)+computeScalarTable scope_table (Let _ _ (Match _ cases body _)) = do+  body_tab <- concatMapM (computeScalarTable $ scope_table <> scopeOf (bodyStms body)) (stmsToList $ bodyStms body)+  cases_tab <-+    concatMapM+      ( \(Case _ b) ->+          concatMapM+            (computeScalarTable $ scope_table <> scopeOf (bodyStms b))+            ( stmsToList $+                bodyStms body+            )+      )+      cases+  pure $ body_tab <> cases_tab+computeScalarTable scope_table (Let _ _ (Op op)) = do+  on_op <- asks scalarTableOnOp+  on_op scope_table op+computeScalarTable _ _ = pure mempty++computeScalarTableGPUMem :: ScopeTab GPUMem -> Op (Aliases GPUMem) -> ScalarTableM GPUMem (M.Map VName (PrimExp VName))+computeScalarTableGPUMem _ (Alloc _ _) = pure mempty+computeScalarTableGPUMem scope_table (Inner (SegOp segop)) = do+  concatMapM+    (computeScalarTable $ scope_table <> scopeOf (kernelBodyStms $ segBody segop) <> scopeOfSegSpace (segSpace segop))+    (stmsToList $ kernelBodyStms $ segBody segop)+computeScalarTableGPUMem _ (Inner (SizeOp _)) = pure mempty+computeScalarTableGPUMem _ (Inner (OtherOp ())) = pure mempty+computeScalarTableGPUMem scope_table (Inner (GPUBody _ body)) =+  concatMapM+    (computeScalarTable $ scope_table <> scopeOf (bodyStms body))+    (stmsToList $ bodyStms body)++filterMapM1 :: (Eq k, Monad m) => (v -> m Bool) -> M.Map k v -> m (M.Map k v)+filterMapM1 f m = fmap M.fromAscList $ filterM (f . snd) $ M.toAscList m

@@ -0,0 +1,420 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.Optimise.ArrayShortCircuiting.DataStructs+  ( Coalesced (..),+    CoalescedKind (..),+    ArrayMemBound (..),+    AllocTab,+    AliasTab,+    LUTabFun,+    CreatesNewArrOp,+    HasMemBlock,+    LUTabPrg,+    ScalarTab,+    CoalsTab,+    ScopeTab,+    CoalsEntry (..),+    FreeVarSubsts,+    LmadRef,+    MemRefs (..),+    AccessSummary (..),+    BotUpEnv (..),+    InhibitTab,+    unionCoalsEntry,+    vnameToPrimExp,+    getArrMemAssocFParam,+    getScopeMemInfo,+    createsNewArrOK,+    getArrMemAssoc,+    getUniqueMemFParam,+    markFailedCoal,+    accessSubtract,+    markSuccessCoal,+  )+where++import Control.Applicative+import Data.Functor ((<&>))+import Data.Map.Strict qualified as M+import Data.Maybe+import Data.Set qualified as S+import Futhark.IR.Aliases+import Futhark.IR.GPUMem+import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.SeqMem+import Futhark.Util.Pretty hiding (line, sep, (</>))+import Prelude++type ScopeTab rep = Scope (Aliases rep)+-- ^ maps array-variable names to various info, including+--   types, memory block and index function, etc.++-- | An LMAD specialized to TPrimExps (a typed primexp)+type LmadRef = IxFun.LMAD (TPrimExp Int64 VName)++-- | Summary of all memory accesses at a given point in the code+data AccessSummary+  = -- | The access summary was statically undeterminable, for instance by+    -- having multiple lmads. In this case, we should conservatively avoid all+    -- coalescing.+    Undeterminable+  | -- | A conservative estimate of the set of accesses up until this point.+    Set (S.Set LmadRef)++instance Semigroup AccessSummary where+  Undeterminable <> _ = Undeterminable+  _ <> Undeterminable = Undeterminable+  (Set a) <> (Set b) =+    Set $ S.union a b++instance Monoid AccessSummary where+  mempty = Set mempty++instance FreeIn AccessSummary where+  freeIn' Undeterminable = mempty+  freeIn' (Set s) = freeIn' s++accessSubtract :: AccessSummary -> AccessSummary -> AccessSummary+accessSubtract Undeterminable _ = Undeterminable+accessSubtract _ Undeterminable = Undeterminable+accessSubtract (Set s1) (Set s2) = Set $ s1 S.\\ s2++data MemRefs = MemRefs+  { -- | The access summary of all references (reads+    -- and writes) to the destination of a coalescing entry+    dstrefs :: AccessSummary,+    -- | The access summary of all writes to the source of a coalescing entry+    srcwrts :: AccessSummary+  }++instance Semigroup MemRefs where+  m1 <> m2 =+    MemRefs (dstrefs m1 <> dstrefs m2) (srcwrts m1 <> srcwrts m2)++instance Monoid MemRefs where+  mempty = MemRefs mempty mempty++data CoalescedKind+  = -- | let x    = copy b^{lu}+    CopyCoal+  | -- | let x[i] = b^{lu}+    InPlaceCoal+  | -- | let x    = concat(a, b^{lu})+    ConcatCoal+  | -- | transitive, i.e., other variables aliased with b.+    TransitiveCoal++-- | Information about a memory block: type, shape, name and ixfun.+data ArrayMemBound = MemBlock+  { primType :: PrimType,+    shape :: Shape,+    memName :: VName,+    ixfun :: IxFun+  }++-- | Free variable substitutions+type FreeVarSubsts = M.Map VName (TPrimExp Int64 VName)++-- | Coalesced Access Entry+data Coalesced+  = Coalesced+      CoalescedKind+      -- ^ the kind of coalescing+      ArrayMemBound+      -- ^ destination mem_block info @f_m_x[i]@ (must be ArrayMem)+      -- (Maybe IxFun) -- the inverse ixfun of a coalesced array, such that+      --                     --  ixfuns can be correctly constructed for aliases;+      FreeVarSubsts+      -- ^ substitutions for free vars in index function++data CoalsEntry = CoalsEntry+  { -- | destination memory block+    dstmem :: VName,+    -- | index function of the destination (used for rebasing)+    dstind :: IxFun,+    -- | aliased destination memory blocks can appear+    --   due to repeated (optimistic) coalescing.+    alsmem :: Names,+    -- | per variable-name coalesced entries+    vartab :: M.Map VName Coalesced,+    -- | keys are variable names, values are memblock names;+    --   it records optimistically added coalesced nodes, e.g.,+    --   in the case of if-then-else expressions. For example:+    --       @x    = map f a@+    --       @.. use of y ..@+    --       @b    = map g a@+    --       @x[i] = b      @+    --       @y[k] = x      @+    --   the coalescing of @b@ in @x[i]@ succeeds, but+    --   is dependent of the success of the coalescing+    --   of @x@ in @y[k]@, which fails in this case+    --   because @y@ is used before the new array creation+    --   of @x = map f@. Hence @optdeps@ of the @m_b@ CoalsEntry+    --   records @x -> m_x@ and at the end of analysis it is removed+    --   from the successfully coalesced table if @m_x@ is+    --   unsuccessful.+    --   Storing @m_x@ would probably be sufficient if memory would+    --     not be reused--e.g., by register allocation on arrays--the+    --     @x@ discriminates between memory being reused across semantically+    --     different arrays (searched in @vartab@ field).+    optdeps :: M.Map VName VName,+    -- | Access summaries of uses and writes of destination and source+    -- respectively.+    memrefs :: MemRefs+  }++type AllocTab = M.Map VName Space+-- ^ the allocatted memory blocks++type AliasTab = M.Map VName Names+-- ^ maps a variable or memory block to its aliases++type LUTabFun = M.Map VName Names+-- ^ maps a name indentifying a stmt to the last uses in that stmt++type LUTabPrg = M.Map Name LUTabFun+-- ^ maps function names to last-use tables++type ScalarTab = M.Map VName (PrimExp VName)+-- ^ maps a variable name to its PrimExp scalar expression++type CoalsTab = M.Map VName CoalsEntry+-- ^ maps a memory-block name to a 'CoalsEntry'. Among other things, it contains+--   @vartab@, a map in which each variable associated to that memory block is+--   bound to its 'Coalesced' info.++type InhibitTab = M.Map VName Names+-- ^ inhibited memory-block mergings from the key (memory block)+--   to the value (set of memory blocks).++data BotUpEnv = BotUpEnv+  { -- | maps scalar variables to theirs PrimExp expansion+    scals :: ScalarTab,+    -- | Optimistic coalescing info. We are currently trying to coalesce these+    -- memory blocks.+    activeCoals :: CoalsTab,+    -- | Committed (successfull) coalescing info. These memory blocks have been+    -- successfully coalesced.+    successCoals :: CoalsTab,+    -- | The coalescing failures from this pass. We will no longer try to merge+    -- these memory blocks.+    inhibit :: InhibitTab+  }++instance Pretty CoalsTab where+  pretty = pretty . M.toList++instance Pretty AccessSummary where+  pretty Undeterminable = "Undeterminable"+  pretty (Set a) = "Access-Set:" <+> pretty (S.toList a) <+> " "++instance Pretty MemRefs where+  pretty (MemRefs a b) = "( Use-Sum:" <+> pretty a <+> "Write-Sum:" <+> pretty b <> ")"++instance Pretty CoalescedKind where+  pretty CopyCoal = "Copy"+  pretty InPlaceCoal = "InPlace"+  pretty ConcatCoal = "Concat"+  pretty TransitiveCoal = "Transitive"++instance Pretty ArrayMemBound where+  pretty (MemBlock ptp shp m_nm ixfn) =+    "{" <> pretty ptp <> "," <+> pretty shp <> "," <+> pretty m_nm <> "," <+> pretty ixfn <> "}"++instance Pretty Coalesced where+  pretty (Coalesced knd mbd _) =+    "(Kind:"+      <+> pretty knd <> ", membds:"+      <+> pretty mbd -- <> ", subs:" <+> pretty subs+        <> ")"+      <+> "\n"++instance Pretty CoalsEntry where+  pretty etry =+    "{"+      <+> "Dstmem:"+      <+> pretty (dstmem etry)+        <> ", AliasMems:"+      <+> pretty (alsmem etry)+      <+> ", optdeps:"+      <+> pretty (M.toList $ optdeps etry)+      <+> ", memrefs:"+      <+> pretty (memrefs etry)+      <+> ", vartab:"+      <+> pretty (M.toList $ vartab etry)+      <+> "}"+      <+> "\n"++-- | Compute the union of two 'CoalsEntry'. If two 'CoalsEntry' do not refer to+-- the same destination memory and use the same index function, the first+-- 'CoalsEntry' is returned.+unionCoalsEntry :: CoalsEntry -> CoalsEntry -> CoalsEntry+unionCoalsEntry etry1 (CoalsEntry dstmem2 dstind2 alsmem2 vartab2 optdeps2 memrefs2) =+  if dstmem etry1 /= dstmem2 || dstind etry1 /= dstind2+    then etry1+    else+      etry1+        { alsmem = alsmem etry1 <> alsmem2,+          optdeps = optdeps etry1 <> optdeps2,+          vartab = vartab etry1 <> vartab2,+          memrefs = memrefs etry1 <> memrefs2+        }++-- | Get the names of array 'PatElem's in a 'Pat' and the corresponding+-- 'ArrayMemBound' information for each array.+getArrMemAssoc :: Pat (aliases, LetDecMem) -> [(VName, ArrayMemBound)]+getArrMemAssoc pat =+  mapMaybe+    ( \patel -> case snd $ patElemDec patel of+        (MemArray tp shp _ (ArrayIn mem_nm indfun)) ->+          Just (patElemName patel, MemBlock tp shp mem_nm indfun)+        MemMem _ -> Nothing+        MemPrim _ -> Nothing+        MemAcc {} -> Nothing+    )+    $ patElems pat++-- | Get the names of arrays in a list of 'FParam' and the corresponding+-- 'ArrayMemBound' information for each array.+getArrMemAssocFParam :: [Param FParamMem] -> [(VName, Uniqueness, ArrayMemBound)]+getArrMemAssocFParam =+  mapMaybe+    ( \param -> case paramDec param of+        (MemArray tp shp u (ArrayIn mem_nm indfun)) ->+          Just (paramName param, u, MemBlock tp shp mem_nm indfun)+        MemMem _ -> Nothing+        MemPrim _ -> Nothing+        MemAcc {} -> Nothing+    )++-- | Get memory blocks in a list of 'FParam' that are used for unique arrays in+-- the same list of 'FParam'.+getUniqueMemFParam :: [Param FParamMem] -> M.Map VName Space+getUniqueMemFParam params =+  let mems = M.fromList $ mapMaybe justMem params+      arrayMems = S.fromList $ mapMaybe (justArrayMem . paramDec) params+   in mems `M.restrictKeys` arrayMems+  where+    justMem (Param _ nm (MemMem sp)) = Just (nm, sp)+    justMem _ = Nothing+    justArrayMem (MemArray _ _ Unique (ArrayIn mem_nm _)) = Just mem_nm+    justArrayMem _ = Nothing++class HasMemBlock rep where+  -- | Looks up 'VName' in the given scope. If it is a 'MemArray', return the+  -- 'ArrayMemBound' information for the array.+  getScopeMemInfo :: VName -> Scope rep -> Maybe ArrayMemBound++instance HasMemBlock (Aliases SeqMem) where+  getScopeMemInfo r scope_env0 =+    case M.lookup r scope_env0 of+      Just (LetName (_, MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)+      Just (FParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)+      Just (LParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)+      _ -> Nothing++instance HasMemBlock (Aliases GPUMem) where+  getScopeMemInfo r scope_env0 =+    case M.lookup r scope_env0 of+      Just (LetName (_, MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)+      Just (FParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)+      Just (LParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)+      _ -> Nothing++-- | @True@ if the expression returns a "fresh" array.+createsNewArrOK :: CreatesNewArrOp (Op rep) => Exp rep -> Bool+createsNewArrOK (BasicOp Replicate {}) = True+createsNewArrOK (BasicOp Iota {}) = True+createsNewArrOK (BasicOp Manifest {}) = True+createsNewArrOK (BasicOp Copy {}) = True+createsNewArrOK (BasicOp Concat {}) = True+createsNewArrOK (BasicOp ArrayLit {}) = True+createsNewArrOK (BasicOp Scratch {}) = True+createsNewArrOK (BasicOp Rotate {}) = True+createsNewArrOK (Op op) = createsNewArrOp op+createsNewArrOK _ = False++class CreatesNewArrOp rep where+  createsNewArrOp :: rep -> Bool++instance CreatesNewArrOp () where+  createsNewArrOp () = False++instance CreatesNewArrOp inner => CreatesNewArrOp (MemOp inner) where+  createsNewArrOp (Alloc _ _) = True+  createsNewArrOp (Inner inner) = createsNewArrOp inner++instance CreatesNewArrOp inner => CreatesNewArrOp (HostOp (Aliases GPUMem) inner) where+  createsNewArrOp (OtherOp op) = createsNewArrOp op+  createsNewArrOp (SegOp (SegMap _ _ _ kbody)) = all isReturns $ kernelBodyResult kbody+  createsNewArrOp (SizeOp _) = False+  createsNewArrOp _ = undefined++isReturns :: KernelResult -> Bool+isReturns Returns {} = True+isReturns _ = False++-- | Memory-block removal from active-coalescing table+--   should only be handled via this function, it is easy+--   to run into infinite execution problem; i.e., the+--   fix-pointed iteration of coalescing transformation+--   assumes that whenever a coalescing fails it is+--   recorded in the @inhibit@ table.+markFailedCoal ::+  (CoalsTab, InhibitTab) ->+  VName ->+  (CoalsTab, InhibitTab)+markFailedCoal (coal_tab, inhb_tab) src_mem =+  case M.lookup src_mem coal_tab of+    Nothing -> (coal_tab, inhb_tab)+    Just coale ->+      let failed_set = oneName $ dstmem coale+          failed_set' = failed_set <> fromMaybe mempty (M.lookup src_mem inhb_tab)+       in ( M.delete src_mem coal_tab,+            M.insert src_mem failed_set' inhb_tab+          )++-- | promotion from active-to-successful coalescing tables+--   should be handled with this function (for clarity).+markSuccessCoal ::+  (CoalsTab, CoalsTab) ->+  VName ->+  CoalsEntry ->+  (CoalsTab, CoalsTab)+markSuccessCoal (actv, succc) m_b info_b =+  ( M.delete m_b actv,+    appendCoalsInfo m_b info_b succc+  )++-- | merges entries in the coalesced table.+appendCoalsInfo :: VName -> CoalsEntry -> CoalsTab -> CoalsTab+appendCoalsInfo mb info_new coalstab =+  case M.lookup mb coalstab of+    Nothing -> M.insert mb info_new coalstab+    Just info_old -> M.insert mb (unionCoalsEntry info_old info_new) coalstab++-- | Attempt to convert a 'VName' to a PrimExp.+--+-- First look in 'ScalarTab' to see if we have recorded the scalar value of the+-- argument. Otherwise look up the type of the argument and return a 'LeafExp'+-- if it is a 'PrimType'.+vnameToPrimExp ::+  (CanBeAliased (Op rep), RepTypes rep) =>+  ScopeTab rep ->+  ScalarTab ->+  VName ->+  Maybe (PrimExp VName)+vnameToPrimExp scopetab scaltab v =+  M.lookup v scaltab+    <|> ( M.lookup v scopetab+            >>= toPrimType . typeOf+            <&> LeafExp v+        )++-- | Attempt to extract the 'PrimType' from a 'TypeBase'.+toPrimType :: TypeBase shp u -> Maybe PrimType+toPrimType (Prim pt) = Just pt+toPrimType _ = Nothing

@@ -0,0 +1,347 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE PartialTypeSignatures #-}+{-# LANGUAGE TypeFamilies #-}++-- | Last use analysis for array short circuiting+--+-- Last-Use analysis of a Futhark program in aliased explicit-memory lore form.+-- Takes as input such a program or a function and produces a 'M.Map VName+-- Names', in which the key identified the let stmt, and the list argument+-- identifies the variables that were lastly used in that stmt.  Note that the+-- results of a body do not have a last use, and neither do a function+-- parameters if it happens to not be used inside function's body.  Such cases+-- are supposed to be treated separately.+--+-- This pass is different from "Futhark.Analysis.LastUse" in that memory blocks+-- are used to alias arrays. For instance, an 'Update' will not result in a last+-- use of the array being updated, because the result lives in the same memory.+module Futhark.Optimise.ArrayShortCircuiting.LastUse (lastUseSeqMem, lastUsePrg, lastUsePrgGPU, lastUseGPUMem) where++import Control.Monad.Reader+import Control.Monad.State.Strict+import Data.Bifunctor (bimap)+import Data.Map.Strict qualified as M+import Data.Maybe+import Data.Sequence (Seq (..))+import Futhark.IR.Aliases+import Futhark.IR.GPUMem+import Futhark.IR.SeqMem+import Futhark.Optimise.ArrayShortCircuiting.DataStructs+import Futhark.Util++-- | 'LastUseReader' allows us to abstract over representations by supplying the+-- 'onOp' function.+newtype LastUseReader rep = LastUseReader+  { onOp :: Op (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)+  }++type LastUseM rep a = StateT AliasTab (Reader (LastUseReader rep)) a++aliasLookup :: VName -> LastUseM rep Names+aliasLookup vname =+  gets $ fromMaybe mempty . M.lookup vname++-- | Perform last-use analysis on a 'Prog' in 'SeqMem'+lastUsePrg :: Prog (Aliases SeqMem) -> LUTabPrg+lastUsePrg prg = M.fromList $ map lastUseSeqMem $ progFuns prg++-- | Perform last-use analysis on a 'Prog' in 'GPUMem'+lastUsePrgGPU :: Prog (Aliases GPUMem) -> LUTabPrg+lastUsePrgGPU prg = M.fromList $ map lastUseGPUMem $ progFuns prg++-- | Perform last-use analysis on a 'FunDef' in 'SeqMem'+lastUseSeqMem :: FunDef (Aliases SeqMem) -> (Name, LUTabFun)+lastUseSeqMem (FunDef _ _ fname _ _ body) =+  let (res, _) =+        runReader+          (evalStateT (lastUseBody body (mempty, mempty)) mempty)+          (LastUseReader lastUseSeqOp)+   in (fname, res)++-- | Perform last-use analysis on a 'FunDef' in 'GPUMem'+lastUseGPUMem :: FunDef (Aliases GPUMem) -> (Name, LUTabFun)+lastUseGPUMem (FunDef _ _ fname _ _ body) =+  let (res, _) =+        runReader+          (evalStateT (lastUseBody body (mempty, mempty)) mempty)+          (LastUseReader lastUseGPUOp)+   in (fname, res)++-- | Performing the last-use analysis on a body.+--+-- The implementation consists of a bottom-up traversal of the body's statements+-- in which the the variables lastly used in a statement are computed as the+-- difference between the free-variables in that stmt and the set of variables+-- known to be used after that statement.+lastUseBody ::+  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>+  -- | The body of statements+  Body (Aliases rep) ->+  -- | The current last-use table, tupled with the known set of already used names+  (LUTabFun, Names) ->+  -- | The result is:+  --      (i) an updated last-use table,+  --     (ii) an updated set of used names (including the binding).+  LastUseM rep (LUTabFun, Names)+lastUseBody bdy@(Body _ stms result) (lutab, used_nms) = do+  -- perform analysis bottom-up in bindings: results are known to be used,+  -- hence they are added to the used_nms set.+  (lutab', _) <-+    lastUseStms stms (lutab, used_nms) $+      namesToList $+        freeIn $+          map resSubExp result+  -- Clean up the used names by recomputing the aliasing transitive-closure+  -- of the free names in body based on the current alias table @alstab@.+  used_in_body <- aliasTransitiveClosure $ freeIn bdy+  pure (lutab', used_nms <> used_in_body)++-- | Performing the last-use analysis on a body.+--+-- The implementation consists of a bottom-up traversal of the body's statements+-- in which the the variables lastly used in a statement are computed as the+-- difference between the free-variables in that stmt and the set of variables+-- known to be used after that statement.+lastUseKernelBody ::+  (CanBeAliased (Op rep), ASTRep rep) =>+  -- | The body of statements+  KernelBody (Aliases rep) ->+  -- | The current last-use table, tupled with the known set of already used names+  (LUTabFun, Names) ->+  -- | The result is:+  --      (i) an updated last-use table,+  --     (ii) an updated set of used names (including the binding).+  LastUseM rep (LUTabFun, Names)+lastUseKernelBody bdy@(KernelBody _ stms result) (lutab, used_nms) = do+  -- perform analysis bottom-up in bindings: results are known to be used,+  -- hence they are added to the used_nms set.+  (lutab', _) <-+    lastUseStms stms (lutab, used_nms) $ namesToList $ freeIn result+  -- Clean up the used names by recomputing the aliasing transitive-closure+  -- of the free names in body based on the current alias table @alstab@.+  used_in_body <- aliasTransitiveClosure $ freeIn bdy+  pure (lutab', used_nms <> used_in_body)++lastUseStms ::+  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>+  Stms (Aliases rep) ->+  (LUTabFun, Names) ->+  [VName] ->+  LastUseM rep (LUTabFun, Names)+lastUseStms Empty (lutab, nms) res_nms = do+  aliases <- concatMapM aliasLookup res_nms+  pure (lutab, nms <> aliases)+lastUseStms (stm@(Let pat _ e) :<| stms) (lutab, nms) res_nms = do+  let extra_alias = case e of+        BasicOp (Update _ old _ _) -> oneName old+        BasicOp (FlatUpdate old _ _) -> oneName old+        _ -> mempty+  -- We build up aliases top-down+  updateAliasing extra_alias pat+  -- But compute last use bottom-up+  (lutab', nms') <- lastUseStms stms (lutab, nms) res_nms+  (lutab'', nms'') <- lastUseStm stm (lutab', nms')+  pure (lutab'', nms'')++lastUseStm ::+  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>+  Stm (Aliases rep) ->+  (LUTabFun, Names) ->+  LastUseM rep (LUTabFun, Names)+lastUseStm (Let pat _ e) (lutab, used_nms) =+  do+    -- analyse the expression and get the+    --  (i)  a new last-use table (in case the @e@ contains bodies of stmts)+    -- (ii) the set of variables lastly used in the current binding.+    -- (iii)  aliased transitive-closure of used names, and+    (lutab', last_uses, used_nms') <- lastUseExp e used_nms+    -- filter-out the binded names from the set of used variables,+    -- since they go out of scope, and update the last-use table.+    let patnms = patNames pat+        used_nms'' = used_nms' `namesSubtract` namesFromList patnms+        lutab'' =+          M.union lutab' $ M.insert (head patnms) last_uses lutab+    pure (lutab'', used_nms'')++--------------------------------++-- | Last-Use Analysis for an expression.+lastUseExp ::+  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>+  -- | The expression to analyse+  Exp (Aliases rep) ->+  -- | The set of used names "after" this expression+  Names ->+  -- | Result:+  --    1. an extra LUTab recording the last use for expression's inner bodies,+  --    2. the set of last-used vars in the expression at this level,+  --    3. the updated used names, now including expression's free vars.+  LastUseM rep (LUTabFun, Names, Names)+lastUseExp (Match _ cases body _) used_nms = do+  -- For an if-then-else, we duplicate the last use at each body level, meaning+  -- we record the last use of the outer statement, and also the last use in the+  -- statement in the inner bodies. We can safely ignore the if-condition as it is+  -- a boolean scalar.+  (lutab_cases, used_cases) <-+    bimap mconcat mconcat . unzip+      <$> mapM (flip lastUseBody (M.empty, used_nms) . caseBody) cases+  (lutab', body_used_nms) <- lastUseBody body (M.empty, used_nms)+  let free_in_body = freeIn body+  let free_in_cases = freeIn cases+  let used_nms' = used_cases <> body_used_nms+  (_, last_used_arrs) <- lastUsedInNames used_nms $ free_in_body <> free_in_cases+  pure (lutab_cases <> lutab', last_used_arrs, used_nms')+lastUseExp (DoLoop var_ses _ body) used_nms0 = do+  free_in_body <- aliasTransitiveClosure $ freeIn body+  -- compute the aliasing transitive closure of initializers that are not last-uses+  var_inis <- catMaybes <$> mapM (initHelper (free_in_body <> used_nms0)) var_ses+  let -- To record last-uses inside the loop body, we call 'lastUseBody' with used-names+      -- being:  (free_in_body - loop-variants-a) + used_nms0. As such we disable cases b)+      -- and c) to produce loop-variant last uses inside the loop, and also we prevent+      -- the free-loop-variables to having last uses inside the loop.+      free_in_body' = free_in_body `namesSubtract` namesFromList (map fst var_inis)+      used_nms = used_nms0 <> free_in_body' <> freeIn (bodyResult body)+  (body_lutab, _) <- lastUseBody body (mempty, used_nms)++  -- add var_inis_a to the body_lutab, i.e., record the last-use of+  -- initializer in the corresponding loop variant.+  let lutab_res = body_lutab <> M.fromList var_inis++      -- the result used names are:+      fpar_nms = namesFromList $ map (identName . paramIdent . fst) var_ses+      used_nms' = (free_in_body <> freeIn (map snd var_ses)) `namesSubtract` fpar_nms+      used_nms_res = used_nms0 <> used_nms' <> freeIn (bodyResult body)++      -- the last-uses at loop-statement level are the loop free variables that+      -- do not belong to @used_nms0@; this includes the initializers of b), @lu_ini_b@+      lu_arrs = used_nms' `namesSubtract` used_nms0+  pure (lutab_res, lu_arrs, used_nms_res)+  where+    initHelper free_and_used (fp, se) = do+      names <- aliasTransitiveClosure $ maybe mempty oneName $ subExpVar se+      if names `namesIntersect` free_and_used+        then pure Nothing+        else pure $ Just (identName $ paramIdent fp, names)+lastUseExp (Op op) used_nms = do+  on_op <- reader onOp+  on_op op used_nms+lastUseExp e used_nms = do+  let free_in_e = freeIn e+  (used_nms', lu_vars) <- lastUsedInNames used_nms free_in_e+  pure (M.empty, lu_vars, used_nms')++lastUseGPUOp :: Op (Aliases GPUMem) -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)+lastUseGPUOp (Alloc se sp) used_nms = do+  let free_in_e = freeIn se <> freeIn sp+  (used_nms', lu_vars) <- lastUsedInNames used_nms free_in_e+  pure (M.empty, lu_vars, used_nms')+lastUseGPUOp (Inner (OtherOp ())) used_nms =+  pure (mempty, mempty, used_nms)+lastUseGPUOp (Inner (SizeOp sop)) used_nms = do+  (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn sop+  pure (mempty, lu_vars, used_nms')+lastUseGPUOp (Inner (SegOp (SegMap _ _ tps kbody))) used_nms = do+  (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn tps+  (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')+  pure (body_lutab, lu_vars, used_nms' <> used_nms'')+lastUseGPUOp (Inner (SegOp (SegRed _ _ sbos tps kbody))) used_nms = do+  (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseSegBinOp sbos used_nms+  (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps+  (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')+  pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'')+lastUseGPUOp (Inner (SegOp (SegScan _ _ sbos tps kbody))) used_nms = do+  (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseSegBinOp sbos used_nms+  (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps+  (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')+  pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'')+lastUseGPUOp (Inner (SegOp (SegHist _ _ hos tps kbody))) used_nms = do+  (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseHistOp hos used_nms+  (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps+  (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')+  pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'')+lastUseGPUOp (Inner (GPUBody tps body)) used_nms = do+  (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn tps+  (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms')+  pure (body_lutab, lu_vars, used_nms' <> used_nms'')++lastUseSegBinOp :: [SegBinOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)+lastUseSegBinOp sbos used_nms = do+  (lutab, lu_vars, used_nms') <- unzip3 <$> mapM helper sbos+  pure (mconcat lutab, mconcat lu_vars, mconcat used_nms')+  where+    helper (SegBinOp _ (Lambda _ body _) neutral shp) = do+      (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn neutral <> freeIn shp+      (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms')+      pure (body_lutab, lu_vars, used_nms'')++lastUseHistOp :: [HistOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)+lastUseHistOp hos used_nms = do+  (lutab, lu_vars, used_nms') <- unzip3 <$> mapM helper hos+  pure (mconcat lutab, mconcat lu_vars, mconcat used_nms')+  where+    helper (HistOp shp rf dest neutral shp' (Lambda _ body _)) = do+      (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn shp <> freeIn rf <> freeIn dest <> freeIn neutral <> freeIn shp'+      (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms')+      pure (body_lutab, lu_vars, used_nms'')++lastUseSeqOp :: Op (Aliases SeqMem) -> Names -> LastUseM SeqMem (LUTabFun, Names, Names)+lastUseSeqOp (Alloc se sp) used_nms = do+  let free_in_e = freeIn se <> freeIn sp+  (used_nms', lu_vars) <- lastUsedInNames used_nms free_in_e+  pure (mempty, lu_vars, used_nms')+lastUseSeqOp (Inner ()) used_nms = do+  pure (mempty, mempty, used_nms)++------------------------------------------------------++-- | Given already used names and newly encountered 'Names', return an updated+-- set used names and the set of names that were last used here.+--+-- For a given name @x@ in the new uses, if neither @x@ nor any of its aliases+-- are present in the set of used names, this is a last use of @x@.+lastUsedInNames ::+  -- | Used names+  Names ->+  -- | New uses+  Names ->+  LastUseM rep (Names, Names)+lastUsedInNames used_nms new_uses = do+  -- a use of an argument x is also a use of any variable in x alias set+  -- so we update the alias-based transitive-closure of used names.+  new_uses_with_aliases <- aliasTransitiveClosure new_uses+  -- if neither a variable x, nor any of its alias set have been used before (in+  -- the backward traversal), then it is a last use of both that variable and+  -- all other variables in its alias set+  last_uses <- filterM isLastUse $ namesToList new_uses+  last_uses' <- aliasTransitiveClosure $ namesFromList last_uses+  pure (used_nms <> new_uses_with_aliases, last_uses')+  where+    isLastUse x = do+      with_aliases <- aliasTransitiveClosure $ oneName x+      pure $ not $ with_aliases `namesIntersect` used_nms++-- | Compute the transitive closure of the aliases of a set of 'Names'.+aliasTransitiveClosure :: Names -> LastUseM rep Names+aliasTransitiveClosure args = do+  res <- foldl (<>) args <$> mapM aliasLookup (namesToList args)+  if res == args+    then pure res+    else aliasTransitiveClosure res++-- | For each 'PatElem' in the 'Pat', add its aliases to the 'AliasTab' in+-- 'LastUseM'. Additionally, 'Names' are added as aliases of all the 'PatElemT'.+updateAliasing ::+  AliasesOf dec =>+  -- | Extra names that all 'PatElem' should alias.+  Names ->+  -- | Pattern to process+  Pat dec ->+  LastUseM rep ()+updateAliasing extra_aliases =+  mapM_ update . patElems+  where+    update :: AliasesOf dec => PatElem dec -> LastUseM rep ()+    update (PatElem name dec) = do+      let aliases = aliasesOf dec+      aliases' <- aliasTransitiveClosure $ extra_aliases <> aliases+      modify $ M.insert name aliases'

@@ -0,0 +1,464 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.Optimise.ArrayShortCircuiting.MemRefAggreg+  ( recordMemRefUses,+    freeVarSubstitutions,+    translateAccessSummary,+    aggSummaryLoopTotal,+    aggSummaryLoopPartial,+    aggSummaryMapPartial,+    aggSummaryMapTotal,+    noMemOverlap,+  )+where++import Control.Monad+import Data.Function ((&))+import Data.List (intersect, partition, uncons)+import Data.List.NonEmpty (NonEmpty (..))+import Data.Map.Strict qualified as M+import Data.Maybe+import Data.Set qualified as S+import Futhark.Analysis.AlgSimplify+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.Aliases+import Futhark.IR.Mem+import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.MonadFreshNames+import Futhark.Optimise.ArrayShortCircuiting.DataStructs+import Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis+import Futhark.Util++-----------------------------------------------------+-- Some translations of Accesses and Ixfuns        --+-----------------------------------------------------++-- | Checks whether the index function can be translated at the current program+-- point and also returns the substitutions.  It comes down to answering the+-- question: "can one perform enough substitutions (from the bottom-up scalar+-- table) until all vars appearing in the index function are defined in the+-- current scope?"+freeVarSubstitutions ::+  FreeIn a =>+  ScopeTab rep ->+  ScalarTab ->+  a ->+  Maybe FreeVarSubsts+freeVarSubstitutions scope0 scals0 indfun =+  freeVarSubstitutions' mempty $ namesToList $ freeIn indfun+  where+    freeVarSubstitutions' :: FreeVarSubsts -> [VName] -> Maybe FreeVarSubsts+    freeVarSubstitutions' subs [] = Just subs+    freeVarSubstitutions' subs0 fvs =+      let fvs_not_in_scope = filter (`M.notMember` scope0) fvs+       in case unzip <$> mapM getSubstitution fvs_not_in_scope of+            -- We require that all free variables can be substituted+            Just (subs, new_fvs) ->+              freeVarSubstitutions' (subs0 <> mconcat subs) $ concat new_fvs+            Nothing -> Nothing+    getSubstitution v+      | Just pe <- M.lookup v scals0,+        IntType _ <- primExpType pe =+          Just (M.singleton v $ TPrimExp pe, namesToList $ freeIn pe)+    getSubstitution _v = Nothing++-- | Translates free variables in an access summary+translateAccessSummary :: ScopeTab rep -> ScalarTab -> AccessSummary -> AccessSummary+translateAccessSummary _ _ Undeterminable = Undeterminable+translateAccessSummary scope0 scals0 (Set slmads)+  | Just subs <- freeVarSubstitutions scope0 scals0 slmads =+      slmads+        & S.map (IxFun.substituteInLMAD subs)+        & Set+translateAccessSummary _ _ _ = Undeterminable++-- | This function computes the written and read memory references for the current statement+getUseSumFromStm ::+  (Op rep ~ MemOp inner, HasMemBlock (Aliases rep)) =>+  TopdownEnv rep ->+  CoalsTab ->+  Stm (Aliases rep) ->+  -- | A pair of written and written+read memory locations, along with their+  -- associated array and the index function used+  Maybe ([(VName, VName, IxFun)], [(VName, VName, IxFun)])+getUseSumFromStm td_env coal_tab (Let _ _ (BasicOp (Index arr (Slice slc))))+  | Just (MemBlock _ shp _ _) <- getScopeMemInfo arr (scope td_env),+    length slc == length (shapeDims shp) && all isFix slc = do+      (mem_b, mem_arr, ixfn_arr) <- getDirAliasedIxfn td_env coal_tab arr+      let new_ixfn = IxFun.slice ixfn_arr $ Slice $ map (fmap pe64) slc+      pure ([], [(mem_b, mem_arr, new_ixfn)])+  where+    isFix DimFix {} = True+    isFix _ = False+getUseSumFromStm _ _ (Let Pat {} _ (BasicOp Index {})) = Just ([], []) -- incomplete slices+getUseSumFromStm _ _ (Let Pat {} _ (BasicOp FlatIndex {})) = Just ([], []) -- incomplete slices+getUseSumFromStm td_env coal_tab (Let (Pat pes) _ (BasicOp (ArrayLit ses _))) =+  let rds = mapMaybe (getDirAliasedIxfn td_env coal_tab) $ mapMaybe seName ses+      wrts = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) pes+   in Just (wrts, wrts ++ rds)+  where+    seName (Var a) = Just a+    seName (Constant _) = Nothing+-- In place update @x[slc] <- a@. In the "in-place update" case,+--   summaries should be added after the old variable @x@ has+--   been added in the active coalesced table.+getUseSumFromStm td_env coal_tab (Let (Pat [x']) _ (BasicOp (Update _ _x (Slice slc) a_se))) = do+  (m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x')+  let x_ixfn_slc = IxFun.slice x_ixfn $ Slice $ map (fmap pe64) slc+      r1 = (m_b, m_x, x_ixfn_slc)+  case a_se of+    Constant _ -> Just ([r1], [r1])+    Var a -> case getDirAliasedIxfn td_env coal_tab a of+      Nothing -> Just ([r1], [r1])+      Just r2 -> Just ([r1], [r1, r2])+getUseSumFromStm td_env coal_tab (Let (Pat [y]) _ (BasicOp (Copy x))) = do+  -- y = copy x+  wrt <- getDirAliasedIxfn td_env coal_tab $ patElemName y+  rd <- getDirAliasedIxfn td_env coal_tab x+  pure ([wrt], [wrt, rd])+getUseSumFromStm _ _ (Let Pat {} _ (BasicOp Copy {})) = error "Impossible"+getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp (Concat _i (a :| bs) _ses))) =+  -- concat+  let ws = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys+      rs = mapMaybe (getDirAliasedIxfn td_env coal_tab) (a : bs)+   in Just (ws, ws ++ rs)+getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp (Manifest _perm x))) =+  let ws = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys+      rs = mapMaybe (getDirAliasedIxfn td_env coal_tab) [x]+   in Just (ws, ws ++ rs)+getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp (Replicate _shp se))) =+  let ws = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys+   in case se of+        Constant _ -> Just (ws, ws)+        Var x -> Just (ws, ws ++ mapMaybe (getDirAliasedIxfn td_env coal_tab) [x])+getUseSumFromStm td_env coal_tab (Let (Pat [x]) _ (BasicOp (FlatUpdate _ (FlatSlice offset slc) v)))+  | Just (m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x) =+      let x_ixfn_slc = IxFun.flatSlice x_ixfn $ FlatSlice (pe64 offset) $ map (fmap pe64) slc+          r1 = (m_b, m_x, x_ixfn_slc)+       in case getDirAliasedIxfn td_env coal_tab v of+            Nothing -> Just ([r1], [r1])+            Just r2 -> Just ([r1], [r1, r2])+getUseSumFromStm _ _ (Let Pat {} _ BasicOp {}) = Just ([], [])+getUseSumFromStm _ _ (Let Pat {} _ (Op (Alloc _ _))) = Just ([], [])+getUseSumFromStm _ _ _ =+  -- if-then-else, loops are supposed to be treated separately,+  -- calls are not supported, and Ops are not yet supported+  Nothing++-- | This function:+--     1. computes the written and read memory references for the current statement+--          (by calling @getUseSumFromStm@)+--     2. fails the entries in active coalesced table for which the write set+--          overlaps the uses of the destination (to that point)+recordMemRefUses ::+  (CanBeAliased (Op rep), RepTypes rep, Op rep ~ MemOp inner, HasMemBlock (Aliases rep)) =>+  TopdownEnv rep ->+  BotUpEnv ->+  Stm (Aliases rep) ->+  (CoalsTab, InhibitTab)+recordMemRefUses td_env bu_env stm =+  let active_tab = activeCoals bu_env+      inhibit_tab = inhibit bu_env+      active_etries = M.toList active_tab+   in case getUseSumFromStm td_env active_tab stm of+        Nothing ->+          M.toList active_tab+            & foldl+              ( \state (m_b, entry) ->+                  if not $ null $ patNames (stmPat stm) `intersect` M.keys (vartab entry)+                    then markFailedCoal state m_b+                    else state+              )+              (active_tab, inhibit_tab)+        Just use_sums ->+          let (mb_wrts, prev_uses, mb_lmads) =+                map (checkOverlapAndExpand use_sums active_tab) active_etries+                  & unzip3++              -- keep only the entries that do not overlap with the memory+              -- blocks defined in @pat@ or @inner_free_vars@.+              -- the others must be recorded in @inhibit_tab@ because+              -- they violate the 3rd safety condition.+              active_tab1 =+                M.fromList+                  $ map+                    ( \(wrts, (uses, prev_use, (k, etry))) ->+                        let mrefs' = (memrefs etry) {dstrefs = prev_use}+                            etry' = etry {memrefs = mrefs'}+                         in (k, addLmads wrts uses etry')+                    )+                  $ mapMaybe (\(x, y) -> (,y) <$> x) -- only keep successful coals+                  $ zip mb_wrts+                  $ zip3 mb_lmads prev_uses active_etries+              failed_tab =+                M.fromList $+                  map snd $+                    filter (isNothing . fst) $+                      zip mb_wrts active_etries+              (_, inhibit_tab1) = foldl markFailedCoal (failed_tab, inhibit_tab) $ M.keys failed_tab+           in (active_tab1, inhibit_tab1)+  where+    checkOverlapAndExpand (stm_wrts, stm_uses) active_tab (m_b, etry) =+      let alias_m_b = getAliases mempty m_b+          stm_uses' = filter ((`notNameIn` alias_m_b) . tupFst) stm_uses+          all_aliases = foldl getAliases mempty $ namesToList $ alsmem etry+          ixfns = map tupThd $ filter ((`nameIn` all_aliases) . tupSnd) stm_uses'+          lmads' = mapMaybe mbLmad ixfns+          lmads'' =+            if length lmads' == length ixfns+              then Set $ S.fromList lmads'+              else Undeterminable+          wrt_ixfns = map tupThd $ filter ((`nameIn` alias_m_b) . tupFst) stm_wrts+          wrt_tmps = mapMaybe mbLmad wrt_ixfns+          prev_use =+            translateAccessSummary (scope td_env) (scalarTable td_env) $+              (dstrefs . memrefs) etry+          wrt_lmads' =+            if length wrt_tmps == length wrt_ixfns+              then Set $ S.fromList wrt_tmps+              else Undeterminable+          original_mem_aliases =+            fmap tupFst stm_uses+              & uncons+              & fmap fst+              & (=<<) (`M.lookup` active_tab)+              & maybe mempty alsmem+          (wrt_lmads'', lmads) =+            if m_b `nameIn` original_mem_aliases+              then (wrt_lmads' <> lmads'', Set mempty)+              else (wrt_lmads', lmads'')+          no_overlap = noMemOverlap td_env prev_use wrt_lmads''+          wrt_lmads =+            if no_overlap+              then Just wrt_lmads''+              else Nothing+       in (wrt_lmads, prev_use, lmads)++    tupFst (a, _, _) = a+    tupSnd (_, b, _) = b+    tupThd (_, _, c) = c+    getAliases acc m =+      oneName m+        <> acc+        <> fromMaybe mempty (M.lookup m (m_alias td_env))+    mbLmad indfun+      | Just subs <- freeVarSubstitutions (scope td_env) (scals bu_env) indfun,+        (IxFun.IxFun (lmad :| []) _ _) <- IxFun.substituteInIxFun subs indfun =+          Just lmad+    mbLmad _ = Nothing+    addLmads wrts uses etry =+      etry {memrefs = MemRefs uses wrts <> memrefs etry}++-- | Check for memory overlap of two access summaries.+--+-- This check is conservative, so unless we can guarantee that there is no+-- overlap, we return 'False'.+noMemOverlap :: (CanBeAliased (Op rep), RepTypes rep) => TopdownEnv rep -> AccessSummary -> AccessSummary -> Bool+noMemOverlap _ _ (Set mr)+  | mr == mempty = True+noMemOverlap _ (Set mr) _+  | mr == mempty = True+noMemOverlap td_env (Set is0) (Set js0)+  | Just non_negs <- mapM (primExpFromSubExpM (vnameToPrimExp (scope td_env) (scalarTable td_env)) . Var) $ namesToList $ nonNegatives td_env =+      let (_, not_disjoints) =+            partition+              ( \i ->+                  all+                    ( \j ->+                        IxFun.disjoint less_thans (nonNegatives td_env) i j+                          || IxFun.disjoint2 () () less_thans (nonNegatives td_env) i j+                          || IxFun.disjoint3 (typeOf <$> scope td_env) asserts less_thans non_negs i j+                    )+                    js+              )+              is+       in null not_disjoints+  where+    less_thans = map (fmap $ fixPoint $ substituteInPrimExp $ scalarTable td_env) $ knownLessThan td_env+    asserts = map (fixPoint (substituteInPrimExp $ scalarTable td_env) . primExpFromSubExp Bool) $ td_asserts td_env+    is = map (fixPoint (IxFun.substituteInLMAD $ TPrimExp <$> scalarTable td_env)) $ S.toList is0+    js = map (fixPoint (IxFun.substituteInLMAD $ TPrimExp <$> scalarTable td_env)) $ S.toList js0+noMemOverlap _ _ _ = False++-- | Computes the total aggregated access summary for a loop by expanding the+-- access summary given according to the iterator variable and bounds of the+-- loop.+--+-- Corresponds to:+--+-- \[+--   \bigcup_{j=0}^{j<n} Access_j+-- \]+aggSummaryLoopTotal ::+  MonadFreshNames m =>+  ScopeTab rep ->+  ScopeTab rep ->+  ScalarTab ->+  Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName)) ->+  AccessSummary ->+  m AccessSummary+aggSummaryLoopTotal _ _ _ _ Undeterminable = pure Undeterminable+aggSummaryLoopTotal _ _ _ _ (Set l)+  | l == mempty = pure $ Set mempty+aggSummaryLoopTotal scope_bef scope_loop scals_loop _ access+  | Set ls <- translateAccessSummary scope_loop scals_loop access,+    nms <- foldl (<>) mempty $ map freeIn $ S.toList ls,+    all inBeforeScope $ namesToList nms = do+      pure $ Set ls+  where+    inBeforeScope v =+      case M.lookup v scope_bef of+        Nothing -> False+        Just _ -> True+aggSummaryLoopTotal _ _ scalars_loop (Just (iterator_var, (lower_bound, upper_bound))) (Set lmads) =+  concatMapM+    ( aggSummaryOne iterator_var lower_bound upper_bound+        . fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars_loop)+    )+    (S.toList lmads)+aggSummaryLoopTotal _ _ _ _ _ = pure Undeterminable++-- | For a given iteration of the loop $i$, computes the aggregated loop access+-- summary of all later iterations.+--+-- Corresponds to:+--+-- \[+--   \bigcup_{j=i+1}^{j<n} Access_j+-- \]+aggSummaryLoopPartial ::+  MonadFreshNames m =>+  ScalarTab ->+  Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName)) ->+  AccessSummary ->+  m AccessSummary+aggSummaryLoopPartial _ _ Undeterminable = pure Undeterminable+aggSummaryLoopPartial _ Nothing _ = pure Undeterminable+aggSummaryLoopPartial scalars_loop (Just (iterator_var, (_, upper_bound))) (Set lmads) = do+  -- map over each index function in the access summary+  --   Substitube a fresh variable k for the loop iterator+  --   if k is in stride or span of ixfun: fall back to total+  --   new_stride = old_offset - old_offset (where k+1 is substituted for k)+  --   new_offset = old_offset where k = lower bound of iteration+  --   new_span = upper bound of iteration+  concatMapM+    ( aggSummaryOne+        iterator_var+        (isInt64 (LeafExp iterator_var $ IntType Int64) + 1)+        (upper_bound - typedLeafExp iterator_var - 1)+        . fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars_loop)+    )+    (S.toList lmads)++-- | For a given map with $k$ dimensions and an index $i$ for each dimension,+-- compute the aggregated access summary of all other threads.+--+-- For the innermost dimension, this corresponds to+--+-- \[+--   \bigcup_{j=0}^{j<i} Access_j \cup \bigcup_{j=i+1}^{j<n} Access_j+-- \]+--+-- where $Access_j$ describes the point accesses in the map. As we move up in+-- dimensionality, the previous access summaries are kept, in addition to the+-- total aggregation of the inner dimensions. For outer dimensions, the equation+-- is the same, the point accesses in $Access_j$ are replaced with the total+-- aggregation of the inner dimensions.+aggSummaryMapPartial :: MonadFreshNames m => ScalarTab -> [(VName, SubExp)] -> LmadRef -> m AccessSummary+aggSummaryMapPartial _ [] = const $ pure mempty+aggSummaryMapPartial scalars dims =+  helper mempty (reverse dims) . Set . S.singleton -- Reverse dims so we work from the inside out+  where+    helper acc [] _ = pure acc+    helper Undeterminable _ _ = pure Undeterminable+    helper _ _ Undeterminable = pure Undeterminable+    helper (Set acc) ((gtid, size) : rest) (Set as) = do+      partial_as <- aggSummaryMapPartialOne scalars (gtid, size) (Set as)+      total_as <-+        concatMapM+          (aggSummaryOne gtid 0 (TPrimExp $ primExpFromSubExp (IntType Int64) size))+          (S.toList as)+      helper (Set acc <> partial_as) rest total_as++-- | Given an access summary $a$, a thread id $i$ and the size $n$ of the+-- dimension, compute the partial map summary.+--+-- Corresponds to+--+-- \[+--   \bigcup_{j=0}^{j<i} a_j \cup \bigcup_{j=i+1}^{j<n} a_j+-- \]+aggSummaryMapPartialOne :: MonadFreshNames m => ScalarTab -> (VName, SubExp) -> AccessSummary -> m AccessSummary+aggSummaryMapPartialOne _ _ Undeterminable = pure Undeterminable+aggSummaryMapPartialOne _ (_, Constant n) (Set _) | oneIsh n = pure mempty+aggSummaryMapPartialOne scalars (gtid, size) (Set lmads0) =+  concatMapM+    helper+    [ (0, isInt64 (LeafExp gtid $ IntType Int64)),+      ( isInt64 (LeafExp gtid $ IntType Int64) + 1,+        isInt64 (primExpFromSubExp (IntType Int64) size) - isInt64 (LeafExp gtid $ IntType Int64) - 1+      )+    ]+  where+    lmads = map (fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars)) $ S.toList lmads0+    helper (x, y) = concatMapM (aggSummaryOne gtid x y) lmads++-- | Computes to total access summary over a multi-dimensional map.+aggSummaryMapTotal :: MonadFreshNames m => ScalarTab -> [(VName, SubExp)] -> AccessSummary -> m AccessSummary+aggSummaryMapTotal _ [] _ = pure mempty+aggSummaryMapTotal _ _ (Set lmads)+  | lmads == mempty = pure mempty+aggSummaryMapTotal _ _ Undeterminable = pure Undeterminable+aggSummaryMapTotal scalars segspace (Set lmads0) =+  foldM+    ( \as' (gtid', size') -> case as' of+        Set lmads' ->+          concatMapM+            ( aggSummaryOne gtid' 0 $+                TPrimExp $+                  primExpFromSubExp (IntType Int64) size'+            )+            (S.toList lmads')+        Undeterminable -> pure Undeterminable+    )+    (Set lmads)+    (reverse segspace)+  where+    lmads =+      S.fromList $+        map (fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars)) $+          S.toList lmads0++-- | Helper function that aggregates the accesses of single LMAD according to a+-- given iterator value, a lower bound and a span.+--+-- If successful, the result is an index function with an extra outer+-- dimension. The stride of the outer dimension is computed by taking the+-- difference between two points in the index function.+--+-- The function returns 'Underterminable' if the iterator is free in the output+-- LMAD or the dimensions of the input LMAD .+aggSummaryOne :: MonadFreshNames m => VName -> TPrimExp Int64 VName -> TPrimExp Int64 VName -> LmadRef -> m AccessSummary+aggSummaryOne iterator_var lower_bound spn lmad@(IxFun.LMAD offset0 dims0)+  | iterator_var `nameIn` freeIn dims0 = pure Undeterminable+  | iterator_var `notNameIn` freeIn offset0 = pure $ Set $ S.singleton lmad+  | otherwise = do+      new_var <- newVName "k"+      let offset = replaceIteratorWith (typedLeafExp new_var) offset0+          offsetp1 = replaceIteratorWith (typedLeafExp new_var + 1) offset0+          new_stride = TPrimExp $ constFoldPrimExp $ simplify $ untyped $ offsetp1 - offset+          new_offset = replaceIteratorWith lower_bound offset0+          new_lmad =+            IxFun.LMAD new_offset $+              IxFun.LMADDim new_stride spn 0 IxFun.Inc : map incPerm dims0+      if new_var `nameIn` freeIn new_lmad+        then pure Undeterminable+        else pure $ Set $ S.singleton new_lmad+  where+    incPerm dim = dim {IxFun.ldPerm = IxFun.ldPerm dim + 1}+    replaceIteratorWith se = TPrimExp . substituteInPrimExp (M.singleton iterator_var $ untyped se) . untyped++-- | Takes a 'VName' and converts it into a 'TPrimExp' with type 'Int64'.+typedLeafExp :: VName -> TPrimExp Int64 VName+typedLeafExp vname = isInt64 $ LeafExp vname (IntType Int64)

@@ -0,0 +1,301 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis+  ( TopdownEnv (..),+    ScopeTab,+    TopDownHelper,+    InhibitTab,+    updateTopdownEnv,+    updateTopdownEnvLoop,+    getDirAliasedIxfn,+    getDirAliasedIxfn',+    addInvAliassesVarTab,+    areAnyAliased,+    isInScope,+    nonNegativesInPat,+  )+where++import Data.Map.Strict qualified as M+import Data.Maybe+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.Aliases+import Futhark.IR.GPUMem+import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.Optimise.ArrayShortCircuiting.DataStructs++type DirAlias = IxFun -> IxFun+-- ^ A direct aliasing transformation++type InvAlias = Maybe (IxFun -> IxFun)+-- ^ An inverse aliasing transformation++type VarAliasTab = M.Map VName (VName, DirAlias, InvAlias)++type MemAliasTab = M.Map VName Names++data TopdownEnv rep = TopdownEnv+  { -- | contains the already allocated memory blocks+    alloc :: AllocTab,+    -- | variable info, including var-to-memblock assocs+    scope :: ScopeTab rep,+    -- | the inherited inhibitions from the previous try+    inhibited :: InhibitTab,+    -- | for statements such as transpose, reshape, index, etc., that alias+    --   an array variable: maps var-names to pair of aliased var name+    --   and index function transformation. For example, for+    --   @let b = a[slc]@ it should add the binding+    --   @ b |-> (a, `slice` slc )@+    v_alias :: VarAliasTab,+    -- | keeps track of memory block aliasing.+    --   this needs to be implemented+    m_alias :: MemAliasTab,+    -- | Contains symbol information about the variables in the program. Used to+    -- determine if a variable is non-negative.+    nonNegatives :: Names,+    scalarTable :: M.Map VName (PrimExp VName),+    -- | A list of known relations of the form 'VName' @<@ 'SubExp', typically+    -- gotten from 'LoopForm' and 'SegSpace'.+    knownLessThan :: [(VName, PrimExp VName)],+    -- | A list of the asserts encountered so far+    td_asserts :: [SubExp]+  }++isInScope :: TopdownEnv rep -> VName -> Bool+isInScope td_env m =+  m `M.member` scope td_env++-- | Get alias and (direct) index function mapping from expression+--+-- For instance, if the expression is a 'Rotate', returns the value being+-- rotated as well as a function for rotating an index function the appropriate+-- amount.+getDirAliasFromExp :: Exp (Aliases rep) -> Maybe (VName, DirAlias)+getDirAliasFromExp (BasicOp (SubExp (Var x))) = Just (x, id)+getDirAliasFromExp (BasicOp (Opaque _ (Var x))) = Just (x, id)+getDirAliasFromExp (BasicOp (Reshape ReshapeCoerce shp x)) =+  Just (x, (`IxFun.coerce` shapeDims (fmap pe64 shp)))+getDirAliasFromExp (BasicOp (Reshape ReshapeArbitrary shp x)) =+  Just (x, (`IxFun.reshape` shapeDims (fmap pe64 shp)))+getDirAliasFromExp (BasicOp (Rearrange _ _)) =+  Nothing+getDirAliasFromExp (BasicOp (Rotate _ _)) =+  Nothing -- Just (x, (`IxFun.rotate` fmap pe64 rs))+getDirAliasFromExp (BasicOp (Index x slc)) =+  Just (x, (`IxFun.slice` (Slice $ map (fmap pe64) $ unSlice slc)))+getDirAliasFromExp (BasicOp (Update _ x _ _elm)) = Just (x, id)+getDirAliasFromExp (BasicOp (FlatIndex x (FlatSlice offset idxs))) =+  Just+    ( x,+      ( `IxFun.flatSlice`+          ( FlatSlice (pe64 offset) $+              map (fmap pe64) idxs+          )+      )+    )+getDirAliasFromExp (BasicOp (FlatUpdate x _ _)) = Just (x, id)+getDirAliasFromExp _ = Nothing++-- | This was former @createsAliasedArrOK@ from DataStructs+--   While Rearrange and Rotate create aliased arrays, we+--   do not yet support them because it would mean we have+--   to "reverse" the index function, for example to support+--   coalescing in the case below,+--       @let a = map f a0   @+--       @let b = transpose a@+--       @let y[4] = copy(b) @+--   we would need to assign to @a@ as index function, the+--   inverse of the transpose, such that, when creating @b@+--   by transposition we get a directly-mapped array, which+--   is expected by the copying in y[4].+--   For the moment we support only transposition and VName-expressions,+--     but rotations and full slices could also be supported.+--+-- This function complements 'getDirAliasFromExp' by returning a function that+-- applies the inverse index function transformation.+getInvAliasFromExp :: Exp (Aliases rep) -> InvAlias+getInvAliasFromExp (BasicOp (SubExp (Var _))) = Just id+getInvAliasFromExp (BasicOp (Opaque _ (Var _))) = Just id+getInvAliasFromExp (BasicOp Update {}) = Just id+getInvAliasFromExp (BasicOp (Rearrange perm _)) =+  let perm' = IxFun.permuteInv perm [0 .. length perm - 1]+   in Just (`IxFun.permute` perm')+getInvAliasFromExp _ = Nothing++class TopDownHelper inner where+  innerNonNegatives :: [VName] -> inner -> Names++  innerKnownLessThan :: inner -> [(VName, PrimExp VName)]++  scopeHelper :: inner -> Scope rep++instance TopDownHelper (HostOp (Aliases GPUMem) ()) where+  innerNonNegatives _ (SegOp seg_op) =+    foldMap (oneName . fst) $ unSegSpace $ segSpace seg_op+  innerNonNegatives [vname] (SizeOp (GetSize _ _)) = oneName vname+  innerNonNegatives [vname] (SizeOp (GetSizeMax _)) = oneName vname+  innerNonNegatives _ _ = mempty++  innerKnownLessThan (SegOp seg_op) =+    map (fmap $ primExpFromSubExp $ IntType Int64) $ unSegSpace $ segSpace seg_op+  innerKnownLessThan _ = mempty++  scopeHelper (SegOp seg_op) = scopeOfSegSpace $ segSpace seg_op+  scopeHelper _ = mempty++instance TopDownHelper () where+  innerNonNegatives _ () = mempty+  innerKnownLessThan () = mempty+  scopeHelper () = mempty++-- | fills in the TopdownEnv table+updateTopdownEnv ::+  (ASTRep rep, Op rep ~ MemOp inner, TopDownHelper (OpWithAliases inner)) =>+  TopdownEnv rep ->+  Stm (Aliases rep) ->+  TopdownEnv rep+updateTopdownEnv env stm@(Let (Pat [pe]) _ (Op (Alloc (Var vname) sp))) =+  env+    { alloc = M.insert (patElemName pe) sp $ alloc env,+      scope = scope env <> scopeOf stm,+      nonNegatives = nonNegatives env <> oneName vname+    }+updateTopdownEnv env stm@(Let pat _ (Op (Inner inner))) =+  env+    { scope = scope env <> scopeOf stm <> scopeHelper inner,+      nonNegatives = nonNegatives env <> innerNonNegatives (patNames pat) inner,+      knownLessThan = knownLessThan env <> innerKnownLessThan inner+    }+updateTopdownEnv env (Let (Pat _) _ (BasicOp (Assert se _ _))) =+  env {td_asserts = se : td_asserts env}+updateTopdownEnv env stm@(Let (Pat [pe]) _ e)+  | Just (x, ixfn) <- getDirAliasFromExp e =+      let ixfn_inv = getInvAliasFromExp e+       in env+            { v_alias = M.insert (patElemName pe) (x, ixfn, ixfn_inv) (v_alias env),+              scope = scope env <> scopeOf stm,+              nonNegatives = nonNegatives env <> nonNegativesInPat (stmPat stm)+            }+updateTopdownEnv env stm =+  env+    { scope = scope env <> scopeOf stm,+      nonNegatives =+        nonNegatives env+          <> nonNegativesInPat (stmPat stm)+    }++nonNegativesInPat :: Typed rep => Pat rep -> Names+nonNegativesInPat (Pat elems) =+  foldMap (namesFromList . mapMaybe subExpVar . arrayDims . typeOf) elems++-- | The topdown handler for loops.+updateTopdownEnvLoop :: TopdownEnv rep -> [(FParam rep, SubExp)] -> LoopForm (Aliases rep) -> TopdownEnv rep+updateTopdownEnvLoop td_env arginis lform =+  let scopetab =+        scope td_env+          <> scopeOfFParams (map fst arginis)+          <> scopeOf lform+      non_negatives =+        nonNegatives td_env <> case lform of+          ForLoop v _ _ _ -> oneName v+          _ -> mempty+      less_than =+        case lform of+          ForLoop v _ b _ -> [(v, primExpFromSubExp (IntType Int64) b)]+          _ -> mempty+   in td_env+        { scope = scopetab,+          nonNegatives = non_negatives,+          knownLessThan = less_than <> knownLessThan td_env+        }++-- | Get direct aliased index function.  Returns a triple of current memory+-- block to be coalesced, the destination memory block and the index function of+-- the access in the space of the destination block.+getDirAliasedIxfn :: HasMemBlock (Aliases rep) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)+getDirAliasedIxfn td_env coals_tab x =+  case getScopeMemInfo x (scope td_env) of+    Just (MemBlock _ _ m_x orig_ixfun) ->+      case M.lookup m_x coals_tab of+        Just coal_etry -> do+          (Coalesced _ (MemBlock _ _ m ixf) _) <- walkAliasTab (v_alias td_env) (vartab coal_etry) x+          pure (m_x, m, ixf)+        Nothing ->+          -- This value is not subject to coalescing at the moment. Just return the+          -- original index function+          Just (m_x, m_x, orig_ixfun)+    Nothing -> Nothing++-- | Like 'getDirAliasedIxfn', but this version returns 'Nothing' if the value+-- is not currently subject to coalescing.+getDirAliasedIxfn' :: HasMemBlock (Aliases rep) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)+getDirAliasedIxfn' td_env coals_tab x =+  case getScopeMemInfo x (scope td_env) of+    Just (MemBlock _ _ m_x _) ->+      case M.lookup m_x coals_tab of+        Just coal_etry -> do+          (Coalesced _ (MemBlock _ _ m ixf) _) <- walkAliasTab (v_alias td_env) (vartab coal_etry) x+          pure (m_x, m, ixf)+        Nothing ->+          -- This value is not subject to coalescing at the moment. Just return the+          -- original index function+          Nothing+    Nothing -> Nothing++-- | Given a 'VName', walk the 'VarAliasTab' until found in the 'Map'.+walkAliasTab ::+  VarAliasTab ->+  M.Map VName Coalesced ->+  VName ->+  Maybe Coalesced+walkAliasTab _ vtab x+  | Just c <- M.lookup x vtab =+      Just c -- @x@ is in @vartab@ together with its new ixfun+walkAliasTab alias_tab vtab x+  | Just (x0, alias0, _) <- M.lookup x alias_tab = do+      Coalesced knd (MemBlock pt shp vname ixf) substs <- walkAliasTab alias_tab vtab x0+      pure $ Coalesced knd (MemBlock pt shp vname $ alias0 ixf) substs+walkAliasTab _ _ _ = Nothing++-- | We assume @x@ is in @vartab@ and we add the variables that @x@ aliases+--   for as long as possible following a chain of direct-aliasing operators,+--   i.e., without considering aliasing of if-then-else, loops, etc. For example:+--     @ x0 = if c then ... else ...@+--     @ x1 = rearrange r1 x0 @+--     @ x2 = reverse x1@+--     @ y[slc] = x2 @+--   We assume @vartab@ constains a binding for @x2@, and calling this function+--     with @x2@ as argument should also insert entries for @x1@ and @x0@ to+--     @vartab@, of course if their aliasing operations are invertible.+--   We assume inverting aliases has been performed by the top-down pass.+addInvAliassesVarTab ::+  HasMemBlock (Aliases rep) =>+  TopdownEnv rep ->+  M.Map VName Coalesced ->+  VName ->+  Maybe (M.Map VName Coalesced)+addInvAliassesVarTab td_env vtab x+  | Just (Coalesced _ (MemBlock _ _ m_y x_ixfun) fv_subs) <- M.lookup x vtab =+      case M.lookup x (v_alias td_env) of+        Nothing -> Just vtab+        Just (_, _, Nothing) -> Nothing -- can't invert ixfun, conservatively fail!+        Just (x0, _, Just inv_alias0) ->+          let x_ixfn0 = inv_alias0 x_ixfun+           in case getScopeMemInfo x0 (scope td_env) of+                Nothing -> error "impossible"+                Just (MemBlock ptp shp _ _) ->+                  let coal = Coalesced TransitiveCoal (MemBlock ptp shp m_y x_ixfn0) fv_subs+                      vartab' = M.insert x0 coal vtab+                   in addInvAliassesVarTab td_env vartab' x0+addInvAliassesVarTab _ _ _ = Nothing++areAliased :: TopdownEnv rep -> VName -> VName -> Bool+areAliased _ m_x m_y =+  -- this is a dummy implementation+  m_x == m_y++areAnyAliased :: TopdownEnv rep -> VName -> [VName] -> Bool+areAnyAliased td_env m_x =+  any (areAliased td_env m_x)

@@ -228,20 +228,20 @@   if any (bad cse_arrays) $ patElems pat     then m [Let pat' (StmAux cs attrs edec) e']     else case M.lookup (edec, e') esubsts of-      Just subpat ->-        local (addNameSubst pat' subpat) $ do+      Just (subcs, subpat) -> do+        let subsumes = all (`elem` unCerts subcs) (unCerts cs)+        -- We can only do a plain name substitution if it doesn't+        -- violate any certificate dependencies.+        local (if subsumes then addNameSubst pat' subpat else id) $ do           let lets =                 [ Let (Pat [patElem']) (StmAux cs attrs edec) $-                    BasicOp $-                      SubExp $-                        Var $-                          patElemName patElem+                    BasicOp (SubExp $ Var $ patElemName patElem)                   | (name, patElem) <- zip (patNames pat') $ patElems subpat,                     let patElem' = patElem {patElemName = name}                 ]           m lets       _ ->-        local (addExpSubst pat' edec e') $+        local (addExpSubst pat' edec cs e') $           m [Let pat' (StmAux cs attrs edec) e']   where     bad cse_arrays pe@@ -253,7 +253,7 @@ type ExpressionSubstitutions rep =   M.Map     (ExpDec rep, Exp rep)-    (Pat (LetDec rep))+    (Certs, Pat (LetDec rep))  type NameSubstitutions = M.Map VName VName @@ -276,11 +276,12 @@   ASTRep rep =>   Pat (LetDec rep) ->   ExpDec rep ->+  Certs ->   Exp rep ->   CSEState rep ->   CSEState rep-addExpSubst pat edec e (CSEState (esubsts, nsubsts) cse_arrays) =-  CSEState (M.insert (edec, e) pat esubsts, nsubsts) cse_arrays+addExpSubst pat edec cs e (CSEState (esubsts, nsubsts) cse_arrays) =+  CSEState (M.insert (edec, e) (cs, pat) esubsts, nsubsts) cse_arrays  -- | The operations that permit CSE. class CSEInOp op where

@@ -1,5 +1,6 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -fno-warn-redundant-constraints #-}  -- | This module defines the concept of a simplification rule for -- bindings.  The intent is that you pass some context (such as symbol@@ -256,7 +257,7 @@ -- of bindings is returned, that bind at least the same names as the -- original binding (and possibly more, for intermediate results). topDownSimplifyStm ::-  (MonadFreshNames m, HasScope rep m) =>+  (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>   RuleBook rep ->   ST.SymbolTable rep ->   Stm rep ->@@ -269,7 +270,7 @@ -- original binding (and possibly more, for intermediate results). -- The first argument is the set of names used after this binding. bottomUpSimplifyStm ::-  (MonadFreshNames m, HasScope rep m) =>+  (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>   RuleBook rep ->   (ST.SymbolTable rep, UT.UsageTable) ->   Stm rep ->@@ -297,7 +298,7 @@   Skip  applyRules ::-  (MonadFreshNames m, HasScope rep m) =>+  (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>   Rules rep a ->   a ->   Stm rep ->

@@ -62,7 +62,7 @@ -- -- This simplistic rule is only valid before we introduce memory. removeUnnecessaryCopy :: BuilderOps rep => BottomUpRuleBasicOp rep-removeUnnecessaryCopy (vtable, used) (Pat [d]) _ (Copy v)+removeUnnecessaryCopy (vtable, used) (Pat [d]) aux (Copy v)   | not (v `UT.isConsumed` used),     -- This two first clauses below are too conservative, but the     -- problem is that 'v' might not look like it has been consumed if@@ -76,7 +76,7 @@       -- used again.       || (v_is_fresh && v_not_used_again),     (v_not_used_again && consumable) || not (patElemName d `UT.isConsumed` used) =-      Simplify $ letBindNames [patElemName d] $ BasicOp $ SubExp $ Var v+      Simplify $ auxing aux $ letBindNames [patElemName d] $ BasicOp $ SubExp $ Var v   where     v_not_used_again = not (v `UT.used` used)     v_is_fresh = v `ST.lookupAliases` vtable == mempty

@@ -154,9 +154,9 @@     defOf = (`ST.lookupExp` vtable)     seType (Var v) = ST.lookupType v vtable     seType (Constant v) = Just $ Prim $ primValueType v-ruleBasicOp vtable pat _ (Update _ src _ (Var v))+ruleBasicOp vtable pat aux (Update _ src _ (Var v))   | Just (BasicOp Scratch {}, _) <- ST.lookupExp v vtable =-      Simplify $ letBind pat $ BasicOp $ SubExp $ Var src+      Simplify $ auxing aux $ letBind pat $ BasicOp $ SubExp $ Var src -- If we are writing a single-element slice from some array, and the -- element of that array can be computed as a PrimExp based on the -- index, let's just write that instead.@@ -170,10 +170,10 @@           letBind pat $             BasicOp $               Update safety src (Slice [DimFix i]) e'-ruleBasicOp vtable pat _ (Update _ dest destis (Var v))+ruleBasicOp vtable pat aux (Update _ dest destis (Var v))   | Just (e, _) <- ST.lookupExp v vtable,     arrayFrom e =-      Simplify $ letBind pat $ BasicOp $ SubExp $ Var dest+      Simplify $ auxing aux $ letBind pat $ BasicOp $ SubExp $ Var dest   where     arrayFrom (BasicOp (Copy copy_v))       | Just (e', _) <- ST.lookupExp copy_v vtable =@@ -187,9 +187,9 @@           True     arrayFrom _ =       False-ruleBasicOp vtable pat _ (Update _ dest is se)+ruleBasicOp vtable pat aux (Update _ dest is se)   | Just dest_t <- ST.lookupType dest vtable,-    isFullSlice (arrayShape dest_t) is = Simplify $+    isFullSlice (arrayShape dest_t) is = Simplify . auxing aux $       case se of         Var v | not $ null $ sliceDims is -> do           v_reshaped <-

@@ -0,0 +1,120 @@+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE NamedFieldPuns #-}+{-# LANGUAGE TypeFamilies #-}++-- | This pass attempts to lift allocations as far towards the top in their body+-- as possible. It does not try to hoist allocations outside across body+-- boundaries.+module Futhark.Pass.LiftAllocations (liftAllocationsSeqMem, liftAllocationsGPUMem) where++import Control.Monad.Reader+import Data.Sequence (Seq (..))+import Futhark.IR.GPUMem+import Futhark.IR.SeqMem+import Futhark.Pass (Pass (..))++liftAllocationsSeqMem :: Pass SeqMem SeqMem+liftAllocationsSeqMem =+  Pass "lift allocations" "lift allocations" $ \prog@Prog {progFuns} ->+    pure $+      prog+        { progFuns =+            fmap+              ( \f@FunDef {funDefBody} ->+                  f {funDefBody = runReader (liftAllocationsInBody funDefBody) (Env pure)}+              )+              progFuns+        }++liftAllocationsGPUMem :: Pass GPUMem GPUMem+liftAllocationsGPUMem =+  Pass "lift allocations gpu" "lift allocations gpu" $ \prog@Prog {progFuns} ->+    pure $+      prog+        { progFuns =+            fmap+              ( \f@FunDef {funDefBody} ->+                  f {funDefBody = runReader (liftAllocationsInBody funDefBody) (Env liftAllocationsInHostOp)}+              )+              progFuns+        }++newtype Env inner = Env+  {onInner :: inner -> LiftM inner inner}++type LiftM inner a = Reader (Env inner) a++liftAllocationsInBody :: (Mem rep inner, LetDec rep ~ LetDecMem) => Body rep -> LiftM inner (Body rep)+liftAllocationsInBody body = do+  stms <- liftAllocationsInStms (bodyStms body) mempty mempty mempty+  pure $ body {bodyStms = stms}++liftAllocationsInStms ::+  (Mem rep inner, LetDec rep ~ LetDecMem) =>+  -- | The input stms+  Stms rep ->+  -- | The lifted allocations and associated statements+  Stms rep ->+  -- | The other statements processed so far+  Stms rep ->+  -- | Names we need to lift+  Names ->+  LiftM inner (Stms rep)+liftAllocationsInStms Empty lifted acc _ = pure $ lifted <> acc+liftAllocationsInStms (stms :|> stm@(Let (Pat [PatElem vname _]) _ (Op (Alloc _ _)))) lifted acc to_lift =+  liftAllocationsInStms stms (stm :<| lifted) acc ((freeIn stm <> to_lift) `namesSubtract` oneName vname)+liftAllocationsInStms (stms :|> stm@(Let pat _ (Op (Inner inner)))) lifted acc to_lift = do+  on_inner <- asks onInner+  inner' <- on_inner inner+  let stm' = stm {stmExp = Op $ Inner inner'}+      pat_names = namesFromList $ patNames pat+  if pat_names `namesIntersect` to_lift+    then liftAllocationsInStms stms (stm' :<| lifted) acc ((to_lift `namesSubtract` pat_names) <> freeIn stm)+    else liftAllocationsInStms stms lifted (stm' :<| acc) to_lift+liftAllocationsInStms (stms :|> stm@(Let pat aux (Match cond_ses cases body dec))) lifted acc to_lift = do+  cases' <- mapM (\(Case p b) -> Case p <$> liftAllocationsInBody b) cases+  body' <- liftAllocationsInBody body+  let stm' = stm {stmExp = Match cond_ses cases' body' dec}+      pat_names = namesFromList $ patNames pat+  if pat_names `namesIntersect` to_lift+    then+      liftAllocationsInStms+        stms+        (stm' :<| lifted)+        acc+        ( (to_lift `namesSubtract` pat_names)+            <> freeIn cond_ses+            <> freeIn cases+            <> freeIn body+            <> freeIn dec+            <> freeIn aux+        )+    else liftAllocationsInStms stms lifted (stm' :<| acc) to_lift+liftAllocationsInStms (stms :|> stm@(Let pat _ (DoLoop params form body))) lifted acc to_lift = do+  body' <- liftAllocationsInBody body+  let stm' = stm {stmExp = DoLoop params form body'}+      pat_names = namesFromList $ patNames pat+  if pat_names `namesIntersect` to_lift+    then liftAllocationsInStms stms (stm' :<| lifted) acc ((to_lift `namesSubtract` pat_names) <> freeIn stm)+    else liftAllocationsInStms stms lifted (stm' :<| acc) to_lift+liftAllocationsInStms (stms :|> stm@(Let pat _ _)) lifted acc to_lift = do+  let pat_names = namesFromList (patNames pat)+  if pat_names `namesIntersect` to_lift+    then liftAllocationsInStms stms (stm :<| lifted) acc ((to_lift `namesSubtract` pat_names) <> freeIn stm)+    else liftAllocationsInStms stms lifted (stm :<| acc) to_lift++liftAllocationsInHostOp :: HostOp GPUMem () -> LiftM (HostOp GPUMem ()) (HostOp GPUMem ())+liftAllocationsInHostOp (SegOp (SegMap lvl sp tps body)) = do+  stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty+  pure $ SegOp $ SegMap lvl sp tps $ body {kernelBodyStms = stms}+liftAllocationsInHostOp (SegOp (SegRed lvl sp binops tps body)) = do+  stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty+  pure $ SegOp $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}+liftAllocationsInHostOp (SegOp (SegScan lvl sp binops tps body)) = do+  stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty+  pure $ SegOp $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}+liftAllocationsInHostOp (SegOp (SegHist lvl sp histops tps body)) = do+  stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty+  pure $ SegOp $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}+liftAllocationsInHostOp op = pure op

@@ -0,0 +1,114 @@+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE NamedFieldPuns #-}+{-# LANGUAGE TypeFamilies #-}++-- | This pass attempts to lower allocations as far towards the bottom of their+-- body as possible.+module Futhark.Pass.LowerAllocations (lowerAllocationsSeqMem, lowerAllocationsGPUMem) where++import Control.Monad.Reader+import Data.Function ((&))+import Data.Map qualified as M+import Data.Sequence (Seq (..))+import Data.Sequence qualified as Seq+import Futhark.IR.GPUMem+import Futhark.IR.SeqMem+import Futhark.Pass (Pass (..))++lowerAllocationsSeqMem :: Pass SeqMem SeqMem+lowerAllocationsSeqMem =+  Pass "lower allocations" "lower allocations" $ \prog@Prog {progFuns} ->+    pure $+      prog+        { progFuns =+            fmap+              ( \f@FunDef {funDefBody} ->+                  f {funDefBody = runReader (lowerAllocationsInBody funDefBody) (Env pure)}+              )+              progFuns+        }++lowerAllocationsGPUMem :: Pass GPUMem GPUMem+lowerAllocationsGPUMem =+  Pass "lower allocations gpu" "lower allocations gpu" $ \prog@Prog {progFuns} ->+    pure $+      prog+        { progFuns =+            fmap+              ( \f@FunDef {funDefBody} ->+                  f {funDefBody = runReader (lowerAllocationsInBody funDefBody) (Env lowerAllocationsInHostOp)}+              )+              progFuns+        }++newtype Env inner = Env+  {onInner :: inner -> LowerM inner inner}++type LowerM inner a = Reader (Env inner) a++lowerAllocationsInBody :: (Mem rep inner, LetDec rep ~ LetDecMem) => Body rep -> LowerM inner (Body rep)+lowerAllocationsInBody body = do+  stms <- lowerAllocationsInStms (bodyStms body) mempty mempty+  pure $ body {bodyStms = stms}++lowerAllocationsInStms ::+  (Mem rep inner, LetDec rep ~ LetDecMem) =>+  -- | The input stms+  Stms rep ->+  -- | The allocations currently being lowered+  M.Map VName (Stm rep) ->+  -- | The other statements processed so far+  Stms rep ->+  LowerM inner (Stms rep)+lowerAllocationsInStms Empty allocs acc = pure $ acc <> Seq.fromList (M.elems allocs)+lowerAllocationsInStms (stm@(Let (Pat [PatElem vname _]) _ (Op (Alloc _ _))) :<| stms) allocs acc =+  lowerAllocationsInStms stms (M.insert vname stm allocs) acc+lowerAllocationsInStms (stm0@(Let _ _ (Op (Inner inner))) :<| stms) alloc0 acc0 = do+  on_inner <- asks onInner+  inner' <- on_inner inner+  let stm = stm0 {stmExp = Op $ Inner inner'}+      (alloc, acc) = insertLoweredAllocs (freeIn stm0) alloc0 acc0+  lowerAllocationsInStms stms alloc (acc :|> stm)+lowerAllocationsInStms (stm@(Let _ _ (Match cond_ses cases body dec)) :<| stms) alloc acc = do+  cases' <- mapM (\(Case pat b) -> Case pat <$> lowerAllocationsInBody b) cases+  body' <- lowerAllocationsInBody body+  let stm' = stm {stmExp = Match cond_ses cases' body' dec}+      (alloc', acc') = insertLoweredAllocs (freeIn stm) alloc acc+  lowerAllocationsInStms stms alloc' (acc' :|> stm')+lowerAllocationsInStms (stm@(Let _ _ (DoLoop params form body)) :<| stms) alloc acc = do+  body' <- lowerAllocationsInBody body+  let stm' = stm {stmExp = DoLoop params form body'}+      (alloc', acc') = insertLoweredAllocs (freeIn stm) alloc acc+  lowerAllocationsInStms stms alloc' (acc' :|> stm')+lowerAllocationsInStms (stm :<| stms) alloc acc = do+  let (alloc', acc') = insertLoweredAllocs (freeIn stm) alloc acc+  lowerAllocationsInStms stms alloc' (acc' :|> stm)++insertLoweredAllocs :: Names -> M.Map VName (Stm rep) -> Stms rep -> (M.Map VName (Stm rep), Stms rep)+insertLoweredAllocs frees alloc acc =+  frees+    `namesIntersection` namesFromList (M.keys alloc)+    & namesToList+    & foldl+      ( \(alloc', acc') name ->+          ( M.delete name alloc',+            acc' :|> alloc' M.! name+          )+      )+      (alloc, acc)++lowerAllocationsInHostOp :: HostOp GPUMem () -> LowerM (HostOp GPUMem ()) (HostOp GPUMem ())+lowerAllocationsInHostOp (SegOp (SegMap lvl sp tps body)) = do+  stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty+  pure $ SegOp $ SegMap lvl sp tps $ body {kernelBodyStms = stms}+lowerAllocationsInHostOp (SegOp (SegRed lvl sp binops tps body)) = do+  stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty+  pure $ SegOp $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}+lowerAllocationsInHostOp (SegOp (SegScan lvl sp binops tps body)) = do+  stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty+  pure $ SegOp $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}+lowerAllocationsInHostOp (SegOp (SegHist lvl sp histops tps body)) = do+  stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty+  pure $ SegOp $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}+lowerAllocationsInHostOp op = pure op

@@ -18,6 +18,7 @@ import Futhark.IR.SOACS (SOACS, usesAD) import Futhark.IR.Seq (Seq) import Futhark.IR.SeqMem (SeqMem)+import Futhark.Optimise.ArrayShortCircuiting qualified as ArrayShortCircuiting import Futhark.Optimise.CSE import Futhark.Optimise.DoubleBuffer import Futhark.Optimise.EntryPointMem@@ -41,6 +42,8 @@ import Futhark.Pass.ExtractMulticore import Futhark.Pass.FirstOrderTransform import Futhark.Pass.KernelBabysitting+import Futhark.Pass.LiftAllocations as LiftAllocations+import Futhark.Pass.LowerAllocations as LowerAllocations import Futhark.Pass.Simplify import Futhark.Pipeline @@ -126,6 +129,14 @@       [ performCSE False,         simplifySeqMem,         entryPointMemSeq,+        simplifySeqMem,+        LiftAllocations.liftAllocationsSeqMem,+        simplifySeqMem,+        ArrayShortCircuiting.optimiseSeqMem,+        simplifySeqMem,+        performCSE False,+        simplifySeqMem,+        LowerAllocations.lowerAllocationsSeqMem,         simplifySeqMem       ] @@ -141,6 +152,16 @@         simplifyGPUMem,         entryPointMemGPU,         doubleBufferGPU,+        simplifyGPUMem,+        performCSE False,+        LiftAllocations.liftAllocationsGPUMem,+        simplifyGPUMem,+        ArrayShortCircuiting.optimiseGPUMem,+        simplifyGPUMem,+        performCSE False,+        simplifyGPUMem,+        LowerAllocations.lowerAllocationsGPUMem,+        performCSE False,         simplifyGPUMem,         MemoryBlockMerging.optimise,         simplifyGPUMem,

@@ -321,6 +321,11 @@         cmdMaybe $ cmdNew server v t vs         pure v +      getField from (f, _) = do+        to <- newVar "field"+        cmdMaybe $ cmdProject server to from $ nameToText f+        pure to+       toVal :: ValOrVar -> m V.Value       toVal (VVal v) = pure v       toVal (VVar v) = readVar server v@@ -347,7 +352,9 @@       interValToVal = traverse scriptValueToVal        -- Apart from type checking, this function also converts-      -- FutharkScript tuples/records to Futhark-level tuples/records.+      -- FutharkScript tuples/records to Futhark-level tuples/records,+      -- as well as maps between different names for the same+      -- tuple/record.       interValToVar :: m VarName -> TypeName -> ExpValue -> m VarName       interValToVar _ t (V.ValueAtom v)         | STValue t == scriptValueType v = scriptValueToVar v@@ -359,6 +366,11 @@         | Just fs <- isRecord t types,           Just vs' <- mapM ((`M.lookup` vs) . nameToText . fst) fs =             mkRecord t =<< zipWithM (interValToVar bad) (map snd fs) vs'+      interValToVar _ t (V.ValueAtom (SValue vt (VVar v)))+        | Just t_fs <- isRecord t types,+          Just vt_fs <- isRecord vt types,+          vt_fs == t_fs =+            mkRecord t =<< mapM (getField v) vt_fs       interValToVar bad _ _ = bad        valToInterVal :: V.CompoundValue -> ExpValue

@@ -19,11 +19,13 @@     takeLast,     dropLast,     mapEither,+    partitionMaybe,     maybeNth,     maybeHead,     splitFromEnd,     splitAt3,     focusNth,+    focusMaybe,     hashText,     unixEnvironment,     isEnvVarAtLeast,@@ -46,6 +48,7 @@     cartesian,     traverseFold,     fixPoint,+    concatMapM,   ) where @@ -60,7 +63,7 @@ import Data.Either import Data.Foldable (fold, toList) import Data.Function ((&))-import Data.List (foldl', genericDrop, genericSplitAt, sortBy)+import Data.List (findIndex, foldl', genericDrop, genericSplitAt, sortBy) import Data.List.NonEmpty qualified as NE import Data.Map qualified as M import Data.Maybe@@ -155,6 +158,16 @@ mapEither :: (a -> Either b c) -> [a] -> ([b], [c]) mapEither f l = partitionEithers $ map f l +-- | A combination of 'partition' and 'mapMaybe'+partitionMaybe :: (a -> Maybe b) -> [a] -> ([b], [a])+partitionMaybe f = helper ([], [])+  where+    helper (acc1, acc2) [] = (reverse acc1, reverse acc2)+    helper (acc1, acc2) (x : xs) =+      case f x of+        Just x' -> helper (x' : acc1, acc2) xs+        Nothing -> helper (acc1, x : acc2) xs+ -- | Return the list element at the given index, if the index is valid. maybeNth :: Integral int => int -> [a] -> Maybe a maybeNth i l@@ -184,8 +197,17 @@   | (bef, x : aft) <- genericSplitAt i xs = Just (bef, x, aft)   | otherwise = Nothing --- | Compute a hash of a pretty that is stable across OS versions.--- Returns the hash as a pretty as well, ready for human consumption.+-- | Return the first list element that satisifes a predicate, along with the+-- elements before and after.+focusMaybe :: (a -> Maybe b) -> [a] -> Maybe ([a], b, [a])+focusMaybe f xs = do+  idx <- findIndex (isJust . f) xs+  (before, focus, after) <- focusNth idx xs+  res <- f focus+  pure (before, res, after)++-- | Compute a hash of a text that is stable across OS versions.+-- Returns the hash as a text as well, ready for human consumption. hashText :: T.Text -> T.Text hashText =   T.decodeUtf8With T.lenientDecode . Base16.encode . MD5.hash . T.encodeUtf8@@ -429,3 +451,6 @@ fixPoint f x =   let x' = f x    in if x' == x then x else fixPoint f x'++concatMapM :: (Monad m, Monoid b) => (a -> m b) -> [a] -> m b+concatMapM f xs = mconcat <$> mapM f xs

@@ -44,6 +44,10 @@ newtype Wrapped a = Wrapped {wrappedValue :: a}   deriving (Eq, Ord, Show) +instance Enum a => Enum (Wrapped a) where+  toEnum a = Wrapped $ toEnum a+  fromEnum (Wrapped a) = fromEnum a+ liftOp ::   (a -> a) ->   Wrapped a ->

@@ -28,8 +28,13 @@ versionString :: String versionString =   showVersion version+    ++ unreleased     ++ gitversion $$tGitInfoCwdTry   where+    unreleased =+      if last (versionBranch version) == 0+        then " (prerelease - include info below when reporting bugs)"+        else mempty     gitversion (Left _) =       case commitIdFromFile of         Nothing -> ""

@@ -300,28 +300,24 @@ instance (Eq vn, IsName vn, Annot f) => Pretty (AppExpBase f vn) where   pretty = prettyAppExp (-1) +prettyInst :: Annot f => f PatType -> Doc a+prettyInst t =+  case unAnnot t of+    Just t'+      | isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 2 ->+          "@" <> parens (align $ pretty t')+    _ -> mempty+ prettyExp :: (Eq vn, IsName vn, Annot f) => Int -> ExpBase f vn -> Doc a-prettyExp _ (Var name t _) = pretty name <> inst-  where-    inst = case unAnnot t of-      Just t'-        | isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 2 ->-            "@" <> parens (align $ pretty t')-      _ -> mempty-prettyExp _ (Hole t _) = "???" <> inst-  where-    inst = case unAnnot t of-      Just t'-        | isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 2 ->-            "@" <> parens (align $ pretty t')-      _ -> mempty+prettyExp _ (Var name t _) = pretty name <> prettyInst t+prettyExp _ (Hole t _) = "???" <> prettyInst t prettyExp _ (Parens e _) = align $ parens $ pretty e prettyExp _ (QualParens (v, _) e _) = pretty v <> "." <> align (parens $ pretty e) prettyExp p (Ascript e t _) =   parensIf (p /= -1) $ prettyExp 0 e <+> ":" <+> align (pretty t) prettyExp _ (Literal v _) = pretty v-prettyExp _ (IntLit v _ _) = pretty v-prettyExp _ (FloatLit v _ _) = pretty v+prettyExp _ (IntLit v t _) = pretty v <> prettyInst t+prettyExp _ (FloatLit v t _) = pretty v <> prettyInst t prettyExp _ (TupLit es _)   | any hasArrayLit es = parens $ commastack $ map pretty es   | otherwise = parens $ commasep $ map pretty es@@ -331,14 +327,8 @@   where     fieldArray (RecordFieldExplicit _ e _) = hasArrayLit e     fieldArray RecordFieldImplicit {} = False-prettyExp _ (ArrayLit es info _) =-  brackets (commasep $ map pretty es) <> info'-  where-    info' = case unAnnot info of-      Just t-        | isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 2 ->-            "@" <> parens (align $ pretty t)-      _ -> mempty+prettyExp _ (ArrayLit es t _) =+  brackets (commasep $ map pretty es) <> prettyInst t prettyExp _ (StringLit s _) =   pretty $ show $ map (chr . fromIntegral) s prettyExp _ (Project k e _ _) = pretty e <> "." <> pretty k

@@ -113,6 +113,7 @@ import Language.Futhark.Syntax import Language.Futhark.Traversals import Language.Futhark.Tuple+import System.FilePath (takeDirectory)  -- | The name of the default program entry point (@main@). defaultEntryPoint :: Name@@ -1172,9 +1173,9 @@     tupInt64 x =       tupleRecord $ replicate x $ Scalar $ Prim $ Signed Int64 --- | Is this file part of the built-in prelude?+-- | Is this include part of the built-in prelude? isBuiltin :: FilePath -> Bool-isBuiltin = ("/prelude/" `isPrefixOf`)+isBuiltin = (== "/prelude") . takeDirectory  -- | Is the position of this thing builtin as per 'isBuiltin'?  Things -- without location are considered not built-in.

@@ -961,55 +961,6 @@       | prev_applied == 1 = "argument"       | otherwise = "arguments" --- | @returnType appres ret_type arg_diet arg_type@ gives result of applying--- an argument the given types to a function with the given return--- type, consuming the argument with the given diet.-returnType ::-  Aliasing ->-  PatType ->-  Diet ->-  PatType ->-  PatType-returnType _ (Array _ Unique et shape) _ _ =-  Array mempty Nonunique et shape -- Intentional!-returnType appres (Array als Nonunique et shape) d arg =-  Array (appres <> als <> arg_als) Nonunique et shape-  where-    arg_als = aliases $ maskAliases arg d-returnType appres (Scalar (Record fs)) d arg =-  Scalar $ Record $ fmap (\et -> returnType appres et d arg) fs-returnType _ (Scalar (Prim t)) _ _ =-  Scalar $ Prim t-returnType _ (Scalar (TypeVar _ Unique t targs)) _ _ =-  Scalar $ TypeVar mempty Nonunique t targs -- Intentional!-returnType appres (Scalar (TypeVar als Nonunique t targs)) d arg =-  Scalar $ TypeVar (appres <> als <> arg_als) Unique t targs-  where-    arg_als = aliases $ maskAliases arg d-returnType _ (Scalar (Arrow old_als v t1 (RetType dims t2))) d arg =-  Scalar $ Arrow als v (t1 `setAliases` mempty) $ RetType dims $ t2 `setAliases` als-  where-    -- Make sure to propagate the aliases of an existing closure.-    als = old_als <> aliases (maskAliases arg d)-returnType appres (Scalar (Sum cs)) d arg =-  Scalar $ Sum $ (fmap . fmap) (\et -> returnType appres et d arg) cs---- @t `maskAliases` d@ removes aliases (sets them to 'mempty') from--- the parts of @t@ that are denoted as consumed by the 'Diet' @d@.-maskAliases ::-  Monoid as =>-  TypeBase shape as ->-  Diet ->-  TypeBase shape as-maskAliases t Consume = t `setAliases` mempty-maskAliases t Observe = t-maskAliases (Scalar (Record ets)) (RecordDiet ds) =-  Scalar $ Record $ M.intersectionWith maskAliases ets ds-maskAliases (Scalar (Sum ets)) (SumDiet ds) =-  Scalar $ Sum $ M.intersectionWith (zipWith maskAliases) ets ds-maskAliases t FuncDiet {} = t-maskAliases _ _ = error "Invalid arguments passed to maskAliases."- consumedByArg :: SrcLoc -> PatType -> Diet -> TermTypeM [Aliasing] consumedByArg loc (Scalar (Record ets)) (RecordDiet ds) =   mconcat . M.elems <$> traverse (uncurry $ consumedByArg loc) (M.intersectionWith (,) ets ds)

@@ -113,7 +113,7 @@             S.toList $               S.map aliasVar (aliases t) `S.intersection` bound_outside =             lift . typeError loop_loc mempty $-              "Return value for loop parameter"+              "Return value for consuming loop parameter"                 <+> dquotes (prettyName pat_v)                 <+> "aliases"                 <+> dquotes (prettyName v) <> "."@@ -189,6 +189,11 @@ type CheckedLoop =   ([VName], Pat, Exp, LoopFormBase Info VName, Exp) +loopReturnType :: Pat -> PatType -> PatType+loopReturnType pat = returnType mempty pat_t (diet pat_t)+  where+    pat_t = patternType pat+ -- | Type-check a @loop@ expression, passing in a function for -- type-checking subexpressions. checkDoLoop ::@@ -227,9 +232,13 @@     -- (There is also a convergence loop for inferring uniqueness, but     -- that's orthogonal to the size handling.) +    -- We don't want the loop parameters to alias their initial+    -- values, so we blank them here.  We will actually check them+    -- properly later.     (merge_t, new_dims_to_initial_dim) <-       -- dim handling (1)-      allDimsFreshInType loc Nonrigid "loop" =<< expTypeFully mergeexp'+      allDimsFreshInType loc Nonrigid "loop" . flip setAliases mempty+        =<< expTypeFully mergeexp'     let new_dims = M.keys new_dims_to_initial_dim      -- dim handling (2)@@ -309,7 +318,7 @@           bound_t <- expTypeFully uboundexp'           bindingIdent i bound_t $ \i' ->             noUnique . bindingPat [] mergepat (Ascribed merge_t) $-              \mergepat' -> onlySelfAliasing . tapOccurrences $ do+              \mergepat' -> tapOccurrences $ do                 loopbody' <- noSizeEscape $ checkExp loopbody                 (sparams, mergepat'') <- checkLoopReturnSize mergepat' loopbody'                 pure@@ -327,7 +336,7 @@               | Just t' <- peelArray 1 t ->                   bindingPat [] xpat (Ascribed t') $ \xpat' ->                     noUnique . bindingPat [] mergepat (Ascribed merge_t) $-                      \mergepat' -> onlySelfAliasing . tapOccurrences $ do+                      \mergepat' -> tapOccurrences $ do                         loopbody' <- noSizeEscape $ checkExp loopbody                         (sparams, mergepat'') <- checkLoopReturnSize mergepat' loopbody'                         pure@@ -342,8 +351,7 @@                       <+> pretty t         While cond ->           noUnique . bindingPat [] mergepat (Ascribed merge_t) $ \mergepat' ->-            onlySelfAliasing-              . tapOccurrences+            tapOccurrences               . sequentially                 ( checkExp cond                     >>= unifies "being the condition of a 'while' loop" (Scalar $ Prim Bool)@@ -363,6 +371,7 @@       convergePat loc mergepat' (allConsumed bodyflow) loopbody_t $         mkUsage (srclocOf loopbody') "being (part of) the result of the loop body" +    merge_t' <- expTypeFully mergeexp'     let consumeMerge (Id _ (Info pt) ploc) mt           | unique pt = consume ploc $ aliases mt         consumeMerge (TuplePat pats _) t@@ -374,14 +383,14 @@           consumeMerge pat t         consumeMerge _ _ =           pure ()-    consumeMerge mergepat'' =<< expTypeFully mergeexp'+    consumeMerge mergepat'' merge_t'      -- dim handling (4)     wellTypedLoopArg Initial sparams mergepat'' mergeexp'      (loopt, retext) <-       freshDimsInType loc (Rigid RigidLoop) "loop" (S.fromList sparams) $-        patternType mergepat''+        loopReturnType mergepat'' merge_t'     -- We set all of the uniqueness to be unique.  This is intentional,     -- and matches what happens for function calls.  Those arrays that     -- really *cannot* be consumed will alias something unconsumable,@@ -396,4 +405,7 @@           second (`S.difference` S.map AliasBound bound_here) $             loopt `setUniqueness` Unique -    pure ((sparams, mergepat'', mergeexp', form', loopbody'), AppRes loopt' retext)+    pure+      ( (sparams, mergepat'', mergeexp', form', loopbody'),+        AppRes loopt' retext+      )

@@ -52,7 +52,6 @@     Names,     Occurrence (..),     Occurrences,-    onlySelfAliasing,     noUnique,     removeSeminullOccurrences,     occur,@@ -67,7 +66,6 @@     allConsumed,      -- * Errors-    useAfterConsume,     unusedSize,     uniqueReturnAliased,     returnAliased,@@ -574,11 +572,11 @@   case tparam of     TypeParamType x _ _ -> do       constrain v . NoConstraint x . mkUsage loc . docText $-        "instantiated type parameter of " <> dquotes (pretty qn) <> "."+        "instantiated type parameter of " <> dquotes (pretty qn)       pure (v, Subst [] $ RetType [] $ Scalar $ TypeVar mempty Nonunique (qualName v) [])     TypeParamDim {} -> do       constrain v . Size Nothing . mkUsage loc . docText $-        "instantiated size parameter of " <> dquotes (pretty qn) <> "."+        "instantiated size parameter of " <> dquotes (pretty qn)       pure (v, SizeSubst $ NamedSize $ qualName v)  checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)@@ -931,16 +929,6 @@ observe (Ident nm (Info t) loc) =   let als = AliasBound nm `S.insert` aliases t    in occur [observation als loc]--onlySelfAliasing :: TermTypeM a -> TermTypeM a-onlySelfAliasing = localScope (\scope -> scope {scopeVtable = M.mapWithKey set $ scopeVtable scope})-  where-    set k (BoundV l tparams t) =-      BoundV l tparams $-        t `addAliases` S.intersection (S.singleton (AliasBound k))-    set _ (OverloadedF ts pts rt) = OverloadedF ts pts rt-    set _ EqualityF = EqualityF-    set _ (WasConsumed loc) = WasConsumed loc  -- | Enter a context where nothing outside can be consumed (i.e. the -- body of a function definition).

@@ -337,7 +337,7 @@        outer_t' <- normTypeFully outer_t       PatAscription-        <$> checkPat' sizes p (Ascribed (addAliasesFromType st outer_t'))+        <$> checkPat' sizes p (Ascribed (addAliasesFromType (fromStruct st) outer_t'))         <*> pure t'         <*> pure loc     NoneInferred ->

@@ -4,6 +4,7 @@     renameRetType,     subtypeOf,     subuniqueOf,+    returnType,     addAliasesFromType,     checkForDuplicateNames,     checkTypeParams,@@ -81,24 +82,67 @@ mustBeExplicitInType :: StructType -> S.Set VName mustBeExplicitInType = snd . determineSizeWitnesses +-- | @returnType appres ret_type arg_diet arg_type@ gives result of applying+-- an argument the given types to a function with the given return+-- type, consuming the argument with the given diet.+returnType :: Aliasing -> PatType -> Diet -> PatType -> PatType+returnType _ (Array _ Unique et shape) _ _ =+  Array mempty Nonunique et shape -- Intentional!+returnType appres (Array als Nonunique et shape) d arg =+  Array (appres <> als <> arg_als) Nonunique et shape+  where+    arg_als = aliases $ maskAliases arg d+returnType appres (Scalar (Record fs)) d arg =+  Scalar $ Record $ fmap (\et -> returnType appres et d arg) fs+returnType _ (Scalar (Prim t)) _ _ =+  Scalar $ Prim t+returnType _ (Scalar (TypeVar _ Unique t targs)) _ _ =+  Scalar $ TypeVar mempty Nonunique t targs -- Intentional!+returnType appres (Scalar (TypeVar als Nonunique t targs)) d arg =+  Scalar $ TypeVar (appres <> als <> arg_als) Unique t targs+  where+    arg_als = aliases $ maskAliases arg d+returnType _ (Scalar (Arrow old_als v t1 (RetType dims t2))) d arg =+  Scalar $ Arrow als v (t1 `setAliases` mempty) $ RetType dims $ t2 `setAliases` als+  where+    -- Make sure to propagate the aliases of an existing closure.+    als = old_als <> aliases (maskAliases arg d)+returnType appres (Scalar (Sum cs)) d arg =+  Scalar $ Sum $ (fmap . fmap) (\et -> returnType appres et d arg) cs++-- @t `maskAliases` d@ removes aliases (sets them to 'mempty') from+-- the parts of @t@ that are denoted as consumed by the 'Diet' @d@.+maskAliases ::+  Monoid as =>+  TypeBase shape as ->+  Diet ->+  TypeBase shape as+maskAliases t Consume = t `setAliases` mempty+maskAliases t Observe = t+maskAliases (Scalar (Record ets)) (RecordDiet ds) =+  Scalar $ Record $ M.intersectionWith maskAliases ets ds+maskAliases (Scalar (Sum ets)) (SumDiet ds) =+  Scalar $ Sum $ M.intersectionWith (zipWith maskAliases) ets ds+maskAliases t FuncDiet {} = t+maskAliases _ _ = error "Invalid arguments passed to maskAliases."+ -- | The two types are assumed to be structurally equal, but not--- necessarily regarding sizes.  Adds aliases from the latter to the--- former.-addAliasesFromType :: StructType -> PatType -> PatType-addAliasesFromType (Array _ u1 et1 shape1) (Array als _ _ _) =-  Array als u1 et1 shape1+-- necessarily regarding sizes.  Combines aliases.+addAliasesFromType :: PatType -> PatType -> PatType+addAliasesFromType (Array als1 u1 et1 shape1) (Array als2 _ _ _) =+  Array (als1 <> als2) u1 et1 shape1 addAliasesFromType-  (Scalar (TypeVar _ u1 tv1 targs1))+  (Scalar (TypeVar als1 u1 tv1 targs1))   (Scalar (TypeVar als2 _ _ _)) =-    Scalar $ TypeVar als2 u1 tv1 targs1+    Scalar $ TypeVar (als1 <> als2) u1 tv1 targs1 addAliasesFromType (Scalar (Record ts1)) (Scalar (Record ts2))   | length ts1 == length ts2,     sort (M.keys ts1) == sort (M.keys ts2) =       Scalar $ Record $ M.intersectionWith addAliasesFromType ts1 ts2 addAliasesFromType-  (Scalar (Arrow _ mn1 pt1 (RetType dims1 rt1)))-  (Scalar (Arrow as2 _ _ (RetType _ rt2))) =-    Scalar (Arrow as2 mn1 pt1 (RetType dims1 rt1'))+  (Scalar (Arrow als1 mn1 pt1 (RetType dims1 rt1)))+  (Scalar (Arrow als2 _ _ (RetType _ rt2))) =+    Scalar (Arrow (als1 <> als2) mn1 pt1 (RetType dims1 rt1'))     where       rt1' = addAliasesFromType rt1 rt2 addAliasesFromType (Scalar (Sum cs1)) (Scalar (Sum cs2))

@@ -268,19 +268,21 @@     . freeInType  typeVarNotes :: MonadUnify m => VName -> m Notes-typeVarNotes v = maybe mempty (aNote . note . snd) . M.lookup v <$> getConstraints+typeVarNotes v = maybe mempty (note . snd) . M.lookup v <$> getConstraints   where     note (HasConstrs cs _) =-      prettyName v-        <+> "="-        <+> mconcat (map ppConstr (M.toList cs))-        <+> "..."+      aNote $+        prettyName v+          <+> "="+          <+> mconcat (map ppConstr (M.toList cs))+          <+> "..."     note (Overloaded ts _) =-      prettyName v <+> "must be one of" <+> mconcat (punctuate ", " (map pretty ts))+      aNote $ prettyName v <+> "must be one of" <+> mconcat (punctuate ", " (map pretty ts))     note (HasFields fs _) =-      prettyName v-        <+> "="-        <+> braces (mconcat (punctuate ", " (map ppField (M.toList fs))))+      aNote $+        prettyName v+          <+> "="+          <+> braces (mconcat (punctuate ", " (map ppField (M.toList fs))))     note _ = mempty      ppConstr (c, _) = "#" <> pretty c <+> "..." <+> "|"@@ -442,9 +444,7 @@           Scalar (Record arg_fs)           )             | M.keys fs == M.keys arg_fs ->-                forM_ (M.toList $ M.intersectionWith (,) fs arg_fs) $ \(k, (k_t1, k_t2)) -> do-                  let bcs' = breadCrumb (MatchingFields [k]) bcs-                  subunify ord bound bcs' k_t1 k_t2+                unifySharedFields onDims usage bound bcs fs arg_fs             | otherwise -> do                 let missing =                       filter (`notElem` M.keys arg_fs) (M.keys fs)@@ -723,30 +723,34 @@                   </> "due to"                   <+> pretty old_usage <> "."     Just (HasFields required_fields old_usage) -> do-      link       case tp of         Scalar (Record tp_fields)           | all (`M.member` tp_fields) $ M.keys required_fields -> do               required_fields' <- mapM normTypeFully required_fields-              let bcs' =-                    breadCrumb-                      ( Matching $-                          prettyName vn-                            <+> "must be a record with at least the fields:"-                            </> indent 2 (pretty (Record required_fields'))-                            </> "due to"-                            <+> pretty old_usage <> "."-                      )-                      bcs-              mapM_ (uncurry $ unifyWith onDims usage bound bcs') $-                M.elems $-                  M.intersectionWith (,) required_fields tp_fields-        Scalar (TypeVar _ _ (QualName [] v) [])-          | not $ isRigid v constraints ->+              let tp' = Scalar $ Record $ required_fields <> tp_fields -- Crucially left-biased.+                  ext = filter (`S.member` freeInType tp') bound               modifyConstraints $-                M.insert-                  v-                  (lvl, HasFields required_fields old_usage)+                M.insert vn (lvl, Constraint (RetType ext tp') usage)+              unifySharedFields onDims usage bound bcs required_fields' tp_fields+        Scalar (TypeVar _ _ (QualName [] v) []) -> do+          case M.lookup v constraints of+            Just (_, HasFields tp_fields _) ->+              unifySharedFields onDims usage bound bcs required_fields tp_fields+            Just (_, NoConstraint {}) -> pure ()+            Just (_, Equality {}) -> pure ()+            _ -> do+              notes <- (<>) <$> typeVarNotes vn <*> typeVarNotes v+              noRecordType notes+          link+          modifyConstraints $+            M.insertWith+              combineFields+              v+              (lvl, HasFields required_fields old_usage)+          where+            combineFields (_, HasFields fs1 usage1) (_, HasFields fs2 _) =+              (lvl, HasFields (M.union fs1 fs2) usage1)+            combineFields hasfs _ = hasfs         _ ->           unifyError usage mempty bcs $             "Cannot instantiate"@@ -771,7 +775,7 @@               unifySharedConstructors onDims usage bound bcs required_cs ts         Scalar (TypeVar _ _ (QualName [] v) []) -> do           case M.lookup v constraints of-            Just (_, HasConstrs v_cs _) -> do+            Just (_, HasConstrs v_cs _) ->               unifySharedConstructors onDims usage bound bcs required_cs v_cs             Just (_, NoConstraint {}) -> pure ()             Just (_, Equality {}) -> pure ()@@ -797,6 +801,12 @@         notes         bcs         "Cannot unify a sum type with a non-sum type"+    noRecordType notes =+      unifyError+        usage+        notes+        bcs+        "Cannot unify a record type with a non-record type"  linkVarToDim ::   MonadUnify m =>@@ -999,6 +1009,19 @@   arrayElemTypeWith usage $ breadCrumb bc noBreadCrumbs   where     bc = Matching $ "When checking" <+> textwrap desc++unifySharedFields ::+  MonadUnify m =>+  UnifyDims m ->+  Usage ->+  [VName] ->+  BreadCrumbs ->+  M.Map Name StructType ->+  M.Map Name StructType ->+  m ()+unifySharedFields onDims usage bound bcs fs1 fs2 =+  forM_ (M.toList $ M.intersectionWith (,) fs1 fs2) $ \(f, (t1, t2)) ->+    unifyWith onDims usage bound (breadCrumb (MatchingFields [f]) bcs) t1 t2  unifySharedConstructors ::   MonadUnify m =>

@@ -0,0 +1,90 @@+{-# LANGUAGE FlexibleInstances #-}+{-# OPTIONS_GHC -fno-warn-orphans #-}+{-# OPTIONS_GHC -fno-warn-unused-imports #-}+{-# OPTIONS_GHC -fno-warn-unused-matches #-}+{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}++module Futhark.Analysis.AlgSimplifyTests+  ( tests,+  )+where++import Control.Monad+import Data.Function ((&))+import Data.List (subsequences)+import Data.Map qualified as M+import Data.Maybe (fromMaybe, mapMaybe)+import Futhark.Analysis.AlgSimplify hiding (add, sub)+import Futhark.Analysis.PrimExp+import Futhark.IR.Syntax.Core+import Test.Tasty+import Test.Tasty.HUnit+import Test.Tasty.QuickCheck++tests :: TestTree+tests =+  testGroup+    "AlgSimplifyTests"+    [ testProperty "simplify is idempotent" $ \(TestableExp e) -> simplify e == simplify (simplify e),+      testProperty "simplify doesn't change exp evalutation result" $+        \(TestableExp e) ->+          evalPrimExp (\_ -> Nothing) e+            == evalPrimExp (\_ -> Nothing) (simplify e)+    ]++eval :: TestableExp -> Int64+eval (TestableExp e) = evalExp e++evalExp :: PrimExp VName -> Int64+evalExp (ValueExp (IntValue (Int64Value i))) = i+evalExp (BinOpExp (Add Int64 OverflowUndef) e1 e2) = evalExp e1 + evalExp e2+evalExp (BinOpExp (Sub Int64 OverflowUndef) e1 e2) = evalExp e1 - evalExp e2+evalExp (BinOpExp (Mul Int64 OverflowUndef) e1 e2) = evalExp e1 * evalExp e2+evalExp _ = undefined++add :: PrimExp VName -> PrimExp VName -> PrimExp VName+add = BinOpExp (Add Int64 OverflowUndef)++sub :: PrimExp VName -> PrimExp VName -> PrimExp VName+sub = BinOpExp (Sub Int64 OverflowUndef)++mul :: PrimExp VName -> PrimExp VName -> PrimExp VName+mul = BinOpExp (Mul Int64 OverflowUndef)++neg :: PrimExp VName -> PrimExp VName+neg = BinOpExp (Sub Int64 OverflowUndef) (val 0)++l :: Int -> PrimExp VName+l i = LeafExp (VName (nameFromString $ show i) i) (IntType Int64)++val :: Int64 -> PrimExp VName+val = ValueExp . IntValue . Int64Value++generateExp :: Gen (PrimExp VName)+generateExp = do+  n <- getSize+  if n <= 1+    then val <$> arbitrary+    else+      oneof+        [ scale (`div` 2) $ generateBinOp add,+          scale (`div` 2) $ generateBinOp sub,+          scale (`div` 2) $ generateBinOp mul,+          scale (`div` 2) generateNeg,+          val <$> arbitrary+        ]++generateBinOp :: (PrimExp VName -> PrimExp VName -> PrimExp VName) -> Gen (PrimExp VName)+generateBinOp op = do+  t1 <- generateExp+  op t1 <$> generateExp++generateNeg :: Gen (PrimExp VName)+generateNeg =+  do neg <$> generateExp++newtype TestableExp = TestableExp (PrimExp VName)+  deriving (Show)++instance Arbitrary TestableExp where+  arbitrary = TestableExp <$> generateExp

@@ -0,0 +1,68 @@+module Futhark.IR.Mem.IntervalTests+  ( tests,+  )+where++import Futhark.Analysis.AlgSimplify+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.Mem.Interval+import Futhark.IR.Syntax+import Futhark.IR.Syntax.Core ()+import Test.Tasty+import Test.Tasty.HUnit++-- Actual tests.+tests :: TestTree+tests =+  testGroup+    "IntervalTests"+    testDistributeOffset++name :: String -> Int -> VName+name s = VName (nameFromString s)++testDistributeOffset :: [TestTree]+testDistributeOffset =+  [ testCase "Stride is (nb-b)" $ do+      let n = TPrimExp $ LeafExp (name "n" 1) $ IntType Int64+          b = TPrimExp $ LeafExp (name "b" 2) $ IntType Int64+      res <-+        distributeOffset+          [Prod False [untyped (n * b - b :: TPrimExp Int64 VName)]]+          [ Interval 0 1 (n * b - b),+            Interval 0 b b,+            Interval 0 b 1+          ]+      res == [Interval 1 1 (n * b - b), Interval 0 b b, Interval 0 b 1] @? "Failed",+    testCase "Stride is 1024r" $ do+      let r = TPrimExp $ LeafExp (name "r" 1) $ IntType Int64+      res <-+        distributeOffset+          [Prod False [untyped (1024 :: TPrimExp Int64 VName), untyped r]]+          [ Interval 0 1 (1024 * r),+            Interval 0 32 32,+            Interval 0 32 1+          ]+      res == [Interval 1 1 (1024 * r), Interval 0 32 32, Interval 0 32 1] @? "Failed. Got " <> show res,+    testCase "Stride is 32, offsets are multples of 32" $ do+      let n = TPrimExp $ LeafExp (name "n" 0) $ IntType Int64+      let g1 = TPrimExp $ LeafExp (name "g" 1) $ IntType Int64+      let g2 = TPrimExp $ LeafExp (name "g" 2) $ IntType Int64+      res <-+        distributeOffset+          [ Prod False [untyped (1024 :: TPrimExp Int64 VName)],+            Prod False [untyped (1024 :: TPrimExp Int64 VName), untyped g1],+            Prod False [untyped (32 :: TPrimExp Int64 VName), untyped g2]+          ]+          [ Interval 0 1 (1024 * n),+            Interval 0 1 32,+            Interval 0 32 1+          ]+      res+        == [ Interval 0 1 (1024 * n),+             Interval (32 + 32 * g1 + g2) 1 32,+             Interval 0 32 1+           ]+        @? "Failed. Got "+          <> show res+  ]

@@ -12,9 +12,13 @@     rebase,     shape,     index,+    disjoint,   ) where +import Data.List qualified as L+import Data.Set qualified as S+import Futhark.IR.Pretty () import Futhark.IR.Prop import Futhark.IR.Syntax   ( DimIndex (..),@@ -27,7 +31,7 @@   ) import Futhark.Util.IntegralExp import Futhark.Util.Pretty-import Prelude hiding (mod)+import Prelude hiding (div, mod, span)  type Shape num = [num] @@ -111,7 +115,7 @@   shape ixfun  index ::-  (IntegralExp num, Eq num) =>+  (Eq num, IntegralExp num) =>   IxFun num ->   Indices num ->   num@@ -164,3 +168,28 @@         r@Rebase {} ->           r    in index fun' is++allPoints :: (IntegralExp num, Enum num) => [num] -> [[num]]+allPoints dims =+  let total = product dims+      strides = drop 1 $ L.reverse $ scanl (*) 1 $ L.reverse dims+   in map (unflatInd strides) [0 .. total - 1]+  where+    unflatInd strides x =+      fst $+        foldl+          ( \(res, acc) span ->+              (res ++ [acc `div` span], acc `mod` span)+          )+          ([], x)+          strides++disjoint :: (IntegralExp num, Ord num, Enum num) => IxFun num -> IxFun num -> Bool+disjoint ixf1 ixf2 =+  let shp1 = shape ixf1+      points1 = S.fromList $ allPoints shp1+      allIdxs1 = S.map (index ixf1) points1+      shp2 = shape ixf2+      points2 = S.fromList $ allPoints shp2+      allIdxs2 = S.map (index ixf2) points2+   in S.disjoint allIdxs1 allIdxs2

@@ -5,11 +5,15 @@   ) where +import Data.Function ((&)) import Data.List qualified as L+import Data.Map qualified as M+import Futhark.Analysis.PrimExp.Convert import Futhark.IR.Mem.IxFun qualified as IxFunLMAD import Futhark.IR.Mem.IxFun.Alg qualified as IxFunAlg import Futhark.IR.Mem.IxFunWrapper import Futhark.IR.Mem.IxFunWrapper qualified as IxFunWrap+import Futhark.IR.Prop import Futhark.IR.Syntax import Futhark.IR.Syntax.Core () import Futhark.Util.IntegralExp qualified as IE@@ -114,6 +118,10 @@         test_flatSlice_flatSlice_iota,         test_flatSlice_slice_iota,         test_flatSlice_transpose_slice_iota+        -- TODO: Without z3, these tests fail. Ideally, our internal simplifier+        -- should be able to handle them:+        --+        -- test_disjoint3       ]  singleton :: TestTree -> [TestTree]@@ -377,3 +385,240 @@         flatSlice (permute (slice (iota [20, 20]) $ Slice [DimSlice 1 5 2, DimSlice 0 5 2]) [1, 0]) flat_slice_1   where     flat_slice_1 = FlatSlice 1 [FlatDimIndex 2 2]++-- test_disjoint2 :: [TestTree]+-- test_disjoint2 =+--   let add_nw64 = (+)++--       mul_nw64 = (*)++--       sub64 = (-)++--       vname s i = VName (nameFromString s) i+--    in [ let gtid_8472 = TPrimExp $ LeafExp (vname "gtid" 8472) $ IntType Int64++--             gtid_8473 = TPrimExp $ LeafExp (vname "gtid" 8473) $ IntType Int64++--             gtid_8474 = TPrimExp $ LeafExp (vname "gtid" 8474) $ IntType Int64++--             num_blocks_8284 = TPrimExp $ LeafExp (vname "num_blocks" 8284) $ IntType Int64++--             nonnegs = freeIn [gtid_8472, gtid_8473, gtid_8474, num_blocks_8284]++--             j_m_i_8287 :: TPrimExp Int64 VName+--             j_m_i_8287 = num_blocks_8284 - 1++--             lessthans :: [(VName, PrimExp VName)]+--             lessthans =+--               [ (head $ namesToList $ freeIn gtid_8472, untyped j_m_i_8287),+--                 (head $ namesToList $ freeIn gtid_8473, untyped j_m_i_8287),+--                 (head $ namesToList $ freeIn gtid_8474, untyped (16 :: TPrimExp Int64 VName))+--               ]++--             lm1 :: IxFunLMAD.LMAD (TPrimExp Int64 VName)+--             lm1 =+--               IxFunLMAD.LMAD+--                 256+--                 [ IxFunLMAD.LMADDim 256 0 (sub64 (num_blocks_8284) 1) 0 IxFunLMAD.Inc,+--                   IxFunLMAD.LMADDim 1 0 16 1 IxFunLMAD.Inc,+--                   IxFunLMAD.LMADDim 16 0 16 2 IxFunLMAD.Inc+--                 ]+--             lm2 :: IxFunLMAD.LMAD (TPrimExp Int64 VName)+--             lm2 =+--               IxFunLMAD.LMAD+--                 (add_nw64 (add_nw64 (add_nw64 (add_nw64 (mul_nw64 (256) (num_blocks_8284)) (256)) (mul_nw64 (gtid_8472) (mul_nw64 (256) (num_blocks_8284)))) (mul_nw64 (gtid_8473) (256))) (mul_nw64 (gtid_8474) (16)))+--                 [IxFunLMAD.LMADDim 1 0 16 0 IxFunLMAD.Inc]+--          in testCase (pretty lm1 <> " and " <> pretty lm2) $ IxFunLMAD.disjoint2 lessthans nonnegs lm1 lm2 @? "Failed"+--       ]++-- test_lessThanish :: [TestTree]+-- test_lessThanish =+--   [testCase "0 < 1" $ IxFunLMAD.lessThanish mempty mempty 0 1 @? "Failed"]++-- test_lessThanOrEqualish :: [TestTree]+-- test_lessThanOrEqualish =+--   [testCase "1 <= 1" $ IxFunLMAD.lessThanOrEqualish mempty mempty 1 1 @? "Failed"]++_test_disjoint3 :: [TestTree]+_test_disjoint3 =+  let foo s = VName (nameFromString s)+      add_nw64 = (+)+      add64 = (+)+      mul_nw64 = (*)+      mul64 = (*)+      sub64 = (-)+      sdiv64 = IE.div+      sub_nw64 = (-)+      disjointTester asserts lessthans lm1 lm2 =+        let nonnegs = map (`LeafExp` IntType Int64) $ namesToList $ freeIn lm1 <> freeIn lm2++            scmap =+              M.fromList $+                map (\x -> (x, Prim $ IntType Int64)) $+                  namesToList $+                    freeIn lm1 <> freeIn lm2 <> freeIn lessthans <> freeIn asserts+         in IxFunLMAD.disjoint3 scmap asserts lessthans nonnegs lm1 lm2+   in [ testCase "lm1 and lm2" $+          let lessthans =+                [ ( i_12214,+                    sdiv64 (sub64 n_blab 1) block_size_12121+                  ),+                  (gtid_12553, add64 1 i_12214)+                ]+                  & map (\(v, p) -> (head $ namesToList $ freeIn v, untyped p))++              asserts =+                [ untyped ((2 * block_size_12121 :: TPrimExp Int64 VName) .<. n_blab :: TPrimExp Bool VName),+                  untyped ((3 :: TPrimExp Int64 VName) .<. n_blab :: TPrimExp Bool VName)+                ]++              block_size_12121 = TPrimExp $ LeafExp (foo "block_size" 12121) $ IntType Int64+              i_12214 = TPrimExp $ LeafExp (foo "i" 12214) $ IntType Int64+              n_blab = TPrimExp $ LeafExp (foo "n" 1337) $ IntType Int64+              gtid_12553 = TPrimExp $ LeafExp (foo "gtid" 12553) $ IntType Int64++              lm1 =+                IxFunLMAD.LMAD+                  (add_nw64 (mul64 block_size_12121 i_12214) (mul_nw64 (add_nw64 gtid_12553 1) (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (add64 1 i_12214) gtid_12553) 1) 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1 IxFunLMAD.Inc+                  ]++              lm2 =+                IxFunLMAD.LMAD+                  (block_size_12121 * i_12214)+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1 IxFunLMAD.Inc+                  ]++              lm_w =+                IxFunLMAD.LMAD+                  (add_nw64 (add64 (add64 1 n_blab) (mul64 block_size_12121 i_12214)) (mul_nw64 gtid_12553 (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))+                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 block_size_12121 1 IxFunLMAD.Inc+                  ]++              lm_blocks =+                IxFunLMAD.LMAD+                  (block_size_12121 * i_12214 + n_blab + 1)+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim n_blab block_size_12121 1 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 block_size_12121 2 IxFunLMAD.Inc+                  ]++              lm_lower_per =+                IxFunLMAD.LMAD+                  (block_size_12121 * i_12214)+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1 IxFunLMAD.Inc+                  ]++              res1 = disjointTester asserts lessthans lm1 lm_w+              res2 = disjointTester asserts lessthans lm2 lm_w+              res3 = disjointTester asserts lessthans lm_lower_per lm_blocks+           in res1 && res2 && res3 @? "Failed",+        testCase "nw second half" $ do+          let lessthans =+                [ ( i_12214,+                    sdiv64 (sub64 n_blab 1) block_size_12121+                  ),+                  (gtid_12553, add64 1 i_12214)+                ]+                  & map (\(v, p) -> (head $ namesToList $ freeIn v, untyped p))++              asserts =+                [ untyped ((2 * block_size_12121 :: TPrimExp Int64 VName) .<. n_blab :: TPrimExp Bool VName),+                  untyped ((3 :: TPrimExp Int64 VName) .<. n_blab :: TPrimExp Bool VName)+                ]++              block_size_12121 = TPrimExp $ LeafExp (foo "block_size" 12121) $ IntType Int64+              i_12214 = TPrimExp $ LeafExp (foo "i" 12214) $ IntType Int64+              n_blab = TPrimExp $ LeafExp (foo "n" 1337) $ IntType Int64+              gtid_12553 = TPrimExp $ LeafExp (foo "gtid" 12553) $ IntType Int64++              lm1 =+                IxFunLMAD.LMAD+                  (add_nw64 (add64 n_blab (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1)) (mul_nw64 (add_nw64 gtid_12553 1) (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim n_blab block_size_12121 1 IxFunLMAD.Inc+                  ]++              lm2 =+                IxFunLMAD.LMAD+                  (add_nw64 (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1) (mul_nw64 (add_nw64 gtid_12553 1) (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1 IxFunLMAD.Inc+                  ]++              lm3 =+                IxFunLMAD.LMAD+                  (add64 n_blab (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1))+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim n_blab block_size_12121 1 IxFunLMAD.Inc+                  ]++              lm4 =+                IxFunLMAD.LMAD+                  (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1)+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1 IxFunLMAD.Inc+                  ]++              lm_w =+                IxFunLMAD.LMAD+                  (add_nw64 (sub64 (mul64 n_blab (add64 2 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) (mul_nw64 gtid_12553 (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))+                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 block_size_12121 1 IxFunLMAD.Inc+                  ]++              res1 = disjointTester asserts lessthans lm1 lm_w+              res2 = disjointTester asserts lessthans lm2 lm_w+              res3 = disjointTester asserts lessthans lm3 lm_w+              res4 = disjointTester asserts lessthans lm4 lm_w+           in res1 && res2 && res3 && res4 @? "Failed " <> show [res1, res2, res3, res4],+        testCase "lud long" $+          let lessthans =+                [ (step, num_blocks - 1 :: TPrimExp Int64 VName)+                ]+                  & map (\(v, p) -> (head $ namesToList $ freeIn v, untyped p))++              step = TPrimExp $ LeafExp (foo "step" 1337) $ IntType Int64++              num_blocks = TPrimExp $ LeafExp (foo "n" 1338) $ IntType Int64++              lm1 =+                IxFunLMAD.LMAD+                  (1024 * num_blocks * (1 + step) + 1024 * step)+                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 32 32 1 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 32 2 IxFunLMAD.Inc+                  ]++              lm_w1 =+                IxFunLMAD.LMAD+                  (1024 * num_blocks * step + 1024 * step)+                  [ IxFunLMAD.LMADDim 32 32 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 32 1 IxFunLMAD.Inc+                  ]++              lm_w2 =+                IxFunLMAD.LMAD+                  ((1 + step) * 1024 * num_blocks + (1 + step) * 1024)+                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1024 (num_blocks - step - 1) 1 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1024 1 2 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 32 1 3 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 128 8 4 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 4 8 5 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 32 4 6 IxFunLMAD.Inc,+                    IxFunLMAD.LMADDim 1 4 7 IxFunLMAD.Inc+                  ]++              asserts =+                [ untyped ((1 :: TPrimExp Int64 VName) .<. num_blocks :: TPrimExp Bool VName)+                ]++              res1 = disjointTester asserts lessthans lm1 lm_w1+              res2 = disjointTester asserts lessthans lm1 lm_w2+           in res1 && res2 @? "Failed"+      ]

@@ -1,7 +1,9 @@ module Main (main) where  import Futhark.AD.DerivativesTests qualified+import Futhark.Analysis.AlgSimplifyTests qualified import Futhark.BenchTests qualified+import Futhark.IR.Mem.IntervalTests qualified import Futhark.IR.Mem.IxFunTests qualified import Futhark.IR.PropTests qualified import Futhark.IR.Syntax.CoreTests qualified@@ -22,9 +24,11 @@       Futhark.IR.PropTests.tests,       Futhark.IR.Syntax.CoreTests.tests,       Futhark.Pkg.SolveTests.tests,+      Futhark.IR.Mem.IntervalTests.tests,       Futhark.IR.Mem.IxFunTests.tests,       Language.Futhark.PrimitiveTests.tests,       Futhark.Optimise.MemoryBlockMerging.GreedyColoringTests.tests,+      Futhark.Analysis.AlgSimplifyTests.tests,       Language.Futhark.TypeCheckerTests.tests     ]