diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -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 diff --git a/src/Futhark/Analysis/AlgSimplify.hs b/src/Futhark/Analysis/AlgSimplify.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/AlgSimplify.hs @@ -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 diff --git a/src/Futhark/Analysis/Interference.hs b/src/Futhark/Analysis/Interference.hs --- a/src/Futhark/Analysis/Interference.hs +++ b/src/Futhark/Analysis/Interference.hs @@ -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 = diff --git a/src/Futhark/Analysis/LastUse.hs b/src/Futhark/Analysis/LastUse.hs --- a/src/Futhark/Analysis/LastUse.hs +++ b/src/Futhark/Analysis/LastUse.hs @@ -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 diff --git a/src/Futhark/CLI/Dev.hs b/src/Futhark/CLI/Dev.hs --- a/src/Futhark/CLI/Dev.hs +++ b/src/Futhark/CLI/Dev.hs @@ -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 diff --git a/src/Futhark/CodeGen/Backends/GenericC/Fun.hs b/src/Futhark/CodeGen/Backends/GenericC/Fun.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Fun.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Fun.hs @@ -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 + }|] + ) diff --git a/src/Futhark/CodeGen/ImpCode.hs b/src/Futhark/CodeGen/ImpCode.hs --- a/src/Futhark/CodeGen/ImpCode.hs +++ b/src/Futhark/CodeGen/ImpCode.hs @@ -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. diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs @@ -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 diff --git a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs @@ -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 diff --git a/src/Futhark/IR/GPU/Op.hs b/src/Futhark/IR/GPU/Op.hs --- a/src/Futhark/IR/GPU/Op.hs +++ b/src/Futhark/IR/GPU/Op.hs @@ -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 diff --git a/src/Futhark/IR/Mem/Interval.hs b/src/Futhark/IR/Mem/Interval.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/Mem/Interval.hs @@ -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 diff --git a/src/Futhark/IR/Mem/IxFun.hs b/src/Futhark/IR/Mem/IxFun.hs --- a/src/Futhark/IR/Mem/IxFun.hs +++ b/src/Futhark/IR/Mem/IxFun.hs @@ -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. -- diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs --- a/src/Futhark/IR/SOACS/Simplify.hs +++ b/src/Futhark/IR/SOACS/Simplify.hs @@ -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. diff --git a/src/Futhark/Internalise/Defunctionalise.hs b/src/Futhark/Internalise/Defunctionalise.hs --- a/src/Futhark/Internalise/Defunctionalise.hs +++ b/src/Futhark/Internalise/Defunctionalise.hs @@ -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. diff --git a/src/Futhark/Internalise/Defunctorise.hs b/src/Futhark/Internalise/Defunctorise.hs --- a/src/Futhark/Internalise/Defunctorise.hs +++ b/src/Futhark/Internalise/Defunctorise.hs @@ -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 = diff --git a/src/Futhark/Optimise/ArrayShortCircuiting.hs b/src/Futhark/Optimise/ArrayShortCircuiting.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ArrayShortCircuiting.hs @@ -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 diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs @@ -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 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 diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs b/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs @@ -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 diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/LastUse.hs b/src/Futhark/Optimise/ArrayShortCircuiting/LastUse.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ArrayShortCircuiting/LastUse.hs @@ -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' diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs b/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs @@ -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 + 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 + 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 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 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) diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs b/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs @@ -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) diff --git a/src/Futhark/Optimise/CSE.hs b/src/Futhark/Optimise/CSE.hs --- a/src/Futhark/Optimise/CSE.hs +++ b/src/Futhark/Optimise/CSE.hs @@ -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 diff --git a/src/Futhark/Optimise/Simplify/Rule.hs b/src/Futhark/Optimise/Simplify/Rule.hs --- a/src/Futhark/Optimise/Simplify/Rule.hs +++ b/src/Futhark/Optimise/Simplify/Rule.hs @@ -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 -> diff --git a/src/Futhark/Optimise/Simplify/Rules.hs b/src/Futhark/Optimise/Simplify/Rules.hs --- a/src/Futhark/Optimise/Simplify/Rules.hs +++ b/src/Futhark/Optimise/Simplify/Rules.hs @@ -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 diff --git a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs --- a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +++ b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs @@ -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 <- diff --git a/src/Futhark/Pass/LiftAllocations.hs b/src/Futhark/Pass/LiftAllocations.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/LiftAllocations.hs @@ -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 diff --git a/src/Futhark/Pass/LowerAllocations.hs b/src/Futhark/Pass/LowerAllocations.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/LowerAllocations.hs @@ -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 diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs --- a/src/Futhark/Passes.hs +++ b/src/Futhark/Passes.hs @@ -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, diff --git a/src/Futhark/Script.hs b/src/Futhark/Script.hs --- a/src/Futhark/Script.hs +++ b/src/Futhark/Script.hs @@ -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 diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs --- a/src/Futhark/Util.hs +++ b/src/Futhark/Util.hs @@ -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 diff --git a/src/Futhark/Util/IntegralExp.hs b/src/Futhark/Util/IntegralExp.hs --- a/src/Futhark/Util/IntegralExp.hs +++ b/src/Futhark/Util/IntegralExp.hs @@ -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 -> diff --git a/src/Futhark/Version.hs b/src/Futhark/Version.hs --- a/src/Futhark/Version.hs +++ b/src/Futhark/Version.hs @@ -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 -> "" diff --git a/src/Language/Futhark/Pretty.hs b/src/Language/Futhark/Pretty.hs --- a/src/Language/Futhark/Pretty.hs +++ b/src/Language/Futhark/Pretty.hs @@ -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 diff --git a/src/Language/Futhark/Prop.hs b/src/Language/Futhark/Prop.hs --- a/src/Language/Futhark/Prop.hs +++ b/src/Language/Futhark/Prop.hs @@ -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. diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -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) diff --git a/src/Language/Futhark/TypeChecker/Terms/DoLoop.hs b/src/Language/Futhark/TypeChecker/Terms/DoLoop.hs --- a/src/Language/Futhark/TypeChecker/Terms/DoLoop.hs +++ b/src/Language/Futhark/TypeChecker/Terms/DoLoop.hs @@ -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 + ) diff --git a/src/Language/Futhark/TypeChecker/Terms/Monad.hs b/src/Language/Futhark/TypeChecker/Terms/Monad.hs --- a/src/Language/Futhark/TypeChecker/Terms/Monad.hs +++ b/src/Language/Futhark/TypeChecker/Terms/Monad.hs @@ -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). diff --git a/src/Language/Futhark/TypeChecker/Terms/Pat.hs b/src/Language/Futhark/TypeChecker/Terms/Pat.hs --- a/src/Language/Futhark/TypeChecker/Terms/Pat.hs +++ b/src/Language/Futhark/TypeChecker/Terms/Pat.hs @@ -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 -> diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs --- a/src/Language/Futhark/TypeChecker/Types.hs +++ b/src/Language/Futhark/TypeChecker/Types.hs @@ -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)) diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs --- a/src/Language/Futhark/TypeChecker/Unify.hs +++ b/src/Language/Futhark/TypeChecker/Unify.hs @@ -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 => diff --git a/unittests/Futhark/Analysis/AlgSimplifyTests.hs b/unittests/Futhark/Analysis/AlgSimplifyTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Analysis/AlgSimplifyTests.hs @@ -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 diff --git a/unittests/Futhark/IR/Mem/IntervalTests.hs b/unittests/Futhark/IR/Mem/IntervalTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/IR/Mem/IntervalTests.hs @@ -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 + ] diff --git a/unittests/Futhark/IR/Mem/IxFun/Alg.hs b/unittests/Futhark/IR/Mem/IxFun/Alg.hs --- a/unittests/Futhark/IR/Mem/IxFun/Alg.hs +++ b/unittests/Futhark/IR/Mem/IxFun/Alg.hs @@ -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 diff --git a/unittests/Futhark/IR/Mem/IxFunTests.hs b/unittests/Futhark/IR/Mem/IxFunTests.hs --- a/unittests/Futhark/IR/Mem/IxFunTests.hs +++ b/unittests/Futhark/IR/Mem/IxFunTests.hs @@ -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" + ] diff --git a/unittests/futhark_tests.hs b/unittests/futhark_tests.hs --- a/unittests/futhark_tests.hs +++ b/unittests/futhark_tests.hs @@ -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 ]